Share the modification registry with multiprocessing workers

alphabase/constants/modification.py

@@ -1,3 +1,10 @@
+"""Modification registry.
+
+`update_all_by_MOD_DF` calculates all other module-level lookups from `MOD_DF`.
+Thus a copy of `MOD_DF` is a complete snapshot of the registry. Change the
+registry only with `update_all_by_MOD_DF`, to keep this true.
+"""
+
 import os
 from typing import Union
 
@@ -463,7 +470,7 @@ def _check_mass_sanity(
     composition_mass = calc_mass_from_formula(composition)
     if not np.allclose(composition_mass, MOD_MASS[mod_name], atol=1e-5):
         raise ValueError(
-            f"Modification mass of {mod_name} is inconsistent with the composition formula: {composition}, df version {MOD_DF.loc[mod_name,['composition']]}"
+            f"Modification mass of {mod_name} is inconsistent with the composition formula: {composition}, df version {MOD_DF.loc[mod_name, ['composition']]}"
             f" calculated_mass={composition_mass}, mod_mass={MOD_MASS[mod_name]}"
         )
 
@@ -533,6 +540,35 @@ def has_custom_mods():
     return len(MOD_DF[MOD_DF["classification"] == _MOD_CLASSIFICATION_USER_ADDED]) > 0
 
 
+def get_modification_state() -> pd.DataFrame:
+    """Snapshot the modification registry, for a multiprocessing worker or a test.
+
+    Returns
+    -------
+    pd.DataFrame
+        A copy of :data:`MOD_DF`, including all run-time changes.
+    """
+    return MOD_DF.copy()
+
+
+def set_modification_state(mod_df: pd.DataFrame) -> None:
+    """Install a registry snapshot and rebuild the derived lookups.
+
+    A multiprocessing worker started with "spawn" imports alphabase afresh and
+    so knows only `modification.tsv`. Installing the parent's snapshot gives it
+    every run-time change: custom modifications, modloss filtering, lower-case
+    amino acids, a custom TSV.
+
+    Parameters
+    ----------
+    mod_df : pd.DataFrame
+        A snapshot from :func:`get_modification_state`.
+    """
+    global MOD_DF
+    MOD_DF = mod_df
+    update_all_by_MOD_DF()
+
+
 def add_new_modifications(new_mods: Union[list, dict]):
     """Add new modifications into :data:`MOD_DF`.
 

alphabase/peptide/precursor.py

@@ -1,10 +1,8 @@
-import multiprocessing as mp
 import typing
 from functools import partial
 
 import numpy as np
 import pandas as pd
-from tqdm import tqdm
 from xxhash import xxh64_intdigest
 
 from alphabase.constants.aa import AA_Composition
@@ -13,6 +11,7 @@
 from alphabase.constants.modification import MOD_Composition, ModificationKeys
 from alphabase.numba_wrapper import numba_njit
 from alphabase.peptide.mass_calc import calc_peptide_masses_for_same_len_seqs
+from alphabase.utils import parallel_apply
 
 
 def _xxh64_str_intdigest(text: str, *, seed: int) -> int:
@@ -480,24 +479,6 @@ def calc_precursor_isotope_info(
     return precursor_df
 
 
-def _batchify_df(df_group, mp_batch_size):
-    """Internal funciton for multiprocessing"""
-    for _, df in df_group:
-        for i in range(0, len(df), mp_batch_size):
-            yield df.iloc[i : i + mp_batch_size, :]
-
-
-def _count_batchify_df(df_group, mp_batch_size):
-    """Internal funciton for multiprocessing"""
-    count = 0
-    for _, df in df_group:
-        for _ in range(0, len(df), mp_batch_size):
-            count += 1
-    return count
-
-
-# `progress_bar` should be replaced by more advanced tqdm wrappers created by Sander
-# I will leave it to alphabase.utils
 def calc_precursor_isotope_info_mp(
     precursor_df: pd.DataFrame,
     processes: int = 8,
@@ -540,23 +521,17 @@ def calc_precursor_isotope_info_mp(
             precursor_df=precursor_df,
             min_right_most_intensity=min_right_most_intensity,
         )
-    df_list = []
-    df_group = precursor_df.groupby("nAA")
-    with mp.get_context("spawn").Pool(processes) as p:
-        processing = p.imap(
-            partial(
-                calc_precursor_isotope_info,
-                min_right_most_intensity=min_right_most_intensity,
-            ),
-            _batchify_df(df_group, mp_batch_size),
-        )
-        if progress_bar:
-            processing = progress_bar(
-                processing, _count_batchify_df(df_group, mp_batch_size)
-            )
-        for df in processing:
-            df_list.append(df)
-    return pd.concat(df_list)
+    return parallel_apply(
+        partial(
+            calc_precursor_isotope_info,
+            min_right_most_intensity=min_right_most_intensity,
+        ),
+        precursor_df,
+        processes=processes,
+        batch_size=mp_batch_size,
+        group_by="nAA",
+        progress=progress_bar,
+    )
 
 
 def calc_precursor_isotope_intensity(
@@ -683,28 +658,20 @@ def calc_precursor_isotope_intensity_mp(
             normalize=normalize,
         )
 
-    df_list = []
-    df_group = precursor_df.groupby("nAA")
-
-    with mp.get_context("spawn").Pool(mp_process_num) as p:
-        processing = p.imap(
-            partial(
-                calc_precursor_isotope_intensity,
-                max_isotope=max_isotope,
-                min_right_most_intensity=min_right_most_intensity,
-                normalize=normalize,
-            ),
-            _batchify_df(df_group, mp_batch_size),
-        )
-
-        if progress_bar:
-            df_list = list(
-                tqdm(processing, total=_count_batchify_df(df_group, mp_batch_size))
-            )
-        else:
-            df_list = list(processing)
-
-    return pd.concat(df_list, ignore_index=True)
+    return parallel_apply(
+        partial(
+            calc_precursor_isotope_intensity,
+            max_isotope=max_isotope,
+            min_right_most_intensity=min_right_most_intensity,
+            normalize=normalize,
+        ),
+        precursor_df,
+        processes=mp_process_num,
+        batch_size=mp_batch_size,
+        group_by="nAA",
+        progress=progress_bar,
+        ignore_index=True,
+    )
 
 
 calc_precursor_isotope = calc_precursor_isotope_intensity

alphabase/psm_reader/sage_reader.py

@@ -1,23 +1,21 @@
 """SageReader for reading Sage output files."""
 
 import logging
-import multiprocessing as mp
 import re
 from abc import ABC
-from collections.abc import Generator
 from functools import partial
 from typing import Optional
 
 import numpy as np
 import pandas as pd
-from tqdm import tqdm
 
 from alphabase.constants.modification import MOD_DF, ModificationKeys
 from alphabase.psm_reader.keys import PsmDfCols
 from alphabase.psm_reader.psm_reader import (
     PSMReaderBase,
     psm_reader_provider,
 )
+from alphabase.utils import parallel_apply
 
 
 class SageModificationTranslator:
@@ -444,27 +442,6 @@ def _apply_translate_modifications(
     return psm_df
 
 
-def _batchify_df(df: pd.DataFrame, mp_batch_size: int) -> Generator:
-    """Internal funciton for applying translation modifications in parallel.
-
-    Parameters
-    ----------
-    df : pd.DataFrame
-        The PSM dataframe.
-
-    mp_batch_size : int
-        The batch size for parallel processing.
-
-    Returns
-    -------
-    typing.Generator
-        A generator for the batchified dataframe.
-
-    """
-    for i in range(0, len(df), mp_batch_size):
-        yield df.iloc[i : i + mp_batch_size, :]
-
-
 def _apply_translate_modifications_mp(
     psm_df: pd.DataFrame,
     mod_translation_df: pd.DataFrame,
@@ -493,22 +470,17 @@ def _apply_translate_modifications_mp(
         Whether to show a progress bar. Defaults to True
 
     """
-    with mp.get_context("spawn").Pool(mp_process_num) as p:
-        processing = p.imap(
-            partial(
-                _apply_translate_modifications,
-                mod_translation_df=mod_translation_df,
-            ),
-            _batchify_df(psm_df, mp_batch_size),
-        )
-        if progress_bar:
-            df_list = list(
-                tqdm(processing, total=int(np.ceil(len(psm_df) / mp_batch_size)))
-            )
-        else:
-            df_list = list(processing)
-
-    return pd.concat(df_list, ignore_index=True)
+    return parallel_apply(
+        partial(
+            _apply_translate_modifications,
+            mod_translation_df=mod_translation_df,
+        ),
+        psm_df,
+        processes=mp_process_num,
+        batch_size=mp_batch_size,
+        progress=progress_bar,
+        ignore_index=True,
+    )
 
 
 def _get_annotated_mod_df() -> pd.DataFrame:

alphabase/spectral_library/base.py

@@ -7,7 +7,6 @@
 import numpy as np
 import pandas as pd
 
-from alphabase.constants.modification import has_custom_mods
 from alphabase.io.hdf import HDF_File
 from alphabase.peptide.fragment import (
     calc_fragment_count,
@@ -429,13 +428,6 @@ def calc_precursor_isotope_intensity(
             mp_process_num > 1 and len(self.precursor_df) > mp_batch_size
         )
 
-        if do_multiprocessing and has_custom_mods():
-            logging.warning(
-                "Multiprocessing not compatible with custom modifications yet, falling back to single process."
-            )
-            do_multiprocessing = False
-            # TODO enable multiprocessing also in this case
-
         if do_multiprocessing:
             (self._precursor_df) = calc_precursor_isotope_intensity_mp(
                 self.precursor_df,
@@ -485,7 +477,8 @@ def calc_precursor_isotope_info(
             (self._precursor_df) = calc_precursor_isotope_info_mp(
                 self.precursor_df,
                 processes=mp_process_num,
-                process_bar=mp_process_bar,
+                mp_batch_size=mp_batch_size,
+                progress_bar=mp_process_bar,
             )
         else:
             (self._precursor_df) = calc_precursor_isotope_info(self.precursor_df)

alphabase/spectral_library/decoy.py

@@ -1,16 +1,10 @@
 import copy
-import multiprocessing as mp
 from typing import Any
 
 import pandas as pd
 
 from alphabase.spectral_library.base import SpecLibBase
-
-
-def _batchify_series(series, mp_batch_size):
-    """Internal funciton for multiprocessing"""
-    for i in range(0, len(series), mp_batch_size):
-        yield series.iloc[i : i + mp_batch_size]
+from alphabase.utils import parallel_apply
 
 
 class BaseDecoyGenerator:
@@ -190,16 +184,13 @@ def decoy_sequence(
             self._remove_target_seqs()
             return
 
-        sequence_batches = list(
-            _batchify_series(self._precursor_df["sequence"], mp_batch_size)
+        self._precursor_df["sequence"] = parallel_apply(
+            self.generator,
+            self._precursor_df["sequence"],
+            processes=mp_process_num,
+            batch_size=mp_batch_size,
+            progress=False,
         )
-
-        series_list = []
-        with mp.get_context("spawn").Pool(mp_process_num) as p:
-            processing = p.imap(self.generator, sequence_batches)
-            for df in processing:
-                series_list.append(df)
-        self._precursor_df["sequence"] = pd.concat(series_list)
         self._remove_target_seqs()
 
     def _remove_target_seqs(self):