Prediction reports: training-set membership, novelty and conformal intervals (v4.4.0)

CHANGELOG.md

@@ -6,6 +6,49 @@ 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.
+
+  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
+  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
 
 ### Added

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",

deeplc/calibration/multihead.py

@@ -28,6 +28,24 @@
 LOGGER = logging.getLogger(__name__)
 
 
+def as_head_matrix(source):
+    """
+    Coerce a source to a two-dimensional float64 matrix, unless it is a head source.
+
+    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.
+    """
+    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):
     """Abstract base class for a calibration that selects its own head(s) from a matrix."""
 
@@ -57,6 +75,17 @@ def transform(self, source: np.ndarray) -> np.ndarray:
         """Transform the source matrix 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. One fitted on a single head has
+        no spread to report and keeps this default.
+        """
+        return None
+
 
 class _SingleHeadCalibration(MultiHeadCalibration):
     """
@@ -120,9 +149,7 @@ 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]
+        source = as_head_matrix(source)
         head = self.selected_model_head
         if source.shape[1] <= head:
             raise CalibrationError(
@@ -132,7 +159,10 @@ def transform(self, source: np.ndarray) -> np.ndarray:
         if source.shape[0] == 0:
             return np.array([])
         return np.asarray(
-            self._inner.transform(source[:, head].astype(np.float32)), dtype=np.float64
+            self._inner.transform(
+                np.asarray(source[:, head], dtype=np.float64).astype(np.float32)
+            ),
+            dtype=np.float64,
         )
 
 
@@ -285,9 +315,7 @@ 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]
+        source = as_head_matrix(source)
         head_idx = cast(np.ndarray, self._head_idx)
         if source.shape[1] <= int(head_idx.max()):
             raise CalibrationError(
@@ -296,13 +324,41 @@ 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:
+        """Give each selected head's own estimate of the retention time, in reference units."""
+        head_idx = cast(np.ndarray, self._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(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[:, i].astype(np.float32)), dtype=np.float64)
+                for i, cal in enumerate(self._head_calibrations)
             ]
         )
-        return np.asarray(self._ridge.predict(calibrated), dtype=np.float64)
+
+    def disagreement(self, source: np.ndarray) -> np.ndarray | None:
+        """
+        Report how far the combined setup heads lie apart for each input, in reference units.
+
+        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 = 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(columns)
+        mean = estimates @ weights
+        return np.sqrt(((estimates - mean[:, None]) ** 2) @ weights)
 
 
 def upgrade_calibration(calibration: Calibration | MultiHeadCalibration) -> MultiHeadCalibration:

deeplc/core.py

@@ -214,6 +214,105 @@ 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, 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.
+
+    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.
+
+    """
+
+    #: 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
+    ):
+        """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 __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 = (
+                wanted,
+                predict(
+                    self._psm_list,
+                    model=self._model,
+                    predict_kwargs={**self._predict_kwargs, "task_idx": list(wanted)},
+                    return_matrix=True,
+                ),
+            )
+        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."""
+        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 +360,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)