From bc3d3092d7098167886251acf62e051a3440df5c Mon Sep 17 00:00:00 2001 From: RobbinBouwmeester Date: Mon, 31 Aug 2026 10:57:18 +0200 Subject: [PATCH 1/2] 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/2] 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