DeepLC 4.5.0: keep the multitask prediction matrix as its low-rank factors

CHANGELOG.md

@@ -6,6 +6,32 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to
 [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [4.5.0] - 2026-09-11
+
+### Changed
+
+- `predict(..., return_matrix=True)` on a multitask model now returns a
+  `FactoredPredictionMatrix` instead of a dense array. `FactorHead` ends with
+  `proj(trunk) @ embedding.T * scale + shift`, so the `(n_peptides, n_tasks)` matrix is
+  determined by `(n_peptides, rank)`: at rank 64 and 6,543 setups that is 102 times less
+  memory, 0.6 GiB instead of 63 GiB for 2.6 million peptides. It reports the same shape and
+  indexes the same way, selecting rows alone gives another factored matrix, and anything it
+  cannot answer from the factors falls through to the dense matrix, so callers that reduce
+  over the result rather than slicing it are unaffected. Pass
+  `predict_kwargs={"factored": False}` for the dense array.
+- `MultiHeadCalibration.fit` no longer expands a lazy source. Head ranking already walked the
+  heads in blocks and every later step reads single columns, so a reference of every
+  confidently identified PSM in a large search no longer has to fit in memory at full width.
+  Block size is now chosen from the row count, holding the peak near 64 MiB whether the
+  reference has a thousand rows or a million.
+
+### Fixed
+
+- Predicting no longer collects every batch in a list and concatenates at the end, which
+  needed the result twice over and only failed once all the work was done. The output is
+  allocated once and filled in place, and an output that does not fit reports its size and
+  points at `task_idx` rather than surfacing a bare allocator failure.
+
 ## [4.4.0] - 2026-09-09
 
 ### Added

deeplc/_architecture.py

@@ -938,6 +938,29 @@ def forward(
         del x_atom_sum  # the fused trunk reads x_atom directly
         return self.head(self.encoder(x_atom, x_global, x_one_hot), task_idx)
 
+    def project(
+        self,
+        x_atom: torch.Tensor,
+        x_atom_sum: torch.Tensor,
+        x_global: torch.Tensor,
+        x_one_hot: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Map the trunk into the head's low-rank space, stopping before the per-setup step.
+
+        ``FactorHead`` finishes with ``projected @ embedding.T * scale + shift``, so this is
+        everything the setups share. Keeping it, rather than the ``(batch, n_tasks)`` matrix
+        it expands to, is what lets a whole run be held in ``(n_peptides, rank)``.
+
+        Returns
+        -------
+        torch.Tensor
+            Shape ``(batch, rank)``.
+
+        """
+        del x_atom_sum  # the fused trunk reads x_atom directly
+        return self.head.proj(self.encoder(x_atom, x_global, x_one_hot))
+
     @property
     def padding_reach(self) -> int | None:
         """

deeplc/_factored.py

@@ -0,0 +1,181 @@
+"""A prediction matrix held as its low-rank factors rather than materialised."""
+
+from __future__ import annotations
+
+import numpy as np
+
+__all__ = ["FactoredPredictionMatrix"]
+
+
+class FactoredPredictionMatrix:
+    """
+    A prediction matrix kept as its low-rank factors.
+
+    Holds the ``(n_peptides, n_tasks)`` output of a
+    :class:`~deeplc._architecture.FactorHead` as the factors it is computed from.
+
+    That head is low rank by construction::
+
+        pred[:, j] = (proj(trunk).embedding[j]) * scale[j] + shift[j]
+
+    so the whole matrix is determined by ``proj(trunk)``, which is ``(n_peptides, rank)``,
+    together with the head's own parameters. At rank 64 and 6,543 setups that is 102 times
+    smaller than the matrix it stands for: 2.6 M peptides need 0.6 GiB of factors instead of
+    63 GiB of predictions.
+
+    It indexes like the matrix. ``m[rows]``, ``m[:, heads]`` and ``m[rows, head]`` each
+    evaluate only what was asked for, so a caller that reads a few thousand rows to choose a
+    head and then one column per run never builds the rest. ``np.asarray(m)`` still gives the
+    dense matrix for code that genuinely needs it, at its full size.
+
+    Parameters
+    ----------
+    projections
+        ``proj(trunk)`` for every peptide, shape ``(n_peptides, rank)``.
+    embedding
+        Per-setup embedding, shape ``(n_tasks, rank)``.
+    scale, shift
+        Per-setup affine map, shape ``(n_tasks,)``.
+
+    """
+
+    #: Marks this as a source a calibration may index instead of a materialised matrix.
+    is_head_source = True
+
+    def __init__(
+        self,
+        projections: np.ndarray,
+        embedding: np.ndarray,
+        scale: np.ndarray,
+        shift: np.ndarray,
+    ) -> None:
+        self._projections = np.asarray(projections)
+        self._embedding = np.asarray(embedding, dtype=self._projections.dtype)
+        self._scale = np.asarray(scale, dtype=self._projections.dtype)
+        self._shift = np.asarray(shift, dtype=self._projections.dtype)
+        if self._embedding.shape[1] != self._projections.shape[1]:
+            raise ValueError(
+                f"embedding rank {self._embedding.shape[1]} does not match projection rank "
+                f"{self._projections.shape[1]}"
+            )
+
+    @property
+    def shape(self) -> tuple[int, int]:
+        """Rows and task count, without evaluating anything."""
+        return (self._projections.shape[0], self._embedding.shape[0])
+
+    @property
+    def ndim(self) -> int:
+        """Always two: this stands in for a matrix."""
+        return 2
+
+    @property
+    def dtype(self) -> np.dtype:
+        """Element type of the matrix it stands for."""
+        return self._projections.dtype
+
+    @property
+    def rank(self) -> int:
+        """Width of the factorisation."""
+        return self._projections.shape[1]
+
+    @property
+    def nbytes(self) -> int:
+        """What the factors occupy."""
+        return int(
+            self._projections.nbytes
+            + self._embedding.nbytes
+            + self._scale.nbytes
+            + self._shift.nbytes
+        )
+
+    @property
+    def dense_nbytes(self) -> int:
+        """What the matrix would occupy if it were materialised."""
+        rows, columns = self.shape
+        return int(rows * columns * self._projections.itemsize)
+
+    def __len__(self) -> int:
+        """Return the number of peptides."""
+        return self.shape[0]
+
+    def _evaluate(self, rows, columns) -> np.ndarray:
+        """Compute just the requested block from the factors."""
+        projections = self._projections if rows is None else self._projections[rows]
+        if projections.ndim == 1:
+            projections = projections[None, :]
+        if columns is None:
+            embedding, scale, shift = self._embedding, self._scale, self._shift
+        else:
+            embedding = np.atleast_2d(self._embedding[columns])
+            scale = np.atleast_1d(self._scale[columns])
+            shift = np.atleast_1d(self._shift[columns])
+        return (projections @ embedding.T) * scale + shift
+
+    def _select_rows(self, rows) -> FactoredPredictionMatrix:
+        """Narrow to a subset of peptides, still as factors."""
+        return FactoredPredictionMatrix(
+            self._projections[rows], self._embedding, self._scale, self._shift
+        )
+
+    def __getitem__(self, key):
+        """
+        Index as the dense matrix would, evaluating only the block that is asked for.
+
+        Selecting rows alone gives another factored matrix rather than a dense one. That is
+        what lets a caller narrow to one run's peptides and still hand the result to a
+        calibration, which reads a few dozen of the thousands of heads: nothing in that path
+        ever builds the wide matrix. A scalar in either position drops that axis, as numpy
+        does, so ``m[rows, head]`` is one dimensional and ``m[i, j]`` is a scalar.
+        """
+        rows, columns = key if isinstance(key, tuple) else (key, None)
+        if isinstance(rows, slice) and rows == slice(None):
+            rows = None
+        if columns is None and rows is not None and not _is_scalar(rows):
+            return self._select_rows(rows)
+        row_scalar = _is_scalar(rows)
+        column_scalar = columns is not None and _is_scalar(columns)
+        block = self._evaluate(rows, columns)
+        if row_scalar:
+            block = block[0]
+            return block[0] if column_scalar else block
+        return block[:, 0] if column_scalar else block
+
+    def __array__(self, dtype=None, copy=None) -> np.ndarray:
+        """Return the whole matrix, for callers that really need it, at its full size."""
+        dense = self._evaluate(None, None)
+        return dense if dtype is None else dense.astype(dtype)
+
+    def __getattr__(self, name: str):
+        """
+        Fall back to the dense matrix for anything not answered from the factors.
+
+        Reductions such as ``.min()`` and ``.mean()`` have no cheap factored form, so code
+        that calls them gets the same answer it always did, at the same cost it always had.
+        Only indexing, which is the path that matters for a wide matrix, stays lazy.
+        """
+        if name.startswith("_"):
+            raise AttributeError(name)
+        return getattr(np.asarray(self), name)
+
+    def __repr__(self) -> str:
+        """Show the shape it stands for and what it actually costs."""
+        rows, columns = self.shape
+        return (
+            f"{type(self).__name__}(shape=({rows}, {columns}), rank={self.rank}, "
+            f"{_human(self.nbytes)} held for a {_human(self.dense_nbytes)} matrix)"
+        )
+
+
+def _is_scalar(index) -> bool:
+    """Whether an index selects one element rather than a subset."""
+    return np.isscalar(index) or (isinstance(index, np.generic) and np.ndim(index) == 0)
+
+
+def _human(n: int) -> str:
+    """Format a byte count in the largest unit that keeps it above one."""
+    for unit in ("B", "KiB", "MiB", "GiB"):
+        if n < 1024 or unit == "GiB":
+            return f"{n:.1f} {unit}" if unit != "B" else f"{n} B"
+        n /= 1024
+    return f"{n:.1f} GiB"

deeplc/_model_ops.py

@@ -19,7 +19,8 @@
 )
 from torch.utils.data import DataLoader, Dataset, Subset
 
-from deeplc._architecture import DeepLCModel, FlexCNNMultitaskModel
+from deeplc._architecture import DeepLCModel, FactorHead, FlexCNNMultitaskModel
+from deeplc._factored import FactoredPredictionMatrix
 from deeplc.data import DeepLCDataset
 
 logger = logging.getLogger(__name__)
@@ -242,6 +243,21 @@ def train(
     return model
 
 
+def supports_factored(model: torch.nn.Module) -> bool:
+    """
+    Whether this model's output can be kept as low-rank factors instead of a matrix.
+
+    True only for a multitask model that still has its ``FactorHead``. A head fine-tuned
+    onto one setup returns that column alone, so there is nothing to factor.
+    """
+    head = getattr(model, "head", None)
+    return (
+        hasattr(model, "project")
+        and isinstance(head, FactorHead)
+        and not getattr(head, "has_new_task", False)
+    )
+
+
 def _output_hint(n_rows: int, columns: int, itemsize: int) -> str:
     """Explain an output that does not fit, and how to make it smaller."""
     return (
@@ -280,10 +296,17 @@ def predict(
     show_progress: bool = True,
     task_idx: Sequence[int] | None = None,
     length_buckets: bool = True,
-) -> torch.Tensor:
+    factored: bool = False,
+) -> torch.Tensor | FactoredPredictionMatrix:
     """
     Predict using the model for the given dataset.
 
+    ``factored`` returns a :class:`~deeplc._factored.FactoredPredictionMatrix` rather than
+    the matrix itself, for the multitask models whose head is low rank. It indexes the same
+    way but holds ``(n_peptides, rank)`` instead of ``(n_peptides, n_tasks)``, which at rank
+    64 and 6,543 setups is 102 times less memory. Ignored when the model cannot be factored
+    or when ``task_idx`` already narrows the output.
+
     ``length_buckets`` runs length-sorted chunks in a window that fits them rather than
     padding every peptide to the model's full window; see :func:`_length_buckets`. It is
     exact for models that report a ``padding_reach`` and ignored for those that do not.
@@ -298,6 +321,8 @@ def predict(
     device = device or ("cuda" if torch.cuda.is_available() else "cpu")
     model = load_model(model, device)
 
+    as_factors = factored and task_idx is None and supports_factored(model)
+
     buckets = _length_buckets(model, data, batch_size) if length_buckets else None
     if buckets is None:
         predictions = _predict_epoch(
@@ -308,8 +333,10 @@ def predict(
             num_workers=num_workers,
             show_progress=show_progress,
             task_idx=task_idx,
+            project=as_factors,
         )
-        return predictions.cpu().detach()
+        result = predictions.cpu().detach()
+        return _as_factored(model, result) if as_factors else result
 
     out: torch.Tensor | None = None
     for indices, subset in buckets:
@@ -321,13 +348,26 @@ def predict(
             num_workers=num_workers,
             show_progress=show_progress,
             task_idx=task_idx,
+            project=as_factors,
         ).cpu()
         if out is None:
             out = _allocate_output(len(data), tuple(part.shape[1:]), part.dtype)
         out[indices] = part
     if out is None:
         raise ValueError("Dataset is empty — nothing to predict.")
-    return out.detach()
+    out = out.detach()
+    return _as_factored(model, out) if as_factors else out
+
+
+def _as_factored(model: torch.nn.Module, projections: torch.Tensor) -> FactoredPredictionMatrix:
+    """Pair the projections with the head parameters that expand them."""
+    head = model.head
+    return FactoredPredictionMatrix(
+        projections.numpy(),
+        head.embedding.detach().cpu().numpy(),
+        head.scale.detach().cpu().numpy(),
+        head.shift.detach().cpu().numpy(),
+    )
 
 
 #: Widest spread of peptide lengths allowed inside one prediction chunk. Small enough that
@@ -503,8 +543,9 @@ def _predict_epoch(
     num_workers: int = 0,
     show_progress: bool = False,
     task_idx: Sequence[int] | None = None,
+    project: bool = False,
 ) -> torch.Tensor:
-    """Predict using the model for one epoch."""
+    """Predict using the model for one epoch, or project into the head's rank if asked."""
     model.eval()
     selected = None
     if task_idx is not None and supports_task_subset(model):
@@ -523,7 +564,12 @@ def _predict_epoch(
             total=total,
         ):
             features = [feature_tensor.to(device) for feature_tensor in features]
-            outputs = model(*features) if selected is None else model(*features, task_idx=selected)
+            if project:
+                outputs = model.project(*features)
+            else:
+                outputs = (
+                    model(*features) if selected is None else model(*features, task_idx=selected)
+                )
             batch = outputs.cpu()
             # The column count is only known once a batch has been through the model, so the
             # output is allocated on the first one and filled from there. A dataset that