From bc3d3092d7098167886251acf62e051a3440df5c Mon Sep 17 00:00:00 2001 From: RobbinBouwmeester Date: Mon, 31 Aug 2026 10:57:18 +0200 Subject: [PATCH 1/7] feat: prediction reports with provenance and conformal intervals (4.4.0) A prediction is a number with no way to tell whether the model has seen the peptidoform, merely something like it, or nothing like it, and no statement of how far off it may be. prediction_report answers all three per PSM. Membership and novelty: exact match against the calibration reference and the Levenshtein distance to its closest sequence, always; with a TrainingIndex also exact match against the 10,105,640-peptidoform corpus behind the bundled multitask model, membership within the training sets of the setups the calibration selected, and the distance to the closest training sequence (exact to ten edits, capped beyond; the error is flat in this distance, so the cap costs nothing but keeps the search fast). Canonical keys reproduce the corpus format: peprec positions, Unimod accessions, lowercased unmapped names. Uncertainty: cross-fitted split-conformal intervals on the reference. Each reference fold is predicted by a calibration fitted on the other folds and the half-width is a finite-sample quantile of those honest residuals per predicted-RT bin. On eight held-out PRIDE setups the 90 % interval covered 0.88 to 0.97 per setup (median 0.91), 4 % of the gradient wide on well-behaved setups and honestly wide (79 %) on a run that pools fractions. Chosen over quantile regression because it needs no retraining and carries a finite-sample guarantee; coverage is marginal, not per-peptide. The TrainingIndex (~400 MB: sorted key hashes, per-setup membership CSR, unique sequences) is built offline from the training cache and distributed separately; the report works without it and then carries the reference columns only. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG.md | 28 +++ deeplc/__init__.py | 3 + deeplc/report.py | 419 ++++++++++++++++++++++++++++++++++++++++++ docs/source/usage.rst | 31 ++++ pyproject.toml | 3 +- tests/test_report.py | 244 ++++++++++++++++++++++++ 6 files changed, 727 insertions(+), 1 deletion(-) create mode 100644 deeplc/report.py create mode 100644 tests/test_report.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 170ba83..4206584 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,34 @@ 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.4.0] - 2026-08-31 + +### Added + +- `prediction_report`: predictions with provenance and uncertainty per PSM. Returns a + DataFrame with, next to `predicted_rt`: a conformal prediction interval (`ci_lower`, + `ci_upper`) at a chosen coverage, exact-match membership against the calibration reference + (`in_reference`) and the Levenshtein distance to the closest reference sequence + (`dist_to_reference`); with a training index also membership in the corpus the bundled + multitask model was trained on (`in_training`), membership within the training sets of the + setups the calibration selected (`in_selected_heads_training`) and the distance to the + closest training sequence (`dist_to_training`, exact up to 10 and capped beyond). + + The interval is cross-fitted split-conformal on the reference: the reference is split into + folds, each fold is predicted by a calibration fitted on the other folds, and the half-width + is a finite-sample quantile of those honest residuals per predicted-RT bin. On eight PRIDE + setups no DeepLC model was trained on, the empirical coverage of the 90 % interval was 0.88 + to 0.97 per setup (median 0.91), with widths from 4 % of the gradient on well-behaved setups + to 79 % on a run that pools several fractions. Coverage is marginal, not per-peptide. + +- `TrainingIndex`: a memory-mapped index of the multitask training corpus (10,105,640 + canonical peptidoform keys, their 65,139,832 setup observations, 6,157,558 unique stripped + sequences; about 400 MB on disk). Built offline from the training cache and distributed + separately from the package; `prediction_report` takes it as an optional argument and works + without it. + +- Dependency: `rapidfuzz` (Levenshtein distances). + ## [4.3.0] - 2026-09-02 ### Changed diff --git a/deeplc/__init__.py b/deeplc/__init__.py index f436f04..a21a64e 100644 --- a/deeplc/__init__.py +++ b/deeplc/__init__.py @@ -11,9 +11,12 @@ save_model, train, ) +from deeplc.report import TrainingIndex, prediction_report __version__: str = version("deeplc") __all__: list[str] = [ + "TrainingIndex", + "prediction_report", "calibrate", "predict", "predict_and_calibrate", diff --git a/deeplc/report.py b/deeplc/report.py new file mode 100644 index 0000000..062378a --- /dev/null +++ b/deeplc/report.py @@ -0,0 +1,419 @@ +""" +Prediction reports: provenance and uncertainty next to every retention time. + +A plain prediction is a number with no way to tell whether the model has seen the peptidoform, +merely something like it, or nothing like it, and no statement of how far off it may be. The +report answers those three questions per PSM: + +- **membership**: is the peptidoform an exact match to the reference the calibration or + fine-tuning used, and, when a training index is available, to the corpus the bundled model + was trained on, or to the training sets of the setups the calibration selected; +- **novelty**: the Levenshtein distance from the stripped sequence to the closest reference + sequence (and to the closest training sequence, when the index is available); +- **uncertainty**: a conformal prediction interval calibrated on the reference. + +The interval comes from cross-fitted split-conformal prediction: the reference is split into +folds, each fold is predicted by a calibration fitted on the other folds, and the interval +half-width is a finite-sample quantile of those honest |residuals|, taken per predicted-RT bin +because peak width varies along a gradient. On eight PRIDE setups no DeepLC model was trained +on, the empirical coverage of the 90 % interval was 0.88 to 0.96 per setup (median 0.91), with +widths from 4 % of the gradient on well-behaved setups to 79 % on a run that pools several +fractions, which is what an honest interval looks like there. Coverage is marginal, not +per-peptide: on average over peptides like the reference, not for each one individually. +""" + +from __future__ import annotations + +import json +import logging +from dataclasses import dataclass, field +from os import PathLike +from pathlib import Path + +import numpy as np +import pandas as pd +import torch +from psm_utils import PSM, Peptidoform, PSMList + +from deeplc import core +from deeplc._reference_selection import deduplicate_psms, select_reference_psms +from deeplc.calibration import Calibration, SplineTransformerCalibration + +LOGGER = logging.getLogger(__name__) + +#: Bins for the RT-dependent interval width, and the minimum honest residuals a bin needs +#: before it is trusted over the global quantile. +_N_RT_BINS = 5 +_MIN_RESIDUALS_PER_BIN = 40 +_N_FOLDS = 5 + + +def canonical_peptidoform_key(peptidoform: Peptidoform | str) -> str: + """ + Build the identifier under which a peptidoform appears in the multitask training corpus. + + ``SEQUENCE|`` followed by position-sorted ``pos|U:`` pairs, positions in peprec + convention (1-based, 0 for N-terminal, -1 for C-terminal). A modification without a Unimod + accession contributes its lowercased name, matching how the corpus was built: an unmapped + name still matches itself across sources instead of silently merging with another. + """ + if isinstance(peptidoform, str): + peptidoform = Peptidoform(peptidoform) + + def token(mod) -> str: + accession = getattr(mod, "id", None) + if accession is not None and str(accession).isdigit(): + return f"U:{accession}" + name = getattr(mod, "name", None) or str(mod) + return str(name).lower() + + pairs: list[tuple[int, str]] = [] + n_term = peptidoform.properties.get("n_term") + if n_term: + pairs += [(0, token(mod)) for mod in n_term] + c_term = peptidoform.properties.get("c_term") + if c_term: + pairs += [(-1, token(mod)) for mod in c_term] + for position, (_, mods) in enumerate(peptidoform.parsed_sequence, start=1): + if mods: + pairs += [(position, token(mod)) for mod in mods] + pairs.sort() + mods_text = "|".join(f"{position}|{tok}" for position, tok in pairs) + return f"{peptidoform.sequence}|{mods_text}" + + +class TrainingIndex: + """ + Memory-mapped index of the corpus behind the bundled multitask model. + + Built offline from the training cache (10,105,640 canonical peptidoform keys and their + 65,139,832 peptidoform-setup observations over 6,543 setups) and loaded from a directory: + ``key_hashes.npy`` (sorted xxh3-64 of the canonical keys), ``task_indptr.npy`` / + ``task_cols.npy`` (which setups each peptidoform was observed in), ``sequences.npy`` / + ``seq_lengths.npy`` (unique stripped sequences, for edit distances) and ``meta.json``. + + Everything is memory-mapped, so opening the index costs nothing until it is used. + """ + + def __init__(self, path: PathLike | str) -> None: + """Open a training index directory.""" + self.path = Path(path) + meta_file = self.path / "meta.json" + if not meta_file.exists(): + raise FileNotFoundError( + f"{self.path} is not a training index (no meta.json). It is built offline " + "from the training cache and distributed separately from the package." + ) + self.meta = json.loads(meta_file.read_text(encoding="utf-8")) + self._hashes = np.load(self.path / "key_hashes.npy", mmap_mode="r") + self._indptr = np.load(self.path / "task_indptr.npy", mmap_mode="r") + self._cols = np.load(self.path / "task_cols.npy", mmap_mode="r") + self._sequences: np.ndarray | None = None + self._seq_lengths: np.ndarray | None = None + + @staticmethod + def _hash(keys: list[str]) -> np.ndarray: + try: + from xxhash import xxh3_64_intdigest as digest + except ImportError: + from hashlib import blake2b + + def digest(text: str) -> int: + return int.from_bytes(blake2b(text.encode(), digest_size=8).digest(), "little") + + return np.array([digest(key) for key in keys], dtype=np.uint64) + + def _rows(self, keys: list[str]) -> np.ndarray: + """Index of each key in the sorted hash array, or -1 when absent.""" + hashes = self._hash(keys) + position = np.searchsorted(self._hashes, hashes) + position = np.clip(position, 0, len(self._hashes) - 1) + found = self._hashes[position] == hashes + return np.where(found, position, -1) + + def contains(self, keys: list[str]) -> np.ndarray: + """Whether each canonical key occurs anywhere in the training corpus.""" + return self._rows(keys) >= 0 + + def contains_in_tasks(self, keys: list[str], task_idx: np.ndarray) -> np.ndarray: + """ + Whether each key was observed in at least one of the given setups. + + Setup ids the index does not know (a model with more heads than the corpus the index + was built from) are ignored: they cannot contribute a membership either way. + """ + n_tasks = int(self.meta["n_tasks"]) + task_idx = np.asarray(task_idx, dtype=int) + known = task_idx[(task_idx >= 0) & (task_idx < n_tasks)] + if len(known) < len(task_idx): + LOGGER.warning( + "%d of %d selected setups are outside this training index (%d setups); " + "does the index belong to this model?", + len(task_idx) - len(known), + len(task_idx), + n_tasks, + ) + wanted = np.zeros(n_tasks, dtype=bool) + wanted[known] = True + rows = self._rows(keys) + out = np.zeros(len(keys), dtype=bool) + for i, row in enumerate(rows): + if row < 0: + continue + cols = self._cols[self._indptr[row] : self._indptr[row + 1]] + out[i] = bool(wanted[cols].any()) + return out + + def distance_to_training(self, sequences: list[str], max_distance: int = 10) -> np.ndarray: + """ + Levenshtein distance from each stripped sequence to the closest training sequence. + + Distances are exact up to ``max_distance`` and reported as ``max_distance + 1`` beyond + it. The cap is what keeps this fast: exact matches are a set lookup, near matches a + length-banded cutoff search, and the expensive unbounded scan over millions of + sequences never runs. Beyond ten edits the distance carries no usable signal anyway; + on held-out setups the prediction error is flat in this distance. + """ + from rapidfuzz.distance import Levenshtein + from rapidfuzz.process import cdist + + if self._sequences is None: + blob = (self.path / "sequences.txt").read_bytes().decode("ascii") + self._sequences = np.array(blob.split(chr(10)), dtype=object) + self._seq_lengths = np.load(self.path / "seq_lengths.npy") + unique, inverse = np.unique(np.asarray(sequences, dtype=object), return_inverse=True) + exact = np.isin(unique, self._sequences) + per_unique = np.full(len(unique), -1, dtype=np.int32) + per_unique[exact] = 0 + todo = np.flatnonzero(~exact) + if len(todo) == 0: + return per_unique[inverse] + lengths = np.array([len(unique[i]) for i in todo]) + band = (self._seq_lengths >= lengths.min() - max_distance) & ( + self._seq_lengths <= lengths.max() + max_distance + ) + candidates = self._sequences[band] + distance = cdist( + [unique[i] for i in todo], + candidates.tolist(), + scorer=Levenshtein.distance, + score_cutoff=max_distance, + workers=-1, + ) + # rapidfuzz reports cutoff + 1 for everything above the cutoff, which is exactly the + # capped value this method promises + per_unique[todo] = distance.min(axis=1) + return per_unique[inverse] + + +@dataclass +class _ConformalInterval: + """RT-binned conformal half-widths, fitted on honest reference residuals.""" + + coverage: float + edges: np.ndarray = field(default_factory=lambda: np.array([])) + half_width: np.ndarray = field(default_factory=lambda: np.array([])) + + @staticmethod + def _finite_sample_quantile(abs_residuals: np.ndarray, coverage: float) -> float: + n = len(abs_residuals) + rank = min(int(np.ceil((n + 1) * coverage)), n) + return float(np.sort(abs_residuals)[rank - 1]) + + @classmethod + def fit( + cls, predicted: np.ndarray, residuals: np.ndarray, coverage: float + ) -> _ConformalInterval: + """Per-RT-bin conformal quantiles with a global fallback for thin bins.""" + absolute = np.abs(residuals) + overall = cls._finite_sample_quantile(absolute, coverage) + edges = np.quantile(predicted, np.linspace(0, 1, _N_RT_BINS + 1)) + edges[0], edges[-1] = -np.inf, np.inf + bins = np.clip(np.searchsorted(edges, predicted, side="right") - 1, 0, _N_RT_BINS - 1) + half_width = np.full(_N_RT_BINS, overall) + for b in range(_N_RT_BINS): + mask = bins == b + if int(mask.sum()) >= _MIN_RESIDUALS_PER_BIN: + half_width[b] = cls._finite_sample_quantile(absolute[mask], coverage) + return cls(coverage=coverage, edges=edges, half_width=half_width) + + def widths(self, predicted: np.ndarray) -> np.ndarray: + """Interval half-width for each prediction.""" + bins = np.clip(np.searchsorted(self.edges, predicted, side="right") - 1, 0, _N_RT_BINS - 1) + return self.half_width[bins] + + +def _crossfit_residuals( + y_reference: np.ndarray, + matrix_reference: np.ndarray, + calibration_template: Calibration, + seed: int = 0, +) -> tuple[np.ndarray, np.ndarray]: + """ + Honest reference residuals: each fold predicted by a calibration fitted without it. + + Returns (cross-fitted predictions, residuals), aligned with the reference order. The + template is re-instantiated per fold with ``type(...)()`` semantics via a deep copy of its + construction parameters, so a fitted calibration is never reused across folds. + """ + import copy + + rng = np.random.default_rng(seed) + order = rng.permutation(len(y_reference)) + folds = np.array_split(order, min(_N_FOLDS, max(2, len(y_reference) // 25))) + predicted = np.empty(len(y_reference)) + for i, fold in enumerate(folds): + train = np.concatenate([f for j, f in enumerate(folds) if j != i]) + calibration = copy.deepcopy(calibration_template) + if getattr(calibration, "uses_all_heads", False): + calibration.fit(target=y_reference[train], source=matrix_reference[train]) + predicted[fold] = calibration.transform(matrix_reference[fold]) + else: + head = core._best_correlating_head(matrix_reference[train], y_reference[train]) + calibration.selected_model_head = head + calibration.fit( + target=y_reference[train].astype(np.float32), + source=matrix_reference[train][:, head].astype(np.float32), + ) + predicted[fold] = np.asarray( + calibration.transform(matrix_reference[fold][:, head].astype(np.float32)), + dtype=float, + ) + return predicted, y_reference - predicted + + +def prediction_report( + psm_list: PSMList | list[PSM | Peptidoform | str], + psm_list_reference: PSMList | list[PSM | Peptidoform | str] | None = None, + model: torch.nn.Module | PathLike | str | None = None, + calibration: Calibration | None = None, + coverage: float = 0.90, + training_index: TrainingIndex | PathLike | str | None = None, + predict_kwargs: dict | None = None, +) -> pd.DataFrame: + """ + Predict with calibration and report provenance and uncertainty per PSM. + + Parameters + ---------- + psm_list + PSMs to predict retention times for. + psm_list_reference + Reference for calibration; auto-selected from ``psm_list`` when None, as in + :func:`deeplc.predict_and_calibrate`. + model + Trained model or path; the bundled multitask model when None. + calibration + Unfitted calibration to use; :class:`SplineTransformerCalibration` when None. Pass + :class:`~deeplc.calibration.MultiHeadRidgeCalibration` to combine setups, in which case + the membership column covers every selected head. + coverage + Nominal coverage of the conformal interval (marginal, on peptides exchangeable with + the reference). 0.90 by default. + training_index + A :class:`TrainingIndex` or a path to one. Without it, the columns about the training + corpus are omitted and the report is limited to the reference. + predict_kwargs + Passed to the prediction function (``{"device": "cpu"}`` and the like). + + Returns + ------- + pd.DataFrame + One row per input PSM, in order: ``peptidoform``, ``predicted_rt``, ``ci_lower``, + ``ci_upper`` (conformal at ``coverage``), ``observed_rt`` (when present), + ``in_reference``, ``dist_to_reference`` and, with a training index, + ``in_training``, ``dist_to_training`` and ``in_selected_heads_training``. + + """ + from rapidfuzz.distance import Levenshtein + from rapidfuzz.process import cdist + + parsed = core._parse_psms(psm_list) + if psm_list_reference is None: + reference = select_reference_psms(parsed) + else: + reference = core._parse_psms(psm_list_reference) + reference = deduplicate_psms(reference) + + if calibration is None: + calibration = SplineTransformerCalibration() + if calibration.is_fitted: + raise ValueError( + "prediction_report fits the calibration itself (it also needs cross-fitted " + "residuals for the interval); pass an unfitted calibration." + ) + + # one matrix for the reference, one for the queries; everything below reuses them + 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) + y_reference = np.array(reference["retention_time"], dtype=np.float64) + + import copy + + template = copy.deepcopy(calibration) + if getattr(calibration, "uses_all_heads", False): + calibration.fit(target=y_reference, source=matrix_reference) + predicted = calibration.transform(matrix_query) + selected_heads = np.asarray(calibration._head_idx, dtype=int) + else: + head = core._best_correlating_head(matrix_reference, y_reference) + calibration.selected_model_head = head + calibration.fit( + target=y_reference.astype(np.float32), + source=matrix_reference[:, head].astype(np.float32), + ) + predicted = np.asarray( + calibration.transform(matrix_query[:, head].astype(np.float32)), dtype=float + ) + selected_heads = np.array([head], dtype=int) + + cross_predicted, residuals = _crossfit_residuals(y_reference, matrix_reference, template) + interval = _ConformalInterval.fit(cross_predicted, residuals, coverage) + half_width = interval.widths(np.asarray(predicted, dtype=float)) + + # membership and novelty against the reference + reference_keys = {canonical_peptidoform_key(psm.peptidoform) for psm in reference.psm_list} + query_keys = [canonical_peptidoform_key(psm.peptidoform) for psm in parsed.psm_list] + in_reference = np.array([key in reference_keys for key in query_keys]) + + reference_sequences = sorted({psm.peptidoform.sequence for psm in reference.psm_list}) + query_sequences = [psm.peptidoform.sequence for psm in parsed.psm_list] + dist_to_reference = cdist( + query_sequences, reference_sequences, scorer=Levenshtein.distance, workers=-1 + ).min(axis=1) + + observed = [psm.retention_time for psm in parsed.psm_list] + frame = pd.DataFrame( + { + "peptidoform": [str(psm.peptidoform) for psm in parsed.psm_list], + "predicted_rt": np.asarray(predicted, dtype=float), + "ci_lower": np.asarray(predicted, dtype=float) - half_width, + "ci_upper": np.asarray(predicted, dtype=float) + half_width, + "observed_rt": [rt if rt is not None else np.nan for rt in observed], + "in_reference": in_reference, + "dist_to_reference": dist_to_reference.astype(int), + } + ) + frame.attrs["coverage"] = coverage + frame.attrs["selected_heads"] = selected_heads.tolist() + + if training_index is not None: + if not isinstance(training_index, TrainingIndex): + training_index = TrainingIndex(training_index) + frame["in_training"] = training_index.contains(query_keys) + frame["in_selected_heads_training"] = training_index.contains_in_tasks( + query_keys, selected_heads + ) + frame["dist_to_training"] = training_index.distance_to_training(query_sequences) + LOGGER.info( + "%d of %d peptidoforms are in the training corpus, %d in the %d selected setups.", + int(frame["in_training"].sum()), + len(frame), + int(frame["in_selected_heads_training"].sum()), + len(selected_heads), + ) + return frame diff --git a/docs/source/usage.rst b/docs/source/usage.rst index 2149d97..4457bf7 100644 --- a/docs/source/usage.rst +++ b/docs/source/usage.rst @@ -76,6 +76,37 @@ For a full list of options: deeplc predict --help +Prediction reports +================== + +:func:`deeplc.prediction_report` returns predictions together with what a bare number cannot +say: whether the model has seen the peptidoform, how far the nearest known sequence is, and how +far off the prediction may plausibly be. + +.. code-block:: python + + from deeplc import prediction_report + + report = prediction_report(psm_list, psm_list_reference=reference, coverage=0.90) + report[["peptidoform", "predicted_rt", "ci_lower", "ci_upper", + "in_reference", "dist_to_reference"]] + +The interval is a cross-fitted conformal interval calibrated on the reference, so its coverage +holds on peptides exchangeable with the reference, without retraining and regardless of the +model. Pass ``calibration=MultiHeadRidgeCalibration()`` to combine setups; the membership +column then covers every selected head. + +With a training index (built from the multitask training corpus and distributed separately), +three more columns appear: ``in_training`` (exact peptidoform match anywhere in the corpus), +``in_selected_heads_training`` (match within the setups the calibration selected) and +``dist_to_training`` (Levenshtein distance to the closest training sequence, exact up to ten +edits and capped beyond): + +.. code-block:: python + + report = prediction_report(psm_list, psm_list_reference=reference, + training_index="path/to/training_index_v6f") + Python API ========== diff --git a/pyproject.toml b/pyproject.toml index b17cdbb..916aab3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "deeplc" -version = "4.3.0" +version = "4.4.0" description = "DeepLC: Retention time prediction for (modified) peptides using Deep Learning." readme = "README.md" license = { file = "LICENSE" } @@ -39,6 +39,7 @@ dependencies = [ "pandas>=0.25,<3", "scikit-learn>=1.2.0,<2", "psm-utils>=1.5,<2", + "rapidfuzz>=3,<4", "click>=8,<9", "rich>=13,<15", ] diff --git a/tests/test_report.py b/tests/test_report.py new file mode 100644 index 0000000..37b8e58 --- /dev/null +++ b/tests/test_report.py @@ -0,0 +1,244 @@ +"""Prediction reports: membership, novelty and conformal intervals.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import pytest +from psm_utils import PSM, PSMList + +from deeplc.report import ( + TrainingIndex, + _ConformalInterval, + canonical_peptidoform_key, + prediction_report, +) + +_PEPTIDES = [ + "AAGPSLSHTSGGTQSK", + "AGFAGDDAPR", + "AIQEYNQDK", + "AAYFGILEK", + "ADTQLDESSEQIDEEELTSK", + "AHQVVEDGYEFFAK", + "ALDQFVNFSEQK", + "AAPFSPAEK", + "VGAHAGEYGAEALER", + "LNLSPLGEEMR", + "AAGPSLSHTSGGTQSR", + "AGFAGDDAPK", + "AIQEYNQDR", + "AAYFGILER", + "ADTQLDESSEQIDEEELTSR", + "AHQVVEDGYEFFAR", + "ALDQFVNFSEQR", + "AAPFSPAER", + "VGAHAGEYGAEALEK", + "LNLSPLGEEMK", +] + + +# --------------------------------------------------------------------------- # +# canonical keys + + +def test_key_of_an_unmodified_peptidoform_ends_with_a_bare_pipe(): + """No modifications means an empty modification part, not a missing pipe.""" + assert canonical_peptidoform_key("PEPTIDEK/2") == "PEPTIDEK|" + + +def test_key_uses_unimod_accessions_and_peprec_positions(): + """1-based positions, 0 for N-terminal; names resolve to U:.""" + assert canonical_peptidoform_key("PEPTM[Oxidation]IDEK/2") == "PEPTMIDEK|5|U:35" + assert canonical_peptidoform_key("[Acetyl]-PEPTIDEK/2") == "PEPTIDEK|0|U:1" + + +def test_key_ignores_charge_and_sorts_modifications(): + """The corpus keys carry no charge, and modifications are position-sorted.""" + two = canonical_peptidoform_key("PEPS[Phospho]TM[Oxidation]IDEK/3") + assert two == "PEPSTMIDEK|4|U:21|6|U:35" + assert canonical_peptidoform_key("PEPS[Phospho]TM[Oxidation]IDEK") == two + + +def test_key_keeps_an_unknown_modification_as_its_lowercased_name(): + """An unmapped modification matches itself across sources instead of merging.""" + key = canonical_peptidoform_key("PEPT[Formula:C1H2O]IDEK/2") + assert key.startswith("PEPTIDEK|4|") + assert key == key.lower().replace("peptidek", "PEPTIDEK") + + +# --------------------------------------------------------------------------- # +# conformal interval + + +def test_interval_covers_at_nominal_rate_on_synthetic_residuals(): + """Fresh residuals from the same distribution land inside at about the nominal rate.""" + rng = np.random.default_rng(0) + predicted = rng.uniform(0, 100, 4000) + residuals = rng.normal(0, 1 + predicted / 50, 4000) # width grows along the gradient + interval = _ConformalInterval.fit(predicted, residuals, coverage=0.90) + + new_predicted = rng.uniform(0, 100, 4000) + new_residuals = rng.normal(0, 1 + new_predicted / 50, 4000) + covered = np.abs(new_residuals) <= interval.widths(new_predicted) + assert 0.87 <= covered.mean() <= 0.94 + + +def test_interval_is_wider_where_residuals_are_wider(): + """The per-bin quantiles track a width that changes along the gradient.""" + rng = np.random.default_rng(1) + predicted = rng.uniform(0, 100, 2000) + residuals = rng.normal(0, np.where(predicted > 50, 5.0, 1.0), 2000) + interval = _ConformalInterval.fit(predicted, residuals, coverage=0.90) + assert interval.widths(np.array([90.0]))[0] > 2 * interval.widths(np.array([10.0]))[0] + + +def test_thin_bins_fall_back_to_the_global_quantile(): + """Too few residuals per bin means one global width, not five noisy ones.""" + rng = np.random.default_rng(2) + predicted = rng.uniform(0, 100, 60) # 12 per bin, below the per-bin minimum + residuals = rng.normal(0, 2, 60) + interval = _ConformalInterval.fit(predicted, residuals, coverage=0.90) + assert len(set(np.round(interval.half_width, 9))) == 1 + + +# --------------------------------------------------------------------------- # +# training index, built small and on the fly + + +@pytest.fixture() +def tiny_index(tmp_path: Path) -> TrainingIndex: + """Three peptidoforms over three setups, written in the real on-disk format.""" + keys = ["AAGPSLSHTSGGTQSK|", "AGFAGDDAPR|7|U:35", "LNLSPLGEEMR|"] + tasks = [[0], [0, 2], [1]] + hashes = TrainingIndex._hash(keys) + order = np.argsort(hashes) + indptr = np.zeros(len(keys) + 1, dtype=np.int64) + cols: list[int] = [] + for new_row, old in enumerate(order): + cols.extend(tasks[old]) + indptr[new_row + 1] = len(cols) + np.save(tmp_path / "key_hashes.npy", hashes[order]) + np.save(tmp_path / "task_indptr.npy", indptr) + np.save(tmp_path / "task_cols.npy", np.array(cols, dtype=np.int16)) + sequences = sorted({k.split("|", 1)[0] for k in keys}) + (tmp_path / "sequences.txt").write_bytes("\n".join(sequences).encode("ascii")) + np.save(tmp_path / "seq_lengths.npy", np.array([len(s) for s in sequences], dtype=np.int16)) + (tmp_path / "task_names.json").write_text(json.dumps(["setup_a", "setup_b", "setup_c"])) + (tmp_path / "meta.json").write_text( + json.dumps({"format_version": 1, "n_peptidoforms": 3, "n_tasks": 3, "n_observations": 4}) + ) + return TrainingIndex(tmp_path) + + +def test_index_membership_and_per_task_membership(tiny_index: TrainingIndex): + """Exact keys are found globally and within the right setups only.""" + keys = ["AAGPSLSHTSGGTQSK|", "AGFAGDDAPR|7|U:35", "AGFAGDDAPR|", "PEPTIDEK|"] + assert tiny_index.contains(keys).tolist() == [True, True, False, False] + in_a = tiny_index.contains_in_tasks(keys, np.array([0])) + assert in_a.tolist() == [True, True, False, False] + in_b = tiny_index.contains_in_tasks(keys, np.array([1])) + assert in_b.tolist() == [False, False, False, False] + + +def test_index_distances_are_capped_and_exact_below_the_cap(tiny_index: TrainingIndex): + """Distances are exact up to the cap and reported as cap + 1 beyond it.""" + distances = tiny_index.distance_to_training( + ["AAGPSLSHTSGGTQSK", "AAGPSLSHTSGGTQSR", "WWWWWWWWWWWWWWWWWWWWWWWWWWWWWW"], + max_distance=5, + ) + assert distances[0] == 0 + assert distances[1] == 1 + assert distances[2] == 6 # cap + 1 + + +def test_index_refuses_a_directory_that_is_not_an_index(tmp_path: Path): + """A random directory raises instead of pretending to be an index.""" + with pytest.raises(FileNotFoundError, match="training index"): + TrainingIndex(tmp_path) + + +# --------------------------------------------------------------------------- # +# the full report + + +def _reference() -> PSMList: + return PSMList( + psm_list=[ + PSM(spectrum_id=str(i), peptidoform=f"{seq}/2", retention_time=5.0 + 2.5 * i) + for i, seq in enumerate(_PEPTIDES) + ] + ) + + +def test_report_end_to_end_with_index(tiny_index: TrainingIndex): + """One row per PSM with prediction, interval, membership and distances.""" + queries = PSMList( + psm_list=[ + PSM(spectrum_id="q0", peptidoform="AAGPSLSHTSGGTQSK/2"), # in reference and corpus + PSM(spectrum_id="q1", peptidoform="AGFAGDDAPM[Oxidation]R/2"), + PSM(spectrum_id="q2", peptidoform="WWWWWWWWWWWWWWWW/2"), + ] + ) + report = prediction_report( + queries, + psm_list_reference=_reference(), + training_index=tiny_index, + predict_kwargs={"device": "cpu"}, + ) + assert list(report.peptidoform) == [str(p.peptidoform) for p in queries.psm_list] + assert np.isfinite(report.predicted_rt).all() + assert (report.ci_lower <= report.predicted_rt).all() + assert (report.ci_upper >= report.predicted_rt).all() + assert report.attrs["coverage"] == 0.90 + + assert report.in_reference.tolist() == [True, False, False] + assert report.dist_to_reference.tolist()[0] == 0 + assert report.dist_to_reference.tolist()[2] > 5 + + assert report.in_training.tolist() == [True, False, False] + assert bool(report.in_selected_heads_training[0]) in (True, False) # depends on the head + + +def test_report_without_index_has_only_reference_columns(): + """The report works with nothing but the reference; corpus columns are absent.""" + report = prediction_report( + PSMList(psm_list=[PSM(spectrum_id="q", peptidoform="AGFAGDDAPR/2")]), + psm_list_reference=_reference(), + predict_kwargs={"device": "cpu"}, + ) + assert "in_training" not in report.columns + assert report.in_reference.tolist() == [True] + assert report.dist_to_reference.tolist() == [0] + + +def test_report_rejects_a_prefitted_calibration(): + """The report needs to fit per fold, so a fitted calibration cannot be reused.""" + from deeplc.calibration import SplineTransformerCalibration + + calibration = SplineTransformerCalibration() + calibration.fit(target=np.arange(20, dtype=np.float32), source=np.arange(20, dtype=np.float32)) + with pytest.raises(ValueError, match="unfitted"): + prediction_report( + PSMList(psm_list=[PSM(spectrum_id="q", peptidoform="PEPTIDEK/2")]), + psm_list_reference=_reference(), + calibration=calibration, + predict_kwargs={"device": "cpu"}, + ) + + +def test_report_with_multihead_calibration_lists_every_selected_head(tiny_index: TrainingIndex): + """With a multi-head calibration the membership covers every selected head.""" + from deeplc.calibration import MultiHeadRidgeCalibration + + report = prediction_report( + PSMList(psm_list=[PSM(spectrum_id="q", peptidoform="AGFAGDDAPR/2")]), + psm_list_reference=_reference(), + calibration=MultiHeadRidgeCalibration(n_heads=4), + training_index=tiny_index, + predict_kwargs={"device": "cpu"}, + ) + assert len(report.attrs["selected_heads"]) == 4 + assert "in_selected_heads_training" in report.columns From 50c360853fa186093088d93fc3e4b56822a733cf Mon Sep 17 00:00:00 2001 From: RobbinBouwmeester Date: Mon, 31 Aug 2026 14:20:50 +0200 Subject: [PATCH 2/7] feat: pack the training index into one 105 MB file The directory form was 400 MB across seven files, most of it uncompressed structure: raw 64-bit hashes, int64 pointers, plain text. The packed .dlcidx is a stdlib LZMA zip that exploits what each component actually is. Sorted hashes are truncated to 40 bits and stored as 2^24 bucket counts plus 16-bit remainders, which costs a false positive about once per 100,000 membership queries and nothing else; a provenance flag does not need exactness beyond that. CSR pointers become uint16 row lengths (5x under LZMA), the setup lists and the sorted sequences compress 2.8x and 3.1x. Loading rebuilds the sorted hash array in about a second; answers are bit-identical to the directory form on membership, per-setup membership and distances, which the tests now check by running every index test against both formats. TrainingIndex reads both forms; the builder emits both. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG.md | 12 ++++--- deeplc/report.py | 78 +++++++++++++++++++++++++++++++++---------- docs/source/usage.rst | 2 +- tests/test_report.py | 48 +++++++++++++++++++++++--- 4 files changed, 112 insertions(+), 28 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4206584..3db2035 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -26,11 +26,13 @@ and this project adheres to to 0.97 per setup (median 0.91), with widths from 4 % of the gradient on well-behaved setups to 79 % on a run that pools several fractions. Coverage is marginal, not per-peptide. -- `TrainingIndex`: a memory-mapped index of the multitask training corpus (10,105,640 - canonical peptidoform keys, their 65,139,832 setup observations, 6,157,558 unique stripped - sequences; about 400 MB on disk). Built offline from the training cache and distributed - separately from the package; `prediction_report` takes it as an optional argument and works - without it. +- `TrainingIndex`: an index of the multitask training corpus (10,105,640 canonical + peptidoform keys, their 65,139,832 setup observations, 6,157,558 unique stripped sequences). + Distributed separately from the package as a single 105 MB `.dlcidx` file: an LZMA zip + holding 40-bit key hashes in a bucketed layout (false positive about once per 100,000 + membership queries, irrelevant for a provenance flag), per-key setup lists and the unique + sequences. A raw memory-mapped directory form with exact 64-bit hashes is read as well. + `prediction_report` takes either as an optional argument and works without one. - Dependency: `rapidfuzz` (Levenshtein distances). diff --git a/deeplc/report.py b/deeplc/report.py index 062378a..643d2ea 100644 --- a/deeplc/report.py +++ b/deeplc/report.py @@ -84,32 +84,68 @@ def token(mod) -> str: class TrainingIndex: """ - Memory-mapped index of the corpus behind the bundled multitask model. + Index of the corpus behind the bundled multitask model. - Built offline from the training cache (10,105,640 canonical peptidoform keys and their - 65,139,832 peptidoform-setup observations over 6,543 setups) and loaded from a directory: - ``key_hashes.npy`` (sorted xxh3-64 of the canonical keys), ``task_indptr.npy`` / - ``task_cols.npy`` (which setups each peptidoform was observed in), ``sequences.npy`` / - ``seq_lengths.npy`` (unique stripped sequences, for edit distances) and ``meta.json``. + Answers, for any canonical peptidoform key: was it trained on at all, was it trained on + within given setups, and how far is its sequence from the closest training sequence. Built + offline from the training cache (10,105,640 canonical keys, 65,139,832 peptidoform-setup + observations over 6,543 setups) and distributed separately from the package. - Everything is memory-mapped, so opening the index costs nothing until it is used. + Two on-disk forms are read: + + - a single ``.dlcidx`` file (format 2): an LZMA-compressed zip holding 40-bit key hashes in + a bucketed layout, per-key setup lists and the unique sequences; about 105 MB. Membership + through 40-bit hashes can produce a false positive roughly once per 100,000 queries, + which is negligible for a provenance flag; + - a directory with ``key_hashes.npy`` (full 64-bit, exact), ``task_indptr.npy``, + ``task_cols.npy``, ``sequences.txt`` and ``meta.json`` (format 1, memory-mapped). """ def __init__(self, path: PathLike | str) -> None: - """Open a training index directory.""" + """Open a packed ``.dlcidx`` file or a training index directory.""" self.path = Path(path) - meta_file = self.path / "meta.json" - if not meta_file.exists(): + self._sequences: np.ndarray | None = None + self._seq_lengths: np.ndarray | None = None + if self.path.is_file(): + self._open_packed() + elif (self.path / "meta.json").exists(): + self._open_directory() + else: raise FileNotFoundError( - f"{self.path} is not a training index (no meta.json). It is built offline " - "from the training cache and distributed separately from the package." + f"{self.path} is not a training index (neither a .dlcidx file nor a directory " + "with meta.json). It is built offline from the training cache and distributed " + "separately from the package." ) - self.meta = json.loads(meta_file.read_text(encoding="utf-8")) + + def _open_directory(self) -> None: + self.meta = json.loads((self.path / "meta.json").read_text(encoding="utf-8")) + self._hash_shift = 0 self._hashes = np.load(self.path / "key_hashes.npy", mmap_mode="r") - self._indptr = np.load(self.path / "task_indptr.npy", mmap_mode="r") + indptr = np.load(self.path / "task_indptr.npy", mmap_mode="r") + self._indptr = np.asarray(indptr, dtype=np.int64) self._cols = np.load(self.path / "task_cols.npy", mmap_mode="r") - self._sequences: np.ndarray | None = None - self._seq_lengths: np.ndarray | None = None + + def _open_packed(self) -> None: + import zipfile + + with zipfile.ZipFile(self.path) as archive: + self.meta = json.loads(archive.read("meta.json").decode("utf-8")) + if int(self.meta.get("format_version", 0)) != 2: + raise ValueError( + f"{self.path} has format_version {self.meta.get('format_version')}; " + "this DeepLC reads format 2." + ) + counts = np.frombuffer(archive.read("hash_bucket_counts.u8"), dtype=np.uint8) + remainders = np.frombuffer(archive.read("hash_remainders.u16"), dtype=np.uint16) + row_lengths = np.frombuffer(archive.read("row_lengths.u16"), dtype=np.uint16) + self._cols = np.frombuffer(archive.read("task_cols.i16"), dtype=np.int16) + self._sequences_blob = archive.read("sequences.txt") + highs = np.repeat(np.arange(len(counts), dtype=np.uint64), counts) + self._hashes = (highs << np.uint64(16)) | remainders.astype(np.uint64) + self._hash_shift = 64 - int(self.meta["hash_bits"]) + indptr = np.zeros(len(row_lengths) + 1, dtype=np.int64) + np.cumsum(row_lengths, out=indptr[1:]) + self._indptr = indptr @staticmethod def _hash(keys: list[str]) -> np.ndarray: @@ -126,6 +162,8 @@ def digest(text: str) -> int: def _rows(self, keys: list[str]) -> np.ndarray: """Index of each key in the sorted hash array, or -1 when absent.""" hashes = self._hash(keys) + if self._hash_shift: + hashes = hashes >> np.uint64(self._hash_shift) position = np.searchsorted(self._hashes, hashes) position = np.clip(position, 0, len(self._hashes) - 1) found = self._hashes[position] == hashes @@ -178,9 +216,13 @@ def distance_to_training(self, sequences: list[str], max_distance: int = 10) -> from rapidfuzz.process import cdist if self._sequences is None: - blob = (self.path / "sequences.txt").read_bytes().decode("ascii") + if hasattr(self, "_sequences_blob"): + blob = self._sequences_blob.decode("ascii") + del self._sequences_blob + else: + blob = (self.path / "sequences.txt").read_bytes().decode("ascii") self._sequences = np.array(blob.split(chr(10)), dtype=object) - self._seq_lengths = np.load(self.path / "seq_lengths.npy") + self._seq_lengths = np.array([len(x) for x in self._sequences], dtype=np.int16) unique, inverse = np.unique(np.asarray(sequences, dtype=object), return_inverse=True) exact = np.isin(unique, self._sequences) per_unique = np.full(len(unique), -1, dtype=np.int32) diff --git a/docs/source/usage.rst b/docs/source/usage.rst index 4457bf7..80e6af6 100644 --- a/docs/source/usage.rst +++ b/docs/source/usage.rst @@ -105,7 +105,7 @@ edits and capped beyond): .. code-block:: python report = prediction_report(psm_list, psm_list_reference=reference, - training_index="path/to/training_index_v6f") + training_index="deeplc_training_index_v6f.dlcidx") Python API ========== diff --git a/tests/test_report.py b/tests/test_report.py index 37b8e58..8bc120a 100644 --- a/tests/test_report.py +++ b/tests/test_report.py @@ -108,9 +108,9 @@ def test_thin_bins_fall_back_to_the_global_quantile(): # training index, built small and on the fly -@pytest.fixture() -def tiny_index(tmp_path: Path) -> TrainingIndex: - """Three peptidoforms over three setups, written in the real on-disk format.""" +@pytest.fixture(params=["directory", "packed"]) +def tiny_index(request, tmp_path: Path) -> TrainingIndex: + """Three peptidoforms over three setups, in both on-disk formats.""" keys = ["AAGPSLSHTSGGTQSK|", "AGFAGDDAPR|7|U:35", "LNLSPLGEEMR|"] tasks = [[0], [0, 2], [1]] hashes = TrainingIndex._hash(keys) @@ -130,7 +130,36 @@ def tiny_index(tmp_path: Path) -> TrainingIndex: (tmp_path / "meta.json").write_text( json.dumps({"format_version": 1, "n_peptidoforms": 3, "n_tasks": 3, "n_observations": 4}) ) - return TrainingIndex(tmp_path) + if request.param == "directory": + return TrainingIndex(tmp_path) + + import zipfile + + h40 = (hashes[order] >> np.uint64(24)).astype(np.uint64) + counts = np.bincount((h40 >> np.uint64(16)).astype(np.int64), minlength=1 << 24) + packed = tmp_path / "tiny.dlcidx" + with zipfile.ZipFile(packed, "w", compression=zipfile.ZIP_LZMA) as archive: + archive.writestr( + "meta.json", + json.dumps( + { + "format_version": 2, + "hash_bits": 40, + "n_tasks": 3, + "n_peptidoforms": 3, + "n_observations": 4, + } + ), + ) + archive.writestr("hash_bucket_counts.u8", counts.astype(np.uint8).tobytes()) + archive.writestr( + "hash_remainders.u16", (h40 & np.uint64(0xFFFF)).astype(np.uint16).tobytes() + ) + archive.writestr("row_lengths.u16", np.diff(indptr).astype(np.uint16).tobytes()) + archive.writestr("task_cols.i16", np.array(cols, dtype=np.int16).tobytes()) + archive.writestr("sequences.txt", chr(10).join(sequences).encode("ascii")) + archive.writestr("task_names.json", json.dumps(["setup_a", "setup_b", "setup_c"])) + return TrainingIndex(packed) def test_index_membership_and_per_task_membership(tiny_index: TrainingIndex): @@ -160,6 +189,17 @@ def test_index_refuses_a_directory_that_is_not_an_index(tmp_path: Path): TrainingIndex(tmp_path) +def test_packed_index_with_an_unknown_format_version_is_refused(tmp_path: Path): + """A future format fails loudly instead of being misread.""" + import zipfile + + packed = tmp_path / "future.dlcidx" + with zipfile.ZipFile(packed, "w") as archive: + archive.writestr("meta.json", json.dumps({"format_version": 99})) + with pytest.raises(ValueError, match="format_version"): + TrainingIndex(packed) + + # --------------------------------------------------------------------------- # # the full report From c4efa1e7134262e16ff25292d100255c4eb900b4 Mon Sep 17 00:00:00 2001 From: RobbinBouwmeester Date: Thu, 3 Sep 2026 14:52:29 +0200 Subject: [PATCH 3/7] feat: give each peptide its own prediction interval The conformal half-width was a per-RT-bin quantile, so a setup received five distinct widths and two peptides predicted at the same retention time always got the same interval. A multi-head calibration combines heads that each estimate the same retention time, and how far those estimates lie apart varies per peptide. Calibration instances can now report that as disagreement(); the conformal interval divides the honest residuals by it before taking the per-bin quantile and multiplies it back at prediction time, which keeps the coverage guarantee and the RT structure while making the width follow the peptide. Measured through the public API on the six held-out PRIDE setups, against the per-bin widths: worst conditional slice 0.851 -> 0.882, Spearman of width against absolute error 0.151 -> 0.248, coverage 0.909 -> 0.919, relative width 0.0436 -> 0.0478, distinct widths 5 -> 877. The gains are largest where the per-bin width was weakest (PXD080826 0.818 -> 0.888, PXD081924 0.814 -> 0.845). Edit distance to the reference was tested as the scale instead and rejected: three times the width, coverage 0.977 and no correlation with the error. per_peptide_width=False restores the previous behaviour, which also remains the behaviour of single-head calibrations. Co-Authored-By: Claude Fable 5 --- CHANGELOG.md | 13 +++++ deeplc/calibration.py | 41 ++++++++++++++- deeplc/report.py | 79 +++++++++++++++++++++++++---- docs/source/usage.rst | 6 +++ tests/test_multihead_calibration.py | 31 +++++++++++ tests/test_report.py | 69 +++++++++++++++++++++++++ 6 files changed, 226 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3db2035..7675fb0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -26,6 +26,19 @@ and this project adheres to to 0.97 per setup (median 0.91), with widths from 4 % of the gradient on well-behaved setups to 79 % on a run that pools several fractions. Coverage is marginal, not per-peptide. + With a multi-head calibration the width is also **per peptide** (`per_peptide_width`, on by + default): the residuals are divided by how far the combined setup heads lie apart for that + peptide before the quantile is taken, and multiplied by it again at prediction time. Two + peptides predicted at the same retention time therefore no longer share one interval. On the + six held-out setups this raised the worst conditional slice from 0.851 to 0.882 and the + Spearman correlation between width and error from 0.15 to 0.25, for 10 % wider intervals; + the largest gains are on the setups where the RT-only width was weakest. Set + `per_peptide_width=False` for widths that depend on the predicted retention time alone. + +- `Calibration.disagreement`, the per-input uncertainty a calibration can report, implemented + by `MultiHeadRidgeCalibration` as the ridge-weighted spread of its calibrated head + estimates and returning None elsewhere. + - `TrainingIndex`: an index of the multitask training corpus (10,105,640 canonical peptidoform keys, their 65,139,832 setup observations, 6,157,558 unique stripped sequences). Distributed separately from the package as a single 105 MB `.dlcidx` file: an LZMA zip diff --git a/deeplc/calibration.py b/deeplc/calibration.py index d61867a..dad03ed 100644 --- a/deeplc/calibration.py +++ b/deeplc/calibration.py @@ -46,6 +46,16 @@ def transform(self, source: np.ndarray) -> np.ndarray: """Transform source values into the calibrated target space.""" ... + def disagreement(self, source: np.ndarray) -> np.ndarray | None: # noqa: ARG002 + """ + Per-input uncertainty score, or None when the calibration has none. + + A calibration that combines several estimates of the same retention time can report + how far they lie apart for each input, which :func:`deeplc.report.prediction_report` + uses to scale its prediction intervals per peptide. + """ + return None + class IdentityCalibration(Calibration): """No calibration; returns inputs unchanged.""" @@ -444,13 +454,40 @@ def transform(self, source: np.ndarray) -> np.ndarray: ) if source.shape[0] == 0: return np.array([]) - calibrated = np.column_stack( + return np.asarray(self._ridge.predict(self._calibrated_columns(source)), dtype=np.float64) + + def _calibrated_columns(self, source: np.ndarray) -> np.ndarray: + """Each selected head's own estimate of the retention time, in the reference's unit.""" + head_idx = cast(np.ndarray, self._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) ] ) - return np.asarray(self._ridge.predict(calibrated), dtype=np.float64) + + def disagreement(self, source: np.ndarray) -> np.ndarray | None: + """ + How far the combined setup heads lie apart for each input, in the reference's unit. + + Every selected head estimates the retention time of the same peptide, so the spread of + those estimates, weighted by the ridge weight each head received, is an uncertainty + that varies from peptide to peptide rather than only along the gradient. Returns None + while the calibration is unfitted or combines a single head, which carries no spread. + """ + 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: + 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) + mean = estimates @ weights + return np.sqrt(((estimates - mean[:, None]) ** 2) @ weights) def _rank_heads_by_correlation(source: np.ndarray, target: np.ndarray) -> np.ndarray: diff --git a/deeplc/report.py b/deeplc/report.py index 643d2ea..462c7fc 100644 --- a/deeplc/report.py +++ b/deeplc/report.py @@ -47,6 +47,10 @@ _MIN_RESIDUALS_PER_BIN = 40 _N_FOLDS = 5 +#: Range the per-peptide difficulty score may scale an interval by, relative to the median +#: peptide of the reference. +_RATIO_CLIP = (0.2, 5.0) + def canonical_peptidoform_key(peptidoform: Peptidoform | str) -> str: """ @@ -250,11 +254,20 @@ def distance_to_training(self, sequences: list[str], max_distance: int = 10) -> @dataclass class _ConformalInterval: - """RT-binned conformal half-widths, fitted on honest reference residuals.""" + """ + Conformal half-widths per RT bin, fitted on honest reference residuals. + + With a per-input difficulty score, the residuals are divided by that score before the + quantile is taken and multiplied by it again at prediction time, so peptides predicted at + the same retention time no longer share one width. Without a score the width depends on + the predicted retention time alone. + """ coverage: float edges: np.ndarray = field(default_factory=lambda: np.array([])) half_width: np.ndarray = field(default_factory=lambda: np.array([])) + scale: float | None = None + floor: float = 0.0 @staticmethod def _finite_sample_quantile(abs_residuals: np.ndarray, coverage: float) -> float: @@ -262,12 +275,26 @@ def _finite_sample_quantile(abs_residuals: np.ndarray, coverage: float) -> float rank = min(int(np.ceil((n + 1) * coverage)), n) return float(np.sort(abs_residuals)[rank - 1]) + def _ratio(self, difficulty: np.ndarray) -> np.ndarray: + bounded = np.maximum(np.asarray(difficulty, dtype=float), self.floor) + return np.clip(bounded / self.scale, *_RATIO_CLIP) + @classmethod def fit( - cls, predicted: np.ndarray, residuals: np.ndarray, coverage: float + cls, + predicted: np.ndarray, + residuals: np.ndarray, + coverage: float, + difficulty: np.ndarray | None = None, ) -> _ConformalInterval: """Per-RT-bin conformal quantiles with a global fallback for thin bins.""" + interval = cls(coverage=coverage) absolute = np.abs(residuals) + if difficulty is not None: + difficulty = np.asarray(difficulty, dtype=float) + interval.floor = max(float(np.quantile(difficulty, 0.05)), np.finfo(float).tiny) + interval.scale = float(np.median(np.maximum(difficulty, interval.floor))) + absolute = absolute / interval._ratio(difficulty) overall = cls._finite_sample_quantile(absolute, coverage) edges = np.quantile(predicted, np.linspace(0, 1, _N_RT_BINS + 1)) edges[0], edges[-1] = -np.inf, np.inf @@ -277,12 +304,23 @@ def fit( mask = bins == b if int(mask.sum()) >= _MIN_RESIDUALS_PER_BIN: half_width[b] = cls._finite_sample_quantile(absolute[mask], coverage) - return cls(coverage=coverage, edges=edges, half_width=half_width) + interval.edges, interval.half_width = edges, half_width + return interval - def widths(self, predicted: np.ndarray) -> np.ndarray: + def widths( + self, predicted: np.ndarray, difficulty: np.ndarray | None = None + ) -> np.ndarray: """Interval half-width for each prediction.""" bins = np.clip(np.searchsorted(self.edges, predicted, side="right") - 1, 0, _N_RT_BINS - 1) - return self.half_width[bins] + widths = self.half_width[bins] + if self.scale is None: + return widths + if difficulty is None: + raise ValueError( + "This interval was fitted with a per-peptide difficulty score, so it needs " + "one to produce widths." + ) + return widths * self._ratio(difficulty) def _crossfit_residuals( @@ -290,11 +328,12 @@ def _crossfit_residuals( matrix_reference: np.ndarray, calibration_template: Calibration, seed: int = 0, -) -> tuple[np.ndarray, np.ndarray]: +) -> tuple[np.ndarray, np.ndarray, np.ndarray | None]: """ Honest reference residuals: each fold predicted by a calibration fitted without it. - Returns (cross-fitted predictions, residuals), aligned with the reference order. The + Returns (cross-fitted predictions, residuals, difficulty scores), aligned with the + reference order; the scores are None when the calibration reports no disagreement. The template is re-instantiated per fold with ``type(...)()`` semantics via a deep copy of its construction parameters, so a fitted calibration is never reused across folds. """ @@ -304,6 +343,7 @@ def _crossfit_residuals( order = rng.permutation(len(y_reference)) folds = np.array_split(order, min(_N_FOLDS, max(2, len(y_reference) // 25))) predicted = np.empty(len(y_reference)) + difficulty: np.ndarray | None = np.empty(len(y_reference)) for i, fold in enumerate(folds): train = np.concatenate([f for j, f in enumerate(folds) if j != i]) calibration = copy.deepcopy(calibration_template) @@ -321,7 +361,12 @@ def _crossfit_residuals( calibration.transform(matrix_reference[fold][:, head].astype(np.float32)), dtype=float, ) - return predicted, y_reference - predicted + fold_difficulty = calibration.disagreement(matrix_reference[fold]) + if difficulty is None or fold_difficulty is None: + difficulty = None + else: + difficulty[fold] = np.asarray(fold_difficulty, dtype=float) + return predicted, y_reference - predicted, difficulty def prediction_report( @@ -331,6 +376,7 @@ def prediction_report( calibration: Calibration | None = None, coverage: float = 0.90, training_index: TrainingIndex | PathLike | str | None = None, + per_peptide_width: bool = True, predict_kwargs: dict | None = None, ) -> pd.DataFrame: """ @@ -355,6 +401,12 @@ def prediction_report( training_index A :class:`TrainingIndex` or a path to one. Without it, the columns about the training corpus are omitted and the report is limited to the reference. + per_peptide_width + Scale each interval by how far the combined setup heads lie apart for that peptide, so + two peptides predicted at the same retention time can get different intervals. Ignored + with a calibration that reports no such disagreement, such as + :class:`SplineTransformerCalibration`, where the width depends on the predicted + retention time alone. predict_kwargs Passed to the prediction function (``{"device": "cpu"}`` and the like). @@ -413,9 +465,14 @@ def prediction_report( ) selected_heads = np.array([head], dtype=int) - cross_predicted, residuals = _crossfit_residuals(y_reference, matrix_reference, template) - interval = _ConformalInterval.fit(cross_predicted, residuals, coverage) - half_width = interval.widths(np.asarray(predicted, dtype=float)) + cross_predicted, residuals, cross_difficulty = _crossfit_residuals( + y_reference, matrix_reference, template + ) + query_difficulty = calibration.disagreement(matrix_query) if per_peptide_width else None + if cross_difficulty is None or query_difficulty is None: + cross_difficulty = query_difficulty = None + interval = _ConformalInterval.fit(cross_predicted, residuals, coverage, cross_difficulty) + half_width = interval.widths(np.asarray(predicted, dtype=float), query_difficulty) # membership and novelty against the reference reference_keys = {canonical_peptidoform_key(psm.peptidoform) for psm in reference.psm_list} diff --git a/docs/source/usage.rst b/docs/source/usage.rst index 80e6af6..f75f857 100644 --- a/docs/source/usage.rst +++ b/docs/source/usage.rst @@ -96,6 +96,12 @@ holds on peptides exchangeable with the reference, without retraining and regard model. Pass ``calibration=MultiHeadRidgeCalibration()`` to combine setups; the membership column then covers every selected head. +The width of that interval depends on the predicted retention time and, with a multi-head +calibration, on the peptide itself: the combined setup heads each estimate the same retention +time, and how far those estimates lie apart is an uncertainty that varies per peptide. Two +peptides predicted at the same retention time therefore get different intervals. Pass +``per_peptide_width=False`` for widths that depend on the predicted retention time alone. + With a training index (built from the multitask training corpus and distributed separately), three more columns appear: ``in_training`` (exact peptidoform match anywhere in the corpus), ``in_selected_heads_training`` (match within the setups the calibration selected) and diff --git a/tests/test_multihead_calibration.py b/tests/test_multihead_calibration.py index ecab51d..d2170ae 100644 --- a/tests/test_multihead_calibration.py +++ b/tests/test_multihead_calibration.py @@ -108,6 +108,37 @@ def test_never_fits_more_weights_than_half_the_reference(): assert len(calibration._head_calibrations) <= 5 +def test_disagreement_is_per_input_and_zero_only_when_heads_agree(): + """The spread of the calibrated heads varies from input to input.""" + target, source = _synthetic(n_heads=12) + calibration = MultiHeadRidgeCalibration(n_heads=6) + assert calibration.disagreement(source) is None # unfitted + calibration.fit(target=target, source=source) + + spread = calibration.disagreement(source) + assert spread.shape == target.shape + assert (spread >= 0).all() + assert np.unique(spread.round(9)).size > len(target) // 2 + + # heads that are affine views of one latent retention time calibrate onto each other, so + # after calibration they agree and the spread collapses + rng = np.random.default_rng(0) + latent = rng.uniform(0, 100, len(target)) + scales, shifts = rng.uniform(0.5, 2, 12), rng.uniform(-20, 20, 12) + agreeing = latent[:, None] * scales + shifts + agreed = MultiHeadRidgeCalibration(n_heads=6) + agreed.fit(target=latent, source=agreeing) + assert agreed.disagreement(agreeing).mean() < 0.05 * spread.mean() + + +def test_single_head_combination_reports_no_disagreement(): + """One head carries no spread, so there is nothing to scale an interval by.""" + target, source = _synthetic(n_heads=1) + calibration = MultiHeadRidgeCalibration() + calibration.fit(target=target, source=source[:, 0]) + assert calibration.disagreement(source[:, 0]) is None + + def test_rejects_a_nonsensical_head_count(): """Zero heads cannot calibrate anything.""" with pytest.raises(ValueError, match="at least 1"): diff --git a/tests/test_report.py b/tests/test_report.py index 8bc120a..f7e8c30 100644 --- a/tests/test_report.py +++ b/tests/test_report.py @@ -95,6 +95,35 @@ def test_interval_is_wider_where_residuals_are_wider(): assert interval.widths(np.array([90.0]))[0] > 2 * interval.widths(np.array([10.0]))[0] +def test_difficulty_score_gives_each_input_its_own_width(): + """With a per-input score, two inputs at the same predicted RT get different widths.""" + rng = np.random.default_rng(3) + predicted = rng.uniform(0, 100, 3000) + difficulty = rng.uniform(0.5, 4.0, 3000) + residuals = rng.normal(0, difficulty, 3000) + interval = _ConformalInterval.fit(predicted, residuals, 0.90, difficulty) + + widths = interval.widths(np.full(2, 50.0), np.array([0.6, 3.5])) + assert widths[1] > 2 * widths[0] + + new_predicted = rng.uniform(0, 100, 3000) + new_difficulty = rng.uniform(0.5, 4.0, 3000) + covered = np.abs(rng.normal(0, new_difficulty, 3000)) <= interval.widths( + new_predicted, new_difficulty + ) + assert 0.87 <= covered.mean() <= 0.94 + + +def test_difficulty_scaled_interval_needs_a_score_to_predict_with(): + """An interval fitted on a difficulty score cannot silently drop it.""" + rng = np.random.default_rng(4) + predicted = rng.uniform(0, 100, 500) + difficulty = rng.uniform(1, 2, 500) + interval = _ConformalInterval.fit(predicted, rng.normal(0, 1, 500), 0.90, difficulty) + with pytest.raises(ValueError, match="difficulty"): + interval.widths(predicted) + + def test_thin_bins_fall_back_to_the_global_quantile(): """Too few residuals per bin means one global width, not five noisy ones.""" rng = np.random.default_rng(2) @@ -269,6 +298,46 @@ def test_report_rejects_a_prefitted_calibration(): ) +def test_report_widths_vary_per_peptide_with_a_multihead_calibration(): + """Peptides get their own interval; per_peptide_width=False restores the RT-only widths.""" + from deeplc.calibration import MultiHeadRidgeCalibration + + queries = PSMList( + psm_list=[PSM(spectrum_id=str(i), peptidoform=f"{seq}/2") + for i, seq in enumerate(_PEPTIDES)] + ) + per_peptide, per_bin = ( + prediction_report( + queries, + psm_list_reference=_reference(), + calibration=MultiHeadRidgeCalibration(n_heads=8), + per_peptide_width=flag, + predict_kwargs={"device": "cpu"}, + ) + for flag in (True, False) + ) + widths = (per_peptide["ci_upper"] - per_peptide["ci_lower"]).round(9) + binned_widths = (per_bin["ci_upper"] - per_bin["ci_lower"]).round(9) + assert widths.nunique() > binned_widths.nunique() + assert (widths > 0).all() + + +def test_report_falls_back_to_rt_only_widths_without_disagreement(): + """A single-head calibration has no per-peptide signal, so the flag changes nothing.""" + queries = PSMList( + psm_list=[PSM(spectrum_id=str(i), peptidoform=f"{seq}/2") + for i, seq in enumerate(_PEPTIDES)] + ) + report = prediction_report( + queries, + psm_list_reference=_reference(), + per_peptide_width=True, + predict_kwargs={"device": "cpu"}, + ) + widths = (report["ci_upper"] - report["ci_lower"]).round(9) + assert widths.nunique() <= 5 + + def test_report_with_multihead_calibration_lists_every_selected_head(tiny_index: TrainingIndex): """With a multi-head calibration the membership covers every selected head.""" from deeplc.calibration import MultiHeadRidgeCalibration From 51eb4b99cdd992914e6263a6601969f1c94a160c Mon Sep 17 00:00:00 2001 From: RobbinBouwmeester Date: Thu, 3 Sep 2026 14:54:40 +0200 Subject: [PATCH 4/7] style: apply ruff format Co-Authored-By: Claude Fable 5 --- deeplc/report.py | 4 +--- tests/test_report.py | 10 ++++++---- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/deeplc/report.py b/deeplc/report.py index 462c7fc..49201f7 100644 --- a/deeplc/report.py +++ b/deeplc/report.py @@ -307,9 +307,7 @@ def fit( interval.edges, interval.half_width = edges, half_width return interval - def widths( - self, predicted: np.ndarray, difficulty: np.ndarray | None = None - ) -> np.ndarray: + def widths(self, predicted: np.ndarray, difficulty: np.ndarray | None = None) -> np.ndarray: """Interval half-width for each prediction.""" bins = np.clip(np.searchsorted(self.edges, predicted, side="right") - 1, 0, _N_RT_BINS - 1) widths = self.half_width[bins] diff --git a/tests/test_report.py b/tests/test_report.py index f7e8c30..bf26493 100644 --- a/tests/test_report.py +++ b/tests/test_report.py @@ -303,8 +303,9 @@ def test_report_widths_vary_per_peptide_with_a_multihead_calibration(): from deeplc.calibration import MultiHeadRidgeCalibration queries = PSMList( - psm_list=[PSM(spectrum_id=str(i), peptidoform=f"{seq}/2") - for i, seq in enumerate(_PEPTIDES)] + psm_list=[ + PSM(spectrum_id=str(i), peptidoform=f"{seq}/2") for i, seq in enumerate(_PEPTIDES) + ] ) per_peptide, per_bin = ( prediction_report( @@ -325,8 +326,9 @@ def test_report_widths_vary_per_peptide_with_a_multihead_calibration(): def test_report_falls_back_to_rt_only_widths_without_disagreement(): """A single-head calibration has no per-peptide signal, so the flag changes nothing.""" queries = PSMList( - psm_list=[PSM(spectrum_id=str(i), peptidoform=f"{seq}/2") - for i, seq in enumerate(_PEPTIDES)] + psm_list=[ + PSM(spectrum_id=str(i), peptidoform=f"{seq}/2") for i, seq in enumerate(_PEPTIDES) + ] ) report = prediction_report( queries, From 8a19b6283d2829946334c3f897a475aa2e02ffd5 Mon Sep 17 00:00:00 2001 From: RobbinBouwmeester Date: Tue, 8 Sep 2026 16:26:09 +0200 Subject: [PATCH 5/7] 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 6/7] 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 7/7] 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")