Refactor parsing of precursor data from spectrum files

ms2rescore/core.py

@@ -3,15 +3,14 @@
 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.feature_generators import FEATURE_GENERATORS
 from ms2rescore.parse_psms import parse_psms
-from ms2rescore.parse_spectra import get_missing_values
+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
@@ -62,20 +61,28 @@ def rescore(configuration: Dict, psm_list: Optional[PSMList] = None) -> None:
  )
 
  # Add missing precursor info from spectrum file if needed
- psm_list = _fill_missing_precursor_info(psm_list, config)
+ available_ms_data = add_precursor_values(
+ psm_list, config["spectrum_path"], config["spectrum_id_pattern"]
+ )
 
  # Add rescoring features
  for fgen_name, fgen_config in config["feature_generators"].items():
- # TODO: Handle this somewhere else, more generally?
- if fgen_name == "maxquant" and not (psm_list["source"] == "msms").all():
- logger.warning(
- "MaxQuant feature generator requires PSMs from a MaxQuant msms.txt file. Skipping "
- "this feature generator."
- )
- continue
+ # Compile configuration
  conf = config.copy()
  conf.update(fgen_config)
  fgen = FEATURE_GENERATORS[fgen_name](**conf)
+
+ # Check if required MS data is available
+ missing_ms_data = fgen.required_ms_data - available_ms_data
+ if missing_ms_data:
+ logger.warning(
+ f"Skipping feature generator {fgen_name} because required MS data is missing: "
+ f"{missing_ms_data}. Ensure that the required MS data is present in the input "
+ "files or disable the feature generator."
+ )
+ continue
+
+ # Add features
  fgen.add_features(psm_list)
  logger.debug(f"Adding features from {fgen_name}: {set(fgen.feature_names)}")
  feature_names[fgen_name] = set(fgen.feature_names)
@@ -102,6 +109,7 @@ def rescore(configuration: Dict, psm_list: Optional[PSMList] = None) -> None:
  # Write feature names to file
  _write_feature_names(feature_names, output_file_root)
 
+ # Rename PSMs to USIs if requested
  if config["rename_to_usi"]:
  logging.debug(f"Creating USIs for {len(psm_list)} PSMs")
  psm_list["spectrum_id"] = [psm.get_usi(as_url=False) for psm in psm_list]
@@ -173,51 +181,6 @@ def rescore(configuration: Dict, psm_list: Optional[PSMList] = None) -> None:
  logger.exception(e)
 
 
-def _fill_missing_precursor_info(psm_list: PSMList, config: Dict) -> PSMList:
- """Fill missing precursor info from spectrum file if needed."""
- # Check if required
- # TODO: avoid hard coding feature generators in some way
- rt_required = ("deeplc" in config["feature_generators"]) and any(
- v is None or v == 0 or np.isnan(v) for v in psm_list["retention_time"]
- )
- im_required = (
- "ionmob" in config["feature_generators"] or "im2deep" in config["feature_generators"]
- ) and any(v is None or v == 0 or np.isnan(v) for v in psm_list["ion_mobility"])
- logger.debug(f"RT required: {rt_required}, IM required: {im_required}")
-
- # Add missing values
- if rt_required or im_required:
- logger.info("Parsing missing retention time and/or ion mobility values from spectra...")
- get_missing_values(psm_list, config, rt_required=rt_required, im_required=im_required)
-
- # Check if values are now present
- for value_name, required in [("retention_time", rt_required), ("ion_mobility", im_required)]:
- if required and (
- 0.0 in psm_list[value_name]
- or None in psm_list[value_name]
- or np.isnan(psm_list[value_name]).any()
- ):
- if all(v is None or v == 0.0 or np.isnan(v) for v in psm_list[value_name]):
- raise exceptions.MissingValuesError(
- f"Could not find any '{value_name}' values in PSM or spectrum files. Disable "
- f"feature generators that require '{value_name}' or ensure that the values are "
- "present in the input files."
- )
- else:
- missing_value_psms = psm_list[
- [v is None or np.isnan(v) for v in psm_list[value_name]]
- ]
- logger.warning(
- f"Found {len(missing_value_psms)} PSMs with missing '{value_name}' values. "
- "These PSMs will be removed."
- )
- psm_list = psm_list[
- [v is not None and not np.isnan(v) for v in psm_list[value_name]]
- ]
-
- return psm_list
-
-
 def _filter_by_rank(psm_list: PSMList, max_rank: int, lower_score_better: bool) -> PSMList:
  """Filter PSMs by rank."""
  psm_list.set_ranks(lower_score_better=lower_score_better)

ms2rescore/feature_generators/base.py

@@ -1,11 +1,17 @@
 from abc import ABC, abstractmethod
+from typing import Set
 
 from psm_utils import PSMList
 
+from ms2rescore.parse_spectra import MSDataType
+
 
 class FeatureGeneratorBase(ABC):
  """Base class from which all feature generators must inherit."""
 
+ # List of required MS data types for feature generation
+ required_ms_data: Set[MSDataType] = set()
+
  def __init__(self, *args, **kwargs) -> None:
  super().__init__()
 

ms2rescore/feature_generators/deeplc.py

@@ -27,6 +27,7 @@
 from psm_utils import PSMList
 
 from ms2rescore.feature_generators.base import FeatureGeneratorBase
+from ms2rescore.parse_spectra import MSDataType
 
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
 logger = logging.getLogger(__name__)
@@ -35,6 +36,8 @@
 class DeepLCFeatureGenerator(FeatureGeneratorBase):
  """DeepLC retention time-based feature generator."""
 
+ required_ms_data = {MSDataType.retention_time}
+
  def __init__(
  self,
  *args,
@@ -138,9 +141,11 @@ def add_features(self, psm_list: PSMList) -> None:
  )
 
  # Disable wild logging to stdout by Tensorflow, unless in debug mode
- with contextlib.redirect_stdout(
- open(os.devnull, "w")
- ) if not self._verbose else contextlib.nullcontext():
+ with (
+ contextlib.redirect_stdout(open(os.devnull, "w"))
+ if not self._verbose
+ else contextlib.nullcontext()
+ ):
  # Make new PSM list for this run (chain PSMs per spectrum to flat list)
  psm_list_run = PSMList(psm_list=list(chain.from_iterable(psms.values())))
 

ms2rescore/feature_generators/im2deep.py

@@ -22,6 +22,7 @@
 from psm_utils import PSMList
 
 from ms2rescore.feature_generators.base import FeatureGeneratorBase
+from ms2rescore.parse_spectra import MSDataType
 
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
 logger = logging.getLogger(__name__)
@@ -30,6 +31,8 @@
 class IM2DeepFeatureGenerator(FeatureGeneratorBase):
  """IM2Deep collision cross section feature generator."""
 
+ required_ms_data = {MSDataType.ion_mobility}
+
  def __init__(
  self,
  *args,

ms2rescore/feature_generators/ionmob.py

@@ -24,6 +24,7 @@
 from psm_utils import Peptidoform, PSMList
 
 from ms2rescore.feature_generators.base import FeatureGeneratorBase, FeatureGeneratorException
+from ms2rescore.parse_spectra import MSDataType
 
 try:
  from ionmob import __file__ as ionmob_file
@@ -55,6 +56,8 @@
 class IonMobFeatureGenerator(FeatureGeneratorBase):
  """Ionmob collisional cross section (CCS)-based feature generator."""
 
+ required_ms_data = {MSDataType.ion_mobility}
+
  def __init__(
  self,
  *args,

ms2rescore/feature_generators/maxquant.py

@@ -23,6 +23,16 @@
 class MaxQuantFeatureGenerator(FeatureGeneratorBase):
  """Generate MaxQuant-derived features."""
 
+ available_features = [
+ "mean_error_top7",
+ "sq_mean_error_top7",
+ "stdev_error_top7",
+ "ln_explained_ion_current",
+ "ln_nterm_ion_current_ratio",
+ "ln_cterm_ion_current_ratio",
+ "ln_ms2_ion_current",
+ ]
+
  def __init__(self, *args, **kwargs) -> None:
  """
  Generate MaxQuant-derived features.
@@ -39,30 +49,30 @@ def __init__(self, *args, **kwargs) -> None:
 
  """
  super().__init__(*args, **kwargs)
+ self._feature_names = self.available_features[:] # Copy list
 
  @property
  def feature_names(self) -> List[str]:
- return [
- "mean_error_top7",
- "sq_mean_error_top7",
- "stdev_error_top7",
- "ln_explained_ion_current",
- "ln_nterm_ion_current_ratio",
- "ln_cterm_ion_current_ratio",
- "ln_ms2_ion_current",
- ]
+ return self._feature_names
 
  def add_features(self, psm_list: PSMList):
  """
- Add MS²PIP-derived features to PSMs.
+ Add MaxQuant-derived features to PSMs.
 
  Parameters
  ----------
  psm_list
  PSMs to add features to.
 
  """
- logger.info("Adding MaxQuant-derived features to PSMs.")
+ # Check if all PSMs are from MaxQuant
+ if not self._all_psms_from_maxquant(psm_list):
+ self._feature_names = [] # Set feature names to empty list to indicate none added
+ logger.warning("Not all PSMs are from MaxQuant. Skipping MaxQuant feature generation.")
+ return
+ else:
+ self._feature_names = self.available_features # Reset feature names
+ logger.info("Adding MaxQuant-derived features to PSMs.")
 
  # Infer mass deviations column name
  for column_name in [
@@ -90,6 +100,11 @@ def add_features(self, psm_list: PSMList):
  for psm in psm_list:
  psm["rescoring_features"].update(self._compute_features(psm["metadata"]))
 
+ @staticmethod
+ def _all_psms_from_maxquant(psm_list):
+ """Check if the PSMs are from MaxQuant."""
+ return (psm_list["source"] == "msms").all()
+
  def _compute_features(self, psm_metadata):
  """Compute features from derived from intensities and mass errors."""
  features = {}

ms2rescore/feature_generators/ms2pip.py

@@ -38,6 +38,7 @@
 from rich.progress import track
 
 from ms2rescore.feature_generators.base import FeatureGeneratorBase, FeatureGeneratorException
+from ms2rescore.parse_spectra import MSDataType
 from ms2rescore.utils import infer_spectrum_path
 
 logger = logging.getLogger(__name__)
@@ -46,6 +47,8 @@
 class MS2PIPFeatureGenerator(FeatureGeneratorBase):
  """Generate MS²PIP-based features."""
 
+ required_ms_data = {MSDataType.ms2_spectra}
+
  def __init__(
  self,
  *args,