Merge pull request #189 from compomics/spectrum-reader-changes

Spectrum reader changes

Author: RalfG

ms2pip/ms2pipC.py

@@ -324,10 +324,11 @@ def process_spectra(
         ces = specdict["ce"]
     else:
         specdict = (
-            data[["spec_id", "peptide", "modifications"]].set_index("spec_id").to_dict()
+            data[["spec_id", "peptide", "modifications", "charge"]].set_index("spec_id").to_dict()
         )
     peptides = specdict["peptide"]
     modifications = specdict["modifications"]
+    charges = specdict["charge"]
 
     # cols contains the names of the computed features
     cols_n = get_feature_names_new()
@@ -377,6 +378,11 @@ def process_spectra(
         peptide = peptide.replace("L", "I")
         mods = modifications[title]
 
+        if spectrum.precursor_charge:
+            charge = spectrum.precursor_charge
+        else:
+            charge = charges[title]  # If charge cannot be parsed from MGF
+
         if "mut" in mods:
             continue
 
@@ -428,15 +434,15 @@ def process_spectra(
                 dvectors.append(
                     np.array(
                         ms2pip_pyx.get_vector_ce(
-                            peptide, modpeptide, spectrum.charge, colen
+                            peptide, modpeptide, charge, colen
                         ),
                         dtype=np.uint16,
                     )
                 )  # SD: added collision energy
             else:
                 dvectors.append(
                     np.array(
-                        ms2pip_pyx.get_vector(peptide, modpeptide, spectrum.charge),
+                        ms2pip_pyx.get_vector(peptide, modpeptide, charge),
                         dtype=np.uint16,
                     )
                 )
@@ -460,7 +466,7 @@ def process_spectra(
             # Predict the b- and y-ion intensities from the peptide
             pepid_buf.append(title)
             peplen_buf.append(len(peptide) - 2)
-            charge_buf.append(spectrum.charge)
+            charge_buf.append(charge)
 
             # get/append ion mzs, targets and predictions
             targets = ms2pip_pyx.get_targets(
@@ -479,13 +485,13 @@ def process_spectra(
             if "xgboost_model_files" in MODELS[model].keys():
                 vector_buf.append(
                     np.array(
-                        ms2pip_pyx.get_vector(peptide, modpeptide, spectrum.charge),
+                        ms2pip_pyx.get_vector(peptide, modpeptide, charge),
                         dtype=np.uint16,
                     )
                 )
             else:
                 predictions = ms2pip_pyx.get_predictions(
-                    peptide, modpeptide, spectrum.charge, model_id, peaks_version, colen
+                    peptide, modpeptide, charge, model_id, peaks_version, colen
                 )
                 prediction_buf.append(
                     [np.array(p, dtype=np.float32) for p in predictions]

ms2pip/spectrum.py

@@ -1,26 +1,28 @@
 """Read MS2 spectra."""
 
 from pathlib import Path
-from typing import Generator, List
+from typing import Generator
 
 import numpy as np
-from pyteomics import mzml, mgf
+from pyteomics import mgf, mzml
 
 from ms2pip.exceptions import (
-    UnsupportedSpectrumFiletypeError,
-    InvalidSpectrumError,
     EmptySpectrumError,
+    InvalidSpectrumError,
+    UnsupportedSpectrumFiletypeError,
 )
 
 
 class Spectrum:
-    def __init__(self, title, charge, pepmass, msms, peaks) -> None:
+    def __init__(
+        self, title, msms, peaks, precursor_charge=None, precursor_mz=None
+    ) -> None:
         """Minimal information on observed MS2 spectrum."""
         self.title = str(title)
-        self.charge = int(charge)
-        self.pepmass = float(pepmass)
         self.msms = np.array(msms, dtype=np.float32)
         self.peaks = np.array(peaks, dtype=np.float32)
+        self.precursor_charge = int(precursor_charge) if precursor_charge else None
+        self.precursor_mz = float(precursor_mz) if precursor_mz else None
 
         self.tic = np.sum(self.peaks)
 
@@ -57,7 +59,9 @@ def remove_reporter_ions(self, label_type=None) -> None:
     def remove_precursor(self, tolerance=0.02) -> None:
         """Remove precursor peak."""
         for mi, mp in enumerate(self.msms):
-            if (mp >= self.pepmass - tolerance) & (mp <= self.pepmass + tolerance):
+            if (mp >= self.precursor_mz - tolerance) & (
+                mp <= self.precursor_mz + tolerance
+            ):
                 self.peaks[mi] = 0
 
     def tic_norm(self) -> None:
@@ -91,10 +95,16 @@ def read_mgf(spec_file) -> Generator[Spectrum, None, None]:
             spec_id = spectrum["params"]["title"]
             peaks = spectrum["intensity array"]
             msms = spectrum["m/z array"]
-            precursor_mz = spectrum["params"]["pepmass"][0]
-            precursor_charge = spectrum["params"]["charge"][0]
+            try:
+                precursor_charge = spectrum["params"]["charge"][0]
+            except KeyError:
+                precursor_charge = None
+            try:
+                precursor_mz = spectrum["params"]["pepmass"][0]
+            except KeyError:
+                precursor_mz = None
             parsed_spectrum = Spectrum(
-                spec_id, precursor_charge, precursor_mz, msms, peaks
+                spec_id, msms, peaks, precursor_charge, precursor_mz
             )
             yield parsed_spectrum
 
@@ -128,7 +138,7 @@ def read_mzml(spec_file) -> Generator[Spectrum, None, None]:
                 precursor_mz = precursor["selected ion m/z"]
                 precursor_charge = precursor["charge state"]
                 parsed_spectrum = Spectrum(
-                    spec_id, precursor_charge, precursor_mz, msms, peaks
+                    spec_id, msms, peaks, precursor_charge, precursor_mz
                 )
                 yield parsed_spectrum