From 1ba246329f752af02bf02eb72293617c482e7d09 Mon Sep 17 00:00:00 2001 From: boopthesnoot Date: Wed, 9 Apr 2025 17:59:52 +0200 Subject: [PATCH 1/9] NEW: annotating real spectra vs calculated one --- alphabase/constants/spectral_library.py | 6 + alphabase/spectral_library/annotate.py | 695 +++++++++++++++++++ alphabase/spectral_library/flat.py | 7 +- tests/unit/spectral_library/test_annotate.py | 497 +++++++++++++ 4 files changed, 1202 insertions(+), 3 deletions(-) create mode 100644 alphabase/constants/spectral_library.py create mode 100644 alphabase/spectral_library/annotate.py create mode 100644 tests/unit/spectral_library/test_annotate.py 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/spectral_library/annotate.py b/alphabase/spectral_library/annotate.py new file mode 100644 index 00000000..fa56f81c --- /dev/null +++ b/alphabase/spectral_library/annotate.py @@ -0,0 +1,695 @@ +"""Module for annotating spectral libraries with raw mass spectrometry data.""" + +from typing import Optional + +import numba as nb +import numpy as np +import pandas as pd +from alpharaw.mzml import MzMLReader +from alpharaw.thermo import ThermoRawData + +from alphabase.constants.spectral_library import LOSS_NUMBER_TO_TYPE +from alphabase.spectral_library.flat import SpecLibFlat + +UNANNOTATED_TYPE = 255 + +REQUIRED_PSM_COLUMNS = [ + "sequence", + "charge", + "raw_name", + "score", + "proteins", + "spec_idx", + "mod_sites", + "mods", + "mod_seq_hash", + "mod_seq_charge_hash", + "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: MzMLReader | ThermoRawData, + 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 : MzMLReader | ThermoRawData + 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: MzMLReader | ThermoRawData, + mass_error_ppm: float = 20, +) -> SpecLibFlat: + """Annotate a spectral library flat with raw data and calculate PIF. + + This function matches observed spectrum peaks to theoretical fragments, + calculates the precursor ion fraction (PIF), and organizes the results + into DataFrames. + + Parameters + ---------- + speclib_flat : SpecLibFlat + A spectral library flat object containing precursor and fragment information. + raw_data : MzMLReader | ThermoRawData + 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 raw_data.peak_df. Please run annotate_fragments_flat() first." + ) + + fragment_df_list = [] + matched_precursor_df = speclib_flat.precursor_df.copy() + matched_precursor_df["pif"] = 0.0 + + 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, "pif"] = calculate_pif(spectrum_peak_df) + 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 + + +def calculate_pif(spectrum_peak_df: pd.DataFrame) -> float: + """Calculate the precursor ion fraction (PIF) of a spectrum. + + This function computes the ratio of the sum of intensities for non-255 type + peaks to the sum of all intensities. It handles edge cases such as empty + DataFrames or arrays, and cases where all intensities are zero. + + Parameters + ---------- + spectrum_peak_df : pandas.DataFrame + A DataFrame containing at least two columns: + - 'intensity': numeric values representing peak intensities + - 'type': integer values where 255 represents a special peak type + + Returns + ------- + float + The calculated Peak Integral Fraction (PIF). + Returns 0 if the DataFrame is empty, if either 'intensity' or 'type' + arrays are empty, or if the sum of all intensities is zero. + + """ + if spectrum_peak_df.empty: + return 0 + + intensities = spectrum_peak_df["intensity"].to_numpy() + types = spectrum_peak_df["type"].to_numpy() + + if len(intensities) == 0 or len(types) == 0: + return 0 + + numerator = np.sum(intensities[types != UNANNOTATED_TYPE]) + denominator = np.sum(intensities) + + if denominator == 0: + return 0 + + return numerator / denominator + + +@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 255 (for integer arrays) or inf (for float array). + + """ + n_peaks = len(spectrum_mz) + annotated_type = np.full(n_peaks, 255, dtype=np.uint8) + annotated_loss_type = np.full(n_peaks, 255, dtype=np.uint8) + annotated_charge = np.full(n_peaks, 255, dtype=np.uint8) + annotated_number = np.full(n_peaks, 255, dtype=np.uint8) + annotated_position = np.full(n_peaks, 255, 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 _get_dense_column( # noqa: PLR0913 + frag_type: int, + loss_type: int, + charge: int, + frag_type_mapping: dict, + loss_type_mapping: dict, + charge_type_mapping: set, +) -> str: + """Convert the fragment type, loss type and charge to a dense column name. + + Parameters + ---------- + frag_type : int + fragment type like ord('b'), ord('y') etc. + loss_type : int + loss type like 17, 18 etc. + charge : int + charge + frag_type_mapping : dict + mapping of fragment type to string + loss_type_mapping : dict + mapping of loss type to string + charge_type_mapping : dict + mapping of charge to string + + Returns + ------- + str + column name + + """ + items = [] + + if frag_type in frag_type_mapping: + items.append(frag_type_mapping[frag_type]) + + if loss_type in loss_type_mapping: + items.append(loss_type_mapping[loss_type]) + + if charge in charge_type_mapping: + items.append(charge_type_mapping[charge]) + + return "_".join(items) + + +def _add_frag_column_annotation( + flatlib: SpecLibFlat, + charged_frag_types: Optional[list[str] | tuple[str, ...]] = None, +) -> None: + """Add the fragment column annotation ['b_z1',...] to the long format fragment dataframe. + + Important + --------- + This function operates in place and modifies the flatlib object. + + Parameters + ---------- + flatlib : SpecLibFlat + The flatlib object to be modified + + charged_frag_types : Union[list, tuple] + 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 + ------- + None + + """ + frag_type_mapping = { + ord(frag_char.split("_")[0]): frag_char.split("_")[0] + for frag_char in charged_frag_types + } + + loss_type_mapping = { + loss_type: LOSS_NUMBER_TO_TYPE[loss_type].replace("_", "") + for loss_type in LOSS_NUMBER_TO_TYPE + if loss_type != 0 + } + + charge_type_mapping = { + int(frag_char.split("_")[-1][1:]): frag_char.split("_")[-1] + for frag_char in charged_frag_types + } + + flatlib.fragment_df["frag_column"] = flatlib.fragment_df.apply( + lambda row: _get_dense_column( + row["type"], + row["loss_type"], + row["charge"], + frag_type_mapping, + loss_type_mapping, + charge_type_mapping, + ), + axis=1, + ) + + +def _assign_to_dense( + flatlib: SpecLibFlat, charged_frag_types: list[str] | tuple[str, ...] +) -> None: + """Assign the fragment intensities to the dense format. + + Important + --------- + This function operates in place and modifies the flatlib object. + + Parameters + ---------- + flatlib : SpecLibFlat + The flatlib object to be modified + + charged_frag_types : Union[list, tuple] + 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 + ------- + None + + """ + flatlib.charged_frag_types = charged_frag_types + flatlib.calc_fragment_mz_df() + flatlib._fragment_intensity_df = pd.DataFrame( # noqa: SLF001 + np.zeros(flatlib.fragment_mz_df.shape), columns=flatlib.fragment_mz_df.columns + ) + + for flat_frag_start_idx, flat_frag_stop_idx, frag_start_idx, frag_stop_idx in zip( + flatlib.precursor_df["flat_frag_start_idx"], + flatlib.precursor_df["flat_frag_stop_idx"], + flatlib.precursor_df["frag_start_idx"], + flatlib.precursor_df["frag_stop_idx"], + ): + current_flat_fragment_df = flatlib.fragment_df.iloc[ + flat_frag_start_idx:flat_frag_stop_idx + ] + for frag_type in charged_frag_types: + current_frag_type_df = current_flat_fragment_df[ + current_flat_fragment_df["frag_column"] == frag_type + ] + + if current_frag_type_df.empty: + continue + + intensity = flatlib._fragment_intensity_df.iloc[ # noqa: SLF001 + frag_start_idx:frag_stop_idx + ][frag_type].to_numpy() + intensity[current_frag_type_df["position"].to_numpy()] = ( + current_frag_type_df["intensity"].to_numpy() + ) + flatlib._fragment_intensity_df.iloc[frag_start_idx:frag_stop_idx][ # noqa: SLF001 + frag_type + ] = intensity + + +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 + ) + + # Vectorized operation to calculate coverage for all precursors at once + 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 + + """ + # Vectorized operation to calculate Gini coefficient for all precursors + 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: + # Handle empty or zero cases + if len(intensities) == 0 or np.sum(intensities) == 0: + return 1.0 + + # Sort intensities in ascending order + sorted_intensities = np.sort(intensities) + n = len(sorted_intensities) + index = np.arange(1, n + 1) + + # Create intermediate values to improve readability and reduce in-place operations + 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. + + """ + # Get all values at once to avoid repeated access + 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() + + # Calculate lengths once + lengths = stop_indices - start_indices + + # Calculate counts using vectorized operations + 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 funmction operates in place and modifies the flatlib object + + """ + # Vectorized operation to calculate coverage for all precursors at once + start_indices = flatlib.precursor_df["flat_frag_start_idx"].to_numpy() + stop_indices = flatlib.precursor_df["flat_frag_stop_idx"].to_numpy() + + flatlib.precursor_df["mass_accuracy"] = np.array( + [ + _nan_median(flatlib._fragment_df["error"].to_numpy()[start:stop]) # noqa: SLF001 + for start, stop in zip(start_indices, stop_indices) + ] + ) + + +def add_dense_lib( + flatlib: SpecLibFlat, charged_frag_types: 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, tuple] + 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 + + """ + outlib_flat = flatlib.copy() + + _add_frag_column_annotation(outlib_flat, charged_frag_types=charged_frag_types) + _assign_to_dense(outlib_flat, charged_frag_types=charged_frag_types) + _sequence_coverage_metric(outlib_flat) + _sequence_gini_metric(outlib_flat) + _normalized_count_metric(outlib_flat) + _mass_accuracy_metric(outlib_flat) + # calculate precursor coverage + + 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/tests/unit/spectral_library/test_annotate.py b/tests/unit/spectral_library/test_annotate.py new file mode 100644 index 00000000..e9ba47fe --- /dev/null +++ b/tests/unit/spectral_library/test_annotate.py @@ -0,0 +1,497 @@ +from unittest.mock import patch + +import numpy as np +import pandas as pd +import pytest +from alpharaw.mzml import MzMLReader + +from alphabase.spectral_library.annotate import ( + _add_frag_column_annotation, + _get_dense_column, + _mass_accuracy_metric, + add_dense_lib, + annotate_fragments_flat, + annotate_precursors_flat, + annotate_spectrum, + calculate_pif, +) +from alphabase.spectral_library.flat import SpecLibFlat + + +@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) + expected = (100 + 200 + 400) / (100 + 200 + 300 + 400 + 500) + assert np.isclose(result, expected) + + +def test_calculate_pif_all_type_255(): + df = pd.DataFrame({"intensity": [100, 200, 300], "type": [255, 255, 255]}) + assert calculate_pif(df) == 0 + + +def test_calculate_pif_no_type_255(): + df = pd.DataFrame({"intensity": [100, 200, 300], "type": [1, 2, 3]}) + assert calculate_pif(df) == 1 + + +def test_calculate_pif_empty_dataframe(): + df = pd.DataFrame({"intensity": [], "type": []}) + assert calculate_pif(df) == 0 + + +def test_calculate_pif_all_zero_intensity(): + df = pd.DataFrame({"intensity": [0, 0, 0], "type": [1, 2, 3]}) + assert calculate_pif(df) == 0 + + +@pytest.fixture +def mock_raw_data(): + raw_data = MzMLReader() + 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 + + +@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): + result = annotate_precursors_flat(mock_speclib_flat, mock_raw_data) + 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_speclib_flat.precursor_df["spec_idx"] = [1, 2] + result = annotate_precursors_flat( + mock_speclib_flat, mock_raw_data, 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): + # 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, spec_idx_offset=1) + + +def test_annotate_precursors_flat_missing_columns(mock_speclib_flat, 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) + + +def test_annotate_precursors_flat_precursor_mz_mismatch( + mock_speclib_flat, 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) + + +# Tests for annotate_fragments_flat + + +def test_annotate_fragments_flat_success(mock_speclib_flat, mock_raw_data): + with ( + patch( + "alphabase.spectral_library.annotate.annotate_spectrum" + ) as mock_annotate_spectrum, + patch( + "alphabase.spectral_library.annotate.calculate_pif" + ) as mock_calculate_pif, + ): + 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), + ) + mock_calculate_pif.return_value = 0.8 + + annotated_speclib = annotate_precursors_flat(mock_speclib_flat, mock_raw_data) + result = annotate_fragments_flat(annotated_speclib, mock_raw_data) + + assert isinstance(result, SpecLibFlat) + assert "pif" in result.precursor_df.columns + assert len(result._fragment_df) == 200 # 100 peaks per spectrum, 2 spectra + assert all(result.precursor_df["pif"] == 0.8) + assert mock_annotate_spectrum.call_count == 2 + assert mock_calculate_pif.call_count == 2 + + +def test_annotate_fragments_flat_missing_columns(mock_speclib_flat, mock_raw_data): + annotated_speclib = annotate_precursors_flat(mock_speclib_flat, mock_raw_data) + 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) + + +def test_annotate_fragments_flat_empty_spectra(mock_speclib_flat, 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) + result = annotate_fragments_flat(annotated_speclib, mock_raw_data) + assert len(result._fragment_df) == 0 + assert all(result.precursor_df["pif"] == 0) + + +@pytest.fixture +def mappings(): + return { + "frag_type_mapping": {98: "b", 121: "y"}, + "loss_type_mapping": {17: "NH3", 18: "H2O"}, + "charge_type_mapping": {1: "z1", 2: "z2"}, + } + + +@pytest.mark.parametrize( + "type, loss_type, charge, expected", + [ + (98, 17, 2, "b_NH3_z2"), # All parameters present + (121, 0, 1, "y_z1"), # No loss type + (99, 18, 2, "H2O_z2"), # Unknown fragment type + (0, 0, 0, ""), # All unknown parameters + (98, 0, 0, "b"), # Only fragment type + (121, 0, 2, "y_z2"), # Fragment and charge, no loss type + ], +) +def test_get_dense_column(type, loss_type, charge, expected, mappings): + result = _get_dense_column(type, loss_type, charge, **mappings) + assert result == expected + + +@pytest.fixture +def mock_speclib_annotation_flat(): + speclib_flat = SpecLibFlat() + speclib_flat.precursor_df = pd.DataFrame( + { + "sequence": ["PEPTI", "SEQUE"], + "charge": [2, 3], + "raw_name": ["file1.raw", "file2.raw"], + "score": [0.9, 0.8], + "proteins": ["ProtA", "ProtB"], + "fdr": [0.01, 0.02], + "spec_idx": [1, 2], + "mod_sites": ["", ""], + "mods": ["", ""], + "mod_seq_hash": ["hash1", "hash2"], + "mod_seq_charge_hash": ["hash1_2", "hash2_3"], + "flat_frag_start_idx": [0, 4], + "flat_frag_stop_idx": [4, 8], + } + ) + speclib_flat._fragment_df = pd.DataFrame( + { + "mz": [100, 200, 300, 400, 500, 600, 700, 800], + "type": [98, 121, 98, 121, 98, 121, 98, 121], + "loss_type": [0, 0, 17, 17, 0, 0, 18, 18], + "charge": [1, 1, 1, 1, 1, 1, 1, 1], + "number": [1, 2, 3, 4, 5, 6, 7, 8], + "intensity": [100, 200, 300, 400, 500, 600, 700, 800], + "position": [0, 1, 2, 3, 0, 1, 2, 3], + "error": [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], + } + ) + return speclib_flat + + +def test_frag_column_annotation(mock_speclib_annotation_flat): + _add_frag_column_annotation( + mock_speclib_annotation_flat, + charged_frag_types=[ + "b_z1", + "y_z1", + "b_NH3_z1", + "y_NH3_z1", + "b_H2O_z1", + "y_H2O_z1", + ], + ) + assert mock_speclib_annotation_flat.fragment_df["frag_column"].tolist() == [ + "b_z1", + "y_z1", + "b_NH3_z1", + "y_NH3_z1", + "b_z1", + "y_z1", + "b_H2O_z1", + "y_H2O_z1", + ] + + +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.all( + outlib._fragment_intensity_df.sum(axis=1) == outlib.fragment_df["intensity"] + ) + assert np.all(outlib.precursor_df["sequence_coverage"] == [1.0, 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 + } + ) + print(speclib_flat.precursor_df) + + _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 + ) From a018f493a5a8da2cba8f8b9cac4bd00ce6b02d55 Mon Sep 17 00:00:00 2001 From: boopthesnoot Date: Thu, 10 Apr 2025 22:26:54 +0200 Subject: [PATCH 2/9] FIX: circular dependencies --- alphabase/spectral_library/annotate.py | 222 ++++--------------- tests/unit/spectral_library/test_annotate.py | 110 ++++----- 2 files changed, 82 insertions(+), 250 deletions(-) diff --git a/alphabase/spectral_library/annotate.py b/alphabase/spectral_library/annotate.py index fa56f81c..8d39ea87 100644 --- a/alphabase/spectral_library/annotate.py +++ b/alphabase/spectral_library/annotate.py @@ -1,14 +1,13 @@ """Module for annotating spectral libraries with raw mass spectrometry data.""" -from typing import Optional - import numba as nb import numpy as np import pandas as pd -from alpharaw.mzml import MzMLReader -from alpharaw.thermo import ThermoRawData -from alphabase.constants.spectral_library import LOSS_NUMBER_TO_TYPE +from alphabase.peptide.fragment import ( + create_dense_matrices, + filter_valid_charged_frag_types, +) from alphabase.spectral_library.flat import SpecLibFlat UNANNOTATED_TYPE = 255 @@ -44,7 +43,7 @@ def annotate_precursors_flat( speclib_flat: SpecLibFlat, - raw_data: MzMLReader | ThermoRawData, + raw_data_spectrum_df: pd.DataFrame, spec_idx_offset: int = 0, ) -> SpecLibFlat: """Annotate precursor information in a spectral library flat with raw data. @@ -57,7 +56,7 @@ def annotate_precursors_flat( ---------- speclib_flat : SpecLibFlat A spectral library flat object containing precursor information. - raw_data : MzMLReader | ThermoRawData + 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. @@ -69,19 +68,19 @@ def annotate_precursors_flat( """ # Define required columns - if "activation" not in raw_data.spectrum_df.columns: - raw_data.spectrum_df["activation"] = "Any" + 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 + 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["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 + 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( @@ -90,14 +89,14 @@ def annotate_precursors_flat( 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): + 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)}" + 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], + raw_data_spectrum_df[REQUIRED_RAW_COLUMNS], left_on="spec_idx", right_on="_spec_idx", how="left", @@ -136,7 +135,7 @@ def annotate_precursors_flat( def annotate_fragments_flat( speclib_flat: SpecLibFlat, - raw_data: MzMLReader | ThermoRawData, + raw_data_peak_df: pd.DataFrame, mass_error_ppm: float = 20, ) -> SpecLibFlat: """Annotate a spectral library flat with raw data and calculate PIF. @@ -149,7 +148,7 @@ def annotate_fragments_flat( ---------- speclib_flat : SpecLibFlat A spectral library flat object containing precursor and fragment information. - raw_data : MzMLReader | ThermoRawData + 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). @@ -165,7 +164,7 @@ def annotate_fragments_flat( for col in ["peak_start_idx", "peak_stop_idx"] ): raise ValueError( - "peak_start_idx and peak_stop_idx must be present in raw_data.peak_df. Please run annotate_fragments_flat() first." + "peak_start_idx and peak_stop_idx must be present in speclib_flat.precursor_df. Please run annotate_precursors_flat() first." ) fragment_df_list = [] @@ -181,7 +180,7 @@ def annotate_fragments_flat( matched_precursor_df["flat_frag_stop_idx"], ) ): - spectrum_peak_df = raw_data.peak_df.iloc[peak_start:peak_stop].copy() + 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] ( @@ -371,159 +370,6 @@ def annotate_spectrum( # noqa: PLR0913 ) -def _get_dense_column( # noqa: PLR0913 - frag_type: int, - loss_type: int, - charge: int, - frag_type_mapping: dict, - loss_type_mapping: dict, - charge_type_mapping: set, -) -> str: - """Convert the fragment type, loss type and charge to a dense column name. - - Parameters - ---------- - frag_type : int - fragment type like ord('b'), ord('y') etc. - loss_type : int - loss type like 17, 18 etc. - charge : int - charge - frag_type_mapping : dict - mapping of fragment type to string - loss_type_mapping : dict - mapping of loss type to string - charge_type_mapping : dict - mapping of charge to string - - Returns - ------- - str - column name - - """ - items = [] - - if frag_type in frag_type_mapping: - items.append(frag_type_mapping[frag_type]) - - if loss_type in loss_type_mapping: - items.append(loss_type_mapping[loss_type]) - - if charge in charge_type_mapping: - items.append(charge_type_mapping[charge]) - - return "_".join(items) - - -def _add_frag_column_annotation( - flatlib: SpecLibFlat, - charged_frag_types: Optional[list[str] | tuple[str, ...]] = None, -) -> None: - """Add the fragment column annotation ['b_z1',...] to the long format fragment dataframe. - - Important - --------- - This function operates in place and modifies the flatlib object. - - Parameters - ---------- - flatlib : SpecLibFlat - The flatlib object to be modified - - charged_frag_types : Union[list, tuple] - 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 - ------- - None - - """ - frag_type_mapping = { - ord(frag_char.split("_")[0]): frag_char.split("_")[0] - for frag_char in charged_frag_types - } - - loss_type_mapping = { - loss_type: LOSS_NUMBER_TO_TYPE[loss_type].replace("_", "") - for loss_type in LOSS_NUMBER_TO_TYPE - if loss_type != 0 - } - - charge_type_mapping = { - int(frag_char.split("_")[-1][1:]): frag_char.split("_")[-1] - for frag_char in charged_frag_types - } - - flatlib.fragment_df["frag_column"] = flatlib.fragment_df.apply( - lambda row: _get_dense_column( - row["type"], - row["loss_type"], - row["charge"], - frag_type_mapping, - loss_type_mapping, - charge_type_mapping, - ), - axis=1, - ) - - -def _assign_to_dense( - flatlib: SpecLibFlat, charged_frag_types: list[str] | tuple[str, ...] -) -> None: - """Assign the fragment intensities to the dense format. - - Important - --------- - This function operates in place and modifies the flatlib object. - - Parameters - ---------- - flatlib : SpecLibFlat - The flatlib object to be modified - - charged_frag_types : Union[list, tuple] - 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 - ------- - None - - """ - flatlib.charged_frag_types = charged_frag_types - flatlib.calc_fragment_mz_df() - flatlib._fragment_intensity_df = pd.DataFrame( # noqa: SLF001 - np.zeros(flatlib.fragment_mz_df.shape), columns=flatlib.fragment_mz_df.columns - ) - - for flat_frag_start_idx, flat_frag_stop_idx, frag_start_idx, frag_stop_idx in zip( - flatlib.precursor_df["flat_frag_start_idx"], - flatlib.precursor_df["flat_frag_stop_idx"], - flatlib.precursor_df["frag_start_idx"], - flatlib.precursor_df["frag_stop_idx"], - ): - current_flat_fragment_df = flatlib.fragment_df.iloc[ - flat_frag_start_idx:flat_frag_stop_idx - ] - for frag_type in charged_frag_types: - current_frag_type_df = current_flat_fragment_df[ - current_flat_fragment_df["frag_column"] == frag_type - ] - - if current_frag_type_df.empty: - continue - - intensity = flatlib._fragment_intensity_df.iloc[ # noqa: SLF001 - frag_start_idx:frag_stop_idx - ][frag_type].to_numpy() - intensity[current_frag_type_df["position"].to_numpy()] = ( - current_frag_type_df["intensity"].to_numpy() - ) - flatlib._fragment_intensity_df.iloc[frag_start_idx:frag_stop_idx][ # noqa: SLF001 - frag_type - ] = intensity - - def _sequence_coverage_metric(flatlib: SpecLibFlat) -> None: """Calculate the sequence coverage metric for the precursors. @@ -652,8 +498,8 @@ def _mass_accuracy_metric(flatlib: SpecLibFlat) -> None: """ # Vectorized operation to calculate coverage for all precursors at once - start_indices = flatlib.precursor_df["flat_frag_start_idx"].to_numpy() - stop_indices = flatlib.precursor_df["flat_frag_stop_idx"].to_numpy() + start_indices = flatlib.precursor_df["frag_start_idx"].to_numpy() + stop_indices = flatlib.precursor_df["frag_stop_idx"].to_numpy() flatlib.precursor_df["mass_accuracy"] = np.array( [ @@ -684,12 +530,26 @@ def add_dense_lib( """ outlib_flat = flatlib.copy() - _add_frag_column_annotation(outlib_flat, charged_frag_types=charged_frag_types) - _assign_to_dense(outlib_flat, charged_frag_types=charged_frag_types) + valid_charged_frag_types = filter_valid_charged_frag_types(charged_frag_types) + + 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 + _sequence_coverage_metric(outlib_flat) _sequence_gini_metric(outlib_flat) _normalized_count_metric(outlib_flat) _mass_accuracy_metric(outlib_flat) - # calculate precursor coverage return outlib_flat diff --git a/tests/unit/spectral_library/test_annotate.py b/tests/unit/spectral_library/test_annotate.py index e9ba47fe..5a368282 100644 --- a/tests/unit/spectral_library/test_annotate.py +++ b/tests/unit/spectral_library/test_annotate.py @@ -3,11 +3,8 @@ import numpy as np import pandas as pd import pytest -from alpharaw.mzml import MzMLReader from alphabase.spectral_library.annotate import ( - _add_frag_column_annotation, - _get_dense_column, _mass_accuracy_metric, add_dense_lib, annotate_fragments_flat, @@ -136,8 +133,7 @@ def test_calculate_pif_all_zero_intensity(): @pytest.fixture def mock_raw_data(): - raw_data = MzMLReader() - raw_data.spectrum_df = pd.DataFrame( + raw_data_spectrum_df = pd.DataFrame( { "spec_idx": [0, 1], "peak_start_idx": [0, 100], @@ -146,10 +142,10 @@ def mock_raw_data(): "precursor_mz": [500.5, 600.6], } ) - raw_data.peak_df = pd.DataFrame( + raw_data_peak_df = pd.DataFrame( {"mz": np.linspace(100, 1000, 200), "intensity": np.random.rand(200)} ) - return raw_data + return raw_data_spectrum_df, raw_data_peak_df @pytest.fixture @@ -190,7 +186,8 @@ def mock_speclib_flat(): def test_annotate_precursors_flat_success(mock_speclib_flat, mock_raw_data): - result = annotate_precursors_flat(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 @@ -199,34 +196,40 @@ def test_annotate_precursors_flat_success(mock_speclib_flat, mock_raw_data): 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, spec_idx_offset=1 + 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 + 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, spec_idx_offset=1) + 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) + 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["precursor_mz"] = [ + 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 @@ -234,13 +237,14 @@ def test_annotate_precursors_flat_precursor_mz_mismatch( 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) + 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" @@ -259,8 +263,10 @@ def test_annotate_fragments_flat_success(mock_speclib_flat, mock_raw_data): ) mock_calculate_pif.return_value = 0.8 - annotated_speclib = annotate_precursors_flat(mock_speclib_flat, mock_raw_data) - result = annotate_fragments_flat(annotated_speclib, mock_raw_data) + 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 "pif" in result.precursor_df.columns @@ -271,18 +277,24 @@ def test_annotate_fragments_flat_success(mock_speclib_flat, mock_raw_data): def test_annotate_fragments_flat_missing_columns(mock_speclib_flat, mock_raw_data): - annotated_speclib = annotate_precursors_flat(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) + 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.peak_df = pd.DataFrame(columns=["mz", "intensity"], dtype=np.float32) - annotated_speclib = annotate_precursors_flat(mock_speclib_flat, mock_raw_data) - result = annotate_fragments_flat(annotated_speclib, 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 assert all(result.precursor_df["pif"] == 0) @@ -296,22 +308,6 @@ def mappings(): } -@pytest.mark.parametrize( - "type, loss_type, charge, expected", - [ - (98, 17, 2, "b_NH3_z2"), # All parameters present - (121, 0, 1, "y_z1"), # No loss type - (99, 18, 2, "H2O_z2"), # Unknown fragment type - (0, 0, 0, ""), # All unknown parameters - (98, 0, 0, "b"), # Only fragment type - (121, 0, 2, "y_z2"), # Fragment and charge, no loss type - ], -) -def test_get_dense_column(type, loss_type, charge, expected, mappings): - result = _get_dense_column(type, loss_type, charge, **mappings) - assert result == expected - - @pytest.fixture def mock_speclib_annotation_flat(): speclib_flat = SpecLibFlat() @@ -347,30 +343,6 @@ def mock_speclib_annotation_flat(): return speclib_flat -def test_frag_column_annotation(mock_speclib_annotation_flat): - _add_frag_column_annotation( - mock_speclib_annotation_flat, - charged_frag_types=[ - "b_z1", - "y_z1", - "b_NH3_z1", - "y_NH3_z1", - "b_H2O_z1", - "y_H2O_z1", - ], - ) - assert mock_speclib_annotation_flat.fragment_df["frag_column"].tolist() == [ - "b_z1", - "y_z1", - "b_NH3_z1", - "y_NH3_z1", - "b_z1", - "y_z1", - "b_H2O_z1", - "y_H2O_z1", - ] - - def test_add_dense_lib(mock_speclib_annotation_flat): outlib = add_dense_lib( mock_speclib_annotation_flat, @@ -394,8 +366,8 @@ 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], + "frag_start_idx": [0, 3, 6], + "frag_stop_idx": [3, 6, 9], "sequence": ["PEP", "PEPT", "PEPTI"], } ) @@ -425,8 +397,8 @@ 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], + "frag_start_idx": [0, 3], + "frag_stop_idx": [3, 6], "sequence": ["PEPTIDE", "SEQUENCE"], } ) @@ -454,8 +426,8 @@ 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], + "frag_start_idx": [0, 3], + "frag_stop_idx": [3, 6], "sequence": ["PEPTIDE", "SEQUENCE"], } ) @@ -483,8 +455,8 @@ 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], + "frag_start_idx": [0, 0], + "frag_stop_idx": [0, 0], "sequence": ["PEPTIDE", "SEQUENCE"], } ) From 231c73d0d859fe1109ee92bc1cd0af09bfd92d12 Mon Sep 17 00:00:00 2001 From: boopthesnoot Date: Thu, 10 Apr 2025 22:41:40 +0200 Subject: [PATCH 3/9] FIX: 3.9 typing --- alphabase/spectral_library/annotate.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/alphabase/spectral_library/annotate.py b/alphabase/spectral_library/annotate.py index 8d39ea87..2c15f2f9 100644 --- a/alphabase/spectral_library/annotate.py +++ b/alphabase/spectral_library/annotate.py @@ -1,5 +1,7 @@ """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 @@ -510,7 +512,7 @@ def _mass_accuracy_metric(flatlib: SpecLibFlat) -> None: def add_dense_lib( - flatlib: SpecLibFlat, charged_frag_types: list[str] | tuple[str, ...] + flatlib: SpecLibFlat, charged_frag_types: Union[List[str], Tuple[str, ...]] ) -> SpecLibFlat: """Add dense format to the flatlib object. @@ -519,7 +521,7 @@ def add_dense_lib( flatlib : SpecLibFlat The flatlib object to be modified - charged_frag_types : Union[list, tuple] + 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 From 1baf584baf98f673a91e561e9a3c872d05e3a835 Mon Sep 17 00:00:00 2001 From: boopthesnoot Date: Thu, 9 Oct 2025 11:48:22 +0200 Subject: [PATCH 4/9] new: annotating a library with seq and spec coverage --- alphabase/spectral_library/annotate.py | 247 ++++--------------- alphabase/spectral_library/metrics.py | 226 +++++++++++++++++ tests/unit/spectral_library/test_annotate.py | 140 ++++++----- 3 files changed, 360 insertions(+), 253 deletions(-) create mode 100644 alphabase/spectral_library/metrics.py diff --git a/alphabase/spectral_library/annotate.py b/alphabase/spectral_library/annotate.py index 2c15f2f9..b1019c5f 100644 --- a/alphabase/spectral_library/annotate.py +++ b/alphabase/spectral_library/annotate.py @@ -8,23 +8,20 @@ from alphabase.peptide.fragment import ( 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 - -UNANNOTATED_TYPE = 255 +from alphabase.spectral_library.metrics import UNANNOTATED_TYPE, apply_precursor_metrics REQUIRED_PSM_COLUMNS = [ "sequence", "charge", - "raw_name", - "score", - "proteins", "spec_idx", "mod_sites", "mods", - "mod_seq_hash", - "mod_seq_charge_hash", "flat_frag_start_idx", "flat_frag_stop_idx", "precursor_mz", @@ -140,11 +137,10 @@ def annotate_fragments_flat( raw_data_peak_df: pd.DataFrame, mass_error_ppm: float = 20, ) -> SpecLibFlat: - """Annotate a spectral library flat with raw data and calculate PIF. + """Annotate a spectral library flat with raw data. - This function matches observed spectrum peaks to theoretical fragments, - calculates the precursor ion fraction (PIF), and organizes the results - into DataFrames. + This function matches observed spectrum peaks to theoretical fragments + and organizes the results into DataFrames. Parameters ---------- @@ -171,7 +167,6 @@ def annotate_fragments_flat( fragment_df_list = [] matched_precursor_df = speclib_flat.precursor_df.copy() - matched_precursor_df["pif"] = 0.0 start_index = 0 for i, (peak_start, peak_stop, flat_frag_start, flat_frag_stop) in enumerate( @@ -203,7 +198,6 @@ def annotate_fragments_flat( mass_error_ppm=mass_error_ppm, ) - matched_precursor_df.loc[i, "pif"] = calculate_pif(spectrum_peak_df) 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 @@ -221,46 +215,6 @@ def annotate_fragments_flat( return outlib_flat -def calculate_pif(spectrum_peak_df: pd.DataFrame) -> float: - """Calculate the precursor ion fraction (PIF) of a spectrum. - - This function computes the ratio of the sum of intensities for non-255 type - peaks to the sum of all intensities. It handles edge cases such as empty - DataFrames or arrays, and cases where all intensities are zero. - - Parameters - ---------- - spectrum_peak_df : pandas.DataFrame - A DataFrame containing at least two columns: - - 'intensity': numeric values representing peak intensities - - 'type': integer values where 255 represents a special peak type - - Returns - ------- - float - The calculated Peak Integral Fraction (PIF). - Returns 0 if the DataFrame is empty, if either 'intensity' or 'type' - arrays are empty, or if the sum of all intensities is zero. - - """ - if spectrum_peak_df.empty: - return 0 - - intensities = spectrum_peak_df["intensity"].to_numpy() - types = spectrum_peak_df["type"].to_numpy() - - if len(intensities) == 0 or len(types) == 0: - return 0 - - numerator = np.sum(intensities[types != UNANNOTATED_TYPE]) - denominator = np.sum(intensities) - - if denominator == 0: - return 0 - - return numerator / denominator - - @nb.njit() def annotate_spectrum( # noqa: PLR0913 spectrum_mz: np.ndarray, @@ -308,15 +262,16 @@ def annotate_spectrum( # noqa: PLR0913 6. Annotated mass errors in ppm Each array corresponds to the input spectrum peaks. Peaks without - a match are annotated with 255 (for integer arrays) or inf (for float array). + 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, 255, dtype=np.uint8) - annotated_loss_type = np.full(n_peaks, 255, dtype=np.uint8) - annotated_charge = np.full(n_peaks, 255, dtype=np.uint8) - annotated_number = np.full(n_peaks, 255, dtype=np.uint8) - annotated_position = np.full(n_peaks, 255, dtype=np.uint8) + 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 @@ -372,143 +327,35 @@ def annotate_spectrum( # noqa: PLR0913 ) -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 +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, ) - # Vectorized operation to calculate coverage for all precursors at once - 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) - ] + precursor_df = template.precursor_df.copy() + fragment_mz_df = init_fragment_by_precursor_dataframe( + precursor_df, charged_frag_types ) - - -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 - - """ - # Vectorized operation to calculate Gini coefficient for all precursors - 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: - # Handle empty or zero cases - if len(intensities) == 0 or np.sum(intensities) == 0: - return 1.0 - - # Sort intensities in ascending order - sorted_intensities = np.sort(intensities) - n = len(sorted_intensities) - index = np.arange(1, n + 1) - - # Create intermediate values to improve readability and reduce in-place operations - 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) - ] + fragment_mz_df = create_fragment_mz_dataframe( + precursor_df, + charged_frag_types, + reference_fragment_df=fragment_mz_df, + inplace_in_reference=True, ) - -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. - - """ - # Get all values at once to avoid repeated access - 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() - - # Calculate lengths once - lengths = stop_indices - start_indices - - # Calculate counts using vectorized operations - counts = np.array( - [ - np.sum(fragment_values[start:stop].flatten() > 0) / length - for start, stop, length in zip(start_indices, stop_indices, lengths) - ] + precursor_df, fragment_df = flatten_fragments( + precursor_df, + fragment_mz_df, + pd.DataFrame(), ) - 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) + theoretical_flat._precursor_df = precursor_df # noqa: SLF001 + theoretical_flat._fragment_df = fragment_df # noqa: SLF001 - -def _mass_accuracy_metric(flatlib: SpecLibFlat) -> None: - """Calculate the mass accuracy metric for the precursors. - - Important: - --------- - This funmction operates in place and modifies the flatlib object - - """ - # Vectorized operation to calculate coverage for all precursors at once - start_indices = flatlib.precursor_df["frag_start_idx"].to_numpy() - stop_indices = flatlib.precursor_df["frag_stop_idx"].to_numpy() - - flatlib.precursor_df["mass_accuracy"] = np.array( - [ - _nan_median(flatlib._fragment_df["error"].to_numpy()[start:stop]) # noqa: SLF001 - for start, stop in zip(start_indices, stop_indices) - ] - ) + return theoretical_flat def add_dense_lib( @@ -530,10 +377,27 @@ def add_dense_lib( The modified flatlib object """ - outlib_flat = flatlib.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( @@ -549,9 +413,6 @@ def add_dense_lib( outlib_flat.precursor_df["frag_start_idx"] = frag_start_idx outlib_flat.precursor_df["frag_stop_idx"] = frag_stop_idx - _sequence_coverage_metric(outlib_flat) - _sequence_gini_metric(outlib_flat) - _normalized_count_metric(outlib_flat) - _mass_accuracy_metric(outlib_flat) + apply_precursor_metrics(outlib_flat) return outlib_flat diff --git a/alphabase/spectral_library/metrics.py b/alphabase/spectral_library/metrics.py new file mode 100644 index 00000000..e16b1752 --- /dev/null +++ b/alphabase/spectral_library/metrics.py @@ -0,0 +1,226 @@ +"""Metrics utilities for spectral library precursor annotation.""" + +from __future__ import annotations + +import numpy as np + +from alphabase.spectral_library.flat import SpecLibFlat + +UNANNOTATED_TYPE = 255 + + +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 funmction operates in place and modifies the flatlib object + + """ + start_indices = flatlib.precursor_df["frag_start_idx"].to_numpy() + stop_indices = flatlib.precursor_df["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 index 5a368282..ed0d32e6 100644 --- a/tests/unit/spectral_library/test_annotate.py +++ b/tests/unit/spectral_library/test_annotate.py @@ -5,14 +5,18 @@ import pytest from alphabase.spectral_library.annotate import ( - _mass_accuracy_metric, add_dense_lib, annotate_fragments_flat, annotate_precursors_flat, annotate_spectrum, - calculate_pif, ) +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 @@ -106,29 +110,63 @@ def sample_df(): def test_calculate_pif_normal_case(sample_df): - result = calculate_pif(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(): - df = pd.DataFrame({"intensity": [100, 200, 300], "type": [255, 255, 255]}) - assert calculate_pif(df) == 0 + 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(): - df = pd.DataFrame({"intensity": [100, 200, 300], "type": [1, 2, 3]}) - assert calculate_pif(df) == 1 + 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_dataframe(): - df = pd.DataFrame({"intensity": [], "type": []}) - assert calculate_pif(df) == 0 +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(): - df = pd.DataFrame({"intensity": [0, 0, 0], "type": [1, 2, 3]}) - assert calculate_pif(df) == 0 + 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 @@ -245,14 +283,9 @@ def test_annotate_precursors_flat_precursor_mz_mismatch( 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, - patch( - "alphabase.spectral_library.annotate.calculate_pif" - ) as mock_calculate_pif, - ): + 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), @@ -261,7 +294,6 @@ def test_annotate_fragments_flat_success(mock_speclib_flat, mock_raw_data): np.zeros(100, dtype=np.uint8), np.zeros(100, dtype=np.float32), ) - mock_calculate_pif.return_value = 0.8 annotated_speclib = annotate_precursors_flat( mock_speclib_flat, mock_raw_data_spectrum_df @@ -269,11 +301,8 @@ def test_annotate_fragments_flat_success(mock_speclib_flat, mock_raw_data): result = annotate_fragments_flat(annotated_speclib, mock_raw_data_peak_df) assert isinstance(result, SpecLibFlat) - assert "pif" in result.precursor_df.columns assert len(result._fragment_df) == 200 # 100 peaks per spectrum, 2 spectra - assert all(result.precursor_df["pif"] == 0.8) assert mock_annotate_spectrum.call_count == 2 - assert mock_calculate_pif.call_count == 2 def test_annotate_fragments_flat_missing_columns(mock_speclib_flat, mock_raw_data): @@ -296,50 +325,39 @@ def test_annotate_fragments_flat_empty_spectra(mock_speclib_flat, mock_raw_data) ) result = annotate_fragments_flat(annotated_speclib, mock_raw_data_peak_df) assert len(result._fragment_df) == 0 - assert all(result.precursor_df["pif"] == 0) - - -@pytest.fixture -def mappings(): - return { - "frag_type_mapping": {98: "b", 121: "y"}, - "loss_type_mapping": {17: "NH3", 18: "H2O"}, - "charge_type_mapping": {1: "z1", 2: "z2"}, - } @pytest.fixture def mock_speclib_annotation_flat(): - speclib_flat = SpecLibFlat() - speclib_flat.precursor_df = pd.DataFrame( + 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], - "raw_name": ["file1.raw", "file2.raw"], - "score": [0.9, 0.8], - "proteins": ["ProtA", "ProtB"], - "fdr": [0.01, 0.02], "spec_idx": [1, 2], - "mod_sites": ["", ""], - "mods": ["", ""], - "mod_seq_hash": ["hash1", "hash2"], - "mod_seq_charge_hash": ["hash1_2", "hash2_3"], - "flat_frag_start_idx": [0, 4], - "flat_frag_stop_idx": [4, 8], - } - ) - speclib_flat._fragment_df = pd.DataFrame( - { - "mz": [100, 200, 300, 400, 500, 600, 700, 800], - "type": [98, 121, 98, 121, 98, 121, 98, 121], - "loss_type": [0, 0, 17, 17, 0, 0, 18, 18], - "charge": [1, 1, 1, 1, 1, 1, 1, 1], - "number": [1, 2, 3, 4, 5, 6, 7, 8], - "intensity": [100, 200, 300, 400, 500, 600, 700, 800], - "position": [0, 1, 2, 3, 0, 1, 2, 3], - "error": [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], + "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 @@ -356,10 +374,12 @@ def test_add_dense_lib(mock_speclib_annotation_flat): ], ) - assert np.all( - outlib._fragment_intensity_df.sum(axis=1) == outlib.fragment_df["intensity"] + 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(): From d54af4f71dee57fc0bf3576a669d585137a96db8 Mon Sep 17 00:00:00 2001 From: boopthesnoot Date: Thu, 9 Oct 2025 11:56:12 +0200 Subject: [PATCH 5/9] fix: ruff --- alphabase/spectral_library/metrics.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/alphabase/spectral_library/metrics.py b/alphabase/spectral_library/metrics.py index e16b1752..39945728 100644 --- a/alphabase/spectral_library/metrics.py +++ b/alphabase/spectral_library/metrics.py @@ -2,9 +2,12 @@ from __future__ import annotations +from typing import TYPE_CHECKING + import numpy as np -from alphabase.spectral_library.flat import SpecLibFlat +if TYPE_CHECKING: + from alphabase.spectral_library.flat import SpecLibFlat UNANNOTATED_TYPE = 255 From c7bd68964e777446fa76aae0b6cdebb53aa22b74 Mon Sep 17 00:00:00 2001 From: boopthesnoot Date: Thu, 9 Oct 2025 14:36:19 +0200 Subject: [PATCH 6/9] enh: moved UNANNOTATED_TYPE to fragment.py --- alphabase/peptide/fragment.py | 2 ++ alphabase/spectral_library/annotate.py | 3 ++- alphabase/spectral_library/metrics.py | 4 ++-- 3 files changed, 6 insertions(+), 3 deletions(-) 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 index b1019c5f..8fa03f56 100644 --- a/alphabase/spectral_library/annotate.py +++ b/alphabase/spectral_library/annotate.py @@ -7,6 +7,7 @@ import pandas as pd from alphabase.peptide.fragment import ( + UNANNOTATED_TYPE, create_dense_matrices, create_fragment_mz_dataframe, filter_valid_charged_frag_types, @@ -14,7 +15,7 @@ init_fragment_by_precursor_dataframe, ) from alphabase.spectral_library.flat import SpecLibFlat -from alphabase.spectral_library.metrics import UNANNOTATED_TYPE, apply_precursor_metrics +from alphabase.spectral_library.metrics import apply_precursor_metrics REQUIRED_PSM_COLUMNS = [ "sequence", diff --git a/alphabase/spectral_library/metrics.py b/alphabase/spectral_library/metrics.py index 39945728..1934cb6c 100644 --- a/alphabase/spectral_library/metrics.py +++ b/alphabase/spectral_library/metrics.py @@ -6,11 +6,11 @@ import numpy as np +from alphabase.peptide.fragment import UNANNOTATED_TYPE + if TYPE_CHECKING: from alphabase.spectral_library.flat import SpecLibFlat -UNANNOTATED_TYPE = 255 - def calculate_pif(intensities: np.ndarray, types: np.ndarray) -> float: """Calculate the precursor ion fraction (PIF) of a spectrum. From 2a3d8b1e98591037bc061d1499fda1a282121de1 Mon Sep 17 00:00:00 2001 From: boopthesnoot Date: Fri, 28 Nov 2025 22:49:57 +0100 Subject: [PATCH 7/9] fix: saving spec_idx --- alphabase/spectral_library/annotate.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/alphabase/spectral_library/annotate.py b/alphabase/spectral_library/annotate.py index 8fa03f56..911bf098 100644 --- a/alphabase/spectral_library/annotate.py +++ b/alphabase/spectral_library/annotate.py @@ -378,6 +378,7 @@ def add_dense_lib( 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"] @@ -416,4 +417,8 @@ def add_dense_lib( apply_precursor_metrics(outlib_flat) + outlib_flat.precursor_df["spec_idx"] = original_spec_idx.iloc[ + outlib_flat.precursor_df["spec_idx"].values + ].values + return outlib_flat From d29634ede00075140942c09eabd643c40e881dc7 Mon Sep 17 00:00:00 2001 From: boopthesnoot Date: Thu, 4 Dec 2025 13:32:26 +0100 Subject: [PATCH 8/9] linting --- alphabase/spectral_library/annotate.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/alphabase/spectral_library/annotate.py b/alphabase/spectral_library/annotate.py index 911bf098..cccabc9e 100644 --- a/alphabase/spectral_library/annotate.py +++ b/alphabase/spectral_library/annotate.py @@ -418,7 +418,7 @@ def add_dense_lib( apply_precursor_metrics(outlib_flat) outlib_flat.precursor_df["spec_idx"] = original_spec_idx.iloc[ - outlib_flat.precursor_df["spec_idx"].values - ].values + outlib_flat.precursor_df["spec_idx"].to_numpy() + ].to_numpy() return outlib_flat From cee8362290f4dbc2d7fec9c312d1f43153425d40 Mon Sep 17 00:00:00 2001 From: boopthesnoot Date: Thu, 4 Dec 2025 13:54:05 +0100 Subject: [PATCH 9/9] fix: _mass_accuracy_metric on flat --- alphabase/spectral_library/metrics.py | 6 +-- tests/unit/spectral_library/test_annotate.py | 50 ++++++++++++++++---- 2 files changed, 44 insertions(+), 12 deletions(-) diff --git a/alphabase/spectral_library/metrics.py b/alphabase/spectral_library/metrics.py index 1934cb6c..67116048 100644 --- a/alphabase/spectral_library/metrics.py +++ b/alphabase/spectral_library/metrics.py @@ -196,11 +196,11 @@ def _mass_accuracy_metric(flatlib: SpecLibFlat) -> None: Important: --------- - This funmction operates in place and modifies the flatlib object + This function operates in place and modifies the flatlib object """ - start_indices = flatlib.precursor_df["frag_start_idx"].to_numpy() - stop_indices = flatlib.precursor_df["frag_stop_idx"].to_numpy() + 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 diff --git a/tests/unit/spectral_library/test_annotate.py b/tests/unit/spectral_library/test_annotate.py index ed0d32e6..1fb071bf 100644 --- a/tests/unit/spectral_library/test_annotate.py +++ b/tests/unit/spectral_library/test_annotate.py @@ -386,8 +386,8 @@ def test_mass_accuracy_metric_normal_case(): speclib_flat = SpecLibFlat() speclib_flat.precursor_df = pd.DataFrame( { - "frag_start_idx": [0, 3, 6], - "frag_stop_idx": [3, 6, 9], + "flat_frag_start_idx": [0, 3, 6], + "flat_frag_stop_idx": [3, 6, 9], "sequence": ["PEP", "PEPT", "PEPTI"], } ) @@ -406,7 +406,6 @@ def test_mass_accuracy_metric_normal_case(): ] # Third precursor } ) - print(speclib_flat.precursor_df) _mass_accuracy_metric(speclib_flat) expected = np.array([2.0, 5.0, 8.0]) # Medians for each group @@ -417,8 +416,8 @@ def test_mass_accuracy_metric_with_nans(): speclib_flat = SpecLibFlat() speclib_flat.precursor_df = pd.DataFrame( { - "frag_start_idx": [0, 3], - "frag_stop_idx": [3, 6], + "flat_frag_start_idx": [0, 3], + "flat_frag_stop_idx": [3, 6], "sequence": ["PEPTIDE", "SEQUENCE"], } ) @@ -446,8 +445,8 @@ def test_mass_accuracy_metric_with_infs(): speclib_flat = SpecLibFlat() speclib_flat.precursor_df = pd.DataFrame( { - "frag_start_idx": [0, 3], - "frag_stop_idx": [3, 6], + "flat_frag_start_idx": [0, 3], + "flat_frag_stop_idx": [3, 6], "sequence": ["PEPTIDE", "SEQUENCE"], } ) @@ -475,8 +474,8 @@ def test_mass_accuracy_metric_empty_fragments(): speclib_flat = SpecLibFlat() speclib_flat.precursor_df = pd.DataFrame( { - "frag_start_idx": [0, 0], - "frag_stop_idx": [0, 0], + "flat_frag_start_idx": [0, 0], + "flat_frag_stop_idx": [0, 0], "sequence": ["PEPTIDE", "SEQUENCE"], } ) @@ -487,3 +486,36 @@ def test_mass_accuracy_metric_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." + )