fix: support DeepLC's multi-head calibration API when available

ms2rescore/feature_generators/deeplc.py

@@ -20,13 +20,23 @@
 from typing import ClassVar
 
 import numpy as np
-from deeplc.calibration import SplineTransformerCalibration
-
-# NOTE: `_best_correlating_head` is a private DeepLC function. Imported here intentionally to
-# reuse DeepLC's multitask head-selection logic. To be replaced once DeepLC exposes a public API.
-from deeplc.core import _best_correlating_head, finetune, predict
+from deeplc.core import finetune, predict
 from psm_utils import PSMList
 
+try:
+    # DeepLC >= 4.3: MultiHeadRidgeCalibration combines several best-correlating heads instead of
+    # calibrating against a single one.
+    from deeplc.calibration import MultiHeadRidgeCalibration
+
+    _HAS_MULTIHEAD_CALIBRATION = True
+except ImportError:
+    # DeepLC < 4.3: no public head-selection API yet. `_best_correlating_head` is a private
+    # DeepLC function, imported here intentionally to reuse DeepLC's head-selection logic.
+    from deeplc.calibration import SplineTransformerCalibration
+    from deeplc.core import _best_correlating_head
+
+    _HAS_MULTIHEAD_CALIBRATION = False
+
 from ms2rescore._utils import get_original_hit_mask
 from ms2rescore.feature_generators.base import FeatureGeneratorBase
 from ms2rescore.parse_spectra import MSDataType
@@ -210,20 +220,25 @@ def add_features(self, psm_list: PSMList) -> None:
             if len(observed_rt_calibration) == 0:
                 raise ValueError(f"Run '{run}' has no target PSMs available for calibration.")
 
-            # Select the head that best correlates with observed RT on the calibration PSMs
             reference_matrix = pred_matrix[calibration_idx]
-            head = _best_correlating_head(reference_matrix, observed_rt_calibration)
-            logger.debug(f"Run '{run}': selected DeepLC head {head}")
+            run_idx = run_df.index.values
 
-            # Fit calibration on the selected head and transform all predictions for this run
-            calibration = SplineTransformerCalibration()
-            calibration.fit(
-                target=observed_rt_calibration,
-                source=reference_matrix[:, head],
-            )
+            if _HAS_MULTIHEAD_CALIBRATION:
+                # MultiHeadRidgeCalibration combines several best-correlating heads and
+                # transforms straight from the full prediction matrix.
+                calibration = MultiHeadRidgeCalibration()
+                calibration.fit(target=observed_rt_calibration, source=reference_matrix)
+                head = calibration.selected_model_head
+                calibrated_rt = calibration.transform(pred_matrix[run_idx])
+            else:
+                # Select the head that best correlates with observed RT on the calibration PSMs,
+                # then fit and transform on that one column.
+                head = _best_correlating_head(reference_matrix, observed_rt_calibration)
+                calibration = SplineTransformerCalibration()
+                calibration.fit(target=observed_rt_calibration, source=reference_matrix[:, head])
+                calibrated_rt = calibration.transform(pred_matrix[run_idx, head])
 
-            run_idx = run_df.index.values
-            calibrated_rt = calibration.transform(pred_matrix[run_idx, head])
+            logger.debug(f"Run '{run}': selected DeepLC head {head}")
 
             # Update predictions with calibrated values
             psm_list_df.loc[psm_list_df["run"] == run, "predicted_retention_time"] = calibrated_rt

tests/test_feature_generators_deeplc.py

@@ -184,9 +184,7 @@ def test_selects_best_correlating_head(monkeypatch):
     mat[:, 0] = np.tile([0.0, 1.0], n)[:n]  # non-zero std, ~zero correlation with obs
     mat[:, 1] = obs + 0.5  # best head
     mat[:, 2] = 5.0  # constant -> zero std -> skipped by head selection
-    monkeypatch.setattr(
-        "ms2rescore.feature_generators.deeplc.predict", _fake_predict_factory(mat)
-    )
+    monkeypatch.setattr("ms2rescore.feature_generators.deeplc.predict", _fake_predict_factory(mat))
 
     DeepLCFeatureGenerator(finetune=False, processes=1).add_features(psm_list)
 
@@ -213,9 +211,7 @@ def test_per_run_head_selection(monkeypatch):
     mat[10:20, 1] = 7.0  # constant for run2 -> skipped
     mat[0:10, 2] = 7.0  # constant for run1 -> skipped
     mat[10:20, 2] = np.array(obs2) + 0.5  # run2 best head
-    monkeypatch.setattr(
-        "ms2rescore.feature_generators.deeplc.predict", _fake_predict_factory(mat)
-    )
+    monkeypatch.setattr("ms2rescore.feature_generators.deeplc.predict", _fake_predict_factory(mat))
 
     DeepLCFeatureGenerator(finetune=False, processes=1).add_features(psm_list)
 
@@ -253,9 +249,7 @@ def test_row_alignment_preserved_across_runs(monkeypatch):
 
     obs = np.array([rt for _, rt in interleaved])
     mat = np.column_stack([obs, obs])  # identity: predicted head == observed
-    monkeypatch.setattr(
-        "ms2rescore.feature_generators.deeplc.predict", _fake_predict_factory(mat)
-    )
+    monkeypatch.setattr("ms2rescore.feature_generators.deeplc.predict", _fake_predict_factory(mat))
 
     DeepLCFeatureGenerator(finetune=False, processes=1).add_features(psm_list)
 
@@ -328,3 +322,54 @@ def test_best_run_by_shared_proteoforms_no_overlap_returns_first():
     proteoforms = ["p1", "p2", "p3"]
     best = DeepLCFeatureGenerator._best_run_by_shared_proteoforms(runs, proteoforms)
     assert best == "A"
+
+
+class _FakeMultiHeadCalibration:
+    """
+    Stand-in for DeepLC's MultiHeadRidgeCalibration.
+
+    Fits on the whole prediction matrix, picks the best-correlating column, and records it as
+    ``selected_model_head``, matching the real class's contract closely enough to exercise the
+    generator's multi-head branch without depending on a DeepLC version that has it.
+    """
+
+    def __init__(self):
+        self.selected_model_head = None
+        self._offset = 0.0
+
+    def fit(self, target, source):
+        correlations = [
+            np.corrcoef(source[:, j], target)[0, 1] if np.std(source[:, j]) > 0 else -1.0
+            for j in range(source.shape[1])
+        ]
+        self.selected_model_head = int(np.argmax(correlations))
+        self._offset = target.mean() - source[:, self.selected_model_head].mean()
+
+    def transform(self, source):
+        return source[:, self.selected_model_head] + self._offset
+
+
+def test_uses_multihead_calibration_when_available(monkeypatch):
+    """
+    When DeepLC exposes MultiHeadRidgeCalibration, the generator uses it, not the private
+    `_best_correlating_head` fallback.
+
+    The installed DeepLC version here predates that class, so both the flag and the class are
+    faked to exercise the branch that would otherwise only run against DeepLC >= 4.3.
+    """
+    psm_list = _make_psm_list()
+    obs = np.array([100.0 + i * 30.0 for i in range(len(_PEPTIDES))])
+    mat = np.column_stack([np.zeros_like(obs), obs + 0.5])  # head 0 noise, head 1 correlates
+
+    monkeypatch.setattr("ms2rescore.feature_generators.deeplc.predict", _fake_predict_factory(mat))
+    monkeypatch.setattr("ms2rescore.feature_generators.deeplc._HAS_MULTIHEAD_CALIBRATION", True)
+    monkeypatch.setattr(
+        "ms2rescore.feature_generators.deeplc.MultiHeadRidgeCalibration",
+        _FakeMultiHeadCalibration,
+        raising=False,
+    )
+
+    DeepLCFeatureGenerator(finetune=False, processes=1).add_features(psm_list)
+
+    preds = np.array([p.rescoring_features["predicted_retention_time"] for p in psm_list])
+    assert np.corrcoef(preds, obs)[0, 1] > 0.99