Fix KeyError on partial/empty rescoring config

A user-provided rescoring dict (e.g. "rescoring": {} or {"model": "lda"})
is passed through as-is by the config cascade -- same as feature_generators
and psm_generator, where an override is taken at face value rather than
merged with defaults (CascadeConfig's max_recursion_depth stops it from
reaching into these nested dicts either way). rescoring.py and parse_psms.py
were indexing config["rescoring"]["train_fdr"]/["model"] directly, which
KeyErrors as soon as either key is missing.

Fixed by using dict.get(..., default) where a value is needed before
calling ristretto, and by passing config["rescoring"] to ristretto.rescore()
via **kwargs instead of naming train_fdr/model explicitly, so ristretto's
own defaults (train_fdr=0.01, model="svm") apply to whatever the user's
config omitted.

Author: RalfG

ms2rescore/parse_psms.py

@@ -40,7 +40,7 @@ def parse_psms(config: Dict, psm_list: Union[PSMList, None]) -> PSMList:
     # Remove invalid AAs and find decoys first, so score direction can be inferred from them
     psm_list = _remove_invalid_aa(psm_list)
     _find_decoys(psm_list, config["id_decoy_pattern"])
-    train_fdr = config["rescoring"]["train_fdr"] if config["rescoring"] else 0.01
+    train_fdr = config["rescoring"].get("train_fdr", 0.01) if config["rescoring"] else 0.01
     lower_score_is_better = infer_score_direction(psm_list, train_fdr)
 
     # Filter by PSM rank

ms2rescore/rescoring.py

@@ -147,7 +147,7 @@ def evaluate_before(psm_list: PSMList, config: Dict) -> RescoreResult:
     original_mask = np.array([_is_original_psm(psm) for psm in psm_list])
     psm_list = psm_list[original_mask]
 
-    train_fdr = config["rescoring"]["train_fdr"] if config["rescoring"] else 0.01
+    train_fdr = config["rescoring"].get("train_fdr", 0.01) if config["rescoring"] else 0.01
     features_df = build_features_dataframe(psm_list, set(), infer_score_direction(psm_list, train_fdr))
     return _trim_and_evaluate(
         features_df,
@@ -197,7 +197,9 @@ def rescore(psm_list: PSMList, config: Dict, output_file_root: str) -> Tuple[PSM
 
     """
     feature_names = {f for psm in psm_list for f in psm.rescoring_features}
-    lower_score_is_better = infer_score_direction(psm_list, config["rescoring"]["train_fdr"])
+    lower_score_is_better = infer_score_direction(
+        psm_list, config["rescoring"].get("train_fdr", 0.01)
+    )
     features_df = build_features_dataframe(psm_list, feature_names, lower_score_is_better)
 
     peptide_col = "peptide"
@@ -212,10 +214,11 @@ def rescore(psm_list: PSMList, config: Dict, output_file_root: str) -> Tuple[PSM
             protein_col=protein_col,
             feature_cols=sorted(feature_names),
             decoy_pattern=decoy_pattern,
-            model=config["rescoring"]["model"],
-            train_fdr=config["rescoring"]["train_fdr"],
             n_jobs=int(config["processes"]),
             multi_rank_rescoring=True,
+            # train_fdr/model; ristretto's own kwarg defaults apply to any key a partial
+            # user-provided rescoring dict omitted.
+            **config["rescoring"],
         )
         final_result = _trim_and_evaluate(
             ml_result.psms,

tests/test_rescoring.py

@@ -1,5 +1,7 @@
 """Tests for ms2rescore.rescoring: the ristretto integration layer."""
 
+from unittest.mock import patch
+
 import numpy as np
 import pandas as pd
 from psm_utils import PSM, PSMList
@@ -193,6 +195,38 @@ def test_rescore_respects_configured_model():
     assert ((q >= 0) & (q <= 1)).all()
 
 
+def test_rescore_handles_partial_rescoring_config():
+    """A partial rescoring dict (e.g. only `model` or only `train_fdr`, as produced by the
+    config cascade for a user override of just one key) must not KeyError on the other."""
+    psm_list = _make_psm_list(n_spectra=30, seed=19)
+
+    for rescoring_config in ({"model": "lda"}, {"train_fdr": 0.1}):
+        config = {**BASE_CONFIG, "rescoring": rescoring_config}
+        new_psm_list, after_result = rescoring.rescore(psm_list, config, "unused-output-root")
+        assert len(new_psm_list) == len(after_result.psms)
+
+
+def test_rescore_empty_rescoring_config_falls_back_to_ristretto_defaults():
+    """An empty rescoring dict (the config cascade's result for `"rescoring": {}`) must not
+    KeyError -- train_fdr/model should simply be omitted from the ristretto.rescore() call, so
+    ristretto's own defaults (train_fdr=0.01, model="svm") apply. Whether ristretto's fit
+    itself converges on this tiny fixture at the (stricter) default train_fdr is irrelevant
+    here -- only that the right kwargs reach it.
+    """
+    psm_list = _make_psm_list(n_spectra=30, seed=19)
+    config = {**BASE_CONFIG, "rescoring": {}}
+
+    with patch.object(rescoring.ristretto, "rescore", wraps=rescoring.ristretto.rescore) as m:
+        try:
+            rescoring.rescore(psm_list, config, "unused-output-root")
+        except rescoring.RescoringError:
+            pass
+
+    passed_kwargs = m.call_args.kwargs
+    assert "train_fdr" not in passed_kwargs
+    assert "model" not in passed_kwargs
+
+
 def test_rescore_multi_rank_output_keeps_multiple_ranks_per_spectrum():
     psm_list = _make_psm_list(n_spectra=30, ranks_per_spectrum=2, seed=6)
     config = {**BASE_CONFIG, "max_psm_rank_output": 2}