diff --git a/CHANGELOG.md b/CHANGELOG.md index 170ba83..3db2035 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,36 @@ 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`: 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). + ## [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..643d2ea --- /dev/null +++ b/deeplc/report.py @@ -0,0 +1,461 @@ +""" +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: + """ + Index of the corpus behind the bundled multitask model. + + 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. + + 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 packed ``.dlcidx`` file or a training index directory.""" + self.path = Path(path) + 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 (neither a .dlcidx file nor a directory " + "with meta.json). It is built offline from the training cache and distributed " + "separately from the package." + ) + + 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") + 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") + + 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: + 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) + 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 + 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: + 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.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) + 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..80e6af6 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="deeplc_training_index_v6f.dlcidx") + 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..8bc120a --- /dev/null +++ b/tests/test_report.py @@ -0,0 +1,284 @@ +"""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(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) + 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}) + ) + 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): + """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) + + +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 + + +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