Merge pull request #126 from wilhelm-lab/feature/change_iontype_order…

…_for_dlmoix

changed order of iontypes in preprocessing

spectrum_fundamentals/fragments.py

@@ -89,7 +89,30 @@ def _xl_sanity_check(noncl_xl: int, peptide_beta_mass: float, xl_pos: float):
 
 def retrieve_ion_types(fragmentation_method: str) -> List[str]:
     """
-    Retrieve the ion types resulting from a fragmentation method.
+    Retrieve the ion types resulting from a fragmentation method in the correct order for dlomix predictions.
+
+    Given the fragmentation method the function returns all ion types that can result from it.
+
+    : param fragmentation_method: fragmentation method used during the MS
+    : raises ValueError: if fragmentation_method is other than one of HCD, CID, ETD, ECD, ETCID, ETHCD, UVPD
+    : return: list of possible ion types
+    """
+    fragmentation_method = fragmentation_method.upper()
+    if fragmentation_method == "HCD" or fragmentation_method == "CID":
+        return ["y", "b"]
+    elif fragmentation_method == "ETD" or fragmentation_method == "ECD":
+        return ["z", "c"]
+    elif fragmentation_method == "ETCID" or fragmentation_method == "ETHCD":
+        return ["y", "b", "z", "c"]
+    elif fragmentation_method == "UVPD":
+        return ["y", "b", "z", "c", "x", "a"]
+    else:
+        raise ValueError(f"Unknown fragmentation method provided: {fragmentation_method}")
+
+
+def retrieve_ion_types_for_peak_initialization(fragmentation_method: str) -> List[str]:
+    """
+    Retrieve the ion types resulting from a fragmentation method in the correct order for peak initialization.
 
     Given the fragmentation method the function returns all ion types that can result from it.
 
@@ -161,7 +184,7 @@ def initialize_peaks(
     _xl_sanity_check(noncl_xl, peptide_beta_mass, xl_pos)
 
     max_charge = min(3, charge)
-    ion_types = retrieve_ion_types(fragmentation_method)
+    ion_types = retrieve_ion_types_for_peak_initialization(fragmentation_method)
     modification_deltas = _get_modifications(sequence, custom_mods=custom_mods)
 
     fragments_meta_data = []

tests/unit_tests/test_fragments.py

@@ -110,12 +110,36 @@ def test_get_ion_types_etd(self):
 
     def test_get_ion_types_etcid(self):
         """Test retrieving ion types for ETCID."""
-        assert fragments.retrieve_ion_types("ETCID") == ["y", "z", "b", "c"]
+        assert fragments.retrieve_ion_types("ETCID") == ["y", "b", "z", "c"]
 
     def test_get_ion_types_lower_case(self):
         """Test lower case fragmentation method."""
-        assert fragments.retrieve_ion_types("uvpd") == ["x", "y", "z", "a", "b", "c"]
+        assert fragments.retrieve_ion_types("uvpd") == ["y", "b", "z", "c", "x", "a"]
 
     def test_invalid_fragmentation_method(self):
         """Test if error is raised for invalid fragmentation method."""
         self.assertRaises(ValueError, fragments.retrieve_ion_types, "XYZ")
+
+
+class TestFragmentationMethodForPeakInitialization(unittest.TestCase):
+    """Class to test the retrieving of the IonTypes for peak initialization."""
+
+    def test_get_ion_types_hcd(self):
+        """Test retrieving ion types for HCD."""
+        assert fragments.retrieve_ion_types_for_peak_initialization("HCD") == ["y", "b"]
+
+    def test_get_ion_types_etd(self):
+        """Test retrieving ion types for ETD."""
+        assert fragments.retrieve_ion_types_for_peak_initialization("ETD") == ["z", "c"]
+
+    def test_get_ion_types_etcid(self):
+        """Test retrieving ion types for ETCID."""
+        assert fragments.retrieve_ion_types_for_peak_initialization("ETCID") == ["y", "z", "b", "c"]
+
+    def test_get_ion_types_lower_case(self):
+        """Test lower case fragmentation method."""
+        assert fragments.retrieve_ion_types_for_peak_initialization("uvpd") == ["x", "y", "z", "a", "b", "c"]
+
+    def test_invalid_fragmentation_method(self):
+        """Test if error is raised for invalid fragmentation method."""
+        self.assertRaises(ValueError, fragments.retrieve_ion_types_for_peak_initialization, "XYZ")