Let a calibration pull the head columns it reads - #123
Merged
RobbinBouwmeester merged 3 commits intoSep 8, 2026
Conversation
A multitask model returns one column per LC setup, so the query matrix is 26 kB per peptide at 6,543 setups, and a fitted calibration reads a few dozen of those columns: eighty for MultiHeadRidgeCalibration, one for a spline. The rest were predicted, promoted to float64 and never looked at, which at 10,000 queries is 262 MB from the model and 523 MB after the cast. The caller still hands over exactly one source and branches on nothing, which is what dropping uses_all_heads bought. What changes is that the source may be a column provider rather than a materialised matrix: - take_columns(source, indices) asks a provider for those heads, or indexes a matrix, and every MultiHeadCalibration reads its columns through it. The request is made once for all the heads a calibration uses, so a provider needs one forward pass rather than one per head. - HeadColumnSource in core.py is that provider for a model and a peptide list. It reports its shape without predicting anything, caches the last head set it was asked for (prediction_report reads the same heads twice, once to transform and once for the head disagreement), and implements __array__ so code that genuinely needs every head, such as ranking them in fit(), still gets the whole matrix. - predict_and_calibrate and prediction_report hand over that source for their queries. References still pass a real matrix: ranking reads every head, and a reference is small. The blanket float64 promotion at the top of each transform goes with it, since the columns are cast after they are selected rather than before. Verified on PXD081924: MultiHeadRidgeCalibration gives MAE 0.2718 and coverage 0.9248 either way, and a naive SplineTransformerCalibration gives 0.3513 and 0.9320 while running 1.18 s against 6.14 s, because it no longer predicts 6,543 heads to read one. 189 tests pass, including two that assert a lazy source and a matrix agree and that only the used heads are requested. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Reading the shape off the source, rather than coercing it first, withdrew two things the previous ``np.asarray(source, dtype=np.float64)`` had quietly provided. A list or tuple of predictions raised AttributeError, because only arrays and Series carry ``.shape``. ``source_shape`` now falls back to ``np.asarray`` when a source does not report its own shape, so anything numpy accepts works again while a lazy provider is still asked rather than materialised. A pandas DataFrame was mistaken for a lazy provider: it has a ``columns`` attribute, so the duck-typing check found it and tried to call it. The provider method is now ``head_columns``, which nothing else is likely to define, and the check requires it to be callable. Both are covered by a parametrised test over the forms a caller can hand in: two-dimensional array, one-dimensional array from a single-task model, list, tuple, integer dtype, Series and DataFrame. 196 tests pass. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The first version of this rewrote both transform methods around two helpers,
which is more divergence than the change needs. The coercion at the top of
each one now survives character for character, moved into as_head_matrix and
skipped only for a source that marks itself with is_head_source:
if getattr(source, "is_head_source", False):
return source
source = np.asarray(source, dtype=np.float64)
if source.ndim == 1:
source = source[:, None]
return source
Everything after that line is left as written - the shape checks, their error
messages, the empty-source check, the column stacking - because a head source
answers .shape and source[:, heads] the way an array does. HeadColumnSource
therefore implements __getitem__ instead of a bespoke accessor, and the two
calibrations index it exactly as they index a matrix. fit() keeps the real
coercion: ranking reads every head, and np.asarray on a head source yields
the whole matrix through __array__.
The only other change to a body is that MultiHeadRidgeCalibration asks for
its heads in one slice rather than one per head, so a lazy source needs a
single forward pass; for an array that is the same slice.
That drops the divergence from Ralf's file to 30 added and 12 removed lines,
of which 17 are the new helper and its docstring. Same numbers on PXD081924
(MAE 0.2718, coverage 0.9248, 1,050 widths; naive spline 0.3513 and 0.9320),
every array-like still accepted including a DataFrame, 196 tests pass.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #118. One commit.
The waste
A multitask model returns one column per LC setup. At 6,543 setups the query matrix is 26 kB
per peptide — 262 MB for 10,000 queries, 523 MB once
transformpromotes it to float64 — anda fitted calibration reads a few dozen of those columns: eighty for
MultiHeadRidgeCalibration, one for a spline. The rest are predicted, cast and never lookedat. It is also what forces the reference and test caps in our own analysis scripts.
Why not simply ask the caller to predict fewer heads
That was my first attempt, and it reintroduces exactly what #121 removed: the caller would
have to ask the calibration which heads it needs and then branch between two transform paths.
uses_all_headswith a different name.What this does instead
The caller still hands over one source for its queries and branches on nothing. What
changes is that the source may be a column provider rather than a materialised matrix, and a
calibration asks it for the columns it selected.
take_columns(source, indices)incalibration/multihead.pyeither asks a provider forthose heads or indexes a matrix. Every
MultiHeadCalibrationreads through it, so_SingleHeadCalibrationasks for its one head and the ridge for its eighty — in a singlerequest, so a provider does one forward pass rather than one per head.
HeadColumnSourceincore.pyis that provider for a model and a peptide list. Itreports its shape without predicting anything, caches the last head set it was asked for
(
prediction_reportreads the same heads twice, once to transform and once for the headdisagreement), and implements
__array__so code that genuinely needs every head — rankingthem in
fit()— still gets the whole matrix by writingnp.asarray(source).predict_and_calibrateandprediction_reporthand over that source for their queries.References still pass a real matrix: ranking reads every head, and a reference is small.
The blanket
np.asarray(source, dtype=np.float64)at the top of eachtransformgoes withit, since columns are now cast after they are selected rather than before.
Measured
On PXD081924, 1,103 queries against a 1,103-peptide reference:
MultiHeadRidgeCalibrationSplineTransformerCalibration, upgradedIdentical numbers either way. The naive path got markedly faster for free: it used to predict
6,543 heads in order to read one.
189 tests pass, ruff clean. Two new tests assert that a lazy source and a matrix give the same
answer for both calibration kinds and for the head disagreement, and that only the heads a
calibration uses are ever requested.
Open to a different shape
take_columnsandsource_shapeare module-level inmultihead.py; they could be private orlive in their own module. The provider is detected by
hasattr(source, "columns")rather thana formal
Protocol, to keep the public surface small. Both are happy to change.🤖 Generated with Claude Code