Merge pull request #264 from CompOmics/fix/predict-library-2

Cherry-pick fixes from fix/predict-library (PR #248)

Author: RalfG

ms2pip/_spectrum_processing.py

@@ -19,6 +19,7 @@
     get_ms2_spectra,  # type: ignore[ty:unresolved-import]
 )
 from psm_utils import PSM, Peptidoform, PSMList
+from pyteomics.proforma import MassModification, ModificationBase
 
 import ms2pip.exceptions as exceptions
 from ms2pip.constants import MODELS
@@ -92,7 +93,11 @@ def proforma_to_mass_shift(peptidoform: Peptidoform) -> str:
         parts.append(aa)
         if mods:
             for mod in mods:
-                parts.append(f"[{mod.mass:+.4f}]")  # type: ignore[ty:unresolved-attribute]
+                if not isinstance(mod, (ModificationBase, MassModification)):
+                    raise ValueError(
+                        f"Unsupported ProForma tag type {type(mod)} in peptidoform {peptidoform}"
+                    )
+                parts.append(f"[{mod.mass:+.4f}]")
     c_term = peptidoform.properties.get("c_term")
     if c_term:
         for mod in c_term:
@@ -224,7 +229,9 @@ def _load_and_match_spectra(
             preprocessed_cache[spec_id] = obs
 
         results.append(
-            MatchedSpectrum(psm_index, psm, preprocessed_cache[spec_id], annotated_spectra[batch_idx])
+            MatchedSpectrum(
+                psm_index, psm, preprocessed_cache[spec_id], annotated_spectra[batch_idx]
+            )
         )
 
     return results

ms2pip/_utils/xgb_models.py

@@ -77,6 +77,7 @@ def load_xgb_models(
         Number of threads for XGBoost prediction. Capped internally.
 
     """
+    os.environ.pop("CUDA_VISIBLE_DEVICES", None)  # Workaround for dmlc/xgboost#11283
     nthread = min(
         processes if processes is not None else (os.cpu_count() or 1),
         _MAX_PREDICTION_THREADS,

ms2pip/result.py

@@ -9,11 +9,6 @@
 from psm_utils import PSM
 from pydantic import BaseModel, ConfigDict
 
-try:
-    import spectrum_utils.plot as sup
-except ImportError:
-    sup = None  # type: ignore[ty:invalid-assignment]
-
 from ms2pip.correlation import pearson
 from ms2pip.spectrum import ObservedSpectrum, PredictedSpectrum
 
@@ -105,6 +100,11 @@ def plot_spectra(self):
         Requires optional dependency ``spectrum_utils`` to be installed.
 
         """
+        try:
+            import spectrum_utils.plot as sup
+        except ImportError as e:
+            raise ImportError("Optional dependency spectrum_utils not installed.") from e
+
         predicted, observed = (
             spec.to_spectrum_utils() if spec else None for spec in self.as_spectra()
         )

ms2pip/search_space.py

@@ -184,8 +184,8 @@ class ProteomeSearchSpace(BaseModel):
     fasta_file: Path
     min_length: int = 8
     max_length: int = 30
-    min_precursor_mz: float | None = 0
-    max_precursor_mz: float | None = np.inf
+    min_precursor_mz: float = 0
+    max_precursor_mz: float = np.inf
     cleavage_rule: str = "trypsin"
     missed_cleavages: int = 2
     semi_specific: bool = False
@@ -194,7 +194,31 @@ class ProteomeSearchSpace(BaseModel):
     max_variable_modifications: int = 3
     charges: list[int] = [2, 3]
 
-    _peptidoform_spaces: list[_PeptidoformSearchSpace] = PrivateAttr(default_factory=list)
+    _peptidoform_spaces: list[_PeptidoformSearchSpace] | None = PrivateAttr(default=None)
+
+    @field_validator("min_precursor_mz", mode="before")
+    @classmethod
+    def _coerce_min_precursor_mz(cls, v):
+        return 0.0 if v is None else v
+
+    @field_validator("max_precursor_mz", mode="before")
+    @classmethod
+    def _coerce_max_precursor_mz(cls, v):
+        return np.inf if v is None else v
+
+    @field_validator("min_length")
+    @classmethod
+    def _validate_min_length(cls, v):
+        if v > 3:
+            return v
+        raise ValueError("Minimum peptide length must be greater than 3.")
+
+    @field_validator("max_length")
+    @classmethod
+    def _validate_max_length(cls, v):
+        if v <= 100:
+            return v
+        raise ValueError("Maximum peptide length must be less than or equal to 100.")
 
     @field_validator("modifications")
     @classmethod
@@ -219,7 +243,7 @@ def _validate_unspecific_cleavage(self):
         return self
 
     def __len__(self):
-        if not self._peptidoform_spaces:
+        if self._peptidoform_spaces is None:
             raise ValueError("Search space must be built before length can be determined.")
         return sum(len(pep_space) for pep_space in self._peptidoform_spaces)
 
@@ -276,7 +300,7 @@ def __iter__(self) -> Generator[PSM, None, None]:  # type: ignore[ty:invalid-met
 
         """
         # Build search space if not already built
-        if not self._peptidoform_spaces:
+        if self._peptidoform_spaces is None:
             raise ValueError("Search space must be built before PSMs can be generated.")
 
         spectrum_id = 0
@@ -295,25 +319,26 @@ def filter_psms_by_mz(self, psms: PSMList) -> PSMList:
             psm_list=[
                 psm
                 for psm in psms
-                if self.min_precursor_mz <= psm.peptidoform.theoretical_mz <= self.max_precursor_mz  # type: ignore[ty:unsupported-operator]
+                if psm.peptidoform.theoretical_mz is not None
+                and self.min_precursor_mz <= psm.peptidoform.theoretical_mz <= self.max_precursor_mz
             ]
         )
 
     def _digest_fasta(self, processes: int = 1):
         """Digest FASTA file to peptides and populate search space."""
         # Convert to string to avoid issues with Path objects
-        self.fasta_file = str(self.fasta_file)  # type: ignore[ty:invalid-assignment]
+        fasta_file = str(self.fasta_file)
         n_proteins = _count_fasta_entries(self.fasta_file)
         if self.add_decoys:
             fasta_db = pyteomics.fasta.decoy_db(
-                self.fasta_file,
+                fasta_file,
                 mode="reverse",
                 decoy_only=False,
                 keep_nterm=True,
             )
             n_proteins *= 2
         else:
-            fasta_db = pyteomics.fasta.FASTA(self.fasta_file)
+            fasta_db = pyteomics.fasta.FASTA(fasta_file)
 
         # Read proteins and digest to peptides
         with _get_pool(processes) as pool:
@@ -335,6 +360,7 @@ def _digest_fasta(self, processes: int = 1):
 
     def _remove_redundancy(self):
         """Remove redundancy in peptides and combine protein lists."""
+        assert self._peptidoform_spaces is not None  # for type checker
         peptide_dict = dict()
         for peptide in track(
             self._peptidoform_spaces,
@@ -351,6 +377,7 @@ def _remove_redundancy(self):
 
     def _add_modifications(self, processes: int = 1):
         """Add modifications to peptides in search space."""
+        assert self._peptidoform_spaces is not None  # for type checker
         modifications_by_target = _restructure_modifications_by_target(self.modifications)
         modification_options = []
         with _get_pool(processes) as pool:
@@ -373,6 +400,7 @@ def _add_modifications(self, processes: int = 1):
 
     def _add_charges(self):
         """Add charge permutations to peptides in search space."""
+        assert self._peptidoform_spaces is not None  # for type checker
         for peptide in track(
             self._peptidoform_spaces,
             description="Adding charge permutations...",

ms2pip/spectrum.py

@@ -3,16 +3,22 @@
 from __future__ import annotations
 
 import warnings
+from typing import Annotated
+
 import numpy as np
 from psm_utils import Peptidoform
-from pydantic import model_validator, field_validator, ConfigDict, BaseModel
+from pydantic import BaseModel, BeforeValidator, ConfigDict, model_validator
+
+
+def _coerce_peptidoform(v):
+    if v is None or isinstance(v, Peptidoform):
+        return v
+    elif isinstance(v, str):
+        return Peptidoform(v)
+    raise ValueError("Peptidoform must be a string, a Peptidoform object, or None.")
+
 
-try:
-    import spectrum_utils.spectrum as sus
-    import spectrum_utils.plot as sup
-except ImportError:
-    sus = None  # type: ignore[ty:invalid-assignment]
-    sup = None  # type: ignore[ty:invalid-assignment]
+_PeptidoformField = Annotated[Peptidoform | None, BeforeValidator(_coerce_peptidoform)]
 
 
 class Spectrum(BaseModel):
@@ -47,7 +53,7 @@ class Spectrum(BaseModel):
     intensity: np.ndarray
     annotations: np.ndarray | None = None
     identifier: str | None = None
-    peptidoform: Peptidoform | str | None = None
+    peptidoform: _PeptidoformField = None
     precursor_mz: float | None = None
     precursor_charge: int | None = None
     retention_time: float | None = None
@@ -73,16 +79,6 @@ def check_array_lengths(cls, data):
                 raise ValueError("Array lengths do not match.")
         return data
 
-    @field_validator("peptidoform")
-    @classmethod
-    def check_peptidoform(cls, value):
-        if not value or isinstance(value, Peptidoform):
-            return value
-        elif isinstance(value, str):
-            return Peptidoform(value)
-        else:
-            raise ValueError("Peptidoform must be a string, a Peptidoform object, or None.")
-
     @property
     def tic(self):
         """Total ion current."""
@@ -138,39 +134,51 @@ def to_spectrum_utils(self):
           Otherwise, ``ValueError`` is raised.
 
         """
-        if not sus:
-            raise ImportError("Optional dependency spectrum_utils not installed.")
+        try:
+            import spectrum_utils.spectrum as sus
+        except ImportError as e:
+            raise ImportError("Optional dependency spectrum_utils not installed.") from e
 
         if self.precursor_charge:
             precursor_charge = self.precursor_charge
         else:
             if not self.peptidoform:
                 raise ValueError("`precursor_charge` or `peptidoform` must be set.")
-            else:
-                precursor_charge = self.peptidoform.precursor_charge  # type: ignore[ty:unresolved-attribute]
+            precursor_charge = self.peptidoform.precursor_charge
+            if precursor_charge is None:
+                raise ValueError("Peptidoform charge state is not set.")
 
         if self.precursor_mz:
-            precursor_mz = self.precursor_mz
+            precursor_mz_float = float(self.precursor_mz)
         else:
             if not self.peptidoform:
                 raise ValueError("`precursor_mz` or `peptidoform` must be set.")
+            elif not self.peptidoform.theoretical_mz:
+                raise ValueError(
+                    "Peptidoform theoretical m/z could not be calculated; ensure the charge state "
+                    " is set."
+                )
             else:
                 warnings.warn("precursor_mz not set, using theoretical precursor m/z.")
-                precursor_mz = self.peptidoform.theoretical_mz  # type: ignore[ty:unresolved-attribute]
+                precursor_mz_float = float(self.peptidoform.theoretical_mz)
 
         spectrum = sus.MsmsSpectrum(
             identifier=self.identifier if self.identifier else "spectrum",
-            precursor_mz=precursor_mz,  # type: ignore[ty:invalid-argument-type]
-            precursor_charge=precursor_charge,  # type: ignore[ty:invalid-argument-type]
+            precursor_mz=precursor_mz_float,
+            precursor_charge=precursor_charge,
             mz=self.mz,
             intensity=self.intensity,
-            retention_time=self.retention_time,  # type: ignore[ty:invalid-argument-type]
+            retention_time=self.retention_time if self.retention_time is not None else 0.0,
         )
-        if self.peptidoform:
+        if (
+            self.peptidoform
+            and self.mass_tolerance is not None
+            and self.mass_tolerance_unit is not None
+        ):
             spectrum.annotate_proforma(
                 str(self.peptidoform),
-                self.mass_tolerance,  # type: ignore[ty:invalid-argument-type]
-                self.mass_tolerance_unit,  # type: ignore[ty:invalid-argument-type]
+                self.mass_tolerance,
+                self.mass_tolerance_unit,
             )
         return spectrum