~5x speed-up and consolidated feature generators

.github/workflows/test.yml

@@ -35,8 +35,15 @@ jobs:
           python-version: ${{ matrix.python-version }}
           enable-cache: true
 
+      # Temporarily skip mumble on 3.14 until rustyms supports it
       - name: Install the project
-        run: uv sync --all-extras --dev
+        run: |
+          if [ "${{ matrix.python-version }}" = "3.14" ]; then
+            uv sync --extra idxml --dev
+          else
+            uv sync --all-extras --dev
+          fi
+
 
       - name: Run tests
         run: uv run pytest

.readthedocs.yml

@@ -1,7 +1,7 @@
 version: 2
 
 build:
-  os: ubuntu-22.04
+  os: ubuntu-lts-latest
   tools:
     python: "3.11"
   jobs:

docs/source/userguide/input-files.rst

@@ -29,6 +29,13 @@ Check the list of :py:mod:`psm_utils` tags in the
 file extension, the file type can also be inferred from the file name. In that case,
 ``psm_file_type`` option can be set to ``infer``.
 
+.. note::
+   If a previous MS²Rescore run crashed during feature generation or rescoring, an intermediate
+   file (``<prefix>.intermediate.psms.tsv``) is automatically saved. This file contains all PSMs
+   with features that were successfully added up to that point. You can resume processing by
+   providing this file as the PSM file (``-p <prefix>.intermediate.psms.tsv -t tsv``) to skip
+   already completed feature generation steps.
+
 
 Spectrum file(s)
 ================

docs/source/userguide/output-files.rst

@@ -31,6 +31,10 @@ Log and configuration files:
 +--------------------------------------+--------------------------------------------------------------------------------------+
 | ``<prefix>.feature_names.tsv``       | List of the features and their descriptions                                          |
 +--------------------------------------+--------------------------------------------------------------------------------------+
+| ``<prefix>.intermediate.psms.tsv``   | Created automatically if the process crashes during feature generation or rescoring. |
+|                                      | Contains all PSMs with successfully added features up to the crash point. Can be     |
+|                                      | used to resume processing with ``-p <prefix>.intermediate.psms.tsv -t tsv``.         |
++--------------------------------------+--------------------------------------------------------------------------------------+
 
 Rescoring engine files:
 

ms2rescore.spec

@@ -16,7 +16,7 @@ project = "ms2rescore"
 bundle_name = "ms2rescore"
 bundle_identifier = f"{bundle_name}.{__version__}"
 
-extra_requirements = {"ionmob"}
+extra_requirements = {}
 
 # Requirements config
 skip_requirements_regex = r"^(?:.*\..*)"

ms2rescore/__main__.py

@@ -6,6 +6,7 @@
 import json
 import logging
 import sys
+from datetime import datetime
 from pathlib import Path
 from typing import Union
 
@@ -196,7 +197,13 @@ def profile(fnc, filepath):
     def inner(*args, **kwargs):
         with cProfile.Profile() as profiler:
             return_value = fnc(*args, **kwargs)
-        profiler.dump_stats(filepath + ".profile.prof")
+
+        # Add timestamp to profiler output filename
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        profile_filename = f"{filepath}.profile_{timestamp}.prof"
+        profiler.dump_stats(profile_filename)
+        LOGGER.info(f"Profile data written to: {profile_filename}")
+
         return return_value
 
     return inner
@@ -248,6 +255,7 @@ def main(tims=False):
     # Run MS²Rescore
     try:
         if config["ms2rescore"]["profile"]:
+            LOGGER.info("Profiling enabled")
             profiled_rescore = profile(rescore, config["ms2rescore"]["output_path"])
             profiled_rescore(configuration=config)
         else:

ms2rescore/constants.py

@@ -0,0 +1,6 @@
+"""Shared constants for ms2rescore."""
+
+import re
+
+# Regex pattern to strip charge state suffix (e.g., "/2") from peptidoform strings
+CHARGE_PATTERN = re.compile(r"(/\d+$)")

ms2rescore/core.py

@@ -3,17 +3,22 @@
 from multiprocessing import cpu_count
 from typing import Dict, Optional
 
+import numpy as np
 import psm_utils.io
 from mokapot.dataset import LinearPsmDataset
 from psm_utils import PSMList
 
 from ms2rescore import exceptions
+from ms2rescore.constants import CHARGE_PATTERN
 from ms2rescore.feature_generators import FEATURE_GENERATORS
 from ms2rescore.parse_psms import parse_psms
 from ms2rescore.parse_spectra import add_precursor_values
 from ms2rescore.report import generate
 from ms2rescore.rescoring_engines import mokapot, percolator
-from ms2rescore.rescoring_engines.mokapot import add_peptide_confidence, add_psm_confidence
+from ms2rescore.rescoring_engines.mokapot import (
+    add_peptide_confidence,
+    add_psm_confidence,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -34,7 +39,9 @@ def rescore(configuration: Dict, psm_list: Optional[PSMList] = None) -> None:
         f"Running MS²Rescore with following configuration: {json.dumps(configuration, indent=4)}"
     )
     config = configuration["ms2rescore"]
-    output_file_root = config["output_path"]
+    output_file_root = config["output_path"].split(".intermediate.")[
+        0
+    ]  # if no intermediate, takes full name
 
     # Write full configuration including defaults to file
     with open(output_file_root + ".full-config.json", "w") as f:
@@ -59,11 +66,24 @@ def rescore(configuration: Dict, psm_list: Optional[PSMList] = None) -> None:
     logger.debug(
         f"PSMs already contain the following rescoring features: {psm_list_feature_names}"
     )
+    # Check if all features are already present; collect generators to skip
+    skip_fgens = set()
+    for fgen_name, fgen_config in config["feature_generators"].items():
+        fgen_features = FEATURE_GENERATORS[fgen_name]().feature_names
+        if set(fgen_features).issubset(psm_list_feature_names):
+            logger.debug(
+                f"Skipping feature generator {fgen_name} because all features are already "
+                "present in the PSM file."
+            )
+            feature_names[fgen_name] = set(fgen_features)
+            feature_names["psm_file"] = psm_list_feature_names - set(fgen_features)
+            skip_fgens.add(fgen_name)
 
     # Add missing precursor info from spectrum file if needed
     required_ms_data = {
         ms_data
         for fgen_name in config["feature_generators"].keys()
+        if fgen_name not in skip_fgens
         for ms_data in FEATURE_GENERATORS[fgen_name].required_ms_data
     }
     available_ms_data = add_precursor_values(
@@ -75,6 +95,8 @@ def rescore(configuration: Dict, psm_list: Optional[PSMList] = None) -> None:
 
     # Add rescoring features
     for fgen_name, fgen_config in config["feature_generators"].items():
+        if fgen_name in skip_fgens:
+            continue
         # Compile configuration
         conf = config.copy()
         conf.update(fgen_config)
@@ -89,12 +111,29 @@ def rescore(configuration: Dict, psm_list: Optional[PSMList] = None) -> None:
                 "files or disable the feature generator."
             )
             continue
-
-        # Add features
-        fgen.add_features(psm_list)
+        try:
+            fgen.add_features(psm_list)
+        except (
+            Exception,
+            KeyboardInterrupt,
+        ) as e:  # Intentionally broad to save intermediate output before re-raising
+            logger.error(
+                f"Error while adding features from {fgen_name}: {e}, writing intermediary output..."
+            )
+            # Write intermediate TSV
+            psm_utils.io.write_file(
+                psm_list, output_file_root + ".intermediate.psms.tsv", filetype="tsv"
+            )
+            raise
         logger.debug(f"Adding features from {fgen_name}: {set(fgen.feature_names)}")
         feature_names[fgen_name] = set(fgen.feature_names)
 
+        # Remove overlapping features from psm_file to avoid duplicates
+        # (e.g., hyperscore can be in both psm_file and ms2pip)
+        overlap = feature_names.get("psm_file", set()) & feature_names[fgen_name]
+        if overlap:
+            feature_names["psm_file"] = feature_names["psm_file"] - overlap
+
     # Filter out psms that do not have all added features
     all_feature_names = {f for fgen in feature_names.values() for f in fgen}
     psms_with_features = [
@@ -114,6 +153,12 @@ def rescore(configuration: Dict, psm_list: Optional[PSMList] = None) -> None:
         )
     psm_list = psm_list[psms_with_features]
 
+    if "mumble" in config["psm_generator"]:
+        from ms2rescore.utils import filter_mumble_psms
+
+        # Remove PSMs where matched_ions_pct drops 25% below the original hit
+        psm_list = filter_mumble_psms(psm_list, threshold=0.75)
+
     # Write feature names to file
     _write_feature_names(feature_names, output_file_root)
 
@@ -160,13 +205,18 @@ def rescore(configuration: Dict, psm_list: Optional[PSMList] = None) -> None:
                 protein_kwargs=protein_kwargs,
                 **config["rescoring_engine"]["mokapot"],
             )
-    except exceptions.RescoringError as e:
+    except (
+        Exception,
+        KeyboardInterrupt,
+    ):  # Intentionally broad to save intermediate output before re-raising
         # Write output
-        logger.info(f"Writing intermediary output to {output_file_root}.psms.tsv...")
-        psm_utils.io.write_file(psm_list, output_file_root + ".psms.tsv", filetype="tsv")
+        logger.info(f"Writing intermediary output to {output_file_root}.intermediate.psms.tsv...")
+        psm_utils.io.write_file(
+            psm_list, output_file_root + ".intermediate.psms.tsv", filetype="tsv"
+        )
 
         # Reraise exception
-        raise e
+        raise
 
     # Post-rescoring processing
     if all(psm_list["pep"] == 1.0):
@@ -219,7 +269,12 @@ def _write_feature_names(feature_names, output_file_root):
 def _log_id_psms_before(psm_list: PSMList, fdr: float = 0.01, max_rank: int = 1) -> int:
     """Log #PSMs identified before rescoring."""
     id_psms_before = (
-        (psm_list["qvalue"] <= 0.01) & (psm_list["rank"] <= max_rank) & (~psm_list["is_decoy"])
+        (psm_list["qvalue"] <= fdr)
+        & (psm_list["rank"] <= max_rank)
+        & (~psm_list["is_decoy"])
+        & np.array(
+            [(metadata or {}).get("original_psm", True) for metadata in psm_list["metadata"]]
+        )
     ).sum()
     logger.info(
         f"Found {id_psms_before} identified PSMs with rank <= {max_rank} at {fdr} FDR before "
@@ -274,7 +329,7 @@ def _calculate_confidence(psm_list: PSMList) -> PSMList:
     psm_df = psm_list.to_dataframe()
     psm_df = psm_df.reset_index(drop=True).reset_index()
     psm_df["peptide"] = (
-        psm_df["peptidoform"].astype(str).str.replace(r"(/\d+$)", "", n=1, regex=True)
+        psm_df["peptidoform"].astype(str).str.replace(CHARGE_PATTERN, "", n=1, regex=True)
     )
     psm_df["is_target"] = ~psm_df["is_decoy"]
     lin_psm_data = LinearPsmDataset(
@@ -285,7 +340,9 @@ def _calculate_confidence(psm_list: PSMList) -> PSMList:
     )
 
     # Recalculate confidence
-    new_confidence = lin_psm_data.assign_confidence(scores=psm_list["score"])
+    new_confidence = lin_psm_data.assign_confidence(
+        scores=list(psm_list["score"])
+    )  # explicity make it a list to avoid TypingError: Failed in nopython mode pipeline (step: nopython frontend) in mokapot
 
     # Add new confidence estimations to PSMList
     add_psm_confidence(psm_list, new_confidence)

ms2rescore/exceptions.py

@@ -41,3 +41,9 @@ class RescoringError(MS2RescoreError):
     """Error while rescoring PSMs."""
 
     pass
+
+
+class ParseSpectrumError(MS2RescoreError):
+    """Error while parsing spectrum files."""
+
+    pass

ms2rescore/feature_generators/__init__.py

@@ -6,15 +6,13 @@
 from ms2rescore.feature_generators.basic import BasicFeatureGenerator
 from ms2rescore.feature_generators.deeplc import DeepLCFeatureGenerator
 from ms2rescore.feature_generators.im2deep import IM2DeepFeatureGenerator
-from ms2rescore.feature_generators.ionmob import IonMobFeatureGenerator
-from ms2rescore.feature_generators.maxquant import MaxQuantFeatureGenerator
+from ms2rescore.feature_generators.ms2 import MS2FeatureGenerator
 from ms2rescore.feature_generators.ms2pip import MS2PIPFeatureGenerator
 
 FEATURE_GENERATORS: dict[str, type[FeatureGeneratorBase]] = {
     "basic": BasicFeatureGenerator,
     "ms2pip": MS2PIPFeatureGenerator,
     "deeplc": DeepLCFeatureGenerator,
-    "maxquant": MaxQuantFeatureGenerator,
-    "ionmob": IonMobFeatureGenerator,
     "im2deep": IM2DeepFeatureGenerator,
+    "ms2": MS2FeatureGenerator,
 }

ms2rescore/feature_generators/base.py

@@ -17,11 +17,11 @@ def __init__(self, *args, **kwargs) -> None:
 
     @property
     @abstractmethod
-    def feature_names(self):
+    def feature_names(self) -> list[str]:
         pass
 
     @abstractmethod
-    def add_features(self, psm_list: PSMList):
+    def add_features(self, psm_list: PSMList) -> None:
         pass