diff --git a/CHANGELOG.md b/CHANGELOG.md index 96ada60..b9f616f 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/deeplc/_architecture.py b/deeplc/_architecture.py index 33b611a..025f995 100644 --- a/deeplc/_architecture.py +++ b/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: """ diff --git a/deeplc/_factored.py b/deeplc/_factored.py new file mode 100644 index 0000000..28e7aa2 --- /dev/null +++ b/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" diff --git a/deeplc/_model_ops.py b/deeplc/_model_ops.py index ad05567..65727fe 100644 --- a/deeplc/_model_ops.py +++ b/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,49 @@ 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 ( + f"Predicting would need an output of {n_rows:,} x {columns:,} values " + f"({n_rows * columns * itemsize / 2**30:.1f} GiB). A multitask model returns one " + f"column per LC setup, so ask for the setups you need with " + f"predict_kwargs={{'task_idx': [...]}}, or predict in smaller groups of peptides." + ) + + +def _allocate_output(n_rows: int, tail: tuple[int, ...], dtype: torch.dtype) -> torch.Tensor: + """ + Allocate the whole prediction output up front. + + Two reasons not to collect batches in a list and concatenate at the end. The + concatenation needs the result twice over, once in the parts and once in the copy, so + the peak is double what the caller ends up holding. And it only fails after every + batch has been predicted, which on a large run means the work is lost. Allocating + first fails immediately and holds one copy. + """ + try: + return torch.empty((n_rows, *tail), dtype=dtype) + except (RuntimeError, MemoryError) as exc: + columns = int(np.prod(tail)) if tail else 1 + itemsize = torch.empty(0, dtype=dtype).element_size() + raise MemoryError(_output_hint(n_rows, columns, itemsize)) from exc + + def predict( model: torch.nn.Module | PathLike | str | None, data: Dataset, @@ -252,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. @@ -270,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( @@ -280,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: @@ -293,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 = torch.empty((len(data), part.shape[1]), dtype=part.dtype) + 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 @@ -475,14 +543,18 @@ 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): selected = torch.as_tensor(list(task_idx), dtype=torch.long, device=device) - predictions = [] - total = int(np.ceil(len(data) / batch_size)) if hasattr(data, "__len__") else None + sized = hasattr(data, "__len__") + total = int(np.ceil(len(data) / batch_size)) if sized else None + out: torch.Tensor | None = None + parts: list[torch.Tensor] = [] + filled = 0 with torch.no_grad(): for features in track( _feature_batches(data, batch_size, num_workers), @@ -492,11 +564,26 @@ 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) - predictions.append(outputs.cpu()) - if not predictions: + 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 + # cannot report its length falls back to collecting the batches. + if sized: + if out is None: + out = _allocate_output(len(data), tuple(batch.shape[1:]), batch.dtype) + out[filled : filled + len(batch)] = batch + else: + parts.append(batch) + filled += len(batch) + if filled == 0: raise ValueError("Dataset is empty — nothing to predict.") - return torch.cat(predictions, dim=0) + return out[:filled] if out is not None else torch.cat(parts, dim=0) def _create_progress(disable: bool = False) -> Progress: diff --git a/deeplc/calibration/multihead.py b/deeplc/calibration/multihead.py index b2b6ea7..8c1a098 100644 --- a/deeplc/calibration/multihead.py +++ b/deeplc/calibration/multihead.py @@ -119,15 +119,22 @@ def fit(self, target: np.ndarray, source: np.ndarray) -> None: # The matrix arrives from the model as float32 and is (n, 6,543) wide; promoting it # here doubled a 276 MB reference to 552 MB for no gain, since every head's column is # cast to float32 again for its spline and the ranking accumulates in float64 itself. - source = np.asarray(source) - if source.ndim == 1: - source = source[:, None] + # + # A lazy source is left lazy. Ranking already walks the heads in blocks and every step + # after it reads single columns, so nothing here needs the matrix whole; forcing it + # would put the reference's full width in memory for no purpose. A reference of every + # confidently identified PSM in a large search is exactly where that bites. + source = as_head_matrix(source) + if not getattr(source, "is_head_source", False): + source = np.asarray(source) + if source.ndim == 1: + source = source[:, None] target = np.asarray(target, dtype=np.float64).ravel() if source.shape[0] != target.shape[0]: raise CalibrationError( f"source has {source.shape[0]} rows and target {target.shape[0]}" ) - finite = np.isfinite(target) & np.isfinite(source).all(axis=1) + finite = np.isfinite(target) & _finite_rows(source) if int(finite.sum()) < 3: raise CalibrationError("Fewer than three reference points with finite values.") source, target = source[finite], target[finite] @@ -267,15 +274,22 @@ def fit(self, target: np.ndarray, source: np.ndarray) -> None: # The matrix arrives from the model as float32 and is (n, 6,543) wide; promoting it # here doubled a 276 MB reference to 552 MB for no gain, since every head's column is # cast to float32 again for its spline and the ranking accumulates in float64 itself. - source = np.asarray(source) - if source.ndim == 1: - source = source[:, None] + # + # A lazy source is left lazy. Ranking already walks the heads in blocks and every step + # after it reads single columns, so nothing here needs the matrix whole; forcing it + # would put the reference's full width in memory for no purpose. A reference of every + # confidently identified PSM in a large search is exactly where that bites. + source = as_head_matrix(source) + if not getattr(source, "is_head_source", False): + source = np.asarray(source) + if source.ndim == 1: + source = source[:, None] target = np.asarray(target, dtype=np.float64).ravel() if source.shape[0] != target.shape[0]: raise CalibrationError( f"source has {source.shape[0]} rows and target {target.shape[0]}" ) - finite = np.isfinite(target) & np.isfinite(source).all(axis=1) + finite = np.isfinite(target) & _finite_rows(source) if int(finite.sum()) < 3: raise CalibrationError("Fewer than three reference points with finite values.") source, target = source[finite], target[finite] @@ -424,6 +438,36 @@ def upgrade_calibration(calibration: Calibration | MultiHeadCalibration) -> Mult return MultiHeadSplineCalibration() +#: Largest block of a lazy source held at once, in bytes. Blocks are sized from the row count +#: so the peak does not grow with the reference: a reference of a million PSMs reads narrower +#: blocks than one of a thousand, and both stay near this figure. +_BLOCK_BYTES = 64 * 2**20 + + +def _head_block(n_rows: int, n_heads: int, itemsize: int = 8) -> int: + """How many heads to read at once so a block stays near :data:`_BLOCK_BYTES`.""" + if n_rows <= 0: + return n_heads + return int(max(1, min(n_heads, _BLOCK_BYTES // max(1, n_rows * itemsize)))) + + +def _finite_rows(source) -> np.ndarray: + """ + Rows whose every head is finite, without holding more than a block of heads at once. + + A dense source is checked in one pass as before. A lazy one is walked in the same blocks + the ranking uses, so the widest thing in memory is a slice rather than the matrix. + """ + if not getattr(source, "is_head_source", False): + return np.isfinite(source).all(axis=1) + n_rows, n_heads = source.shape + finite = np.ones(n_rows, dtype=bool) + block = _head_block(n_rows, n_heads, itemsize=4) + for start in range(0, n_heads, block): + finite &= np.isfinite(np.asarray(source[:, start : start + block])).all(axis=1) + return finite + + def _rank_heads_by_correlation(source: np.ndarray, target: np.ndarray) -> np.ndarray: """ Head indices by decreasing Pearson correlation with the target, in one pass. @@ -442,7 +486,7 @@ def _rank_heads_by_correlation(source: np.ndarray, target: np.ndarray) -> np.nda # product per block and the variance follows from the block's own sums. Blocks keep the # accumulation in float64 without ever holding more than a slice of the matrix. correlation = np.empty(n_heads, dtype=np.float64) - block = 512 + block = _head_block(n_rows, n_heads) with np.errstate(invalid="ignore", divide="ignore"): for start in range(0, n_heads, block): chunk = np.asarray(source[:, start : start + block], dtype=np.float64) diff --git a/deeplc/core.py b/deeplc/core.py index 2105c8f..ca8a69c 100644 --- a/deeplc/core.py +++ b/deeplc/core.py @@ -12,6 +12,7 @@ from torch.utils.data import DataLoader from deeplc import _model_ops +from deeplc._factored import FactoredPredictionMatrix from deeplc._reference_selection import deduplicate_psms, select_reference_psms from deeplc.calibration import ( Calibration, @@ -76,7 +77,11 @@ def predict( model Trained model or path to model file. If None, the default DeepLC model is used. predict_kwargs - Additional keyword arguments to pass to the prediction function. + Additional keyword arguments to pass to the prediction function. Pass + ``{"factored": False}`` to force a dense ``ndarray`` from ``return_matrix=True``. + The default hands back a + :class:`~deeplc._factored.FactoredPredictionMatrix`, which holds the head's low-rank + factors and indexes identically at a fraction of the memory. return_matrix If True, return the full prediction matrix of shape ``(n, n_heads)`` when using a multitask model. If False (default), return a 1D array of shape ``(n,)`` for the @@ -114,6 +119,15 @@ def predict( ): kwargs["task_idx"] = [_default_task_idx(loaded_model)] + # The matrix a multitask model returns is its head's low-rank factors expanded out, so the + # factors are handed back instead. They report the same shape and index the same way, but + # hold (n_peptides, rank) rather than (n_peptides, n_tasks): at rank 64 and 6,543 setups + # that is 102 times less, the difference between 0.6 GiB and 63 GiB on a 2.6 M peptide run. + # Anything the factors cannot answer falls through to the dense matrix, so a caller that + # reduces over the result rather than slicing it behaves exactly as before. + if return_matrix and "task_idx" not in kwargs and _model_ops.supports_factored(loaded_model): + kwargs.setdefault("factored", True) + result = _model_ops.predict( model=loaded_model, data=DeepLCDataset.from_psm_list( @@ -122,7 +136,8 @@ def predict( **_feature_kwargs_from_spec(feature_spec), ), **kwargs, - ).numpy() + ) + result = result if isinstance(result, FactoredPredictionMatrix) else result.numpy() if not return_matrix: return result[:, 0 if "task_idx" in kwargs else _default_task_idx(loaded_model)] return result diff --git a/pyproject.toml b/pyproject.toml index 916aab3..3d0e722 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "deeplc" -version = "4.4.0" +version = "4.5.0" description = "DeepLC: Retention time prediction for (modified) peptides using Deep Learning." readme = "README.md" license = { file = "LICENSE" } diff --git a/tests/test_factored.py b/tests/test_factored.py new file mode 100644 index 0000000..32fb3f6 --- /dev/null +++ b/tests/test_factored.py @@ -0,0 +1,149 @@ +"""The factored prediction matrix must equal the matrix it stands for.""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch + +from deeplc import _model_ops, core +from deeplc._architecture import DeepLCModel +from deeplc._factored import FactoredPredictionMatrix +from deeplc.core import DEFAULT_MODEL + +PREDICT_KWARGS = {"device": "cpu", "show_progress": False} +PEPTIDES = [ + "PEPTIDEK", + "LVVVGAGGVGK", + "GPNGPWSVMK", + "YPLQLAELLK", + "VVEEAVDLFK", + "AAELALR", + "SLIDLLQK", + "ELVISLIVESK", + "AAAAAAAAAAK", + "WWWWK", +] + + +def _matrices(): + """Return the same predictions as factors and as a dense matrix.""" + lazy = core.predict(PEPTIDES, model=None, predict_kwargs=PREDICT_KWARGS, return_matrix=True) + dense = core.predict( + PEPTIDES, + model=None, + predict_kwargs={**PREDICT_KWARGS, "factored": False}, + return_matrix=True, + ) + return lazy, dense + + +requires_bundled_model = pytest.mark.skipif( + not DEFAULT_MODEL.exists(), reason="multitask model not bundled" +) + + +@requires_bundled_model +def test_factored_is_the_default_and_dense_is_available(): + """A multitask matrix comes back factored; `factored: False` forces the dense one.""" + lazy, dense = _matrices() + assert isinstance(lazy, FactoredPredictionMatrix) + assert isinstance(dense, np.ndarray) + assert lazy.shape == dense.shape + assert lazy.ndim == dense.ndim == 2 + assert len(lazy) == len(PEPTIDES) + + +@requires_bundled_model +@pytest.mark.parametrize( + "index", + [ + np.s_[:], + np.s_[np.array([0, 3, 7])], + np.s_[2], + np.s_[:, 2016], + np.s_[:, [12, 2016, 4000]], + np.s_[np.array([0, 3, 7]), 2016], + np.s_[3, 2016], + ], + ids=["all", "rows", "one-row", "one-column", "columns", "rows-and-column", "scalar"], +) +def test_every_indexing_form_matches_the_dense_matrix(index): + """Shape and values must agree, so a caller cannot tell which it was handed.""" + lazy, dense = _matrices() + from_factors = np.asarray(lazy[index]) + from_dense = np.asarray(dense[index]) + assert from_factors.shape == from_dense.shape + np.testing.assert_allclose(from_factors, from_dense, atol=1e-3) + + +@requires_bundled_model +def test_selecting_rows_stays_factored(): + """A run's rows must remain factors, or a calibration would build every head for them.""" + lazy, dense = _matrices() + rows = np.array([0, 3, 7]) + view = lazy[rows] + assert isinstance(view, FactoredPredictionMatrix) + assert view.shape == (len(rows), lazy.shape[1]) + np.testing.assert_allclose(np.asarray(view), dense[rows], atol=1e-3) + + +@requires_bundled_model +def test_array_methods_fall_back_to_the_dense_matrix(): + """Reductions have no factored form, so they must still give the same answer.""" + lazy, dense = _matrices() + assert np.isfinite(lazy).all() + np.testing.assert_allclose(lazy.min(), dense.min(), atol=1e-3) + np.testing.assert_allclose(lazy.mean(), dense.mean(), atol=1e-3) + + +@requires_bundled_model +def test_asarray_gives_the_dense_matrix(): + """Code that genuinely needs the matrix still gets it.""" + lazy, dense = _matrices() + np.testing.assert_allclose(np.asarray(lazy), dense, atol=1e-3) + + +@requires_bundled_model +def test_factors_are_smaller_than_the_matrix_they_stand_for(): + """The point of the exercise: a run's factors cost a fraction of its predictions.""" + lazy, _ = _matrices() + # Per peptide, rank floats instead of n_tasks floats. + per_peptide_factored = lazy.rank * lazy.dtype.itemsize + per_peptide_dense = lazy.shape[1] * lazy.dtype.itemsize + assert per_peptide_dense / per_peptide_factored > 100 + + +@requires_bundled_model +def test_calibration_accepts_it_as_a_head_source(): + """`as_head_matrix` must pass it through rather than densifying it.""" + from deeplc.calibration.multihead import as_head_matrix + + lazy, _ = _matrices() + assert lazy.is_head_source is True + assert as_head_matrix(lazy) is lazy + + +def test_single_task_models_are_not_factorable(): + """Only a multitask FactorHead has factors to keep.""" + assert _model_ops.supports_factored(DeepLCModel(n_heads=1)) is False + + +@requires_bundled_model +def test_a_head_finetuned_onto_one_setup_is_not_factorable(): + """`add_task` makes the head return one column, so there is nothing to factor.""" + model = _model_ops.load_model(DEFAULT_MODEL, device="cpu") + assert _model_ops.supports_factored(model) is True + model.head.add_task(torch.zeros(4)) + assert _model_ops.supports_factored(model) is False + + +def test_rank_mismatch_is_rejected(): + """Factors that cannot multiply are caught at construction, not at indexing.""" + with pytest.raises(ValueError, match="rank"): + FactoredPredictionMatrix( + np.zeros((4, 8), dtype=np.float32), + np.zeros((3, 16), dtype=np.float32), + np.ones(3, dtype=np.float32), + np.zeros(3, dtype=np.float32), + ) diff --git a/tests/test_model_ops.py b/tests/test_model_ops.py index e3f187d..8c81a7b 100644 --- a/tests/test_model_ops.py +++ b/tests/test_model_ops.py @@ -81,3 +81,45 @@ def test_load_multitask_model_without_prior_shim(): assert out.ndim == 2 assert out.shape[0] == 2 assert out.shape[1] > 1 # multiple heads + + +def test_predict_output_matches_batchwise_concatenation(): + """The preallocated output must equal what concatenating the batches produced.""" + model, dataset = DeepLCModel(n_heads=1), _TinyDeepLCDataset(length=37) + batched = _model_ops.predict( + model, dataset, device="cpu", batch_size=8, show_progress=False, length_buckets=False + ) + single = _model_ops.predict( + model, dataset, device="cpu", batch_size=1000, show_progress=False, length_buckets=False + ) + assert batched.shape == single.shape == (37, 1) + torch.testing.assert_close(batched, single) + + +def test_allocate_output_reports_the_size_when_it_does_not_fit(): + """An output that cannot be allocated says how large it was and how to shrink it.""" + with pytest.raises(MemoryError, match=r"task_idx"): + _model_ops._allocate_output(2**40, (6543,), torch.float32) + + +def test_predict_propagates_the_allocation_failure(monkeypatch): + """The helper's message reaches the caller of predict() rather than a raw allocator error.""" + def _refuse(*args, **kwargs): + raise MemoryError(_model_ops._output_hint(2**40, 6543, 4)) + + monkeypatch.setattr(_model_ops, "_allocate_output", _refuse) + with pytest.raises(MemoryError, match=r"task_idx"): + _model_ops.predict( + model=DeepLCModel(n_heads=1), + data=_TinyDeepLCDataset(length=4), + device="cpu", + show_progress=False, + length_buckets=False, + ) + + +def test_output_hint_mentions_the_size_in_gigabytes(): + """The hint names the size that failed, so a log line is enough to diagnose it.""" + hint = _model_ops._output_hint(2_587_932, 6543, 4) + assert "63.1 GiB" in hint + assert "task_idx" in hint