revert use a different impl if processes=1

MannLabs · Nov 20, 2024 · 6e58644 · 6e58644
1 parent 3e3bbc5
commit 6e58644
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 18 deletions.
diff --git a/alphabase/psm_reader/sage_reader.py b/alphabase/psm_reader/sage_reader.py
@@ -489,25 +489,20 @@ def _apply_translate_modifications_mp(
         Whether to show a progress bar. Defaults to True
 
     """
-    df_list = []
-    if mp_process_num == 1:
-        for batch_df in _batchify_df(psm_df, mp_batch_size):
-            df_list.append(_apply_translate_modifications(batch_df, mod_translation_df))  # noqa: PERF401
-    else:
-        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),
+    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)))
             )
-            if progress_bar:
-                df_list = list(
-                    tqdm(processing, total=int(np.ceil(len(psm_df) / mp_batch_size)))
-                )
-            else:
-                df_list = list(processing)
+        else:
+            df_list = list(processing)
 
     return pd.concat(df_list, ignore_index=True)
 

diff --git a/tests/integration/test_psm_readers.py b/tests/integration/test_psm_readers.py
@@ -213,6 +213,9 @@ def _assert_reference_df_equal(psm_df: pd.DataFrame, test_case_name: str) -> Non
     if out_file_path.exists():
         expected_df = pd.read_parquet(out_file_path)
 
+        expected_df.sort_values(by=["rt_norm", "precursor_mz"], inplace=True)
+        psm_df.sort_values(by=["rt_norm", "precursor_mz"], inplace=True)
+
         pd.testing.assert_frame_equal(psm_df, expected_df)
     else:
         psm_df.to_parquet(out_file_path)