From 8a19b6283d2829946334c3f897a475aa2e02ffd5 Mon Sep 17 00:00:00 2001 From: RobbinBouwmeester Date: Tue, 8 Sep 2026 16:26:09 +0200 Subject: [PATCH 1/3] feat: let a calibration pull the head columns it reads 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) --- deeplc/calibration/multihead.py | 72 ++++++++++++++-------- deeplc/core.py | 92 +++++++++++++++++++++++++++-- deeplc/report.py | 7 ++- tests/test_multihead_calibration.py | 62 +++++++++++++++++++ 4 files changed, 201 insertions(+), 32 deletions(-) diff --git a/deeplc/calibration/multihead.py b/deeplc/calibration/multihead.py index 3382447..ea4836d 100644 --- a/deeplc/calibration/multihead.py +++ b/deeplc/calibration/multihead.py @@ -13,6 +13,7 @@ import logging from abc import ABC, abstractmethod +from collections.abc import Sequence from typing import cast import numpy as np @@ -28,6 +29,33 @@ LOGGER = logging.getLogger(__name__) +def take_columns(source, indices: Sequence[int]) -> np.ndarray: + """ + Take the named head columns from a source, as float64 of shape ``(n, len(indices))``. + + The source is normally the ``(n, n_heads)`` matrix a model returned. It may instead be an + object offering ``columns(indices)``, such as :class:`deeplc.core.HeadColumnSource`, which + evaluates only the heads asked for: a calibration reads a few dozen of the thousands a + multitask model has, and at 6,543 setups the unread columns are 26 kB per peptide. Which + of the two it is makes no difference to a calibration, and the caller hands over the same + thing either way. + """ + if hasattr(source, "columns"): + taken = source.columns(indices) + else: + matrix = np.asarray(source) + if matrix.ndim == 1: + matrix = matrix[:, None] + taken = matrix[:, list(indices)] + return np.asarray(taken, dtype=np.float64) + + +def source_shape(source) -> tuple[int, int]: + """Rows and head count of a source, without materialising a lazy one.""" + shape = tuple(source.shape) + return (shape[0], shape[1] if len(shape) > 1 else 1) + + class MultiHeadCalibration(ABC): """Abstract base class for a calibration that selects its own head(s) from a matrix.""" @@ -131,20 +159,17 @@ def transform(self, source: np.ndarray) -> np.ndarray: """ if not self.is_fitted: raise CalibrationError("The model has not been fitted yet. Call fit() first.") - source = np.asarray(source, dtype=np.float64) - if source.ndim == 1: - source = source[:, None] head = self.selected_model_head - if source.shape[1] <= head: + rows, n_heads = source_shape(source) + if n_heads <= head: raise CalibrationError( - f"source has {source.shape[1]} heads, but the calibration was fitted on a model " + f"source has {n_heads} heads, but the calibration was fitted on a model " f"with at least {head + 1}." ) - if source.shape[0] == 0: + if rows == 0: return np.array([]) - return np.asarray( - self._inner.transform(source[:, head].astype(np.float32)), dtype=np.float64 - ) + column = take_columns(source, [head])[:, 0] + return np.asarray(self._inner.transform(column.astype(np.float32)), dtype=np.float64) class MultiHeadPiecewiseLinearCalibration(_SingleHeadCalibration): @@ -296,26 +321,29 @@ def transform(self, source: np.ndarray) -> np.ndarray: """ if not self.is_fitted: raise CalibrationError("The model has not been fitted yet. Call fit() first.") - source = np.asarray(source, dtype=np.float64) - if source.ndim == 1: - source = source[:, None] head_idx = cast(np.ndarray, self._head_idx) - if source.shape[1] <= int(head_idx.max()): + rows, n_heads = source_shape(source) + if n_heads <= int(head_idx.max()): raise CalibrationError( - f"source has {source.shape[1]} heads, but the calibration was fitted on a model " + f"source has {n_heads} heads, but the calibration was fitted on a model " f"with at least {int(head_idx.max()) + 1}." ) - if source.shape[0] == 0: + if rows == 0: return np.array([]) return np.asarray(self._ridge.predict(self._calibrated_columns(source)), dtype=np.float64) - def _calibrated_columns(self, source: np.ndarray) -> np.ndarray: + def _calibrated_columns(self, source) -> np.ndarray: """Give each selected head's own estimate of the retention time, in reference units.""" head_idx = cast(np.ndarray, self._head_idx) + # One request for every selected head, so a lazy source evaluates them in a single + # pass rather than once per head. + columns = take_columns(source, head_idx) return np.column_stack( [ - np.asarray(cal.transform(source[:, head].astype(np.float32)), dtype=np.float64) - for cal, head in zip(self._head_calibrations, head_idx, strict=True) + np.asarray( + cal.transform(columns[:, position].astype(np.float32)), dtype=np.float64 + ) + for position, cal in enumerate(self._head_calibrations) ] ) @@ -330,15 +358,13 @@ def disagreement(self, source: np.ndarray) -> np.ndarray | None: """ if not self.is_fitted: return None - columns = np.asarray(source, dtype=np.float64) - if columns.ndim == 1: - columns = columns[:, None] - if columns.shape[0] == 0 or len(cast(np.ndarray, self._head_idx)) < 2: + rows, _ = source_shape(source) + if rows == 0 or len(cast(np.ndarray, self._head_idx)) < 2: return None weights = np.abs(np.asarray(self._ridge.coef_, dtype=np.float64).ravel()) total = weights.sum() weights = weights / total if total > 0 else np.full(len(weights), 1 / len(weights)) - estimates = self._calibrated_columns(columns) + estimates = self._calibrated_columns(source) mean = estimates @ weights return np.sqrt(((estimates - mean[:, None]) ** 2) @ weights) diff --git a/deeplc/core.py b/deeplc/core.py index f3e1714..155dc63 100644 --- a/deeplc/core.py +++ b/deeplc/core.py @@ -214,6 +214,89 @@ def calibrate( return calibration +class HeadColumnSource: + """ + A model's predictions for whichever heads are asked for, evaluated on demand. + + Stands in for the ``(n, n_heads)`` matrix wherever a calibration is given its source. A + multitask model has one head per LC setup, so that matrix is 26 kB per peptide at 6,543 + setups, and a fitted calibration reads a few dozen columns of it; asking the model for + those columns instead costs 320 bytes per peptide and skips the rest of the head layer. + + Passing this or a real matrix makes no difference to the calibration, and none to the + caller, which hands over one source either way. ``np.asarray`` on it still yields the + whole matrix, so code that genuinely needs every head, such as ranking them during + ``fit``, keeps working. + + Parameters + ---------- + psm_list + The peptides to predict. + model + Model or path, as :func:`predict` takes it. + predict_kwargs + Extra arguments for the prediction, such as the device and batch size. + n_heads + How many heads the model has, so the shape is known without predicting anything. + + """ + + def __init__( + self, psm_list, model=None, predict_kwargs: dict | None = None, n_heads: int | None = None + ): + """Initialize the source; nothing is predicted until a column is asked for.""" + self._psm_list = _parse_psms(psm_list) + self._model = model + self._predict_kwargs = dict(predict_kwargs or {}) + loaded = _model_ops.load_model( + model or DEFAULT_MODEL, device=self._predict_kwargs.get("device") + ) + self._n_heads = int(n_heads if n_heads is not None else getattr(loaded, "n_tasks", 1)) + self._cache: tuple[tuple[int, ...], np.ndarray] | None = None + + @property + def shape(self) -> tuple[int, int]: + """Rows and head count, without evaluating anything.""" + return (len(self._psm_list), self._n_heads) + + @property + def ndim(self) -> int: + """Always two: this stands in for a matrix.""" + return 2 + + def columns(self, indices) -> np.ndarray: + """Predictions for the given heads, shape ``(n, len(indices))``, in that order.""" + wanted = tuple(int(i) for i in indices) + # Callers ask for the same heads more than once - prediction_report transforms the + # queries and then asks the same calibration for its head disagreement - and each ask + # would otherwise repeat the forward pass. + if self._cache is None or self._cache[0] != wanted: + self._cache = ( + wanted, + predict( + self._psm_list, + model=self._model, + predict_kwargs={**self._predict_kwargs, "task_idx": list(wanted)}, + return_matrix=True, + ), + ) + return self._cache[1] + + def __array__(self, dtype=None, copy=None) -> np.ndarray: + """Every head, for the callers that really need the whole matrix.""" + matrix = predict( + self._psm_list, + model=self._model, + predict_kwargs=self._predict_kwargs, + return_matrix=True, + ) + return matrix if dtype is None else matrix.astype(dtype) + + def __len__(self) -> int: + """Return the number of peptides.""" + return len(self._psm_list) + + def predict_and_calibrate( psm_list: PSMList | list[PSM | Peptidoform | str], psm_list_reference: PSMList | list[PSM | Peptidoform | str] | None = None, @@ -261,12 +344,9 @@ def predict_and_calibrate( # Predict initial retention times LOGGER.info("Predicting retention times...") - predicted_rt = predict( - psm_list=parsed_psm_list, - model=model, - predict_kwargs=predict_kwargs, - return_matrix=True, - ) + # A source rather than a matrix: the calibration pulls the heads it reads, which for a + # multitask model is a few dozen of thousands. + predicted_rt = HeadColumnSource(parsed_psm_list, model=model, predict_kwargs=predict_kwargs) if calibration is not None: calibration = upgrade_calibration(calibration) diff --git a/deeplc/report.py b/deeplc/report.py index 74424ae..190c4c1 100644 --- a/deeplc/report.py +++ b/deeplc/report.py @@ -434,9 +434,10 @@ def prediction_report( matrix_reference = core.predict( reference, model=model, predict_kwargs=predict_kwargs, return_matrix=True ).astype(np.float64) - matrix_query = core.predict( - parsed, model=model, predict_kwargs=predict_kwargs, return_matrix=True - ).astype(np.float64) + # The queries are handed over as a source rather than a matrix: the calibration asks it + # for the heads it reads, which for a fitted MultiHeadRidgeCalibration is eighty of 6,543. + # Materialising all of them costs 26 kB per peptide and none of it is read. + matrix_query = core.HeadColumnSource(parsed, model=model, predict_kwargs=predict_kwargs) y_reference = np.array(reference["retention_time"], dtype=np.float64) import copy diff --git a/tests/test_multihead_calibration.py b/tests/test_multihead_calibration.py index 9655dea..b6ece5b 100644 --- a/tests/test_multihead_calibration.py +++ b/tests/test_multihead_calibration.py @@ -380,3 +380,65 @@ def test_core_rejects_a_fitted_naive_calibration(): calibration=naive, predict_kwargs={"device": "cpu"}, ) + + +class _CountingSource: + """A column source that records which heads were asked for.""" + + def __init__(self, matrix: np.ndarray): + self._matrix = matrix + self.requests: list[tuple[int, ...]] = [] + + @property + def shape(self): + return self._matrix.shape + + @property + def ndim(self): + return 2 + + def columns(self, indices) -> np.ndarray: + wanted = tuple(int(i) for i in indices) + self.requests.append(wanted) + return self._matrix[:, list(wanted)] + + def __array__(self, dtype=None, copy=None): + raise AssertionError("the whole matrix should not be materialised") + + +def test_column_source_matches_a_matrix(): + """ + A calibration must not care whether its source is a matrix or a column provider. + + That equivalence is what lets the caller hand over one source and never branch on the + calibration, while a multitask model evaluates only the heads that get read. + """ + rng = np.random.RandomState(0) + source = rng.randn(200, 300) * 5 + 40 + target = source[:, 11] * 1.1 + 2 + rng.randn(200) * 0.1 + query = rng.randn(40, 300) * 5 + 40 + + for calibration in (MultiHeadRidgeCalibration(n_heads=12), SplineTransformerCalibration()): + fitted = upgrade_calibration(calibration) + fitted.fit(target, source) + lazy = _CountingSource(query) + np.testing.assert_allclose(fitted.transform(lazy), fitted.transform(query), atol=1e-8) + # every read is one request for all the heads that calibration uses + assert len(lazy.requests) == 1 + assert len(lazy.requests[0]) == len(getattr(fitted, "_head_idx", [0])) + + +def test_column_source_serves_the_disagreement_too(): + """The per-peptide spread reads the same columns, so it works off a lazy source as well.""" + rng = np.random.RandomState(1) + source = rng.randn(200, 120) * 5 + 40 + target = source[:, 3] * 0.9 + 1 + rng.randn(200) * 0.2 + query = rng.randn(30, 120) * 5 + 40 + + calibration = MultiHeadRidgeCalibration(n_heads=10) + calibration.fit(target, source) + np.testing.assert_allclose( + calibration.disagreement(_CountingSource(query)), + calibration.disagreement(query), + atol=1e-8, + ) From 9ad7f56ea94f25d2e59016b0778a99a71c81af4c Mon Sep 17 00:00:00 2001 From: RobbinBouwmeester Date: Tue, 8 Sep 2026 16:44:38 +0200 Subject: [PATCH 2/3] fix: keep taking every array-like transform used to take 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) --- deeplc/calibration/multihead.py | 31 +++++++++++++++++++++------- deeplc/core.py | 2 +- tests/test_multihead_calibration.py | 32 ++++++++++++++++++++++++++++- 3 files changed, 56 insertions(+), 9 deletions(-) diff --git a/deeplc/calibration/multihead.py b/deeplc/calibration/multihead.py index ea4836d..61bb462 100644 --- a/deeplc/calibration/multihead.py +++ b/deeplc/calibration/multihead.py @@ -34,14 +34,18 @@ def take_columns(source, indices: Sequence[int]) -> np.ndarray: Take the named head columns from a source, as float64 of shape ``(n, len(indices))``. The source is normally the ``(n, n_heads)`` matrix a model returned. It may instead be an - object offering ``columns(indices)``, such as :class:`deeplc.core.HeadColumnSource`, which - evaluates only the heads asked for: a calibration reads a few dozen of the thousands a - multitask model has, and at 6,543 setups the unread columns are 26 kB per peptide. Which + object offering ``head_columns(indices)``, such as :class:`deeplc.core.HeadColumnSource`, + which evaluates only the heads asked for: a calibration reads a few dozen of the thousands + a multitask model has, and at 6,543 setups the unread columns are 26 kB per peptide. Which of the two it is makes no difference to a calibration, and the caller hands over the same thing either way. + + The method is called ``head_columns`` rather than ``columns`` because a pandas DataFrame + has a ``columns`` attribute, and a caller passing one deserves to have it read as a matrix + rather than mistaken for a lazy provider. """ - if hasattr(source, "columns"): - taken = source.columns(indices) + if callable(getattr(source, "head_columns", None)): + taken = source.head_columns(indices) else: matrix = np.asarray(source) if matrix.ndim == 1: @@ -51,8 +55,21 @@ def take_columns(source, indices: Sequence[int]) -> np.ndarray: def source_shape(source) -> tuple[int, int]: - """Rows and head count of a source, without materialising a lazy one.""" - shape = tuple(source.shape) + """ + Give the rows and head count of a source, without materialising a lazy one. + + Anything array-like is accepted, a list of predictions included: ``transform`` used to + coerce its argument with ``np.asarray`` before reading a shape off it, and that let + callers pass whatever numpy would take. A source that reports its own shape, such as a + lazy column provider, is asked rather than converted. A one-dimensional source is one + head, which is what a single-task model returns. + """ + shape = getattr(source, "shape", None) + if shape is None: + shape = np.asarray(source).shape + shape = tuple(shape) + if not shape: + raise CalibrationError("source has no rows to calibrate") return (shape[0], shape[1] if len(shape) > 1 else 1) diff --git a/deeplc/core.py b/deeplc/core.py index 155dc63..1485bf7 100644 --- a/deeplc/core.py +++ b/deeplc/core.py @@ -264,7 +264,7 @@ def ndim(self) -> int: """Always two: this stands in for a matrix.""" return 2 - def columns(self, indices) -> np.ndarray: + def head_columns(self, indices) -> np.ndarray: """Predictions for the given heads, shape ``(n, len(indices))``, in that order.""" wanted = tuple(int(i) for i in indices) # Callers ask for the same heads more than once - prediction_report transforms the diff --git a/tests/test_multihead_calibration.py b/tests/test_multihead_calibration.py index b6ece5b..2607cbf 100644 --- a/tests/test_multihead_calibration.py +++ b/tests/test_multihead_calibration.py @@ -3,6 +3,7 @@ from __future__ import annotations import numpy as np +import pandas as pd import pytest from psm_utils import PSM, PSMList @@ -397,7 +398,7 @@ def shape(self): def ndim(self): return 2 - def columns(self, indices) -> np.ndarray: + def head_columns(self, indices) -> np.ndarray: wanted = tuple(int(i) for i in indices) self.requests.append(wanted) return self._matrix[:, list(wanted)] @@ -442,3 +443,32 @@ def test_column_source_serves_the_disagreement_too(): calibration.disagreement(query), atol=1e-8, ) + + +@pytest.mark.parametrize("wrap", [ + pytest.param(lambda column: column[:, None], id="2-D array"), + pytest.param(lambda column: column, id="1-D array"), + pytest.param(lambda column: [float(v) for v in column], id="list"), + pytest.param(lambda column: tuple(float(v) for v in column), id="tuple"), + pytest.param(lambda column: column.astype(int), id="int array"), + pytest.param(lambda column: pd.Series(column), id="pandas Series"), + pytest.param(lambda column: pd.DataFrame({"head": column}), id="pandas DataFrame"), +]) +def test_transform_takes_whatever_numpy_takes(wrap): + """ + Every array-like a caller could hand to transform keeps working. + + ``transform`` used to coerce its argument with ``np.asarray(source, dtype=np.float64)`` + before touching it, which quietly accepted a list, a tuple, a one-dimensional array from a + single-task model, or an integer dtype. Reading the shape off the source directly, so a + lazy provider is not materialised, must not withdraw that. + """ + rng = np.random.RandomState(3) + source = rng.randn(120, 1) * 5 + 40 + calibration = MultiHeadRidgeCalibration(n_heads=1) + calibration.fit(source[:, 0] * 1.1 + 2, source) + + column = rng.randn(10) * 5 + 40 + out = calibration.transform(wrap(column)) + assert np.shape(out) == (10,) + assert np.isfinite(out).all() From 1633e3eb52b3b8310f53b22d1535ca7ff8a74075 Mon Sep 17 00:00:00 2001 From: RobbinBouwmeester Date: Tue, 8 Sep 2026 16:55:09 +0200 Subject: [PATCH 3/3] refactor: keep the existing transform bodies as they were 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) --- deeplc/calibration/multihead.py | 97 +++++++++++------------------ deeplc/core.py | 30 ++++++--- tests/test_multihead_calibration.py | 13 ++-- 3 files changed, 68 insertions(+), 72 deletions(-) diff --git a/deeplc/calibration/multihead.py b/deeplc/calibration/multihead.py index 61bb462..24bd02a 100644 --- a/deeplc/calibration/multihead.py +++ b/deeplc/calibration/multihead.py @@ -13,7 +13,6 @@ import logging from abc import ABC, abstractmethod -from collections.abc import Sequence from typing import cast import numpy as np @@ -29,48 +28,22 @@ LOGGER = logging.getLogger(__name__) -def take_columns(source, indices: Sequence[int]) -> np.ndarray: +def as_head_matrix(source): """ - Take the named head columns from a source, as float64 of shape ``(n, len(indices))``. - - The source is normally the ``(n, n_heads)`` matrix a model returned. It may instead be an - object offering ``head_columns(indices)``, such as :class:`deeplc.core.HeadColumnSource`, - which evaluates only the heads asked for: a calibration reads a few dozen of the thousands - a multitask model has, and at 6,543 setups the unread columns are 26 kB per peptide. Which - of the two it is makes no difference to a calibration, and the caller hands over the same - thing either way. - - The method is called ``head_columns`` rather than ``columns`` because a pandas DataFrame - has a ``columns`` attribute, and a caller passing one deserves to have it read as a matrix - rather than mistaken for a lazy provider. - """ - if callable(getattr(source, "head_columns", None)): - taken = source.head_columns(indices) - else: - matrix = np.asarray(source) - if matrix.ndim == 1: - matrix = matrix[:, None] - taken = matrix[:, list(indices)] - return np.asarray(taken, dtype=np.float64) - + Coerce a source to a two-dimensional float64 matrix, unless it is a head source. -def source_shape(source) -> tuple[int, int]: + Anything array-like is converted exactly as before, so a list, a tuple or the + one-dimensional output of a single-task model all keep working. An object that marks + itself with ``is_head_source``, such as :class:`deeplc.core.HeadColumnSource`, is passed + through: it answers ``.shape`` and ``source[:, heads]`` like an array but evaluates only + the heads that are asked for, which for a multitask model is a few dozen of thousands. """ - Give the rows and head count of a source, without materialising a lazy one. - - Anything array-like is accepted, a list of predictions included: ``transform`` used to - coerce its argument with ``np.asarray`` before reading a shape off it, and that let - callers pass whatever numpy would take. A source that reports its own shape, such as a - lazy column provider, is asked rather than converted. A one-dimensional source is one - head, which is what a single-task model returns. - """ - shape = getattr(source, "shape", None) - if shape is None: - shape = np.asarray(source).shape - shape = tuple(shape) - if not shape: - raise CalibrationError("source has no rows to calibrate") - return (shape[0], shape[1] if len(shape) > 1 else 1) + 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 class MultiHeadCalibration(ABC): @@ -176,17 +149,21 @@ def transform(self, source: np.ndarray) -> np.ndarray: """ if not self.is_fitted: raise CalibrationError("The model has not been fitted yet. Call fit() first.") + source = as_head_matrix(source) head = self.selected_model_head - rows, n_heads = source_shape(source) - if n_heads <= head: + if source.shape[1] <= head: raise CalibrationError( - f"source has {n_heads} heads, but the calibration was fitted on a model " + f"source has {source.shape[1]} heads, but the calibration was fitted on a model " f"with at least {head + 1}." ) - if rows == 0: + if source.shape[0] == 0: return np.array([]) - column = take_columns(source, [head])[:, 0] - return np.asarray(self._inner.transform(column.astype(np.float32)), dtype=np.float64) + return np.asarray( + self._inner.transform( + np.asarray(source[:, head], dtype=np.float64).astype(np.float32) + ), + dtype=np.float64, + ) class MultiHeadPiecewiseLinearCalibration(_SingleHeadCalibration): @@ -338,29 +315,27 @@ def transform(self, source: np.ndarray) -> np.ndarray: """ if not self.is_fitted: raise CalibrationError("The model has not been fitted yet. Call fit() first.") + source = as_head_matrix(source) head_idx = cast(np.ndarray, self._head_idx) - rows, n_heads = source_shape(source) - if n_heads <= int(head_idx.max()): + if source.shape[1] <= int(head_idx.max()): raise CalibrationError( - f"source has {n_heads} heads, but the calibration was fitted on a model " + f"source has {source.shape[1]} heads, but the calibration was fitted on a model " f"with at least {int(head_idx.max()) + 1}." ) - if rows == 0: + if source.shape[0] == 0: return np.array([]) return np.asarray(self._ridge.predict(self._calibrated_columns(source)), dtype=np.float64) - def _calibrated_columns(self, source) -> np.ndarray: + def _calibrated_columns(self, source: np.ndarray) -> np.ndarray: """Give each selected head's own estimate of the retention time, in reference units.""" head_idx = cast(np.ndarray, self._head_idx) - # One request for every selected head, so a lazy source evaluates them in a single - # pass rather than once per head. - columns = take_columns(source, head_idx) + # Asked for in one go rather than head by head: a lazy source then evaluates them in + # a single pass, and for an array this is the same slice. + columns = np.asarray(source[:, head_idx], dtype=np.float64) return np.column_stack( [ - np.asarray( - cal.transform(columns[:, position].astype(np.float32)), dtype=np.float64 - ) - for position, cal in enumerate(self._head_calibrations) + np.asarray(cal.transform(columns[:, i].astype(np.float32)), dtype=np.float64) + for i, cal in enumerate(self._head_calibrations) ] ) @@ -375,13 +350,13 @@ def disagreement(self, source: np.ndarray) -> np.ndarray | None: """ if not self.is_fitted: return None - rows, _ = source_shape(source) - if rows == 0 or len(cast(np.ndarray, self._head_idx)) < 2: + columns = as_head_matrix(source) + if columns.shape[0] == 0 or len(cast(np.ndarray, self._head_idx)) < 2: return None weights = np.abs(np.asarray(self._ridge.coef_, dtype=np.float64).ravel()) total = weights.sum() weights = weights / total if total > 0 else np.full(len(weights), 1 / len(weights)) - estimates = self._calibrated_columns(source) + estimates = self._calibrated_columns(columns) mean = estimates @ weights return np.sqrt(((estimates - mean[:, None]) ** 2) @ weights) diff --git a/deeplc/core.py b/deeplc/core.py index 1485bf7..c3d7d74 100644 --- a/deeplc/core.py +++ b/deeplc/core.py @@ -218,7 +218,8 @@ class HeadColumnSource: """ A model's predictions for whichever heads are asked for, evaluated on demand. - Stands in for the ``(n, n_heads)`` matrix wherever a calibration is given its source. A + Stands in for the ``(n, n_heads)`` matrix wherever a calibration is given its source, and + is indexed the same way: ``source[:, heads]`` predicts those heads and nothing else. A multitask model has one head per LC setup, so that matrix is 26 kB per peptide at 6,543 setups, and a fitted calibration reads a few dozen columns of it; asking the model for those columns instead costs 320 bytes per peptide and skips the rest of the head layer. @@ -241,6 +242,9 @@ class HeadColumnSource: """ + #: Marks this as a source a calibration may index instead of a materialised matrix. + is_head_source = True + def __init__( self, psm_list, model=None, predict_kwargs: dict | None = None, n_heads: int | None = None ): @@ -264,11 +268,22 @@ def ndim(self) -> int: """Always two: this stands in for a matrix.""" return 2 - def head_columns(self, indices) -> np.ndarray: - """Predictions for the given heads, shape ``(n, len(indices))``, in that order.""" - wanted = tuple(int(i) for i in indices) - # Callers ask for the same heads more than once - prediction_report transforms the - # queries and then asks the same calibration for its head disagreement - and each ask + def __getitem__(self, key) -> np.ndarray: + """ + Predict the heads a ``[:, heads]`` slice asks for, and nothing else. + + Only the column part of the key is read; the row part must be everything, because a + calibration slices heads and not peptides. This is what lets the calibrations index a + source exactly as they index a matrix. + """ + rows, heads = key if isinstance(key, tuple) else (key, None) + if heads is None: + raise TypeError("a head source is indexed as source[:, heads]") + if not (isinstance(rows, slice) and rows == slice(None)): + raise TypeError("a head source cannot slice peptides, only heads") + wanted = (int(heads),) if np.isscalar(heads) else tuple(int(i) for i in heads) + # The same heads are asked for more than once - prediction_report transforms the + # queries and then asks the calibration for its head disagreement - and each ask # would otherwise repeat the forward pass. if self._cache is None or self._cache[0] != wanted: self._cache = ( @@ -280,7 +295,8 @@ def head_columns(self, indices) -> np.ndarray: return_matrix=True, ), ) - return self._cache[1] + matrix = self._cache[1] + return matrix[:, 0] if np.isscalar(heads) else matrix def __array__(self, dtype=None, copy=None) -> np.ndarray: """Every head, for the callers that really need the whole matrix.""" diff --git a/tests/test_multihead_calibration.py b/tests/test_multihead_calibration.py index 2607cbf..af6f5de 100644 --- a/tests/test_multihead_calibration.py +++ b/tests/test_multihead_calibration.py @@ -384,7 +384,9 @@ def test_core_rejects_a_fitted_naive_calibration(): class _CountingSource: - """A column source that records which heads were asked for.""" + """A head source that records which heads were asked for, and refuses to be materialised.""" + + is_head_source = True def __init__(self, matrix: np.ndarray): self._matrix = matrix @@ -398,10 +400,13 @@ def shape(self): def ndim(self): return 2 - def head_columns(self, indices) -> np.ndarray: - wanted = tuple(int(i) for i in indices) + def __getitem__(self, key) -> np.ndarray: + rows, heads = key + assert isinstance(rows, slice) and rows == slice(None) + wanted = (int(heads),) if np.isscalar(heads) else tuple(int(i) for i in heads) self.requests.append(wanted) - return self._matrix[:, list(wanted)] + taken = self._matrix[:, list(wanted)] + return taken[:, 0] if np.isscalar(heads) else taken def __array__(self, dtype=None, copy=None): raise AssertionError("the whole matrix should not be materialised")