diff --git a/alphabase/constants/modification.py b/alphabase/constants/modification.py index bf944bf5..f38c83dc 100644 --- a/alphabase/constants/modification.py +++ b/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`. diff --git a/alphabase/peptide/precursor.py b/alphabase/peptide/precursor.py index ce3684cf..c26ed27c 100644 --- a/alphabase/peptide/precursor.py +++ b/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 diff --git a/alphabase/psm_reader/sage_reader.py b/alphabase/psm_reader/sage_reader.py index 67eab2c6..a99eb650 100644 --- a/alphabase/psm_reader/sage_reader.py +++ b/alphabase/psm_reader/sage_reader.py @@ -1,16 +1,13 @@ """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 @@ -18,6 +15,7 @@ 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: diff --git a/alphabase/spectral_library/base.py b/alphabase/spectral_library/base.py index 1cb5e786..a1b0fab4 100644 --- a/alphabase/spectral_library/base.py +++ b/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) diff --git a/alphabase/spectral_library/decoy.py b/alphabase/spectral_library/decoy.py index 2389ad23..d4000735 100644 --- a/alphabase/spectral_library/decoy.py +++ b/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): diff --git a/alphabase/utils.py b/alphabase/utils.py index 82f9405d..07fa1a77 100644 --- a/alphabase/utils.py +++ b/alphabase/utils.py @@ -1,10 +1,16 @@ import io import itertools +import multiprocessing as mp import warnings import pandas as pd import tqdm +from alphabase.constants.modification import ( + get_modification_state, + set_modification_state, +) + class AlphabaseDeprecationWarning(DeprecationWarning): pass @@ -118,3 +124,63 @@ def _sanitize_missing_values(df: pd.DataFrame) -> pd.DataFrame: df[column] = df[column].fillna("") return df + + +def _spawn_pool(processes: int): + """Create a pool whose workers start with the parent's modification registry. + + See :func:`set_modification_state` for why a spawned worker needs it. + """ + registry = get_modification_state() + return mp.get_context("spawn").Pool( + processes, initializer=set_modification_state, initargs=(registry,) + ) + + +def _batchify(obj, batch_size: int, group_by=None) -> list: + """Give the row batches of a DataFrame or Series, each within one group.""" + groups = [group for _, group in obj.groupby(group_by)] if group_by else [obj] + return [ + group.iloc[i : i + batch_size] + for group in groups + for i in range(0, len(group), batch_size) + ] + + +def _with_progress(iterator, total, progress): + """True gives a tqdm bar, a callable `progress(iterator, total)` its own, falsy none.""" + if progress is True: + return tqdm.tqdm(iterator, total=total) + return progress(iterator, total) if callable(progress) else iterator + + +def parallel_apply( + func, + obj, + *, + processes: int, + batch_size: int, + group_by=None, + progress=True, + ignore_index: bool = False, +): + """Apply `func` to row batches of `obj` in spawned workers, then join the results. + + Parameters + ---------- + obj : pd.DataFrame or pd.Series + The object to divide into batches. + + processes : int + The number of worker processes. + + group_by : optional + A column to group by. Each batch then stays in one group. + + progress : bool or callable, optional + See :func:`_with_progress`. + """ + batches = _batchify(obj, batch_size, group_by) + with _spawn_pool(processes) as pool: + results = _with_progress(pool.imap(func, batches), len(batches), progress) + return pd.concat(list(results), ignore_index=ignore_index) diff --git a/tests/unit/peptide/test_precursor_mp.py b/tests/unit/peptide/test_precursor_mp.py new file mode 100644 index 00000000..609a989b --- /dev/null +++ b/tests/unit/peptide/test_precursor_mp.py @@ -0,0 +1,158 @@ +"""A spawned worker must end up with the parent's modification registry.""" + +import os + +import pandas as pd +import pytest + +from alphabase.constants._const import CONST_FILE_FOLDER +from alphabase.constants.modification import ( + MOD_MASS, + add_modifications_for_lower_case_AA, + add_new_modifications, + get_modification_state, + keep_modloss_by_importance, + load_mod_df, + set_modification_state, +) +from alphabase.peptide.precursor import ( + calc_precursor_isotope_intensity, + calc_precursor_isotope_intensity_mp, + update_precursor_mz, +) +from alphabase.spectral_library.base import SpecLibBase +from alphabase.utils import _spawn_pool + +CUSTOM_MOD = "TestCustomMod@K" +CUSTOM_MOD_COMPOSITION = "H(4)O(2)" + + +@pytest.fixture +def restore_registry(): + """Undo the changes that a test makes to the global registry.""" + snapshot = get_modification_state() + yield + set_modification_state(snapshot) + + +def _worker_registry(_): + """Give the worker's registry, and one lookup that is derived from it.""" + return get_modification_state(), dict(MOD_MASS) + + +def _add_custom_mod(): + add_new_modifications({CUSTOM_MOD: {"composition": CUSTOM_MOD_COMPOSITION}}) + + +def _filter_modloss(): + # level 0 differs from the level `load_mod_df` uses at import + keep_modloss_by_importance(0.0) + + +def _add_lower_case_aa(): + add_modifications_for_lower_case_AA() + + +def _load_custom_tsv(): + """Load a TSV that is different from the default TSV of alphabase.""" + default_tsv = os.path.join(CONST_FILE_FOLDER, "modification.tsv") + mod_df = pd.read_table(default_tsv, keep_default_na=False) + extra = mod_df.iloc[[0]].copy() + extra["mod_name"] = "TestTsvMod@K" + extra["composition"] = CUSTOM_MOD_COMPOSITION + custom_tsv = os.path.join( + os.path.dirname(default_tsv), "_test_modification_tmp.tsv" + ) + pd.concat([mod_df, extra], ignore_index=True).to_csv( + custom_tsv, sep="\t", index=False + ) + try: + load_mod_df(custom_tsv) + finally: + os.remove(custom_tsv) + + +@pytest.mark.parametrize( + "mutate", + [_add_custom_mod, _filter_modloss, _add_lower_case_aa, _load_custom_tsv], + ids=["custom_mods", "modloss_filtering", "lower_case_AA", "custom_tsv"], +) +def test_worker_registry_matches_parent(mutate, restore_registry): + # Given a registry that changed at run time + mutate() + expected_df, expected_mass = get_modification_state(), dict(MOD_MASS) + + # When each worker reports its registry and a lookup derived from it + with _spawn_pool(2) as pool: + reports = pool.map(_worker_registry, [None, None]) + + # Then both match the parent, so the worker rebuilt its derived lookups too + for mod_df, mod_mass in reports: + pd.testing.assert_frame_equal(mod_df, expected_df) + assert mod_mass == expected_mass + + +def _precursor_df(n_precursors=40, mod=CUSTOM_MOD): + df = pd.DataFrame( + { + "sequence": ["PEPTIDEK"] * n_precursors, + "mods": [mod] * n_precursors, + "mod_sites": ["8" if mod else ""] * n_precursors, + "charge": [2] * n_precursors, + } + ) + df["nAA"] = df["sequence"].str.len() + return update_precursor_mz(df) + + +@pytest.mark.requires_numba +def test_isotope_intensity_mp_matches_single_process(restore_registry): + # Given a library with a custom modification on its precursors + _add_custom_mod() + isotope_cols = [f"i_{i}" for i in range(6)] + + # When the isotope intensities are calculated with and without workers + single = calc_precursor_isotope_intensity(_precursor_df(), max_isotope=6) + multi = calc_precursor_isotope_intensity_mp( + _precursor_df(), max_isotope=6, mp_batch_size=10, mp_process_num=2 + ) + + # Then the two results are the same + pd.testing.assert_frame_equal( + single.sort_index()[isotope_cols], multi.sort_index()[isotope_cols] + ) + + +@pytest.mark.requires_numba +def test_caller_supplied_progress_bar_is_used(): + # Given a caller that supplies its own progress bar + seen_totals = [] + + def progress(iterator, total): + seen_totals.append(total) + return iterator + + # When the isotope intensities are calculated with workers + calc_precursor_isotope_intensity_mp( + _precursor_df(40, mod=""), + max_isotope=6, + mp_batch_size=10, + mp_process_num=2, + progress_bar=progress, + ) + + # Then the code uses the bar, and does not use it as a flag + assert seen_totals == [4] + + +@pytest.mark.requires_numba +def test_speclib_isotope_info_runs_with_multiprocessing(): + # Given a library that is large enough to use the workers + lib = SpecLibBase() + lib._precursor_df = _precursor_df(20_000, mod="") + + # When the isotope info is calculated + lib.calc_precursor_isotope_info(mp_process_num=2, mp_batch_size=1000) + + # Then it completes. Before, an unknown keyword caused a TypeError. + assert "isotope_apex_offset" in lib.precursor_df.columns