diff --git a/deeplc/calibration/multihead.py b/deeplc/calibration/multihead.py index 3382447..24bd02a 100644 --- a/deeplc/calibration/multihead.py +++ b/deeplc/calibration/multihead.py @@ -28,6 +28,24 @@ LOGGER = logging.getLogger(__name__) +def as_head_matrix(source): + """ + Coerce a source to a two-dimensional float64 matrix, unless it is a head source. + + 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. + """ + 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): """Abstract base class for a calibration that selects its own head(s) from a matrix.""" @@ -131,9 +149,7 @@ 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] + source = as_head_matrix(source) head = self.selected_model_head if source.shape[1] <= head: raise CalibrationError( @@ -143,7 +159,10 @@ def transform(self, source: np.ndarray) -> np.ndarray: if source.shape[0] == 0: return np.array([]) return np.asarray( - self._inner.transform(source[:, head].astype(np.float32)), dtype=np.float64 + self._inner.transform( + np.asarray(source[:, head], dtype=np.float64).astype(np.float32) + ), + dtype=np.float64, ) @@ -296,9 +315,7 @@ 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] + source = as_head_matrix(source) head_idx = cast(np.ndarray, self._head_idx) if source.shape[1] <= int(head_idx.max()): raise CalibrationError( @@ -312,10 +329,13 @@ def transform(self, source: np.ndarray) -> 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) + # 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(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[:, i].astype(np.float32)), dtype=np.float64) + for i, cal in enumerate(self._head_calibrations) ] ) @@ -330,9 +350,7 @@ 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] + 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()) diff --git a/deeplc/core.py b/deeplc/core.py index f3e1714..c3d7d74 100644 --- a/deeplc/core.py +++ b/deeplc/core.py @@ -214,6 +214,105 @@ 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, 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. + + 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. + + """ + + #: 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 + ): + """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 __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 = ( + wanted, + predict( + self._psm_list, + model=self._model, + predict_kwargs={**self._predict_kwargs, "task_idx": list(wanted)}, + return_matrix=True, + ), + ) + 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.""" + 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 +360,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..af6f5de 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 @@ -380,3 +381,99 @@ def test_core_rejects_a_fitted_naive_calibration(): calibration=naive, predict_kwargs={"device": "cpu"}, ) + + +class _CountingSource: + """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 + self.requests: list[tuple[int, ...]] = [] + + @property + def shape(self): + return self._matrix.shape + + @property + def ndim(self): + return 2 + + 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) + 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") + + +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, + ) + + +@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()