simplify code and add z● to ethcd/etcid

spectrum_fundamentals/constants.py

@@ -17,6 +17,42 @@
 VEC_LENGTH_CMS2 = (SEQ_LEN - 1) * 2 * 3 * 2
 # peptide of length 30 can have 29 b, y, b_short, y_short, b_long and y_long ions, each with charge 1+, 2+ and 3+
 # we do not annotate fragments wth charge 3+. All fragmets with charge 3+ convert to -1
+
+
+#############
+# ION TYPES #
+#############
+FORWARD_IONS = ["a", "b", "c"]
+BACKWARDS_IONS = ["x", "y", "z", "z●"]
+IONS = FORWARD_IONS + BACKWARDS_IONS
+
+FRAGMENTATION_TO_IONS_BY_PAIRS = {
+    "HCD": [BACKWARDS_IONS[1], FORWARD_IONS[1]],  # y,b
+    "CID": [BACKWARDS_IONS[1], FORWARD_IONS[1]],  # y,b
+    "ETD": [BACKWARDS_IONS[-1], FORWARD_IONS[2]],  # z●,c
+    "ECD": [BACKWARDS_IONS[-1], FORWARD_IONS[2]],  # z●,c
+    "ETHCD": [BACKWARDS_IONS[1], FORWARD_IONS[1], BACKWARDS_IONS[-1], FORWARD_IONS[2]],  # y,b,z●,c
+    "ETCID": [BACKWARDS_IONS[1], FORWARD_IONS[1], BACKWARDS_IONS[-1], FORWARD_IONS[2]],  # y,b,z●,c
+    "UVPD": [
+        BACKWARDS_IONS[0],
+        FORWARD_IONS[0],
+        BACKWARDS_IONS[1],
+        FORWARD_IONS[1],
+        BACKWARDS_IONS[2],
+        FORWARD_IONS[2],
+    ],  # y,b,z,c,x,a
+}
+
+FRAGMENTATION_TO_IONS_BY_DIRECTION = {
+    "HCD": [BACKWARDS_IONS[1], FORWARD_IONS[1]],  # y,b
+    "CID": [BACKWARDS_IONS[1], FORWARD_IONS[1]],  # y,b
+    "ETD": [BACKWARDS_IONS[-1], FORWARD_IONS[2]],  # z●,c
+    "ECD": [BACKWARDS_IONS[-1], FORWARD_IONS[2]],  # z●,c
+    "ETHCD": [BACKWARDS_IONS[1], BACKWARDS_IONS[-1]] + FORWARD_IONS[1:],  # y,z●,b,c
+    "ETCID": [BACKWARDS_IONS[1], BACKWARDS_IONS[-1]] + FORWARD_IONS[1:],  # y,z●,b,c
+    "UVPD": BACKWARDS_IONS[:-1] + FORWARD_IONS,  # y,z,x,b,c,a
+}
+
 #############
 # ALPHABETS #
 #############

spectrum_fundamentals/fragments.py

@@ -6,7 +6,14 @@
 import numpy as np
 import pandas as pd
 
-from .constants import AA_MASSES, ATOM_MASSES, MOD_MASSES, PARTICLE_MASSES
+from .constants import (
+    AA_MASSES,
+    ATOM_MASSES,
+    FRAGMENTATION_TO_IONS_BY_DIRECTION,
+    FRAGMENTATION_TO_IONS_BY_PAIRS,
+    MOD_MASSES,
+    PARTICLE_MASSES,
+)
 from .mod_string import internal_without_mods
 
 logger = logging.getLogger(__name__)
@@ -95,16 +102,10 @@ def retrieve_ion_types(fragmentation_method: str) -> List[str]:
     : 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:
+    ions = FRAGMENTATION_TO_IONS_BY_PAIRS.get(fragmentation_method, [])
+    if not ions:
         raise ValueError(f"Unknown fragmentation method provided: {fragmentation_method}")
+    return ions
 
 
 def retrieve_ion_types_for_peak_initialization(fragmentation_method: str) -> List[str]:
@@ -118,16 +119,10 @@ def retrieve_ion_types_for_peak_initialization(fragmentation_method: str) -> Lis
     : 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", "z", "b", "c"]
-    elif fragmentation_method == "UVPD":
-        return ["x", "y", "z", "a", "b", "c"]
-    else:
+    ions = FRAGMENTATION_TO_IONS_BY_DIRECTION.get(fragmentation_method, [])
+    if not ions:
         raise ValueError(f"Unknown fragmentation method provided: {fragmentation_method}")
+    return ions
 
 
 def get_ion_delta(ion_types: List[str]) -> np.ndarray: