updated and fixed failed tests

casanovo/data/ms_io.py

@@ -41,6 +41,9 @@ class PepSpecMatch:
     aa_scores : Iterable[float]
         A list of scores for individual amino acids in the peptide
         sequence, where len(aa_scores) == len(sequence)
+    protein : str
+        For DB-search mode, the protein from which the peptide
+        in the PSM was derived. Default value is "null".
     """
 
     sequence: str
@@ -50,6 +53,7 @@ class PepSpecMatch:
     calc_mz: float
     exp_mz: float
     aa_scores: Iterable[float]
+    protein: str = "null"
 
 
 class MztabWriter:
@@ -228,7 +232,7 @@ def save(self) -> None:
                         "PSM",
                         psm.sequence,  # sequence
                         i,  # PSM_ID
-                        "null" if len(psm) < 8 else psm[7],  # accession
+                        psm.protein,  # accession
                         "null",  # unique
                         "null",  # database
                         "null",  # database_version

casanovo/denovo/model.py

@@ -1088,16 +1088,16 @@ def on_predict_batch_end(
             protein,
         ) in outputs:
             self.out_writer.psms.append(
-                (
-                    peptide,
-                    tuple(spectrum_i),
-                    peptide_score,
-                    charge,
-                    precursor_mz,
-                    self.peptide_mass_calculator.mass(peptide, charge),
-                    ",".join(list(map("{:.5f}".format, aa_scores))),
-                    protein,
-                ),
+                ms_io.PepSpecMatch(
+                    sequence=peptide,
+                    spectrum_id=tuple(spectrum_i),
+                    peptide_score=peptide_score,
+                    charge=int(charge),
+                    calc_mz=precursor_mz,
+                    exp_mz=self.peptide_mass_calculator.mass(peptide, charge),
+                    aa_scores=aa_scores,
+                    protein=protein,
+                )
             )
 
 

tests/conftest.py

@@ -42,14 +42,17 @@ def mgf_medium(tmp_path):
     return _create_mgf(peptides, mgf_file, mod_aa_mass={"C": 160.030649})
 
 
+@pytest.fixture
 def mgf_small_unannotated(tmp_path):
     """An MGF file with 2 unannotated spectra."""
     peptides = ["LESLIEK", "PEPTIDEK"]
     mgf_file = tmp_path / "small_unannotated.mgf"
     return _create_mgf(peptides, mgf_file, annotate=False)
 
 
-def _create_mgf(peptides, mgf_file, random_state=42, mod_aa_mass=None, annotate=True):
+def _create_mgf(
+    peptides, mgf_file, random_state=42, mod_aa_mass=None, annotate=True
+):
     """
     Create a fake MGF file from one or more peptides.
 
@@ -73,7 +76,10 @@ def _create_mgf(peptides, mgf_file, random_state=42, mod_aa_mass=None, annotate=
     """
     rng = np.random.default_rng(random_state)
     entries = [
-        _create_mgf_entry(p, rng.choice([2, 3]), mod_aa_mass=mod_aa_mass, annotate=annotate) for p in peptides
+        _create_mgf_entry(
+            p, rng.choice([2, 3]), mod_aa_mass=mod_aa_mass, annotate=annotate
+        )
+        for p in peptides
     ]
     with mgf_file.open("w+") as mgf_ref:
         mgf_ref.write("\n".join(entries))