Merge pull request #98 from hgb-bin-proteomics/feat-exculde_redundant_annotations

Excude redundant internal ions in fragannot

Author: levitsky

internal_ions/fragannot/fragannot_call.py

@@ -1,25 +1,26 @@
 from .fragannot_numba import FragannotNumba
 
-from typing import Dict
-from typing import List
 from ..util.spectrumio import SpectrumFile
 from psm_utils.psm_list import PSMList
 
 
 def fragannot_call(spectrum_file: SpectrumFile,
                    psms: PSMList,
                    tolerance: float,
-                   fragment_types: List[str],
-                   charges: List[str],
-                   losses: List[str],
+                   nterm_fragment_types: list[str],
+                   cterm_fragment_types: list[str],
+                   charges: list[str],
+                   losses: list[str],
                    deisotope: bool,
-                   verbose: bool = False) -> Dict:
+                   verbose: bool = False) -> dict:
 
     frag = FragannotNumba(do_parallel = True)
+
     fragannot_dict = frag.fragment_annotation(psms,
                                               spectrum_file,
                                               tolerance,
-                                              fragment_types,
+                                              nterm_fragment_types,
+                                              cterm_fragment_types,
                                               charges,
                                               losses,
                                               deisotope,

internal_ions/fragannot/fragannot_numba.py

@@ -2,7 +2,7 @@
 import json
 import re
 from pyteomics import mass, parser
-from itertools import tee
+from itertools import tee, product
 import ms_deisotope
 import logging
 
@@ -38,7 +38,8 @@ def fragment_annotation(
             psms: PSMList,
             spectra_file: SpectrumFile,
             tolerance: float,
-            fragment_types: list[str],
+            nterm_fragment_types: list[str],
+            cterm_fragment_types: list[str],
             charges: list[str],
             losses: list[str],
             deisotope: bool,
@@ -48,7 +49,8 @@ def fragment_annotation(
             psms=psms,
             spectra_file=spectra_file,
             tolerance=tolerance,
-            fragment_types=fragment_types,
+            nterm_fragment_types=nterm_fragment_types,
+            cterm_fragment_types=cterm_fragment_types,
             charges=charges,
             losses=losses,
             deisotope=deisotope,
@@ -62,7 +64,8 @@ def fragment_annotation(
         psms: PSMList,
         spectra_file: SpectrumFile,
         tolerance: float,
-        fragment_types: list[str],
+        nterm_fragment_types: list[str],
+        cterm_fragment_types: list[str],
         charges: list[str] | str,
         losses: list[str],
         deisotope: bool,
@@ -112,6 +115,8 @@ def fragment_annotation(
 
     print(f"\nAnnotating spectra in {'parallel' if do_parallel else 'serial'}...\n")
 
+    internal_fragment_types = get_internal_ion_types(nterm_fragment_types, cterm_fragment_types)
+
     if micro_batch:
         i = 0
         still_spectra_available = True
@@ -125,17 +130,17 @@ def fragment_annotation(
                 current_batch = psms[i:len(psms)]
                 still_spectra_available = False
             if do_parallel:
-                p_result = Parallel(n_jobs=nr_used_cores)(delayed(calculate_ions_for_psms)(psm, tolerance, fragment_types, charges, losses, deisotope) for psm in current_batch)
+                p_result = Parallel(n_jobs=nr_used_cores)(delayed(calculate_ions_for_psms)(psm, tolerance, nterm_fragment_types, cterm_fragment_types, internal_fragment_types, charges, losses, deisotope) for psm in current_batch)
             else:
-                p_result = [calculate_ions_for_psms(psm, tolerance, fragment_types, charges, losses, deisotope) for psm in current_batch]
+                p_result = [calculate_ions_for_psms(psm, tolerance, nterm_fragment_types, cterm_fragment_types, internal_fragment_types, charges, losses, deisotope) for psm in current_batch]
             psms_json += list(p_result)
             i += batch_size
     else:
         p_psms = tqdm(psms)  # tqdm is good for cli but bad for streamlit progress
         if do_parallel:
-            p_result = Parallel(n_jobs=nr_used_cores)(delayed(calculate_ions_for_psms)(psm, tolerance, fragment_types, charges, losses, deisotope) for psm in p_psms)
+            p_result = Parallel(n_jobs=nr_used_cores)(delayed(calculate_ions_for_psms)(psm, tolerance, nterm_fragment_types, cterm_fragment_types, internal_fragment_types, charges, losses, deisotope) for psm in p_psms)
         else:
-            p_result = [calculate_ions_for_psms(psm, tolerance, fragment_types, charges, losses, deisotope) for psm in p_psms]
+            p_result = [calculate_ions_for_psms(psm, tolerance, nterm_fragment_types, cterm_fragment_types, internal_fragment_types, charges, losses, deisotope) for psm in p_psms]
         psms_json = list(p_result)
 
     for psm in psms_json:
@@ -160,9 +165,25 @@ def fragment_annotation(
     return psms_dict
 
 
+def get_internal_ion_types(nterm_fragment_types: list[str], cterm_fragment_types: list[str]) -> list[str]:
+    """Get non-redundant internal ion types based on provided n-term and c-term fragment types."""
+    internal_ions = []
+    equiv_internal_ions = set()
+    for n_frag, c_frag in product(nterm_fragment_types, cterm_fragment_types):
+        equiv = constants.IDENTICAL_INTERNAL_IONS[(n_frag, c_frag)]
+        if equiv not in equiv_internal_ions:
+            equiv_internal_ions.add(equiv)
+            internal_ions.append(f"{n_frag}:{c_frag}")
+        else:
+            logger.debug("Skipping redundant internal ion: %s", f"{n_frag}:{c_frag}")
+    return internal_ions
+
+
 def calculate_ions_for_psms(psm,
                             tolerance: float,
-                            fragment_types: list[str],
+                            nterm_fragment_types: list[str],
+                            cterm_fragment_types: list[str],
+                            internal_fragment_types: list[str],
                             charges: list[str] | str,
                             losses: list[str],
                             deisotope: bool) -> dict[str, Any]:
@@ -182,11 +203,12 @@ def calculate_ions_for_psms(psm,
 
     theoretical_fragment_code = compute_theoretical_fragments(
         sequence_length=len(psm.peptidoform.sequence),
-        fragment_types=typed.List(fragment_types),
+        n_term_ions=typed.List(nterm_fragment_types),
+        c_term_ions=typed.List(cterm_fragment_types),
         charges=typed.List([int(c) for c in charges_used]),
         neutral_losses=typed.List(losses),
-        ion_directions=ion_directions,
-        internal=True
+        # ion_directions=ion_directions,
+        internal_ions=typed.List(internal_fragment_types),
     )
 
     theoretical_fragment_dict = {
@@ -234,14 +256,16 @@ def deisotope_peak_list(mzs: list[float], intensities: list[float]) -> tuple[lis
 @jit(nopython=True, cache=True)
 def compute_theoretical_fragments(
         sequence_length: int,
-        fragment_types: list[str],
+        n_term_ions: list[str],
+        c_term_ions: list[str],
         charges: list[int] = [1],
         neutral_losses: list[str] = [],
-        ion_directions: dict[str, str] = {},
-        internal: bool = True) -> list[str]:
+        # ion_directions: dict[str, str] = {},
+        internal_ions: list[str] = [],
+        ) -> list[str]:
 
-    n_term_ions = [ion_type for ion_type in fragment_types if ion_directions[ion_type] == "n-term"]
-    c_term_ions = [ion_type for ion_type in fragment_types if ion_directions[ion_type] == "c-term"]
+    # n_term_ions = [ion_type for ion_type in fragment_types if ion_directions[ion_type] == "n-term"]
+    # c_term_ions = [ion_type for ion_type in fragment_types if ion_directions[ion_type] == "c-term"]
 
     n_term = ["t:" + ion_type for ion_type in n_term_ions]
     c_term = [ion_type + ":t" for ion_type in c_term_ions]
@@ -276,20 +300,18 @@ def compute_theoretical_fragments(
 
     internal_frags_with_nl = ["" for _ in range(0)]  # typed empty list for Numba
 
-    if internal:
+    if internal_ions:
         # internal fragments
-        internal_frags = [
-            f"{n_term_ion}:{c_term_ion}" for n_term_ion in n_term_ions for c_term_ion in c_term_ions
-        ]
+
         internal_pos = [
             f"{i}:{j}"
             for i in range(2, sequence_length)
             for j in range(2, sequence_length)
             if i <= j
         ]
         internal_frags = [
-            f"{internal_ions}@{internal_positions}"
-            for internal_ions in internal_frags
+            f"{ions}@{internal_positions}"
+            for ions in internal_ions
             for internal_positions in internal_pos
         ]
 

internal_ions/tab1.py

@@ -72,7 +72,7 @@ def main(argv=None) -> None:
         else:
             st.warning("Please upload a spectrum file before the identifications file!", icon="⚠️")
 
-    st.session_state["fragannot_call_ion_selection"] = st.session_state["selected_ions_nterm"] + st.session_state["selected_ions_cterm"]
+    # st.session_state["fragannot_call_ion_selection"] = st.session_state["selected_ions_nterm"] + st.session_state["selected_ions_cterm"]
 
     charges_str = st.text_input("Charges to consider [comma delimited]:",
                                 value="+1",
@@ -119,7 +119,7 @@ def main(argv=None) -> None:
             run_info_str = f"\tSpectrum file name: {spectrum_file.name}\n" + \
                            f"\tIdentifications file name: {st.session_state.identifications_file.name}\n" + \
                            f"\tTolerance: {st.session_state.tolerance}\n" + \
-                           f"\tSelected ions: {', '.join(st.session_state['fragannot_call_ion_selection'])}\n" + \
+                           f"\tSelected ions: {', '.join(st.session_state['selected_ions_nterm'] + st.session_state['selected_ions_cterm'])}\n" + \
                            f"\tCharges: {', '.join(st.session_state['charges'])}\n" + \
                            f"\tLosses: {', '.join(st.session_state['losses'])}\n" + \
                            f"\tDeisotope: {st.session_state.deisotope}"
@@ -130,7 +130,8 @@ def main(argv=None) -> None:
                         result = fragannot_call(spectrum_file,
                                                 psm_list,
                                                 float(st.session_state.tolerance),
-                                                st.session_state["fragannot_call_ion_selection"],
+                                                st.session_state["selected_ions_nterm"],
+                                                st.session_state["selected_ions_cterm"],
                                                 st.session_state["charges"],
                                                 st.session_state["losses"],
                                                 st.session_state.deisotope)

internal_ions/tab2.py

@@ -105,9 +105,10 @@ def main(argv=None) -> None:
                 st.markdown("**Figure 1:** Histogram illustrating the frequency distribution of ion types present in the dataset.")
 
             with frag_center_plot_col1_2:
-                st.markdown("**Ion Type Proportions**")
-                st.plotly_chart(common_type_pie(st.session_state["frag_center_filtered"]), width="stretch")
-                st.markdown("**Figure 2:** Pie chart displaying the proportional composition of ion types within the dataset.")
+                if st.session_state["frag_center_filtered"].shape[0]:
+                    st.markdown("**Ion Type Proportions**")
+                    st.plotly_chart(common_type_pie(st.session_state["frag_center_filtered"]), width="stretch")
+                    st.markdown("**Figure 2:** Pie chart displaying the proportional composition of ion types within the dataset.")
 
             st.markdown("**Distribution of m/z Values Across Ion Types**")
             st.plotly_chart(mz_dist_ion_type(st.session_state["frag_center_filtered"]), width="stretch")

internal_ions/util/constants.py

@@ -1,4 +1,5 @@
 from pyteomics import mass
+import logging
 
 from psm_utils.io import FILETYPES
 SUPPORTED_FILETYPES = list(FILETYPES)
@@ -8,11 +9,20 @@
 DIV_COLOR = "rainbow"
 FRAGANNOT_ION_NAMES = ["a", "b", "c", "cdot", "c-1", "c+1", "x", "y", "z", "zdot", "z+1", "z+2", "z+3"]
 
-ion_comp = mass.std_ion_comp.copy()
+logger = logging.getLogger(__name__)
+
+
+class HashableComp(mass.Composition):
+    def __hash__(self):
+        return hash(tuple(sorted(self.items())))
+
+
+ion_comp = {key: HashableComp(val) for (key, val) in mass.std_ion_comp.items()}
+
 for k in list(ion_comp):
     if k.endswith('dot'):
         ion_comp[k.replace('-', '')] = ion_comp.pop(k)
-ion_comp['t'] = {}
+ion_comp['t'] = HashableComp({})
 
 ion_cap_delta_mass = {
     name: mass.calculate_mass(composition=comp, absolute=False) for (name, comp) in ion_comp.items()
@@ -35,6 +45,28 @@
 }
 
 
+def _identical_internal_ions() -> dict[tuple[str, str], tuple[str, str]]:
+    nterm = [key for key, val in ion_direction.items() if val == "n-term"]
+    cterm = [key for key, val in ion_direction.items() if val == "c-term"]
+    internal_ion_comps = {}
+    ion_mapping = {}
+    for n in nterm:
+        for c in cterm:
+            internal_ion_comps.setdefault((ion_comp[n] + ion_comp[c]), []).append((n, c))
+    for values in internal_ion_comps.values():
+        values.sort(key=lambda x: len(x[0]) + len(x[1]))
+        for v in values:
+            ion_mapping[v] = values[0]
+    return ion_mapping
+
+IDENTICAL_INTERNAL_IONS = _identical_internal_ions()
+
+for key, val in IDENTICAL_INTERNAL_IONS.items():
+    if key != val:
+        logger.debug(f"Identical internal ions: {key} -> {val}")
+
+
+
 # ---------------------------------------------------------------------------- #
 #                               For visualization                              #
 # ---------------------------------------------------------------------------- #