Merge pull request #450 from MannLabs/outputaccumlator-unittest

Add additional unittests for outputaccumulator

tests/unit_tests/conftest.py

@@ -34,7 +34,7 @@ def mock_precursor_df(
     precursor_df : pd.DataFrame
         A mock precursor dataframe
     """
-
+    np.random.seed(42)  # for reproducibility
     precursor_mz = np.random.rand(n_precursor) * 2000 + 500
     precursor_charge = np.random.choice([2, 3], size=n_precursor)
 
@@ -107,7 +107,7 @@ def mock_fragment_df(n_fragments: int = 10, n_precursor: int = 20):
     fragment_df : pd.DataFrame
         A mock fragment dataframe
     """
-
+    np.random.seed(42)  # for reproducibility
     precursor_intensity = np.random.rand(n_precursor, 1)
 
     # create a dataframe with n_precursor * n_fragments rows

tests/unit_tests/test_outputaccumulator.py

@@ -9,6 +9,7 @@
 
 from alphadia import outputtransform
 from alphadia.constants.keys import SearchStepFiles
+from alphadia.outputaccumulator import ms2_quality_control
 from alphadia.workflow.base import QUANT_FOLDER_NAME
 
 
@@ -252,3 +253,80 @@ def test_default_column_assignment():
             assert built_lib.precursor_df[f"{col}"].equals(
                 built_lib.precursor_df[f"{col}_library"]
             ), f"{col} != {col}_library"
+
+
+def test_non_nan_fragments():
+    """
+    Test that the accumulated fragments data frame has no nan values
+    """
+    # Given:
+    config, temp_folder, raw_folders, psm_dfs, fragment_dfs = prepare_input_data()
+    keep_top = 2
+    config["transfer_library"]["top_k_samples"] = keep_top
+
+    # When:
+    output = outputtransform.SearchPlanOutput(config, temp_folder)
+    _ = output.build_transfer_library(raw_folders, save=True)
+    built_lib = SpecLibBase()
+    built_lib.load_hdf(
+        os.path.join(temp_folder, f"{output.TRANSFER_OUTPUT}.hdf"), load_mod_seq=True
+    )
+
+    # Then: The fragment dataframe should have no nan values
+    assert (
+        not built_lib.fragment_intensity_df.isnull().values.any()
+    ), "There are nan values in the fragment dataframe"
+
+
+def test_use_for_ms2():
+    """
+    Test that the ms2 quality control is correctly applied by checking the use_for_ms2 column in the precursor_df
+    """
+    # Given:
+    precursor_correlation_cutoff = 0.6
+    fragment_correlation_ratio = 0.9
+
+    # dummy precursor data (3 rows)
+    precursor_df = pd.DataFrame(
+        {
+            "frag_start_idx": [0, 2, 4],
+            "frag_stop_idx": [2, 4, 6],
+            "sequence": ["aa", "bb", "cc"],  # Dummy sequences
+        }
+    )
+
+    # Define dummy fragment intensity data (6 rows)
+    fragment_intensity_df = pd.DataFrame(
+        {"intensity_1": [10, 20, 30, 40, 0, 0], "intensity_2": [5, 15, 25, 35, 0, 0]}
+    )
+
+    # Define dummy fragment correlation data (6 rows)
+    fragment_correlation_df = pd.DataFrame(
+        {
+            "intensity_1": [
+                0.4,
+                0.5,  # Median < 0.6 (False)
+                0.8,
+                0.9,  # Median > 0.6 (True)
+                0.8,
+                0.9,
+            ],  # Median > 0.6 but all zero intensities (False)
+            "intensity_2": [0.3, 0.5, 0.9, 0.7, 0.8, 0.9],
+        }
+    )
+
+    speclib = SpecLibBase()
+    speclib._precursor_df = precursor_df
+    speclib._fragment_intensity_df = fragment_intensity_df
+    speclib._fragment_correlation_df = fragment_correlation_df
+
+    expected_use_for_ms2 = [False, True, False]
+    # When:
+    ms2_quality_control(
+        speclib, precursor_correlation_cutoff, fragment_correlation_ratio
+    )
+
+    # Then:
+    np.testing.assert_array_equal(
+        speclib.precursor_df["use_for_ms2"].values, expected_use_for_ms2
+    )