Let a calibration pull the head columns it reads

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())

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)

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

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()