diff --git a/alphabase/constants/spectral_library.py b/alphabase/constants/spectral_library.py new file mode 100644 index 00000000..ca5ac2d3 --- /dev/null +++ b/alphabase/constants/spectral_library.py @@ -0,0 +1,6 @@ +"""Constants for spectral library related functionality.""" + +# Loss type mappings - these map loss type numeric codes to their string representations +# Used in annotate.py and flat.py + +LOSS_NUMBER_TO_TYPE = {0: "", 17: "_NH3", 18: "_H2O", 98: "_modloss"} diff --git a/alphabase/peptide/fragment.py b/alphabase/peptide/fragment.py index dbc1d58a..71c3cce9 100644 --- a/alphabase/peptide/fragment.py +++ b/alphabase/peptide/fragment.py @@ -83,6 +83,8 @@ class Series: SERIES_MAPPING_INV = {v: k for k, v in SERIES_MAPPING.items()} +UNANNOTATED_TYPE = 255 + @dataclass(frozen=True) class FragmentType: diff --git a/alphabase/spectral_library/annotate.py b/alphabase/spectral_library/annotate.py new file mode 100644 index 00000000..cccabc9e --- /dev/null +++ b/alphabase/spectral_library/annotate.py @@ -0,0 +1,424 @@ +"""Module for annotating spectral libraries with raw mass spectrometry data.""" + +from typing import List, Tuple, Union + +import numba as nb +import numpy as np +import pandas as pd + +from alphabase.peptide.fragment import ( + UNANNOTATED_TYPE, + create_dense_matrices, + create_fragment_mz_dataframe, + filter_valid_charged_frag_types, + flatten_fragments, + init_fragment_by_precursor_dataframe, +) +from alphabase.spectral_library.flat import SpecLibFlat +from alphabase.spectral_library.metrics import apply_precursor_metrics + +REQUIRED_PSM_COLUMNS = [ + "sequence", + "charge", + "spec_idx", + "mod_sites", + "mods", + "flat_frag_start_idx", + "flat_frag_stop_idx", + "precursor_mz", +] + +REQUIRED_RAW_COLUMNS = [ + "_spec_idx", + "peak_start_idx", + "peak_stop_idx", + "precursor_mz", + "nce", + "rt", + "rt_max", + "rt_norm", + "activation", +] + + +def annotate_precursors_flat( + speclib_flat: SpecLibFlat, + raw_data_spectrum_df: pd.DataFrame, + spec_idx_offset: int = 0, +) -> SpecLibFlat: + """Annotate precursor information in a spectral library flat with raw data. + + This function merges precursor information from the spectral library flat + with raw data, performs necessary data cleaning, and updates the precursor + DataFrame in the spectral library flat. + + Parameters + ---------- + speclib_flat : SpecLibFlat + A spectral library flat object containing precursor information. + raw_data_spectrum_df : pd.DataFrame + Raw data object containing spectrum information. + spec_idx_offset : int, optional + Offset to apply to spec_idx in raw_data, by default 0. + + Returns + ------- + SpecLibFlat + An updated SpecLibFlat object with annotated precursor information. + + """ + # Define required columns + if "activation" not in raw_data_spectrum_df.columns: + raw_data_spectrum_df["activation"] = "Any" + + if "nce" not in raw_data_spectrum_df.columns: + raw_data_spectrum_df["nce"] = 0 + + raw_data_spectrum_df["rt_max"] = raw_data_spectrum_df["rt"].max() + raw_data_spectrum_df["rt_norm"] = ( + raw_data_spectrum_df["rt"] / raw_data_spectrum_df["rt_max"] + ) + + raw_data_spectrum_df["_spec_idx"] = ( + raw_data_spectrum_df["spec_idx"] + spec_idx_offset + ) + # Check if all required columns are present + if not all( + col in speclib_flat.precursor_df.columns for col in REQUIRED_PSM_COLUMNS + ): + raise ValueError( + f"The following columns are missing from speclib_flat.precursor_df: {set(REQUIRED_PSM_COLUMNS) - set(speclib_flat.precursor_df.columns)}" + ) + if not all(col in raw_data_spectrum_df.columns for col in REQUIRED_RAW_COLUMNS): + raise ValueError( + f"The following columns are missing from raw_data_spectrum_df: {set(REQUIRED_RAW_COLUMNS) - set(raw_data_spectrum_df.columns)}" + ) + + # Merge precursor data with raw data + merged_precursor_df = speclib_flat.precursor_df[REQUIRED_PSM_COLUMNS].merge( + raw_data_spectrum_df[REQUIRED_RAW_COLUMNS], + left_on="spec_idx", + right_on="_spec_idx", + how="left", + suffixes=("", "_observed"), + ) + + # sanity check precursor_mz + if not np.all( + np.isclose( + merged_precursor_df["precursor_mz"], + merged_precursor_df["precursor_mz_observed"], + atol=5, + ) + ): + raise ValueError( + "precursor_mz and precursor_mz_observed deviate by more than 5 m/z units. Please check the input data." + ) + + # Clean up the merged data + merged_precursor_df = merged_precursor_df.dropna( + subset=["peak_start_idx", "peak_stop_idx"] + ) + merged_precursor_df["peak_start_idx"] = merged_precursor_df[ + "peak_start_idx" + ].astype(int) + merged_precursor_df["peak_stop_idx"] = merged_precursor_df["peak_stop_idx"].astype( + int + ) + + # Create a new SpecLibFlat object with the updated precursor DataFrame + annotated_speclib_flat = speclib_flat.copy() + annotated_speclib_flat.precursor_df = merged_precursor_df + + return annotated_speclib_flat + + +def annotate_fragments_flat( + speclib_flat: SpecLibFlat, + raw_data_peak_df: pd.DataFrame, + mass_error_ppm: float = 20, +) -> SpecLibFlat: + """Annotate a spectral library flat with raw data. + + This function matches observed spectrum peaks to theoretical fragments + and organizes the results into DataFrames. + + Parameters + ---------- + speclib_flat : SpecLibFlat + A spectral library flat object containing precursor and fragment information. + raw_data_peak_df : pd.DataFrame + Raw data object containing peak information. + mass_error_ppm : float, optional + Mass error tolerance in parts per million (default: 20.0). + + Returns + ------- + SpecLibFlat + An updated SpecLibFlat object with annotated fragment information. + + """ + if not all( + col in speclib_flat.precursor_df.columns + for col in ["peak_start_idx", "peak_stop_idx"] + ): + raise ValueError( + "peak_start_idx and peak_stop_idx must be present in speclib_flat.precursor_df. Please run annotate_precursors_flat() first." + ) + + fragment_df_list = [] + matched_precursor_df = speclib_flat.precursor_df.copy() + + start_index = 0 + for i, (peak_start, peak_stop, flat_frag_start, flat_frag_stop) in enumerate( + zip( + matched_precursor_df["peak_start_idx"], + matched_precursor_df["peak_stop_idx"], + matched_precursor_df["flat_frag_start_idx"], + matched_precursor_df["flat_frag_stop_idx"], + ) + ): + spectrum_peak_df = raw_data_peak_df.iloc[peak_start:peak_stop].copy() + fragment_df = speclib_flat.fragment_df.iloc[flat_frag_start:flat_frag_stop] + + ( + spectrum_peak_df["type"], + spectrum_peak_df["loss_type"], + spectrum_peak_df["charge"], + spectrum_peak_df["number"], + spectrum_peak_df["position"], + spectrum_peak_df["error"], + ) = annotate_spectrum( + spectrum_peak_df["mz"].to_numpy(), + fragment_df["mz"].to_numpy(), + fragment_df["type"].to_numpy(), + fragment_df["loss_type"].to_numpy(), + fragment_df["charge"].to_numpy(), + fragment_df["number"].to_numpy(), + fragment_df["position"].to_numpy(), + mass_error_ppm=mass_error_ppm, + ) + + matched_precursor_df.loc[i, "flat_frag_start_idx"] = start_index + matched_precursor_df.loc[i, "flat_frag_stop_idx"] = start_index + len( + spectrum_peak_df + ) + start_index += len(spectrum_peak_df) + + fragment_df_list.append(spectrum_peak_df) + + outlib_flat = SpecLibFlat() + outlib_flat.precursor_df = matched_precursor_df.drop( + columns=["peak_start_idx", "peak_stop_idx"] + ) + outlib_flat._fragment_df = pd.concat(fragment_df_list, ignore_index=True) # noqa: SLF001 + + return outlib_flat + + +@nb.njit() +def annotate_spectrum( # noqa: PLR0913 + spectrum_mz: np.ndarray, + fragment_mz: np.ndarray, + fragment_type: np.ndarray, + fragment_loss_type: np.ndarray, + fragment_charge: np.ndarray, + fragment_number: np.ndarray, + fragment_position: np.ndarray, + mass_error_ppm: float = 20.0, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Annotate a spectrum with theoretical fragment information. + + This function matches observed spectrum peaks to theoretical fragments + based on m/z values within a specified mass error tolerance. + + Parameters + ---------- + spectrum_mz : np.ndarray + M/Z values of the observed spectrum peaks. Shape: (n_peaks,) + fragment_mz : np.ndarray + M/Z values of the theoretical fragments. Shape: (n_fragments,) + fragment_type : np.ndarray + Types of the theoretical fragments. Shape: (n_fragments,) + fragment_loss_type : np.ndarray + Loss types of the theoretical fragments. Shape: (n_fragments,) + fragment_charge : np.ndarray + Charges of the theoretical fragments. Shape: (n_fragments,) + fragment_number : np.ndarray + Numbers of the theoretical fragments. Shape: (n_fragments,) + fragment_position : np.ndarray + Positions of the theoretical fragments. Shape: (n_fragments,) + mass_error_ppm : float, optional + Mass error tolerance in parts per million (default: 20.0). + + Returns + ------- + tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray] + A tuple containing six numpy arrays, each with shape (n_peaks,): + 1. Annotated fragment types + 2. Annotated fragment loss types + 3. Annotated fragment charges + 4. Annotated fragment numbers + 5. Annotated fragment positions + 6. Annotated mass errors in ppm + + Each array corresponds to the input spectrum peaks. Peaks without + a match are annotated with ``UNANNOTATED_TYPE`` (for integer arrays) + or ``np.inf`` (for the float array). + + """ + n_peaks = len(spectrum_mz) + annotated_type = np.full(n_peaks, UNANNOTATED_TYPE, dtype=np.uint8) + annotated_loss_type = np.full(n_peaks, UNANNOTATED_TYPE, dtype=np.uint8) + annotated_charge = np.full(n_peaks, UNANNOTATED_TYPE, dtype=np.uint8) + annotated_number = np.full(n_peaks, UNANNOTATED_TYPE, dtype=np.uint8) + annotated_position = np.full(n_peaks, UNANNOTATED_TYPE, dtype=np.uint8) + annotated_error = np.full(n_peaks, np.inf, dtype=np.float32) + + # Sort fragment arrays based on m/z values + sort_idx = np.argsort(fragment_mz) + fragment_mz_sorted = fragment_mz[sort_idx] + fragment_type_sorted = fragment_type[sort_idx] + fragment_loss_type_sorted = fragment_loss_type[sort_idx] + fragment_charge_sorted = fragment_charge[sort_idx] + fragment_number_sorted = fragment_number[sort_idx] + fragment_position_sorted = fragment_position[sort_idx] + + for i, mz in enumerate(spectrum_mz): + # Find the index of the closest theoretical fragment + idx = np.searchsorted(fragment_mz_sorted, mz) + + # Check the error for this index and the previous one (if it exists) + error = np.inf + best_idx = -1 + + if idx < len(fragment_mz_sorted): + error_right = ( + abs(mz - fragment_mz_sorted[idx]) / fragment_mz_sorted[idx] * 1e6 + ) + if error_right <= mass_error_ppm: + error = error_right + best_idx = idx + + if idx > 0: + error_left = ( + abs(mz - fragment_mz_sorted[idx - 1]) + / fragment_mz_sorted[idx - 1] + * 1e6 + ) + if error_left <= mass_error_ppm and error_left < error: + error = error_left + best_idx = idx - 1 + + if best_idx != -1: + annotated_type[i] = fragment_type_sorted[best_idx] + annotated_loss_type[i] = fragment_loss_type_sorted[best_idx] + annotated_charge[i] = fragment_charge_sorted[best_idx] + annotated_number[i] = fragment_number_sorted[best_idx] + annotated_position[i] = fragment_position_sorted[best_idx] + annotated_error[i] = error + + return ( + annotated_type, + annotated_loss_type, + annotated_charge, + annotated_number, + annotated_position, + annotated_error, + ) + + +def _build_theoretical_flatlib( + template: SpecLibFlat, charged_frag_types: List[str] +) -> SpecLibFlat: + """Generate a flat spectral library containing theoretical fragments.""" + theoretical_flat = SpecLibFlat( + charged_frag_types=charged_frag_types, + ) + + precursor_df = template.precursor_df.copy() + fragment_mz_df = init_fragment_by_precursor_dataframe( + precursor_df, charged_frag_types + ) + fragment_mz_df = create_fragment_mz_dataframe( + precursor_df, + charged_frag_types, + reference_fragment_df=fragment_mz_df, + inplace_in_reference=True, + ) + + precursor_df, fragment_df = flatten_fragments( + precursor_df, + fragment_mz_df, + pd.DataFrame(), + ) + + theoretical_flat._precursor_df = precursor_df # noqa: SLF001 + theoretical_flat._fragment_df = fragment_df # noqa: SLF001 + + return theoretical_flat + + +def add_dense_lib( + flatlib: SpecLibFlat, charged_frag_types: Union[List[str], Tuple[str, ...]] +) -> SpecLibFlat: + """Add dense format to the flatlib object. + + Parameters + ---------- + flatlib : SpecLibFlat + The flatlib object to be modified + + charged_frag_types : Union[List[str], Tuple[str, ...]] + A list of fragment types that should be considered as charged. for example: `get_charged_frag_types(["b", "y", "b_NH3", "y_NH3","b_H2O", "y_H2O"], 2)` + + Returns + ------- + SpecLibFlat + The modified flatlib object + + """ + original_spec_idx = flatlib.precursor_df["spec_idx"].copy() + flatlib.precursor_df["spec_idx"] = np.arange(len(flatlib.precursor_df)) + flatlib.precursor_df["peak_start_idx"] = flatlib.precursor_df["flat_frag_start_idx"] + flatlib.precursor_df["peak_stop_idx"] = flatlib.precursor_df["flat_frag_stop_idx"] + + valid_charged_frag_types = filter_valid_charged_frag_types(charged_frag_types) + + theoretical_flat = _build_theoretical_flatlib(flatlib, valid_charged_frag_types) + raw_precursor_df = flatlib.precursor_df.copy() + raw_fragment_df = flatlib.fragment_df.copy() + + outlib_flat = annotate_precursors_flat( + theoretical_flat, + raw_precursor_df, + spec_idx_offset=0, + ) + outlib_flat = annotate_fragments_flat( + outlib_flat, + raw_fragment_df, + mass_error_ppm=20, + ) + + outlib_flat.charged_frag_types = valid_charged_frag_types + + df_collection, frag_start_idx, frag_stop_idx = create_dense_matrices( + outlib_flat.precursor_df, + outlib_flat.fragment_df, + valid_charged_frag_types, + flat_columns=["intensity", "error"], + ) + + outlib_flat._fragment_mz_df = df_collection["mz"] # noqa: SLF001 + outlib_flat._fragment_intensity_df = df_collection["intensity"] # noqa: SLF001 + + outlib_flat.precursor_df["frag_start_idx"] = frag_start_idx + outlib_flat.precursor_df["frag_stop_idx"] = frag_stop_idx + + apply_precursor_metrics(outlib_flat) + + outlib_flat.precursor_df["spec_idx"] = original_spec_idx.iloc[ + outlib_flat.precursor_df["spec_idx"].to_numpy() + ].to_numpy() + + return outlib_flat diff --git a/alphabase/spectral_library/flat.py b/alphabase/spectral_library/flat.py index bf3e6573..be95c98d 100644 --- a/alphabase/spectral_library/flat.py +++ b/alphabase/spectral_library/flat.py @@ -3,6 +3,7 @@ import pandas as pd +from alphabase.constants.spectral_library import LOSS_NUMBER_TO_TYPE from alphabase.io.hdf import HDF_File from alphabase.peptide.fragment import ( create_dense_matrices, @@ -261,16 +262,16 @@ def get_full_charged_types(self, frag_df: pd.DataFrame) -> list: # Now if we have a fragment type that is a,b,c we should have the corresponding x,y,z corresponding = {"a": "x", "b": "y", "c": "z", "x": "a", "y": "b", "z": "c"} - loss_number_to_type = {0: "", 18: "_H2O", 17: "_NH3", 98: "_modloss"} + for type, loss, max_charge in unique_charge_type_pairs.values: for possible_charge in range(1, max_charge + 1): # Add the string for this pair charged_frag_types.add( - f"{self.frag_types_as_char[type]}{loss_number_to_type[loss]}_z{possible_charge}" + f"{self.frag_types_as_char[type]}{LOSS_NUMBER_TO_TYPE[loss]}_z{possible_charge}" ) # Add the string for the corresponding pair charged_frag_types.add( - f"{corresponding[self.frag_types_as_char[type]]}{loss_number_to_type[loss]}_z{possible_charge}" + f"{corresponding[self.frag_types_as_char[type]]}{LOSS_NUMBER_TO_TYPE[loss]}_z{possible_charge}" ) return list(charged_frag_types) diff --git a/alphabase/spectral_library/metrics.py b/alphabase/spectral_library/metrics.py new file mode 100644 index 00000000..67116048 --- /dev/null +++ b/alphabase/spectral_library/metrics.py @@ -0,0 +1,229 @@ +"""Metrics utilities for spectral library precursor annotation.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np + +from alphabase.peptide.fragment import UNANNOTATED_TYPE + +if TYPE_CHECKING: + from alphabase.spectral_library.flat import SpecLibFlat + + +def calculate_pif(intensities: np.ndarray, types: np.ndarray) -> float: + """Calculate the precursor ion fraction (PIF) of a spectrum. + + This function computes the ratio of the sum of intensities for annotated + fragment peaks to the sum of all intensities. It handles edge cases such as + empty arrays and cases where the sum of all intensities is zero. + + Parameters + ---------- + intensities : np.ndarray + Array of peak intensities. + types : np.ndarray + Array of fragment types where ``UNANNOTATED_TYPE`` represents a special peak type. + + Returns + ------- + float + The calculated Peak Integral Fraction (PIF). + Returns 0 if the arrays are empty or if the sum of all intensities is zero. + + """ + if len(intensities) == 0 or len(types) == 0: + return 0.0 + + numerator = np.sum(intensities[types != UNANNOTATED_TYPE]) + denominator = np.sum(intensities) + + if denominator == 0: + return 0.0 + + return numerator / denominator + + +def _pif_metric(flatlib: SpecLibFlat) -> None: + """Calculate the precursor ion fraction (PIF) metric for the precursors. + + Important + --------- + This function operates in place and modifies the flatlib object. + The PIF is the ratio of annotated fragment intensity to total intensity. + + Parameters + ---------- + flatlib : SpecLibFlat + The flatlib object to be modified + + Returns + ------- + None + + """ + start_indices = flatlib.precursor_df["flat_frag_start_idx"].to_numpy() + stop_indices = flatlib.precursor_df["flat_frag_stop_idx"].to_numpy() + + intensity_array = flatlib._fragment_df["intensity"].to_numpy() # noqa: SLF001 + type_array = flatlib._fragment_df["type"].to_numpy() # noqa: SLF001 + + flatlib.precursor_df["pif"] = np.array( + [ + calculate_pif( + intensity_array[start:stop], + type_array[start:stop], + ) + for start, stop in zip(start_indices, stop_indices) + ] + ) + + +def _sequence_coverage_metric(flatlib: SpecLibFlat) -> None: + """Calculate the sequence coverage metric for the precursors. + + Important + --------- + This function operates in place and modifies the flatlib object + + Parameters + ---------- + flatlib : SpecLibFlat + The flatlib object to be modified + + Returns + ------- + None + + """ + fragment_coverage = ( + (flatlib._fragment_intensity_df.sum(axis=1) > 0).astype(int).to_numpy() # noqa: SLF001 + ) + + start_indices = flatlib.precursor_df["frag_start_idx"].to_numpy() + stop_indices = flatlib.precursor_df["frag_stop_idx"].to_numpy() + + flatlib.precursor_df["sequence_coverage"] = np.array( + [ + fragment_coverage[start:stop].mean() + for start, stop in zip(start_indices, stop_indices) + ] + ) + + +def _sequence_gini_metric(flatlib: SpecLibFlat) -> None: + """Calculate the sequence gini coefficient metric for the precursors. + + Important + --------- + This function operates in place and modifies the flatlib object. + The Gini coefficient ranges from 0 (perfect equality) to 1 (perfect inequality). + For fragment intensities, a lower Gini coefficient indicates more evenly + distributed fragment intensities. + + Parameters + ---------- + flatlib : SpecLibFlat + The flatlib object to be modified + + Returns + ------- + None + + """ + start_indices = flatlib.precursor_df["frag_start_idx"].to_numpy() + stop_indices = flatlib.precursor_df["frag_stop_idx"].to_numpy() + + def gini(intensities: np.ndarray) -> float: + if len(intensities) == 0 or np.sum(intensities) == 0: + return 1.0 + + sorted_intensities = np.sort(intensities) + n = len(sorted_intensities) + index = np.arange(1, n + 1) + + intensity_sum = np.sum(sorted_intensities) + weighted_sum = np.sum(index * sorted_intensities) + + return (2 * weighted_sum) / (n * intensity_sum) - ((n + 1) / n) + + flatlib.precursor_df["sequence_gini"] = np.array( + [ + gini(flatlib._fragment_intensity_df.to_numpy()[start:stop].flatten()) # noqa: SLF001 + for start, stop in zip(start_indices, stop_indices) + ] + ) + + +def _normalized_count_metric(flatlib: SpecLibFlat) -> None: + """Calculate the number of fragments per precursor normalized by the precursor length. + + Important: + --------- + This function operates in place and modifies the flatlib object. + + """ + fragment_values = flatlib._fragment_intensity_df.to_numpy() # noqa: SLF001 + start_indices = flatlib.precursor_df["frag_start_idx"].to_numpy() + stop_indices = flatlib.precursor_df["frag_stop_idx"].to_numpy() + + lengths = stop_indices - start_indices + + counts = np.array( + [ + np.sum(fragment_values[start:stop].flatten() > 0) / length + for start, stop, length in zip(start_indices, stop_indices, lengths) + ] + ) + + flatlib.precursor_df["normalized_count"] = counts + + +def _nan_median(series: np.ndarray) -> float: + """Calculate the median of a numpy array, ignoring NaNs and infs. + + defaults to inf if there are no valid values + """ + valid_values = series[~np.isnan(series) & ~np.isinf(series)] + if len(valid_values) == 0: + return np.inf + return np.median(valid_values) + + +def _mass_accuracy_metric(flatlib: SpecLibFlat) -> None: + """Calculate the mass accuracy metric for the precursors. + + Important: + --------- + This function operates in place and modifies the flatlib object + + """ + start_indices = flatlib.precursor_df["flat_frag_start_idx"].to_numpy() + stop_indices = flatlib.precursor_df["flat_frag_stop_idx"].to_numpy() + + error_array = flatlib._fragment_df["error"].to_numpy() # noqa: SLF001 + + flatlib.precursor_df["mass_accuracy"] = np.array( + [ + _nan_median(error_array[start:stop]) + for start, stop in zip(start_indices, stop_indices) + ] + ) + + +def apply_precursor_metrics(flatlib: SpecLibFlat) -> SpecLibFlat: + """Populate derived precursor metrics on the provided spectral library.""" + _pif_metric(flatlib) + _sequence_coverage_metric(flatlib) + _sequence_gini_metric(flatlib) + _normalized_count_metric(flatlib) + _mass_accuracy_metric(flatlib) + return flatlib + + +__all__ = [ + "UNANNOTATED_TYPE", + "apply_precursor_metrics", + "calculate_pif", +] diff --git a/tests/unit/spectral_library/test_annotate.py b/tests/unit/spectral_library/test_annotate.py new file mode 100644 index 00000000..1fb071bf --- /dev/null +++ b/tests/unit/spectral_library/test_annotate.py @@ -0,0 +1,521 @@ +from unittest.mock import patch + +import numpy as np +import pandas as pd +import pytest + +from alphabase.spectral_library.annotate import ( + add_dense_lib, + annotate_fragments_flat, + annotate_precursors_flat, + annotate_spectrum, +) +from alphabase.spectral_library.base import SpecLibBase +from alphabase.spectral_library.flat import SpecLibFlat +from alphabase.spectral_library.metrics import ( + _mass_accuracy_metric, + _pif_metric, + calculate_pif, +) + + +@pytest.fixture +def sample_data(): + return { + "spectrum_mz": np.array( + [50.0, 100.0, 150.0, 200.0, 250.0, 300.0, 350.0, 500.0] + ), + "fragment_mz": np.array([100.0, 200.001, 300.001, 500.0]), + "fragment_type": np.array([1, 2, 3, 4], dtype=np.uint8), + "fragment_loss_type": np.array([10, 20, 30, 40], dtype=np.uint8), + "fragment_charge": np.array([1, 1, 2, 2], dtype=np.uint8), + "fragment_number": np.array([1, 2, 3, 4], dtype=np.uint8), + "fragment_position": np.array([10, 20, 30, 40], dtype=np.uint8), + } + + +def test_annotate_spectrum_basic(sample_data): + result = annotate_spectrum(**sample_data, mass_error_ppm=20.0) + assert len(result) == 6 + assert all(len(arr) == len(sample_data["spectrum_mz"]) for arr in result) + assert result[0].dtype == np.uint8 + assert result[5].dtype == np.float32 + + +def test_annotate_spectrum_matching(sample_data): + result = annotate_spectrum(**sample_data, mass_error_ppm=20.0) + # Check matching peaks + expected_types = np.array([255, 1, 255, 2, 255, 3, 255, 4], dtype=np.uint8) + assert np.array_equal(result[0], expected_types) + + expected_loss_types = np.array([255, 10, 255, 20, 255, 30, 255, 40], dtype=np.uint8) + assert np.array_equal(result[1], expected_loss_types) + + expected_charges = np.array([255, 1, 255, 1, 255, 2, 255, 2], dtype=np.uint8) + assert np.array_equal(result[2], expected_charges) + + expected_numbers = np.array([255, 1, 255, 2, 255, 3, 255, 4], dtype=np.uint8) + assert np.array_equal(result[3], expected_numbers) + + expected_positions = np.array([255, 10, 255, 20, 255, 30, 255, 40], dtype=np.uint8) + assert np.array_equal(result[4], expected_positions) + + # Check errors + assert np.all( + result[5][[1, 3, 5, 7]] < 20.0 + ) # Errors for matching peaks should be less than 20 ppm + assert np.all( + np.isinf(result[5][[0, 2, 4, 6]]) + ) # Errors for non-matching peaks should be inf + + +def test_annotate_spectrum_no_match(sample_data): + sample_data["spectrum_mz"] = np.array([1000.0, 2000.0, 3000.0, 4000.0]) + result = annotate_spectrum(**sample_data, mass_error_ppm=20.0) + assert np.all(result[0] == 255) + assert np.all(np.isinf(result[5])) + + +def test_annotate_spectrum_mass_error_threshold(sample_data): + # Test with 1 ppm (should match only the exact 100.0 m/z) + result = annotate_spectrum(**sample_data, mass_error_ppm=1.0) + expected_types = np.array([255, 1, 255, 255, 255, 255, 255, 4], dtype=np.uint8) + assert np.array_equal(result[0], expected_types) + assert np.all(np.isinf(result[5][[0, 2, 3, 4, 5, 6]])) + assert np.all(result[5][[1, 7]] < 1.0) + + +def test_annotate_spectrum_empty_input(): + empty_data = { + key: np.array([]) + for key in [ + "spectrum_mz", + "fragment_mz", + "fragment_type", + "fragment_loss_type", + "fragment_charge", + "fragment_number", + "fragment_position", + ] + } + result = annotate_spectrum(**empty_data, mass_error_ppm=20.0) + assert all(len(arr) == 0 for arr in result) + + +@pytest.fixture +def sample_df(): + return pd.DataFrame( + {"intensity": [100, 200, 300, 400, 500], "type": [1, 2, 255, 3, 255]} + ) + + +def test_calculate_pif_normal_case(sample_df): + result = calculate_pif( + sample_df["intensity"].to_numpy(), sample_df["type"].to_numpy() + ) + expected = (100 + 200 + 400) / (100 + 200 + 300 + 400 + 500) + assert np.isclose(result, expected) + + +def test_calculate_pif_all_type_255(): + intensities = np.array([100, 200, 300]) + types = np.array([255, 255, 255], dtype=np.uint8) + assert calculate_pif(intensities, types) == 0 + + +def test_calculate_pif_no_type_255(): + intensities = np.array([100, 200, 300]) + types = np.array([1, 2, 3], dtype=np.uint8) + assert calculate_pif(intensities, types) == 1 + + +def test_calculate_pif_empty_input(): + intensities = np.array([], dtype=np.float64) + types = np.array([], dtype=np.uint8) + assert calculate_pif(intensities, types) == 0 + + +def test_calculate_pif_all_zero_intensity(): + intensities = np.array([0, 0, 0]) + types = np.array([1, 2, 3], dtype=np.uint8) + assert calculate_pif(intensities, types) == 0 + + +def test_pif_metric_updates_precursor_df(pif_speclib_flat): + _pif_metric(pif_speclib_flat) + expected = np.array([0.5, 4 / 9]) + assert np.allclose(pif_speclib_flat.precursor_df["pif"], expected) + + +@pytest.fixture +def pif_speclib_flat(): + flat = SpecLibFlat() + flat.precursor_df = pd.DataFrame( + { + "sequence": ["PEP", "PEPT"], + "mods": ["", ""], + "mod_sites": ["", ""], + "charge": [2, 2], + "flat_frag_start_idx": [0, 3], + "flat_frag_stop_idx": [3, 5], + } + ) + flat._fragment_df = pd.DataFrame( + { + "intensity": [100.0, 200.0, 300.0, 400.0, 500.0], + "type": np.array([1, 2, 255, 3, 255], dtype=np.uint8), + } + ) + return flat + + +@pytest.fixture +def mock_raw_data(): + raw_data_spectrum_df = pd.DataFrame( + { + "spec_idx": [0, 1], + "peak_start_idx": [0, 100], + "peak_stop_idx": [100, 200], + "rt": [10.5, 15.7], + "precursor_mz": [500.5, 600.6], + } + ) + raw_data_peak_df = pd.DataFrame( + {"mz": np.linspace(100, 1000, 200), "intensity": np.random.rand(200)} + ) + return raw_data_spectrum_df, raw_data_peak_df + + +@pytest.fixture +def mock_speclib_flat(): + speclib_flat = SpecLibFlat() + speclib_flat.precursor_df = pd.DataFrame( + { + "sequence": ["PEPTIDE", "SEQUENCE"], + "charge": [2, 3], + "raw_name": ["file1.raw", "file2.raw"], + "score": [0.9, 0.8], + "proteins": ["ProtA", "ProtB"], + "fdr": [0.01, 0.02], + "spec_idx": [0, 1], + "mod_sites": ["", ""], + "mods": ["", ""], + "mod_seq_hash": ["hash1", "hash2"], + "mod_seq_charge_hash": ["hash1_2", "hash2_3"], + "flat_frag_start_idx": [0, 10], + "flat_frag_stop_idx": [10, 20], + "precursor_mz": [500.5, 600.6], + } + ) + speclib_flat._fragment_df = pd.DataFrame( + { + "mz": np.linspace(100, 1000, 20), + "type": np.random.randint(0, 5, 20), + "loss_type": np.random.randint(0, 3, 20), + "charge": np.random.randint(1, 4, 20), + "number": np.random.randint(1, 10, 20), + "position": np.random.randint(0, 7, 20), + } + ) + return speclib_flat + + +# Tests for annotate_precursors_flat + + +def test_annotate_precursors_flat_success(mock_speclib_flat, mock_raw_data): + mock_raw_data_spectrum_df, _mock_raw_data_peak_df = mock_raw_data + result = annotate_precursors_flat(mock_speclib_flat, mock_raw_data_spectrum_df) + assert isinstance(result, SpecLibFlat) + assert "peak_start_idx" in result.precursor_df.columns + assert "peak_stop_idx" in result.precursor_df.columns + assert "rt" in result.precursor_df.columns + assert len(result.precursor_df) == 2 + + +def test_annotate_precursors_flat_spec_idx_offset(mock_speclib_flat, mock_raw_data): + mock_raw_data_spectrum_df, _mock_raw_data_peak_df = mock_raw_data + mock_speclib_flat.precursor_df["spec_idx"] = [1, 2] + result = annotate_precursors_flat( + mock_speclib_flat, mock_raw_data_spectrum_df, spec_idx_offset=1 + ) + assert len(result.precursor_df) == 2 + assert np.all( + result.precursor_df["spec_idx"] == mock_raw_data_spectrum_df["spec_idx"] + 1 + ) + + +def test_annotate_precursors_flat_spec_idx_no_offset(mock_speclib_flat, mock_raw_data): + mock_raw_data_spectrum_df, _mock_raw_data_peak_df = mock_raw_data + # Use a spec_idx offset that doesn't match the raw data + with pytest.raises(ValueError): + annotate_precursors_flat( + mock_speclib_flat, mock_raw_data_spectrum_df, spec_idx_offset=1 + ) + + +def test_annotate_precursors_flat_missing_columns(mock_speclib_flat, mock_raw_data): + mock_raw_data_spectrum_df, _mock_raw_data_peak_df = mock_raw_data + mock_speclib_flat.precursor_df = mock_speclib_flat.precursor_df.drop( + columns=["sequence"] + ) + with pytest.raises(ValueError): + annotate_precursors_flat(mock_speclib_flat, mock_raw_data_spectrum_df) + + +def test_annotate_precursors_flat_precursor_mz_mismatch( + mock_speclib_flat, mock_raw_data +): + mock_raw_data_spectrum_df, _mock_raw_data_peak_df = mock_raw_data + mock_raw_data_spectrum_df["precursor_mz"] = [ + 510.5, + 610.6, + ] # Deviate by more than 5 m/z + with pytest.raises( + ValueError, + match="precursor_mz and precursor_mz_observed deviate by more than 5 m/z units", + ): + annotate_precursors_flat(mock_speclib_flat, mock_raw_data_spectrum_df) + + +# Tests for annotate_fragments_flat + + +def test_annotate_fragments_flat_success(mock_speclib_flat, mock_raw_data): + mock_raw_data_spectrum_df, mock_raw_data_peak_df = mock_raw_data + with patch( + "alphabase.spectral_library.annotate.annotate_spectrum" + ) as mock_annotate_spectrum: + mock_annotate_spectrum.return_value = ( + np.zeros(100, dtype=np.uint8), + np.zeros(100, dtype=np.uint8), + np.zeros(100, dtype=np.uint8), + np.zeros(100, dtype=np.uint8), + np.zeros(100, dtype=np.uint8), + np.zeros(100, dtype=np.float32), + ) + + annotated_speclib = annotate_precursors_flat( + mock_speclib_flat, mock_raw_data_spectrum_df + ) + result = annotate_fragments_flat(annotated_speclib, mock_raw_data_peak_df) + + assert isinstance(result, SpecLibFlat) + assert len(result._fragment_df) == 200 # 100 peaks per spectrum, 2 spectra + assert mock_annotate_spectrum.call_count == 2 + + +def test_annotate_fragments_flat_missing_columns(mock_speclib_flat, mock_raw_data): + mock_raw_data_spectrum_df, mock_raw_data_peak_df = mock_raw_data + annotated_speclib = annotate_precursors_flat( + mock_speclib_flat, mock_raw_data_spectrum_df + ) + annotated_speclib.precursor_df = annotated_speclib.precursor_df.drop( + columns=["peak_start_idx"] + ) + with pytest.raises(ValueError): + annotate_fragments_flat(annotated_speclib, mock_raw_data_peak_df) + + +def test_annotate_fragments_flat_empty_spectra(mock_speclib_flat, mock_raw_data): + mock_raw_data_spectrum_df, mock_raw_data_peak_df = mock_raw_data + mock_raw_data_peak_df = pd.DataFrame(columns=["mz", "intensity"], dtype=np.float32) + annotated_speclib = annotate_precursors_flat( + mock_speclib_flat, mock_raw_data_spectrum_df + ) + result = annotate_fragments_flat(annotated_speclib, mock_raw_data_peak_df) + assert len(result._fragment_df) == 0 + + +@pytest.fixture +def mock_speclib_annotation_flat(): + charged_frag_types = [ + "b_z1", + "y_z1", + "b_NH3_z1", + "y_NH3_z1", + "b_H2O_z1", + "y_H2O_z1", + ] + + speclib_base = SpecLibBase(charged_frag_types=charged_frag_types) + speclib_base.precursor_df = pd.DataFrame( + { + "sequence": ["PEPTI", "SEQUE"], + "mods": ["", ""], + "mod_sites": ["", ""], + "charge": [2, 3], + "spec_idx": [1, 2], + "precursor_mz": [500.5, 600.6], + "rt": [10.0, 20.0], + "nce": [30.0, 32.0], + "activation": ["HCD", "HCD"], + } + ) + speclib_base.calc_fragment_mz_df() + speclib_base._fragment_intensity_df = speclib_base.fragment_mz_df.copy() + + speclib_flat = SpecLibFlat(charged_frag_types=charged_frag_types) + speclib_flat.parse_base_library(speclib_base) + + return speclib_flat + + +def test_add_dense_lib(mock_speclib_annotation_flat): + outlib = add_dense_lib( + mock_speclib_annotation_flat, + charged_frag_types=[ + "b_z1", + "y_z1", + "b_NH3_z1", + "y_NH3_z1", + "b_H2O_z1", + "y_H2O_z1", + ], + ) + + assert np.isclose( + outlib._fragment_intensity_df.values.sum(), + outlib.fragment_df["intensity"].sum(), + ) + assert np.all(outlib.precursor_df["sequence_coverage"] == [1.0, 1.0]) + assert np.allclose(outlib.precursor_df["pif"], 1.0) + + +def test_mass_accuracy_metric_normal_case(): + speclib_flat = SpecLibFlat() + speclib_flat.precursor_df = pd.DataFrame( + { + "flat_frag_start_idx": [0, 3, 6], + "flat_frag_stop_idx": [3, 6, 9], + "sequence": ["PEP", "PEPT", "PEPTI"], + } + ) + speclib_flat._fragment_df = pd.DataFrame( + { + "error": [ + 1.0, + 2.0, + 3.0, # First precursor + 4.0, + 5.0, + 6.0, # Second precursor + 7.0, + 8.0, + 9.0, + ] # Third precursor + } + ) + + _mass_accuracy_metric(speclib_flat) + expected = np.array([2.0, 5.0, 8.0]) # Medians for each group + assert np.allclose(speclib_flat.precursor_df["mass_accuracy"], expected) + + +def test_mass_accuracy_metric_with_nans(): + speclib_flat = SpecLibFlat() + speclib_flat.precursor_df = pd.DataFrame( + { + "flat_frag_start_idx": [0, 3], + "flat_frag_stop_idx": [3, 6], + "sequence": ["PEPTIDE", "SEQUENCE"], + } + ) + speclib_flat._fragment_df = pd.DataFrame( + { + "error": [ + 1.0, + np.nan, + 3.0, # First precursor + np.nan, + np.nan, + np.nan, + ] # Second precursor + } + ) + + _mass_accuracy_metric(speclib_flat) + expected = np.array([2.0, np.inf]) # Median of [1.0, 3.0] and inf for all NaNs + assert np.allclose( + speclib_flat.precursor_df["mass_accuracy"], expected, equal_nan=True + ) + + +def test_mass_accuracy_metric_with_infs(): + speclib_flat = SpecLibFlat() + speclib_flat.precursor_df = pd.DataFrame( + { + "flat_frag_start_idx": [0, 3], + "flat_frag_stop_idx": [3, 6], + "sequence": ["PEPTIDE", "SEQUENCE"], + } + ) + speclib_flat._fragment_df = pd.DataFrame( + { + "error": [ + 1.0, + np.inf, + 3.0, # First precursor + np.inf, + np.inf, + np.inf, + ] # Second precursor + } + ) + + _mass_accuracy_metric(speclib_flat) + expected = np.array([2.0, np.inf]) # Median of [1.0, 3.0] and inf for all infs + assert np.allclose( + speclib_flat.precursor_df["mass_accuracy"], expected, equal_nan=True + ) + + +def test_mass_accuracy_metric_empty_fragments(): + speclib_flat = SpecLibFlat() + speclib_flat.precursor_df = pd.DataFrame( + { + "flat_frag_start_idx": [0, 0], + "flat_frag_stop_idx": [0, 0], + "sequence": ["PEPTIDE", "SEQUENCE"], + } + ) + speclib_flat._fragment_df = pd.DataFrame({"error": []}) + + _mass_accuracy_metric(speclib_flat) + expected = np.array([np.inf, np.inf]) # inf for empty fragments + assert np.allclose( + speclib_flat.precursor_df["mass_accuracy"], expected, equal_nan=True + ) + + +def test_mass_accuracy_metric_index_mismatch(): + """Verify that mass_accuracy uses correct indices (flat vs dense). + + This test creates a scenario where flat and dense indices differ significantly. + If _mass_accuracy_metric incorrectly uses dense indices (frag_start_idx) to + slice the flat _fragment_df, the results will be wrong. + """ + speclib_flat = SpecLibFlat() + speclib_flat.precursor_df = pd.DataFrame( + { + "frag_start_idx": [0, 4], + "frag_stop_idx": [4, 8], + "flat_frag_start_idx": [0, 100], + "flat_frag_stop_idx": [100, 200], + "sequence": ["PEPTIDE", "SEQUENCE"], + } + ) + + # Flat fragment_df has 200 rows total + # First 100 (precursor 0's peaks) have error=1.0 + # Next 100 (precursor 1's peaks) have error=2.0 + speclib_flat._fragment_df = pd.DataFrame({"error": [1.0] * 100 + [2.0] * 100}) + + _mass_accuracy_metric(speclib_flat) + + expected = np.array([1.0, 2.0]) + actual = speclib_flat.precursor_df["mass_accuracy"].to_numpy() + assert np.allclose(actual, expected), ( + f"mass_accuracy uses wrong indices! Expected {expected}, got {actual}. " + "Should use flat_frag_start_idx/flat_frag_stop_idx, not frag_start_idx/frag_stop_idx." + )