Calibrate against several LC setups at once (v4.3.0)

CHANGELOG.md

@@ -6,6 +6,38 @@ 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.3.0] - 2026-09-02
+
+### Changed
+
+- Calibration of a multitask model now combines several LC-setup heads instead of keeping only
+  the best-correlating one. `calibrate` and `predict_and_calibrate` default to the new
+  `MultiHeadRidgeCalibration`: heads are ranked by Pearson correlation to the reference as
+  before, the 80 best are each calibrated with `SplineTransformerCalibration`, and a ridge
+  regression maps the calibrated estimates onto the observed retention times.
+
+  On the eight PRIDE setups that no DeepLC model was trained on this lowered the held-out error
+  on all eight, by a median of 13 % relative to the observed gradient (0.01248 to 0.01090
+  MAE/span). Fitting is faster than the previous path because the head ranking is vectorised
+  (median 1.0 s against 2.3 s), and prediction is unchanged since the full head matrix is
+  computed either way.
+
+  Single-task models keep the previous default (`SplineTransformerCalibration`), and passing
+  a calibration instance restores the old behaviour on any model:
+
+  ```python
+  from deeplc import predict_and_calibrate
+  from deeplc.calibration import SplineTransformerCalibration
+
+  rt = predict_and_calibrate(psms, psm_list_reference=reference,
+                             calibration=SplineTransformerCalibration())
+  ```
+
+### Added
+
+- `Calibration.uses_all_heads`, telling `calibrate` and `predict_and_calibrate` to hand a
+  calibration the whole `(n, n_heads)` prediction matrix rather than a single column.
+
 ## [4.2.0] - 2026-08-28
 
 ### Changed

CLAUDE.md

@@ -0,0 +1,23 @@
+# Instructions for AI assistants
+
+## Documentation style
+
+Documentation must be self-contained, not conversation-dependent. Comments, docstrings, and
+PR descriptions must be understandable to a reader with no access to the conversation or
+reasoning that produced the code.
+
+- Comments and docstrings describe the current state and behavior of the code: what it does
+  and how it is used. They must not read as a log of the development process.
+- Historical context (why a previous approach was replaced, what motivated a design choice)
+  belongs in commit messages and PR descriptions, not in comments or docstrings, unless it is
+  strictly necessary to understand the current implementation, in which case keep it brief.
+- Benchmark results belong in the PR description and the changelog, not in docstrings.
+- Use inline comments only where the code is not self-explanatory.
+- PR descriptions are the place for process: reasoning behind the change, alternatives
+  considered, and measurements.
+
+## Project conventions
+
+- Follow [Keep a Changelog](https://keepachangelog.com/) in `CHANGELOG.md` and semantic
+  versioning in `pyproject.toml`.
+- Run the test suite with `pytest tests` before pushing.

deeplc/calibration.py

@@ -7,7 +7,7 @@
 from typing import cast
 
 import numpy as np
-from sklearn.linear_model import LinearRegression  # type: ignore[import]
+from sklearn.linear_model import LinearRegression, RidgeCV  # type: ignore[import]
 from sklearn.pipeline import Pipeline, make_pipeline  # type: ignore[import]
 from sklearn.preprocessing import SplineTransformer  # type: ignore[import]
 
@@ -21,6 +21,10 @@ class Calibration(ABC):
 
     selected_model_head: int | None = None
 
+    #: Whether ``fit`` and ``transform`` take the full ``(n, n_heads)`` prediction matrix
+    #: instead of a single series.
+    uses_all_heads: bool = False
+
     @abstractmethod
     def __init__(self, *args, **kwargs):
         """Initialize the calibration model."""
@@ -325,6 +329,146 @@ def transform(self, source: np.ndarray) -> np.ndarray:
         return np.array(cal_preds)
 
 
+class MultiHeadRidgeCalibration(Calibration):
+    """
+    Calibrate a multitask model against several of its LC-setup heads at once.
+
+    Heads are ranked by Pearson correlation to the reference, the ``n_heads`` best are each
+    calibrated with :class:`SplineTransformerCalibration`, and a ridge regression maps the
+    calibrated estimates onto the observed retention times. Never fits more head weights than
+    half the reference size. For a single-task model (one head) this reduces to a spline
+    followed by a linear rescaling.
+
+    Parameters
+    ----------
+    n_heads
+        How many of the best-correlating heads to combine.
+    alphas
+        Ridge strengths offered to the internal cross-validation.
+
+    """
+
+    uses_all_heads = True
+
+    def __init__(self, n_heads: int = 80, alphas: np.ndarray | None = None) -> None:
+        """Initialize MultiHeadRidgeCalibration."""
+        super().__init__()
+        if n_heads < 1:
+            raise ValueError(f"n_heads must be at least 1, got {n_heads}")
+        self.n_heads = n_heads
+        self.alphas = np.logspace(-3, 6, 19) if alphas is None else np.asarray(alphas)
+        self._head_idx: np.ndarray | None = None
+        self._head_calibrations: list[SplineTransformerCalibration] = []
+        self._ridge = None
+
+    @property
+    def is_fitted(self) -> bool:
+        """True once the heads are selected, calibrated and weighted."""
+        return self._head_idx is not None and self._ridge is not None
+
+    def fit(self, target: np.ndarray, source: np.ndarray) -> None:
+        """
+        Select, calibrate and weight the heads.
+
+        Parameters
+        ----------
+        target
+            Observed retention times of the reference, shape ``(n,)``.
+        source
+            Reference predictions for every head, shape ``(n, n_heads_total)``. A 1-D array is
+            accepted and treated as a single head, so a single-task model still works.
+
+        """
+        source = np.asarray(source, dtype=np.float64)
+        if source.ndim == 1:
+            source = source[:, None]
+        target = np.asarray(target, dtype=np.float64).ravel()
+        if source.shape[0] != target.shape[0]:
+            raise CalibrationError(
+                f"source has {source.shape[0]} rows and target {target.shape[0]}"
+            )
+        finite = np.isfinite(target) & np.isfinite(source).all(axis=1)
+        if int(finite.sum()) < 3:
+            raise CalibrationError("Fewer than three reference points with finite values.")
+        source, target = source[finite], target[finite]
+
+        order = _rank_heads_by_correlation(source, target)
+        # never fit more weights than half the reference: a 230-peptide reference cannot support
+        # eighty of them, and the ridge would be extrapolating its own regularisation
+        n_heads = int(min(self.n_heads, source.shape[1], max(1, len(target) // 2)))
+        self._head_idx = order[:n_heads]
+        self.selected_model_head = int(order[0])
+
+        calibrated = np.empty((len(target), n_heads), dtype=np.float64)
+        self._head_calibrations = []
+        for position, head in enumerate(self._head_idx):
+            head_calibration = SplineTransformerCalibration()
+            column = source[:, head].astype(np.float32)
+            head_calibration.fit(target=target.astype(np.float32), source=column)
+            calibrated[:, position] = np.asarray(
+                head_calibration.transform(column), dtype=np.float64
+            )
+            self._head_calibrations.append(head_calibration)
+
+        n_splits = int(min(5, max(2, len(target) // 20)))
+        self._ridge = RidgeCV(alphas=self.alphas, cv=n_splits).fit(calibrated, target)
+        LOGGER.info(
+            "Calibrated on %d of %d heads with ridge strength %.4g; head %d correlates best.",
+            n_heads,
+            source.shape[1],
+            float(getattr(self._ridge, "alpha_", float("nan"))),
+            self.selected_model_head,
+        )
+
+    def transform(self, source: np.ndarray) -> np.ndarray:
+        """
+        Calibrate predictions of the model this calibration was fitted with.
+
+        Parameters
+        ----------
+        source
+            Predictions for every head, shape ``(n, n_heads_total)``, as returned by
+            ``predict(..., return_matrix=True)``.
+
+        """
+        if not self.is_fitted:
+            raise CalibrationError("The model has not been fitted yet. Call fit() first.")
+        source = np.asarray(source, dtype=np.float64)
+        if source.ndim == 1:
+            source = source[:, None]
+        head_idx = cast(np.ndarray, self._head_idx)
+        if source.shape[1] <= int(head_idx.max()):
+            raise CalibrationError(
+                f"source has {source.shape[1]} heads, but the calibration was fitted on a model "
+                f"with at least {int(head_idx.max()) + 1}."
+            )
+        if source.shape[0] == 0:
+            return np.array([])
+        calibrated = np.column_stack(
+            [
+                np.asarray(cal.transform(source[:, head].astype(np.float32)), dtype=np.float64)
+                for cal, head in zip(self._head_calibrations, head_idx, strict=True)
+            ]
+        )
+        return np.asarray(self._ridge.predict(calibrated), dtype=np.float64)
+
+
+def _rank_heads_by_correlation(source: np.ndarray, target: np.ndarray) -> np.ndarray:
+    """
+    Head indices by decreasing Pearson correlation with the target, in one pass.
+
+    The same criterion as :func:`deeplc.core._best_correlating_head`, which takes the first
+    element of this order, but vectorised because thousands of heads are ranked at once.
+    """
+    centred = source - source.mean(axis=0)
+    target_centred = target - target.mean()
+    with np.errstate(invalid="ignore", divide="ignore"):
+        denominator = np.sqrt((centred**2).sum(axis=0) * (target_centred**2).sum())
+        correlation = (centred * target_centred[:, None]).sum(axis=0) / denominator
+    correlation = np.where(np.isfinite(correlation), correlation, -np.inf)
+    return np.argsort(-correlation)
+
+
 def _prepare_series(
     target: np.ndarray,
     source: np.ndarray,

deeplc/core.py

@@ -15,6 +15,7 @@
 from deeplc._reference_selection import deduplicate_psms, select_reference_psms
 from deeplc.calibration import (
     Calibration,
+    MultiHeadRidgeCalibration,
     SplineTransformerCalibration,
 )
 from deeplc.data import DeepLCDataset, split_datasets
@@ -154,7 +155,9 @@ def calibrate(
     model
         Trained model or path to model file.
     calibration
-        Calibration instance to use. If None, SplineTransformerCalibration is used.
+        Calibration instance to use. If None, a multitask model gets
+        MultiHeadRidgeCalibration (combining the best-correlating setup heads) and a
+        single-task model gets SplineTransformerCalibration.
     predict_kwargs
         Additional keyword arguments to pass to the prediction function.
 
@@ -171,15 +174,11 @@ def calibrate(
     # Calibration itself and passes it in already fitted.
     psm_list_reference = deduplicate_psms(psm_list_reference)
 
-    # Get calibration
-    if calibration is None:
-        LOGGER.debug("No calibration provided, using SplineTransformerCalibration by default.")
-        calibration = SplineTransformerCalibration()
-    elif not isinstance(calibration, Calibration):
+    if calibration is not None and not isinstance(calibration, Calibration):
         raise ValueError(
             f"Expected calibration to be of type `Calibration`, got {type(calibration)}"
         )
-    if calibration.is_fitted:
+    if calibration is not None and calibration.is_fitted:
         LOGGER.warning(
             "Provided Calibration is already fitted. Refitting will overwrite existing fit."
         )
@@ -199,10 +198,25 @@ def calibrate(
         return_matrix=True,
     )
 
+    # The default depends on the model: a multitask model is calibrated against its
+    # best-correlating setup heads combined, a single-task model against its one output.
+    if calibration is None:
+        if source_rt_cal.shape[1] > 1:
+            calibration = MultiHeadRidgeCalibration()
+        else:
+            calibration = SplineTransformerCalibration()
+        LOGGER.debug("No calibration provided, using %s.", type(calibration).__name__)
+
     # Fit calibration
     LOGGER.debug("Fitting calibration...")
     target_rt_cal = np.array(psm_list_reference["retention_time"], dtype=np.float32)
 
+    # A calibration that combines heads is given the whole matrix and picks its own; it sets
+    # selected_model_head itself, for callers that want to know which setup came out on top.
+    if getattr(calibration, "uses_all_heads", False):
+        calibration.fit(target=target_rt_cal, source=source_rt_cal)
+        return calibration
+
     # Select the best head for calibration if the model predicts for multiple LC setups
     if source_rt_cal.shape[1] > 1:
         calibration.selected_model_head = _best_correlating_head(source_rt_cal, target_rt_cal)
@@ -235,7 +249,9 @@ def predict_and_calibrate(
     model
         Trained model or path to model file.
     calibration
-        Calibration instance to use. If None, SplineTransformerCalibration is used.
+        Calibration instance to use. If None, a multitask model gets
+        MultiHeadRidgeCalibration (combining the best-correlating setup heads) and a
+        single-task model gets SplineTransformerCalibration.
     predict_kwargs
         Additional keyword arguments to pass to the prediction function.
 
@@ -278,6 +294,10 @@ def predict_and_calibrate(
     else:
         LOGGER.info("Calibration is already fitted, skipping fitting step.")
 
+    if getattr(calibration, "uses_all_heads", False):
+        # the calibration combines several heads, so it takes the matrix as it is
+        return calibration.transform(predicted_rt)
+
     if predicted_rt.shape[1] > 1:
         if calibration.selected_model_head is None:
             raise ValueError(

docs/source/models.rst

@@ -22,6 +22,31 @@ The 4.0 default, ``multitask_model.pt`` (shared trunk, one head per setup), stay
 bundled as :data:`deeplc.core.LEGACY_MULTITASK_MODEL` and can be passed as
 ``model=`` to any core function to reproduce 4.0 and 4.1.0 predictions.
 
+Calibrating against several setups at once
+==========================================
+
+Since 4.3.0 a multitask model is calibrated with
+:class:`~deeplc.calibration.MultiHeadRidgeCalibration` by default: every head is ranked by
+Pearson correlation to the reference, the 80 best are calibrated individually, and a ridge
+regression maps those calibrated estimates onto the observed retention times, so several setups
+contribute. The number of heads is the one parameter worth changing: 80 sits on a flat optimum
+between roughly 40 and 320, and the class never fits more weights than half the reference allows.
+Prediction costs nothing extra, because the full head matrix is computed either way.
+
+The previous behaviour, a spline on the single best-correlating head, remains available by
+passing the calibration explicitly (it is also still the default for single-task models):
+
+.. code-block:: python
+
+   from deeplc import predict_and_calibrate
+   from deeplc.calibration import SplineTransformerCalibration
+
+   calibrated_rt = predict_and_calibrate(
+       psm_list,
+       psm_list_reference=reference,
+       calibration=SplineTransformerCalibration(),
+   )
+
 Training a model from scratch
 ==============================
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "deeplc"
-version = "4.2.0"
+version = "4.3.0"
 description = "DeepLC: Retention time prediction for (modified) peptides using Deep Learning."
 readme = "README.md"
 license = { file = "LICENSE" }

tests/test_deduplication.py

@@ -164,7 +164,10 @@ def test_calibrate_always_uses_the_first_observations():
 
     calibration = core.calibrate(reference, predict_kwargs={"device": "cpu"})
     predicted = core.predict(targets, return_matrix=True)
-    calibrated = calibration.transform(predicted[:, calibration.selected_model_head or 0])
+    if calibration.uses_all_heads:
+        calibrated = calibration.transform(predicted)
+    else:
+        calibrated = calibration.transform(predicted[:, calibration.selected_model_head or 0])
 
     assert np.isfinite(calibrated).all()
     clean_low, clean_high = 5.0, 5.0 + 3.0 * (len(_PEPTIDES) - 1)