From bd03ca30fe95dfb0a7984011590669a98619eed6 Mon Sep 17 00:00:00 2001 From: GeorgWa Date: Mon, 31 Aug 2026 15:56:43 +0200 Subject: [PATCH 01/14] Characterize flatten_fragments before the memory refactor flatten_fragments had no unit test. Add black-box tests that pin its contract, so the refactoring that follows can be checked step by step. `_dense_library` builds dense frames with mz == 0 padding from both sources that produce it: unmodified precursors carry no modloss fragments, and some fragments fall outside the mz range. Intensities are distinct, so the top-k selection is unambiguous. `_expected_keep_mask` derives the expected mask per precursor, without flatten_fragments. The tests cover the kept fragments over five combinations of `keep_top_k_fragments` and `min_fragment_intensity`, every annotation column, the reannotated precursor pointers, `custom_df` and `custom_columns`, the path without intensities, an empty library, and a precursor long enough to need more than a uint8 fragment number. Co-Authored-By: Claude Opus 5 (1M context) --- tests/unit/peptide/test_fragment.py | 274 ++++++++++++++++++++++++++++ 1 file changed, 274 insertions(+) diff --git a/tests/unit/peptide/test_fragment.py b/tests/unit/peptide/test_fragment.py index cd59b3d3..39de19a0 100644 --- a/tests/unit/peptide/test_fragment.py +++ b/tests/unit/peptide/test_fragment.py @@ -1,9 +1,15 @@ from unittest.mock import patch +import numpy as np +import pandas as pd import pytest from alphabase.peptide.fragment import ( + FRAGMENT_TYPES, + PEAK_INTENSITY_DTYPE, + PEAK_MZ_DTYPE, filter_valid_charged_frag_types, + flatten_fragments, parse_charged_frag_type, ) @@ -51,3 +57,271 @@ def test_filter_valid_charged_frag_types(mock_parse): ) assert result == ["b_z1", "y_z2"] assert len(recorded_warnings) == 1 # Should have 2 warning messages + + +CHARGED_FRAG_TYPES = [ + "b_z1", + "b_z2", + "y_z1", + "y_z2", + "b_modloss_z1", + "y_modloss_z1", +] +ROWS_PER_PRECURSOR = 4 +N_PRECURSORS = 5 + + +def _dense_library(): + """Make dense fragment frames with padding, and the matching precursor pointers.""" + rng = np.random.default_rng(0) + n_rows = N_PRECURSORS * ROWS_PER_PRECURSOR + n_types = len(CHARGED_FRAG_TYPES) + + mz = (rng.random((n_rows, n_types)) * 1000 + 100).astype(PEAK_MZ_DTYPE) + # unmodified precursors have no modloss fragments, and some fragments fall + # outside the mz range. Both cases give mz == 0 padding. + mz[: 2 * ROWS_PER_PRECURSOR, 4:] = 0 + mz[3, 0] = 0 + # distinct intensities make the top-k selection unambiguous + intensity = rng.permutation(n_rows * n_types).reshape(n_rows, n_types) + intensity = (intensity / intensity.max()).astype(PEAK_INTENSITY_DTYPE) + + frag_start_idx = np.arange(N_PRECURSORS) * ROWS_PER_PRECURSOR + precursor_df = pd.DataFrame( + { + "frag_start_idx": frag_start_idx, + "frag_stop_idx": frag_start_idx + ROWS_PER_PRECURSOR, + } + ) + return ( + precursor_df, + pd.DataFrame(mz, columns=CHARGED_FRAG_TYPES), + pd.DataFrame(intensity, columns=CHARGED_FRAG_TYPES), + ) + + +def _expected_keep_mask( + precursor_df, mz_df, intensity_df, keep_top_k_fragments, min_fragment_intensity +): + """Give the keep mask over all dense slots. This does not use flatten_fragments.""" + n_types = mz_df.shape[1] + mz = mz_df.values.reshape(-1) + intensity = None if len(intensity_df) == 0 else intensity_df.values.reshape(-1) + + mask = np.zeros(mz.size, dtype=bool) + for start, stop in zip(precursor_df.frag_start_idx, precursor_df.frag_stop_idx): + block = slice(start * n_types, stop * n_types) + keep = mz[block] != 0 + if intensity is not None: + keep &= intensity[block] >= min_fragment_intensity + n_slots = (stop - start) * n_types + if keep_top_k_fragments < n_slots: + in_top_k = np.zeros(n_slots, dtype=bool) + in_top_k[np.argsort(intensity[block])[-keep_top_k_fragments:]] = True + keep &= in_top_k + mask[block] = keep + return mask + + +@pytest.mark.requires_numba +@pytest.mark.parametrize( + "keep_top_k_fragments, min_fragment_intensity", + [(1000, -1), (1000, 0.5), (8, -1), (5, 0.2), (1, -1)], +) +def test_flatten_fragments_retains_expected_fragments( + keep_top_k_fragments, min_fragment_intensity +): + """The flat library keeps only the slots that pass all filters.""" + precursor_df, mz_df, intensity_df = _dense_library() + expected_mask = _expected_keep_mask( + precursor_df, + mz_df, + intensity_df, + keep_top_k_fragments, + min_fragment_intensity, + ) + + _, frag_df = flatten_fragments( + precursor_df, + mz_df, + intensity_df, + min_fragment_intensity=min_fragment_intensity, + keep_top_k_fragments=keep_top_k_fragments, + ) + + assert len(frag_df) == expected_mask.sum() + np.testing.assert_array_equal( + frag_df["mz"].values, mz_df.values.reshape(-1)[expected_mask] + ) + np.testing.assert_array_equal( + frag_df["intensity"].values, intensity_df.values.reshape(-1)[expected_mask] + ) + + +@pytest.mark.requires_numba +@pytest.mark.parametrize("keep_top_k_fragments", [1000, 5]) +def test_flatten_fragments_annotates_retained_fragments(keep_top_k_fragments): + """Each annotation column describes the dense slot of its fragment.""" + precursor_df, mz_df, intensity_df = _dense_library() + n_types = mz_df.shape[1] + expected_mask = _expected_keep_mask( + precursor_df, mz_df, intensity_df, keep_top_k_fragments, -1 + ) + kept_slots = np.flatnonzero(expected_mask) + kept_columns = kept_slots % n_types + kept_rows = kept_slots // n_types + + _, frag_df = flatten_fragments( + precursor_df, mz_df, intensity_df, keep_top_k_fragments=keep_top_k_fragments + ) + + frag_types = [parse_charged_frag_type(col) for col in CHARGED_FRAG_TYPES] + np.testing.assert_array_equal( + frag_df["charge"].values, + np.array([charge for _, charge in frag_types])[kept_columns], + ) + np.testing.assert_array_equal( + frag_df["type"].values, + np.array([FRAGMENT_TYPES[name].series_id for name, _ in frag_types])[ + kept_columns + ], + ) + np.testing.assert_array_equal( + frag_df["loss_type"].values, + np.array([FRAGMENT_TYPES[name].loss_id for name, _ in frag_types])[ + kept_columns + ], + ) + + # position counts the fragment rows of a precursor. number counts the ion + # series in the direction of the fragment type. + expected_position = kept_rows % ROWS_PER_PRECURSOR + np.testing.assert_array_equal(frag_df["position"].values, expected_position) + directions = np.array( + [FRAGMENT_TYPES[name].direction_id for name, _ in frag_types] + )[kept_columns] + expected_number = np.where( + directions == 1, expected_position + 1, ROWS_PER_PRECURSOR - expected_position + ) + np.testing.assert_array_equal(frag_df["number"].values, expected_number) + + +@pytest.mark.requires_numba +@pytest.mark.parametrize("keep_top_k_fragments", [1000, 5]) +def test_flatten_fragments_reannotates_precursor_pointers(keep_top_k_fragments): + """The flat pointers of a precursor address only its own fragments.""" + precursor_df, mz_df, intensity_df = _dense_library() + n_types = mz_df.shape[1] + mz = mz_df.values.reshape(-1) + expected_mask = _expected_keep_mask( + precursor_df, mz_df, intensity_df, keep_top_k_fragments, -1 + ) + + precursor_df, frag_df = flatten_fragments( + precursor_df, mz_df, intensity_df, keep_top_k_fragments=keep_top_k_fragments + ) + + for row in precursor_df.itertuples(): + block = slice(row.frag_start_idx * n_types, row.frag_stop_idx * n_types) + np.testing.assert_array_equal( + frag_df["mz"].values[row.flat_frag_start_idx : row.flat_frag_stop_idx], + mz[block][expected_mask[block]], + ) + + # the pointers must cover the fragment dataframe with no gap and no overlap + assert precursor_df.flat_frag_start_idx.iloc[0] == 0 + assert precursor_df.flat_frag_stop_idx.iloc[-1] == len(frag_df) + np.testing.assert_array_equal( + precursor_df.flat_frag_start_idx.values[1:], + precursor_df.flat_frag_stop_idx.values[:-1], + ) + + +@pytest.mark.requires_numba +def test_flatten_fragments_filters_custom_df_columns(): + """flatten_fragments filters a custom_df column like the mz column.""" + precursor_df, mz_df, intensity_df = _dense_library() + cardinality_df = pd.DataFrame( + np.arange(mz_df.size, dtype=np.uint8).reshape(mz_df.shape), + columns=CHARGED_FRAG_TYPES, + ) + expected_mask = _expected_keep_mask(precursor_df, mz_df, intensity_df, 5, -1) + + _, frag_df = flatten_fragments( + precursor_df, + mz_df, + intensity_df, + keep_top_k_fragments=5, + custom_df={"cardinality": cardinality_df}, + ) + + assert "cardinality" in frag_df.columns + np.testing.assert_array_equal( + frag_df["cardinality"].values, + cardinality_df.values.reshape(-1)[expected_mask], + ) + + +@pytest.mark.requires_numba +def test_flatten_fragments_without_intensity(): + """Without intensities, flatten_fragments removes only the mz == 0 padding.""" + precursor_df, mz_df, _ = _dense_library() + expected_mask = _expected_keep_mask(precursor_df, mz_df, pd.DataFrame(), 1000, -1) + + _, frag_df = flatten_fragments(precursor_df, mz_df, pd.DataFrame()) + + assert "intensity" not in frag_df.columns + np.testing.assert_array_equal( + frag_df["mz"].values, mz_df.values.reshape(-1)[expected_mask] + ) + + +@pytest.mark.requires_numba +def test_flatten_fragments_selects_custom_columns(): + """flatten_fragments creates only the requested annotation columns.""" + precursor_df, mz_df, intensity_df = _dense_library() + + _, frag_df = flatten_fragments( + precursor_df, mz_df, intensity_df, custom_columns=["number", "charge"] + ) + + assert list(frag_df.columns) == ["mz", "intensity", "charge", "number"] + + +@pytest.mark.requires_numba +def test_flatten_fragments_empty_precursor_df(): + """An empty library gives an empty fragment dataframe.""" + _, mz_df, intensity_df = _dense_library() + + precursor_df, frag_df = flatten_fragments( + pd.DataFrame({"frag_start_idx": [], "frag_stop_idx": []}), mz_df, intensity_df + ) + + assert len(precursor_df) == 0 + assert len(frag_df) == 0 + + +@pytest.mark.requires_numba +def test_flatten_fragments_long_precursor(): + """A long precursor keeps fragment numbers above 255, and the columns stay uint32.""" + n_rows = 260 # close to the max_frag_per_peptide limit of _fill_in_indices + n_types = len(CHARGED_FRAG_TYPES) + rng = np.random.default_rng(0) + mz_df = pd.DataFrame( + (rng.random((n_rows, n_types)) * 1000 + 100).astype(PEAK_MZ_DTYPE), + columns=CHARGED_FRAG_TYPES, + ) + intensity_df = pd.DataFrame( + rng.random((n_rows, n_types)).astype(PEAK_INTENSITY_DTYPE), + columns=CHARGED_FRAG_TYPES, + ) + precursor_df = pd.DataFrame({"frag_start_idx": [0], "frag_stop_idx": [n_rows]}) + + _, frag_df = flatten_fragments( + precursor_df, mz_df, intensity_df, keep_top_k_fragments=n_rows * n_types + ) + + assert frag_df["position"].max() == n_rows - 1 + assert frag_df["number"].max() == n_rows + assert frag_df["position"].dtype == np.uint32 + assert frag_df["number"].dtype == np.uint32 From 9a6154163e77d1d822061ec6f092df5f3cace115 Mon Sep 17 00:00:00 2001 From: GeorgWa Date: Mon, 31 Aug 2026 15:57:24 +0200 Subject: [PATCH 02/14] Make the fragment parsing helpers private `parse_fragment`, `fill_in_indices` and `calculate_fragment_numbers` only serve `flatten_fragments`, and nothing in alphabase, alphadia or peptdeep calls them. A leading underscore states that, and frees their signatures for the refactoring that follows. Co-Authored-By: Claude Opus 5 (1M context) --- alphabase/peptide/fragment.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/alphabase/peptide/fragment.py b/alphabase/peptide/fragment.py index d0c3e0ab..ae9d0bb8 100644 --- a/alphabase/peptide/fragment.py +++ b/alphabase/peptide/fragment.py @@ -802,7 +802,7 @@ def mask_fragments_for_charge_greater_than_precursor_charge( @numba_njit(parallel=True) -def fill_in_indices( +def _fill_in_indices( frag_start_idxes: np.ndarray, frag_stop_idxes: np.ndarray, indices: np.ndarray, @@ -877,7 +877,7 @@ def fill_in_indices( @numba_vectorize( [nb_.uint32(nb_.int8, nb_.uint32, nb_.uint32, nb_.uint32)], target="parallel" ) -def calculate_fragment_numbers( +def _calculate_fragment_numbers( frag_direction: np.int8, frag_number: np.uint32, index: np.uint32, @@ -907,7 +907,7 @@ def calculate_fragment_numbers( return frag_number -def parse_fragment( +def _parse_fragment( frag_directions: np.ndarray, frag_start_idxes: np.ndarray, frag_stop_idxes: np.ndarray, @@ -954,7 +954,7 @@ def parse_fragment( ) # Fill in indices, max indices and excluded indices - fill_in_indices( + _fill_in_indices( frag_start_idxes, frag_stop_idxes, indices, @@ -966,7 +966,7 @@ def parse_fragment( ) # Calculate fragment numbers - frag_numbers = calculate_fragment_numbers( + frag_numbers = _calculate_fragment_numbers( frag_directions, frag_numbers, indices, max_indices ) return frag_numbers, indices, excluded_indices @@ -1081,7 +1081,7 @@ def flatten_fragments( np.tile(frag_directions, (len(fragment_mz_df), 1)), dtype=np.int8 ) - numbers, positions, excluded_indices = parse_fragment( + numbers, positions, excluded_indices = _parse_fragment( frag_directions, precursor_df.frag_start_idx.values, precursor_df.frag_stop_idx.values, From 45870d6671d8fff2bb4740d4c1b99400b77900de Mon Sep 17 00:00:00 2001 From: GeorgWa Date: Mon, 31 Aug 2026 15:59:19 +0200 Subject: [PATCH 03/14] Extract the charged fragment type annotation and the pointer reannotation Move two concerns out of the body of `flatten_fragments`: the per-column annotation of the charged fragment types, and the reannotation of the precursor pointers. The function now reads as its steps, and the memory refactoring that follows changes a named helper instead of a long function body. `_annotate_charged_frag_types` returns typed arrays rather than four parallel lists of Python ints. Tiling a typed array keeps its dtype, so the dense type/loss_type/charge/direction arrays no longer pass through an int64 intermediate of 8 bytes per dense slot. No functional change. Output stays bit-identical to main across 5400 generated configurations of library shape, padding, top-k, intensity threshold and custom_columns. Co-Authored-By: Claude Opus 5 (1M context) --- alphabase/peptide/fragment.py | 122 ++++++++++++++++++++++------------ 1 file changed, 81 insertions(+), 41 deletions(-) diff --git a/alphabase/peptide/fragment.py b/alphabase/peptide/fragment.py index ae9d0bb8..4a6f39e3 100644 --- a/alphabase/peptide/fragment.py +++ b/alphabase/peptide/fragment.py @@ -972,6 +972,76 @@ def _parse_fragment( return frag_numbers, indices, excluded_indices +def _annotate_charged_frag_types( + charged_frag_types: np.ndarray, +) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Give the series id, loss id, charge and direction of every charged fragment type. + + Parameters + ---------- + charged_frag_types : np.ndarray + names of the charged fragment types, in the order of the dense columns + + Returns + ------- + Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray] + series ids, loss ids, charges and directions, one value per charged fragment type + + """ + series_ids = [] + loss_ids = [] + charges = [] + directions = [] # 'abc': direction=1, 'xyz': direction=-1, otherwise 0 + + for charged_frag_type in charged_frag_types: + frag_type, charge = parse_charged_frag_type(charged_frag_type) + series_ids.append(FRAGMENT_TYPES[frag_type].series_id) + loss_ids.append(FRAGMENT_TYPES[frag_type].loss_id) + charges.append(charge) + directions.append(FRAGMENT_TYPES[frag_type].direction_id) + + return ( + np.array(series_ids, dtype=np.int8), + np.array(loss_ids, dtype=np.int16), + np.array(charges, dtype=np.int8), + np.array(directions, dtype=np.int8), + ) + + +def _reannotate_precursor_pointers( + precursor_df: pd.DataFrame, excluded: np.ndarray, n_fragment_types: int +) -> None: + """Add the flat fragment pointers to `precursor_df`, in place. + + Parameters + ---------- + precursor_df : pd.DataFrame + precursor dataframe with the `frag_start_idx` and `frag_stop_idx` columns + + excluded : np.ndarray + whether each dense slot is excluded from the flat fragment dataframe + + n_fragment_types : int + number of charged fragment types, that is the number of dense columns + + """ + precursor_df["flat_frag_start_idx"] = precursor_df.frag_start_idx + precursor_df["flat_frag_stop_idx"] = precursor_df.frag_stop_idx + precursor_df[["flat_frag_start_idx", "flat_frag_stop_idx"]] *= n_fragment_types + + # cumulative sum counts the number of fragments before the given fragment which were removed. + # This sum does not include the fragment at the index position and has therefore len N +1 + cum_sum_tresh = np.zeros(shape=len(excluded) + 1, dtype=np.int64) + cum_sum_tresh[1:] = np.cumsum(excluded) + + precursor_df["flat_frag_start_idx"] -= cum_sum_tresh[ + precursor_df.flat_frag_start_idx.values + ] + precursor_df["flat_frag_stop_idx"] -= cum_sum_tresh[ + precursor_df.flat_frag_stop_idx.values + ] + + def flatten_fragments( precursor_df: pd.DataFrame, fragment_mz_df: pd.DataFrame, @@ -1052,37 +1122,23 @@ def flatten_fragments( for col_name, df in custom_df.items(): frag_df[col_name] = df.values.reshape(-1) - frag_types = [] - frag_loss_types = [] - frag_charges = [] - frag_directions = [] # 'abc': direction=1, 'xyz': direction=-1, otherwise 0 - - for charged_frag_type in fragment_mz_df.columns.values: - frag_type, charge = parse_charged_frag_type(charged_frag_type) - frag_charges.append(charge) - frag_types.append(FRAGMENT_TYPES[frag_type].series_id) - frag_loss_types.append(FRAGMENT_TYPES[frag_type].loss_id) - frag_directions.append(FRAGMENT_TYPES[frag_type].direction_id) + series_ids, loss_ids, charges, directions = _annotate_charged_frag_types( + fragment_mz_df.columns.values + ) + # tiling the typed arrays keeps their dtype. A tiled list of Python ints gave + # a dense int64 array first, which cost 8 bytes per dense slot. if "type" in custom_columns: - frag_df["type"] = np.array( - np.tile(frag_types, len(fragment_mz_df)), dtype=np.int8 - ) + frag_df["type"] = np.tile(series_ids, len(fragment_mz_df)) if "loss_type" in custom_columns: - frag_df["loss_type"] = np.array( - np.tile(frag_loss_types, len(fragment_mz_df)), dtype=np.int16 - ) + frag_df["loss_type"] = np.tile(loss_ids, len(fragment_mz_df)) if "charge" in custom_columns: - frag_df["charge"] = np.array( - np.tile(frag_charges, len(fragment_mz_df)), dtype=np.int8 - ) + frag_df["charge"] = np.tile(charges, len(fragment_mz_df)) - frag_directions = np.array( - np.tile(frag_directions, (len(fragment_mz_df), 1)), dtype=np.int8 - ) + dense_directions = np.tile(directions, (len(fragment_mz_df), 1)) numbers, positions, excluded_indices = _parse_fragment( - frag_directions, + dense_directions, precursor_df.frag_start_idx.values, precursor_df.frag_stop_idx.values, keep_top_k_fragments, @@ -1096,12 +1152,6 @@ def flatten_fragments( if "position" in custom_columns: frag_df["position"] = positions.reshape(-1) - precursor_df["flat_frag_start_idx"] = precursor_df.frag_start_idx - precursor_df["flat_frag_stop_idx"] = precursor_df.frag_stop_idx - precursor_df[["flat_frag_start_idx", "flat_frag_stop_idx"]] *= len( - fragment_mz_df.columns - ) - if use_intensity: frag_df["intensity"][frag_df["mz"] == 0.0] = 0.0 @@ -1117,17 +1167,7 @@ def flatten_fragments( frag_df = frag_df[~excluded] frag_df = frag_df.reset_index(drop=True) - # cumulative sum counts the number of fragments before the given fragment which were removed. - # This sum does not include the fragment at the index position and has therefore len N +1 - cum_sum_tresh = np.zeros(shape=len(excluded) + 1, dtype=np.int64) - cum_sum_tresh[1:] = np.cumsum(excluded) - - precursor_df["flat_frag_start_idx"] -= cum_sum_tresh[ - precursor_df.flat_frag_start_idx.values - ] - precursor_df["flat_frag_stop_idx"] -= cum_sum_tresh[ - precursor_df.flat_frag_stop_idx.values - ] + _reannotate_precursor_pointers(precursor_df, excluded, len(fragment_mz_df.columns)) return precursor_df, frag_df From dbea5d6536b8d94065cbbf456af79167821e7972 Mon Sep 17 00:00:00 2001 From: GeorgWa Date: Mon, 31 Aug 2026 16:01:37 +0200 Subject: [PATCH 04/14] Build the flat fragment columns from the kept indices Building the flat fragment dataframe materialised every column at the full dense length of n_fragment_rows * n_fragment_types, assembled them into a dataframe and then copied the retained subset out of it. Reannotating the precursor pointers allocated an int64 cumulative sum over all dense slots on top of that. Only mz and intensity are needed at dense length, because they give the keep mask. `_select_dense_fragments` accumulates that mask in place and returns the ascending indices of the kept slots. `_annotate_kept_fragments` then reads the type, the loss type and the charge straight off those indices, because the index of a slot gives its dense column. The dataframe is built from the kept values alone, so the filtering copy is gone. `_reannotate_precursor_pointers` finds a new pointer with a binary search on the kept indices, which is equivalent to subtracting the cumulative number of removed fragments. Output stays bit-identical to main across 5400 generated configurations of library shape, padding, top-k, intensity threshold and custom_columns: same fragment count, column order, dtypes, values and precursor pointers. At 48M dense slots, peak RSS drops from 3.49 GB to 1.77 GB and the call is ~27% faster. Co-Authored-By: Claude Opus 5 (1M context) --- alphabase/peptide/fragment.py | 270 +++++++++++++++++++++++++--------- 1 file changed, 201 insertions(+), 69 deletions(-) diff --git a/alphabase/peptide/fragment.py b/alphabase/peptide/fragment.py index 4a6f39e3..77257039 100644 --- a/alphabase/peptide/fragment.py +++ b/alphabase/peptide/fragment.py @@ -1008,38 +1008,180 @@ def _annotate_charged_frag_types( ) +def _select_dense_fragments( + precursor_df: pd.DataFrame, + mz: np.ndarray, + intensity: Union[np.ndarray, None], + directions: np.ndarray, + n_fragment_rows: int, + n_fragment_types: int, + keep_top_k_fragments: int, + min_fragment_intensity: float, +) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """Select the dense fragment slots to keep, and number every dense slot. + + A slot is excluded if it holds padding (`mz == 0`), if its intensity is below + `min_fragment_intensity`, or if it is not one of the `keep_top_k_fragments` + most intense slots of its precursor. + + Parameters + ---------- + precursor_df : pd.DataFrame + precursor dataframe with the `frag_start_idx` and `frag_stop_idx` columns + + mz : np.ndarray + flattened fragment mz values, of length n_fragment_rows * n_fragment_types + + intensity : np.ndarray or None + flattened fragment intensities, or None to filter on mz only. + Padding slots are set to zero intensity. + + directions : np.ndarray + direction of every charged fragment type + + n_fragment_rows : int + number of rows of the dense fragment dataframes + + n_fragment_types : int + number of charged fragment types, that is the number of dense columns + + keep_top_k_fragments : int + number of most intense slots to keep per precursor + + min_fragment_intensity : float + minimum intensity of a slot to keep + + Returns + ------- + Tuple[np.ndarray, np.ndarray, np.ndarray] + ascending indices of the kept slots, and the fragment number and position + of every dense slot + + """ + dense_directions = np.tile(directions, (n_fragment_rows, 1)) + + numbers, positions, not_top_k = _parse_fragment( + dense_directions, + precursor_df.frag_start_idx.values, + precursor_df.frag_stop_idx.values, + keep_top_k_fragments, + intensity, + n_fragment_types, + ) + del dense_directions + + if intensity is None: + excluded = mz == 0 + else: + is_padding = mz == 0 + intensity[is_padding] = 0.0 + # in-place operations prevent more dense boolean arrays + excluded = intensity < min_fragment_intensity + excluded |= is_padding + excluded |= not_top_k + del is_padding + del not_top_k + + # The in-place inversion prevents one more dense array. The indices stay + # ascending, so the fragments keep their dense order. + np.logical_not(excluded, out=excluded) + return np.flatnonzero(excluded), numbers, positions + + +def _annotate_kept_fragments( + kept_indices: np.ndarray, + numbers: np.ndarray, + positions: np.ndarray, + series_ids: np.ndarray, + loss_ids: np.ndarray, + charges: np.ndarray, + n_fragment_types: int, + custom_columns: list, +) -> Dict[str, np.ndarray]: + """Give the requested annotation columns of the kept fragments. + + The index of a kept slot gives its charged fragment type, so the type, the + loss type and the charge need no dense array. + + Parameters + ---------- + kept_indices : np.ndarray + ascending indices of the kept dense slots + + numbers : np.ndarray + fragment number of every dense slot + + positions : np.ndarray + fragment position of every dense slot + + series_ids : np.ndarray + series id of every charged fragment type + + loss_ids : np.ndarray + loss id of every charged fragment type + + charges : np.ndarray + charge of every charged fragment type + + n_fragment_types : int + number of charged fragment types, that is the number of dense columns + + custom_columns : list + names of the annotation columns to create + + Returns + ------- + Dict[str, np.ndarray] + the requested columns, in the column order of the flat fragment dataframe + + """ + needs_frag_type = bool({"type", "loss_type", "charge"}.intersection(custom_columns)) + kept_frag_types = kept_indices % n_fragment_types if needs_frag_type else None + + columns = {} + if "type" in custom_columns: + columns["type"] = series_ids[kept_frag_types] + if "loss_type" in custom_columns: + columns["loss_type"] = loss_ids[kept_frag_types] + if "charge" in custom_columns: + columns["charge"] = charges[kept_frag_types] + if "number" in custom_columns: + columns["number"] = numbers.reshape(-1)[kept_indices] + if "position" in custom_columns: + columns["position"] = positions.reshape(-1)[kept_indices] + + return columns + + def _reannotate_precursor_pointers( - precursor_df: pd.DataFrame, excluded: np.ndarray, n_fragment_types: int + precursor_df: pd.DataFrame, kept_indices: np.ndarray, n_fragment_types: int ) -> None: """Add the flat fragment pointers to `precursor_df`, in place. + A new pointer is the count of kept fragments before the dense position. A + binary search on the ascending `kept_indices` gives this count and replaces a + dense cumulative sum. + Parameters ---------- precursor_df : pd.DataFrame precursor dataframe with the `frag_start_idx` and `frag_stop_idx` columns - excluded : np.ndarray - whether each dense slot is excluded from the flat fragment dataframe + kept_indices : np.ndarray + ascending indices of the kept dense slots n_fragment_types : int number of charged fragment types, that is the number of dense columns """ - precursor_df["flat_frag_start_idx"] = precursor_df.frag_start_idx - precursor_df["flat_frag_stop_idx"] = precursor_df.frag_stop_idx - precursor_df[["flat_frag_start_idx", "flat_frag_stop_idx"]] *= n_fragment_types - - # cumulative sum counts the number of fragments before the given fragment which were removed. - # This sum does not include the fragment at the index position and has therefore len N +1 - cum_sum_tresh = np.zeros(shape=len(excluded) + 1, dtype=np.int64) - cum_sum_tresh[1:] = np.cumsum(excluded) - - precursor_df["flat_frag_start_idx"] -= cum_sum_tresh[ - precursor_df.flat_frag_start_idx.values - ] - precursor_df["flat_frag_stop_idx"] -= cum_sum_tresh[ - precursor_df.flat_frag_stop_idx.values - ] + dense_start_idx = precursor_df.frag_start_idx.values.astype(np.int64) + dense_stop_idx = precursor_df.frag_stop_idx.values.astype(np.int64) + precursor_df["flat_frag_start_idx"] = np.searchsorted( + kept_indices, dense_start_idx * n_fragment_types + ) + precursor_df["flat_frag_stop_idx"] = np.searchsorted( + kept_indices, dense_stop_idx * n_fragment_types + ) def flatten_fragments( @@ -1107,69 +1249,59 @@ def flatten_fragments( """ if len(precursor_df) == 0: return precursor_df, pd.DataFrame() - # new dataframes for fragments and precursors are created - frag_df = {} - frag_df["mz"] = fragment_mz_df.values.reshape(-1) - if len(fragment_intensity_df) > 0: - frag_df["intensity"] = fragment_intensity_df.values.astype( - PEAK_INTENSITY_DTYPE - ).reshape(-1) - use_intensity = True - else: - use_intensity = False - # add additional columns to the fragment dataframe - # each column in the flat fragment dataframe is a whole pandas dataframe in the dense representation - for col_name, df in custom_df.items(): - frag_df[col_name] = df.values.reshape(-1) + + n_fragment_types = len(fragment_mz_df.columns) + + # Only mz and intensity need the full dense length, because they give the + # keep mask. Dense arrays for the other columns and a copy from a filtered + # dataframe increase the peak memory. For top-k libraries, mz == 0 padding + # fills most dense slots. + mz = fragment_mz_df.values.reshape(-1) + use_intensity = len(fragment_intensity_df) > 0 + intensity = ( + fragment_intensity_df.values.astype(PEAK_INTENSITY_DTYPE).reshape(-1) + if use_intensity + else None + ) series_ids, loss_ids, charges, directions = _annotate_charged_frag_types( fragment_mz_df.columns.values ) - # tiling the typed arrays keeps their dtype. A tiled list of Python ints gave - # a dense int64 array first, which cost 8 bytes per dense slot. - if "type" in custom_columns: - frag_df["type"] = np.tile(series_ids, len(fragment_mz_df)) - if "loss_type" in custom_columns: - frag_df["loss_type"] = np.tile(loss_ids, len(fragment_mz_df)) - if "charge" in custom_columns: - frag_df["charge"] = np.tile(charges, len(fragment_mz_df)) - - dense_directions = np.tile(directions, (len(fragment_mz_df), 1)) - - numbers, positions, excluded_indices = _parse_fragment( - dense_directions, - precursor_df.frag_start_idx.values, - precursor_df.frag_stop_idx.values, + kept_indices, numbers, positions = _select_dense_fragments( + precursor_df, + mz, + intensity, + directions, + len(fragment_mz_df), + n_fragment_types, keep_top_k_fragments, - frag_df["intensity"] if use_intensity else None, - len(fragment_mz_df.columns), + min_fragment_intensity, ) - if "number" in custom_columns: - frag_df["number"] = numbers.reshape(-1) - - if "position" in custom_columns: - frag_df["position"] = positions.reshape(-1) - + frag_df = {} + frag_df["mz"] = mz[kept_indices] if use_intensity: - frag_df["intensity"][frag_df["mz"] == 0.0] = 0.0 - - excluded = ( - frag_df["mz"] == 0 - if not use_intensity - else (frag_df["intensity"] < min_fragment_intensity) - | (frag_df["mz"] == 0) - | (excluded_indices) + frag_df["intensity"] = intensity[kept_indices] + # each custom_df value is a dense dataframe + for col_name, df in custom_df.items(): + frag_df[col_name] = df.values.reshape(-1)[kept_indices] + frag_df.update( + _annotate_kept_fragments( + kept_indices, + numbers, + positions, + series_ids, + loss_ids, + charges, + n_fragment_types, + custom_columns, + ) ) - frag_df = pd.DataFrame(frag_df) - frag_df = frag_df[~excluded] - frag_df = frag_df.reset_index(drop=True) - - _reannotate_precursor_pointers(precursor_df, excluded, len(fragment_mz_df.columns)) + _reannotate_precursor_pointers(precursor_df, kept_indices, n_fragment_types) - return precursor_df, frag_df + return precursor_df, pd.DataFrame(frag_df, copy=False) @numba_njit From d07f56105c1b98ef8f37306c7e81ed080a4150c0 Mon Sep 17 00:00:00 2001 From: GeorgWa Date: Tue, 15 Sep 2026 09:57:30 +0200 Subject: [PATCH 05/14] Pin flatten_fragments against a hand-written library _expected_keep_mask reimplemented the filtering logic, so it could agree with a broken flatten_fragments. Add a minimal_library fixture of two precursors over a forward and a reverse fragment type, with the kept slots written out per filter setting, and assert both the helper and flatten_fragments against it. Derive the long precursor row count from the uint8 boundary it tests instead of picking a number near the fill_in_indices limit, and guard the limit explicitly so exceeding it fails on an assertion rather than on a broadcast error inside numba. Co-Authored-By: Claude Opus 5 (1M context) --- tests/unit/peptide/test_fragment.py | 146 +++++++++++++++++++++++++++- 1 file changed, 145 insertions(+), 1 deletion(-) diff --git a/tests/unit/peptide/test_fragment.py b/tests/unit/peptide/test_fragment.py index 39de19a0..1c484c17 100644 --- a/tests/unit/peptide/test_fragment.py +++ b/tests/unit/peptide/test_fragment.py @@ -69,6 +69,8 @@ def test_filter_valid_charged_frag_types(mock_parse): ] ROWS_PER_PRECURSOR = 4 N_PRECURSORS = 5 +# mirrors the max_frag_per_peptide default of fill_in_indices +MAX_FRAG_PER_PEPTIDE = 300 def _dense_library(): @@ -123,6 +125,144 @@ def _expected_keep_mask( return mask +# 2 precursors x 2 fragment rows x 2 types, b forward and y reverse +MINIMAL_FRAG_TYPES = ["b_z1", "y_z1"] +MINIMAL_N_PRECURSORS = 2 +MINIMAL_ROWS_PER_PRECURSOR = 2 +MINIMAL_MZ = [ + # b_z1, y_z1 + [100.0, 200.0], # precursor 0, row 0 + [0.0, 300.0], # precursor 0, row 1, b_z1 padding + [400.0, 500.0], # precursor 1, row 0 + [600.0, 0.0], # precursor 1, row 1, y_z1 padding +] +MINIMAL_INTENSITY = [ + [0.10, 0.80], + [0.00, 0.50], + [0.90, 0.35], + [0.40, 0.00], +] +# keep_top_k_fragments, min_fragment_intensity, kept dense slots. +# Slots are row-major, per precursor [row 0 b, row 0 y, row 1 b, row 1 y]. +MINIMAL_FILTER_CASES = [ + pytest.param( + 1000, + -1, + [True, True, False, True, True, True, True, False], + id="padding_only", + ), + pytest.param( + 2, + -1, + [False, True, False, True, True, False, True, False], + id="top_2_per_precursor", + ), + pytest.param( + 1000, + 0.3, + [False, True, False, True, True, True, True, False], + id="min_intensity", + ), +] + + +@pytest.fixture +def minimal_library(): + """Give MINIMAL_MZ and MINIMAL_INTENSITY with the matching precursor pointers.""" + frag_start_idx = np.arange(MINIMAL_N_PRECURSORS) * MINIMAL_ROWS_PER_PRECURSOR + precursor_df = pd.DataFrame( + { + "frag_start_idx": frag_start_idx, + "frag_stop_idx": frag_start_idx + MINIMAL_ROWS_PER_PRECURSOR, + } + ) + return ( + precursor_df, + pd.DataFrame( + np.array(MINIMAL_MZ, dtype=PEAK_MZ_DTYPE), columns=MINIMAL_FRAG_TYPES + ), + pd.DataFrame( + np.array(MINIMAL_INTENSITY, dtype=PEAK_INTENSITY_DTYPE), + columns=MINIMAL_FRAG_TYPES, + ), + ) + + +@pytest.mark.parametrize( + "keep_top_k_fragments, min_fragment_intensity, expected_mask", MINIMAL_FILTER_CASES +) +def test_expected_keep_mask_matches_hand_written_mask( + minimal_library, keep_top_k_fragments, min_fragment_intensity, expected_mask +): + """_expected_keep_mask reproduces the masks of MINIMAL_FILTER_CASES.""" + # Given + precursor_df, mz_df, intensity_df = minimal_library + + # When + mask = _expected_keep_mask( + precursor_df, mz_df, intensity_df, keep_top_k_fragments, min_fragment_intensity + ) + + # Then + np.testing.assert_array_equal(mask, expected_mask) + + +@pytest.mark.requires_numba +@pytest.mark.parametrize( + "keep_top_k_fragments, min_fragment_intensity, expected_mask", MINIMAL_FILTER_CASES +) +def test_flatten_fragments_on_minimal_library( + minimal_library, keep_top_k_fragments, min_fragment_intensity, expected_mask +): + """flatten_fragments keeps exactly the slots of MINIMAL_FILTER_CASES.""" + # Given + precursor_df, mz_df, intensity_df = minimal_library + expected_mask = np.array(expected_mask) + + # When + _, frag_df = flatten_fragments( + precursor_df, + mz_df, + intensity_df, + min_fragment_intensity=min_fragment_intensity, + keep_top_k_fragments=keep_top_k_fragments, + ) + + # Then + np.testing.assert_array_equal( + frag_df["mz"].values, + np.array(MINIMAL_MZ, dtype=PEAK_MZ_DTYPE).reshape(-1)[expected_mask], + ) + np.testing.assert_array_equal( + frag_df["intensity"].values, + np.array(MINIMAL_INTENSITY, dtype=PEAK_INTENSITY_DTYPE).reshape(-1)[ + expected_mask + ], + ) + + +@pytest.mark.requires_numba +def test_flatten_fragments_annotates_minimal_library(minimal_library): + """flatten_fragments annotates each kept slot with its dense position.""" + # Given + precursor_df, mz_df, intensity_df = minimal_library + + # When + _, frag_df = flatten_fragments(precursor_df, mz_df, intensity_df) + + # Then + # kept slots: p0 r0 b, p0 r0 y, p0 r1 y, p1 r0 b, p1 r0 y, p1 r1 b + np.testing.assert_array_equal( + frag_df["mz"].values, [100.0, 200.0, 300.0, 400.0, 500.0, 600.0] + ) + np.testing.assert_array_equal(frag_df["position"].values, [0, 0, 1, 0, 0, 1]) + # b: number = position + 1, y: number = MINIMAL_ROWS_PER_PRECURSOR - position + np.testing.assert_array_equal(frag_df["number"].values, [1, 2, 1, 1, 2, 2]) + np.testing.assert_array_equal(frag_df["type"].values, [98, 121, 121, 98, 121, 98]) + np.testing.assert_array_equal(frag_df["charge"].values, [1, 1, 1, 1, 1, 1]) + np.testing.assert_array_equal(frag_df["loss_type"].values, [0, 0, 0, 0, 0, 0]) + + @pytest.mark.requires_numba @pytest.mark.parametrize( "keep_top_k_fragments, min_fragment_intensity", @@ -304,7 +444,11 @@ def test_flatten_fragments_empty_precursor_df(): @pytest.mark.requires_numba def test_flatten_fragments_long_precursor(): """A long precursor keeps fragment numbers above 255, and the columns stay uint32.""" - n_rows = 260 # close to the max_frag_per_peptide limit of _fill_in_indices + # position (n_rows - 1) and number (n_rows) must both exceed the uint8 range + n_rows = np.iinfo(np.uint8).max + 2 + assert ( + n_rows <= MAX_FRAG_PER_PEPTIDE + ), "fill_in_indices cannot index a precursor this long" n_types = len(CHARGED_FRAG_TYPES) rng = np.random.default_rng(0) mz_df = pd.DataFrame( From 8be793d7823d33fa5f1f9adc7971c38a455dde96 Mon Sep 17 00:00:00 2001 From: GeorgWa Date: Tue, 15 Sep 2026 10:02:28 +0200 Subject: [PATCH 06/14] Document the fragment numbering convention "fragment number" read as ambiguous between a count and an enumeration. Spell out that it is the MS series numbering, and contrast it with the fragment position that all series share. Co-Authored-By: Claude Opus 5 (1M context) --- alphabase/peptide/fragment.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/alphabase/peptide/fragment.py b/alphabase/peptide/fragment.py index ae9d0bb8..26f1d6e7 100644 --- a/alphabase/peptide/fragment.py +++ b/alphabase/peptide/fragment.py @@ -886,6 +886,10 @@ def _calculate_fragment_numbers( """ Calculate fragment numbers for each fragment based on the fragment direction. + The fragment number is the MS series numbering, i.e. the 1-based index within the + ion series (the n of b2 or y7). Forward series are numbered from the N-terminus, + reverse series from the C-terminus. + Parameters ---------- frag_direction : np.int8 @@ -919,6 +923,10 @@ def _parse_fragment( Parse fragments to get fragment numbers, fragment positions and not top k excluded indices in one hit faster than doing each operation individually, and makes the most of the operations that are done in parallel. + The fragment number is the MS series numbering, i.e. the 1-based index within the ion + series (the n of b2 or y7). The fragment position is the 0-based index of the fragment + in the sequence, from left to right, and is shared by all series. + Parameters ---------- frag_directions : np.ndarray From 8b43cf677539b228a8877c1ad526f1b727211ddb Mon Sep 17 00:00:00 2001 From: GeorgWa Date: Tue, 15 Sep 2026 10:16:43 +0200 Subject: [PATCH 07/14] Assert against a hand-written library The library was generated from a seeded rng and the expectations came from _expected_keep_mask, which reimplemented the filtering logic and so could agree with a broken flatten_fragments. Replace both with one library fixture over hand-written mz and intensity grids, and hand-written keep masks per filter setting. Two fragment types carry the whole annotation surface: y_z1 is reverse without loss and b_modloss_z1 is forward with loss, so direction, series and loss all vary over two columns. Co-Authored-By: Claude Opus 5 (1M context) --- tests/unit/peptide/test_fragment.py | 361 +++++++++------------------- 1 file changed, 110 insertions(+), 251 deletions(-) diff --git a/tests/unit/peptide/test_fragment.py b/tests/unit/peptide/test_fragment.py index 1c484c17..0332ab3d 100644 --- a/tests/unit/peptide/test_fragment.py +++ b/tests/unit/peptide/test_fragment.py @@ -5,7 +5,6 @@ import pytest from alphabase.peptide.fragment import ( - FRAGMENT_TYPES, PEAK_INTENSITY_DTYPE, PEAK_MZ_DTYPE, filter_valid_charged_frag_types, @@ -59,165 +58,98 @@ def test_filter_valid_charged_frag_types(mock_parse): assert len(recorded_warnings) == 1 # Should have 2 warning messages -CHARGED_FRAG_TYPES = [ - "b_z1", - "b_z2", - "y_z1", - "y_z2", - "b_modloss_z1", - "y_modloss_z1", -] -ROWS_PER_PRECURSOR = 4 -N_PRECURSORS = 5 +# one reverse series without loss and one forward series with loss, so that direction, +# series and loss all vary over two columns +CHARGED_FRAG_TYPES = ["y_z1", "b_modloss_z1"] +N_PRECURSORS = 2 +ROWS_PER_PRECURSOR = 2 # mirrors the max_frag_per_peptide default of fill_in_indices MAX_FRAG_PER_PEPTIDE = 300 - - -def _dense_library(): - """Make dense fragment frames with padding, and the matching precursor pointers.""" - rng = np.random.default_rng(0) - n_rows = N_PRECURSORS * ROWS_PER_PRECURSOR - n_types = len(CHARGED_FRAG_TYPES) - - mz = (rng.random((n_rows, n_types)) * 1000 + 100).astype(PEAK_MZ_DTYPE) - # unmodified precursors have no modloss fragments, and some fragments fall - # outside the mz range. Both cases give mz == 0 padding. - mz[: 2 * ROWS_PER_PRECURSOR, 4:] = 0 - mz[3, 0] = 0 - # distinct intensities make the top-k selection unambiguous - intensity = rng.permutation(n_rows * n_types).reshape(n_rows, n_types) - intensity = (intensity / intensity.max()).astype(PEAK_INTENSITY_DTYPE) - - frag_start_idx = np.arange(N_PRECURSORS) * ROWS_PER_PRECURSOR - precursor_df = pd.DataFrame( - { - "frag_start_idx": frag_start_idx, - "frag_stop_idx": frag_start_idx + ROWS_PER_PRECURSOR, - } - ) - return ( - precursor_df, - pd.DataFrame(mz, columns=CHARGED_FRAG_TYPES), - pd.DataFrame(intensity, columns=CHARGED_FRAG_TYPES), - ) - - -def _expected_keep_mask( - precursor_df, mz_df, intensity_df, keep_top_k_fragments, min_fragment_intensity -): - """Give the keep mask over all dense slots. This does not use flatten_fragments.""" - n_types = mz_df.shape[1] - mz = mz_df.values.reshape(-1) - intensity = None if len(intensity_df) == 0 else intensity_df.values.reshape(-1) - - mask = np.zeros(mz.size, dtype=bool) - for start, stop in zip(precursor_df.frag_start_idx, precursor_df.frag_stop_idx): - block = slice(start * n_types, stop * n_types) - keep = mz[block] != 0 - if intensity is not None: - keep &= intensity[block] >= min_fragment_intensity - n_slots = (stop - start) * n_types - if keep_top_k_fragments < n_slots: - in_top_k = np.zeros(n_slots, dtype=bool) - in_top_k[np.argsort(intensity[block])[-keep_top_k_fragments:]] = True - keep &= in_top_k - mask[block] = keep - return mask - - -# 2 precursors x 2 fragment rows x 2 types, b forward and y reverse -MINIMAL_FRAG_TYPES = ["b_z1", "y_z1"] -MINIMAL_N_PRECURSORS = 2 -MINIMAL_ROWS_PER_PRECURSOR = 2 -MINIMAL_MZ = [ - # b_z1, y_z1 - [100.0, 200.0], # precursor 0, row 0 - [0.0, 300.0], # precursor 0, row 1, b_z1 padding - [400.0, 500.0], # precursor 1, row 0 - [600.0, 0.0], # precursor 1, row 1, y_z1 padding +# each mz encodes its own slot as 100 + row * 10 + column +MZ = [ + # y_z1, b_modloss_z1 + [100.0, 0.0], # precursor 0, row 0, unmodified: no modloss fragment + [110.0, 0.0], # precursor 0, row 1, unmodified: no modloss fragment + [120.0, 121.0], # precursor 1, row 0 + [0.0, 131.0], # precursor 1, row 1, y_z1 outside the mz range +] +# distinct over the whole library, so the top k selection is unambiguous +INTENSITY = [ + [0.11, 0.01], + [0.82, 0.02], + [0.93, 0.48], + [0.03, 0.61], ] -MINIMAL_INTENSITY = [ - [0.10, 0.80], - [0.00, 0.50], - [0.90, 0.35], - [0.40, 0.00], +# slots kept when only the mz == 0 padding is dropped +KEEP_UNFILTERED = [ + [1, 0], + [1, 0], + [1, 1], + [0, 1], ] -# keep_top_k_fragments, min_fragment_intensity, kept dense slots. -# Slots are row-major, per precursor [row 0 b, row 0 y, row 1 b, row 1 y]. -MINIMAL_FILTER_CASES = [ +# keep_top_k_fragments, min_fragment_intensity, kept slots per dense row +FILTER_CASES = [ + pytest.param(1000, -1, KEEP_UNFILTERED, id="padding_only"), pytest.param( 1000, - -1, - [True, True, False, True, True, True, True, False], - id="padding_only", + 0.3, + [ + [0, 0], + [1, 0], + [1, 1], + [0, 1], + ], + id="min_intensity", ), pytest.param( - 2, + 1, -1, - [False, True, False, True, True, False, True, False], - id="top_2_per_precursor", - ), - pytest.param( - 1000, - 0.3, - [False, True, False, True, True, True, True, False], - id="min_intensity", + [ + [0, 0], + [1, 0], + [1, 0], + [0, 0], + ], + id="top_1_per_precursor", ), ] @pytest.fixture -def minimal_library(): - """Give MINIMAL_MZ and MINIMAL_INTENSITY with the matching precursor pointers.""" - frag_start_idx = np.arange(MINIMAL_N_PRECURSORS) * MINIMAL_ROWS_PER_PRECURSOR +def library(): + """Give MZ and INTENSITY as fragment frames with the matching precursor pointers.""" + frag_start_idx = np.arange(N_PRECURSORS) * ROWS_PER_PRECURSOR precursor_df = pd.DataFrame( { "frag_start_idx": frag_start_idx, - "frag_stop_idx": frag_start_idx + MINIMAL_ROWS_PER_PRECURSOR, + "frag_stop_idx": frag_start_idx + ROWS_PER_PRECURSOR, } ) return ( precursor_df, + pd.DataFrame(np.array(MZ, dtype=PEAK_MZ_DTYPE), columns=CHARGED_FRAG_TYPES), pd.DataFrame( - np.array(MINIMAL_MZ, dtype=PEAK_MZ_DTYPE), columns=MINIMAL_FRAG_TYPES - ), - pd.DataFrame( - np.array(MINIMAL_INTENSITY, dtype=PEAK_INTENSITY_DTYPE), - columns=MINIMAL_FRAG_TYPES, + np.array(INTENSITY, dtype=PEAK_INTENSITY_DTYPE), columns=CHARGED_FRAG_TYPES ), ) -@pytest.mark.parametrize( - "keep_top_k_fragments, min_fragment_intensity, expected_mask", MINIMAL_FILTER_CASES -) -def test_expected_keep_mask_matches_hand_written_mask( - minimal_library, keep_top_k_fragments, min_fragment_intensity, expected_mask -): - """_expected_keep_mask reproduces the masks of MINIMAL_FILTER_CASES.""" - # Given - precursor_df, mz_df, intensity_df = minimal_library - - # When - mask = _expected_keep_mask( - precursor_df, mz_df, intensity_df, keep_top_k_fragments, min_fragment_intensity - ) - - # Then - np.testing.assert_array_equal(mask, expected_mask) +def _flat(values, dtype): + """Flatten a dense grid to the slot order of the flat fragment dataframe.""" + return np.array(values, dtype=dtype).reshape(-1) @pytest.mark.requires_numba @pytest.mark.parametrize( - "keep_top_k_fragments, min_fragment_intensity, expected_mask", MINIMAL_FILTER_CASES + "keep_top_k_fragments, min_fragment_intensity, expected_keep", FILTER_CASES ) -def test_flatten_fragments_on_minimal_library( - minimal_library, keep_top_k_fragments, min_fragment_intensity, expected_mask +def test_flatten_fragments_retains_expected_fragments( + library, keep_top_k_fragments, min_fragment_intensity, expected_keep ): - """flatten_fragments keeps exactly the slots of MINIMAL_FILTER_CASES.""" + """The flat library keeps only the slots listed in FILTER_CASES.""" # Given - precursor_df, mz_df, intensity_df = minimal_library - expected_mask = np.array(expected_mask) + precursor_df, mz_df, intensity_df = library + keep = _flat(expected_keep, bool) # When _, frag_df = flatten_fragments( @@ -229,59 +161,51 @@ def test_flatten_fragments_on_minimal_library( ) # Then + assert len(frag_df) == keep.sum() + np.testing.assert_array_equal(frag_df["mz"].values, _flat(MZ, PEAK_MZ_DTYPE)[keep]) np.testing.assert_array_equal( - frag_df["mz"].values, - np.array(MINIMAL_MZ, dtype=PEAK_MZ_DTYPE).reshape(-1)[expected_mask], - ) - np.testing.assert_array_equal( - frag_df["intensity"].values, - np.array(MINIMAL_INTENSITY, dtype=PEAK_INTENSITY_DTYPE).reshape(-1)[ - expected_mask - ], + frag_df["intensity"].values, _flat(INTENSITY, PEAK_INTENSITY_DTYPE)[keep] ) @pytest.mark.requires_numba -def test_flatten_fragments_annotates_minimal_library(minimal_library): - """flatten_fragments annotates each kept slot with its dense position.""" +def test_flatten_fragments_annotates_retained_fragments(library): + """Each annotation column describes the dense slot of its fragment.""" # Given - precursor_df, mz_df, intensity_df = minimal_library + precursor_df, mz_df, intensity_df = library # When _, frag_df = flatten_fragments(precursor_df, mz_df, intensity_df) # Then - # kept slots: p0 r0 b, p0 r0 y, p0 r1 y, p1 r0 b, p1 r0 y, p1 r1 b np.testing.assert_array_equal( - frag_df["mz"].values, [100.0, 200.0, 300.0, 400.0, 500.0, 600.0] + frag_df["mz"].values, _flat(MZ, PEAK_MZ_DTYPE)[_flat(KEEP_UNFILTERED, bool)] ) - np.testing.assert_array_equal(frag_df["position"].values, [0, 0, 1, 0, 0, 1]) - # b: number = position + 1, y: number = MINIMAL_ROWS_PER_PRECURSOR - position - np.testing.assert_array_equal(frag_df["number"].values, [1, 2, 1, 1, 2, 2]) - np.testing.assert_array_equal(frag_df["type"].values, [98, 121, 121, 98, 121, 98]) - np.testing.assert_array_equal(frag_df["charge"].values, [1, 1, 1, 1, 1, 1]) - np.testing.assert_array_equal(frag_df["loss_type"].values, [0, 0, 0, 0, 0, 0]) + # position is the dense row of the fragment within its precursor + np.testing.assert_array_equal(frag_df["position"].values, [0, 1, 0, 0, 1]) + # b_modloss counts forward, y counts backward over ROWS_PER_PRECURSOR rows + np.testing.assert_array_equal(frag_df["number"].values, [2, 1, 2, 1, 2]) + np.testing.assert_array_equal(frag_df["type"].values, [121, 121, 121, 98, 98]) + np.testing.assert_array_equal(frag_df["loss_type"].values, [0, 0, 0, 98, 98]) + np.testing.assert_array_equal(frag_df["charge"].values, [1, 1, 1, 1, 1]) @pytest.mark.requires_numba @pytest.mark.parametrize( - "keep_top_k_fragments, min_fragment_intensity", - [(1000, -1), (1000, 0.5), (8, -1), (5, 0.2), (1, -1)], + "keep_top_k_fragments, min_fragment_intensity, expected_keep", FILTER_CASES ) -def test_flatten_fragments_retains_expected_fragments( - keep_top_k_fragments, min_fragment_intensity +def test_flatten_fragments_reannotates_precursor_pointers( + library, keep_top_k_fragments, min_fragment_intensity, expected_keep ): - """The flat library keeps only the slots that pass all filters.""" - precursor_df, mz_df, intensity_df = _dense_library() - expected_mask = _expected_keep_mask( - precursor_df, - mz_df, - intensity_df, - keep_top_k_fragments, - min_fragment_intensity, - ) + """The flat pointers of a precursor address only its own fragments.""" + # Given + precursor_df, mz_df, intensity_df = library + n_types = len(CHARGED_FRAG_TYPES) + mz = _flat(MZ, PEAK_MZ_DTYPE) + keep = _flat(expected_keep, bool) - _, frag_df = flatten_fragments( + # When + precursor_df, frag_df = flatten_fragments( precursor_df, mz_df, intensity_df, @@ -289,83 +213,12 @@ def test_flatten_fragments_retains_expected_fragments( keep_top_k_fragments=keep_top_k_fragments, ) - assert len(frag_df) == expected_mask.sum() - np.testing.assert_array_equal( - frag_df["mz"].values, mz_df.values.reshape(-1)[expected_mask] - ) - np.testing.assert_array_equal( - frag_df["intensity"].values, intensity_df.values.reshape(-1)[expected_mask] - ) - - -@pytest.mark.requires_numba -@pytest.mark.parametrize("keep_top_k_fragments", [1000, 5]) -def test_flatten_fragments_annotates_retained_fragments(keep_top_k_fragments): - """Each annotation column describes the dense slot of its fragment.""" - precursor_df, mz_df, intensity_df = _dense_library() - n_types = mz_df.shape[1] - expected_mask = _expected_keep_mask( - precursor_df, mz_df, intensity_df, keep_top_k_fragments, -1 - ) - kept_slots = np.flatnonzero(expected_mask) - kept_columns = kept_slots % n_types - kept_rows = kept_slots // n_types - - _, frag_df = flatten_fragments( - precursor_df, mz_df, intensity_df, keep_top_k_fragments=keep_top_k_fragments - ) - - frag_types = [parse_charged_frag_type(col) for col in CHARGED_FRAG_TYPES] - np.testing.assert_array_equal( - frag_df["charge"].values, - np.array([charge for _, charge in frag_types])[kept_columns], - ) - np.testing.assert_array_equal( - frag_df["type"].values, - np.array([FRAGMENT_TYPES[name].series_id for name, _ in frag_types])[ - kept_columns - ], - ) - np.testing.assert_array_equal( - frag_df["loss_type"].values, - np.array([FRAGMENT_TYPES[name].loss_id for name, _ in frag_types])[ - kept_columns - ], - ) - - # position counts the fragment rows of a precursor. number counts the ion - # series in the direction of the fragment type. - expected_position = kept_rows % ROWS_PER_PRECURSOR - np.testing.assert_array_equal(frag_df["position"].values, expected_position) - directions = np.array( - [FRAGMENT_TYPES[name].direction_id for name, _ in frag_types] - )[kept_columns] - expected_number = np.where( - directions == 1, expected_position + 1, ROWS_PER_PRECURSOR - expected_position - ) - np.testing.assert_array_equal(frag_df["number"].values, expected_number) - - -@pytest.mark.requires_numba -@pytest.mark.parametrize("keep_top_k_fragments", [1000, 5]) -def test_flatten_fragments_reannotates_precursor_pointers(keep_top_k_fragments): - """The flat pointers of a precursor address only its own fragments.""" - precursor_df, mz_df, intensity_df = _dense_library() - n_types = mz_df.shape[1] - mz = mz_df.values.reshape(-1) - expected_mask = _expected_keep_mask( - precursor_df, mz_df, intensity_df, keep_top_k_fragments, -1 - ) - - precursor_df, frag_df = flatten_fragments( - precursor_df, mz_df, intensity_df, keep_top_k_fragments=keep_top_k_fragments - ) - + # Then for row in precursor_df.itertuples(): block = slice(row.frag_start_idx * n_types, row.frag_stop_idx * n_types) np.testing.assert_array_equal( frag_df["mz"].values[row.flat_frag_start_idx : row.flat_frag_stop_idx], - mz[block][expected_mask[block]], + mz[block][keep[block]], ) # the pointers must cover the fragment dataframe with no gap and no overlap @@ -378,65 +231,71 @@ def test_flatten_fragments_reannotates_precursor_pointers(keep_top_k_fragments): @pytest.mark.requires_numba -def test_flatten_fragments_filters_custom_df_columns(): +def test_flatten_fragments_filters_custom_df_columns(library): """flatten_fragments filters a custom_df column like the mz column.""" - precursor_df, mz_df, intensity_df = _dense_library() - cardinality_df = pd.DataFrame( - np.arange(mz_df.size, dtype=np.uint8).reshape(mz_df.shape), - columns=CHARGED_FRAG_TYPES, - ) - expected_mask = _expected_keep_mask(precursor_df, mz_df, intensity_df, 5, -1) + # Given + precursor_df, mz_df, intensity_df = library + cardinality = np.arange(mz_df.size, dtype=np.uint8).reshape(mz_df.shape) + cardinality_df = pd.DataFrame(cardinality, columns=CHARGED_FRAG_TYPES) + # When _, frag_df = flatten_fragments( precursor_df, mz_df, intensity_df, - keep_top_k_fragments=5, custom_df={"cardinality": cardinality_df}, ) - assert "cardinality" in frag_df.columns + # Then np.testing.assert_array_equal( frag_df["cardinality"].values, - cardinality_df.values.reshape(-1)[expected_mask], + cardinality.reshape(-1)[_flat(KEEP_UNFILTERED, bool)], ) @pytest.mark.requires_numba -def test_flatten_fragments_without_intensity(): +def test_flatten_fragments_without_intensity(library): """Without intensities, flatten_fragments removes only the mz == 0 padding.""" - precursor_df, mz_df, _ = _dense_library() - expected_mask = _expected_keep_mask(precursor_df, mz_df, pd.DataFrame(), 1000, -1) + # Given + precursor_df, mz_df, _ = library + # When _, frag_df = flatten_fragments(precursor_df, mz_df, pd.DataFrame()) + # Then assert "intensity" not in frag_df.columns np.testing.assert_array_equal( - frag_df["mz"].values, mz_df.values.reshape(-1)[expected_mask] + frag_df["mz"].values, _flat(MZ, PEAK_MZ_DTYPE)[_flat(KEEP_UNFILTERED, bool)] ) @pytest.mark.requires_numba -def test_flatten_fragments_selects_custom_columns(): +def test_flatten_fragments_selects_custom_columns(library): """flatten_fragments creates only the requested annotation columns.""" - precursor_df, mz_df, intensity_df = _dense_library() + # Given + precursor_df, mz_df, intensity_df = library + # When _, frag_df = flatten_fragments( precursor_df, mz_df, intensity_df, custom_columns=["number", "charge"] ) + # Then assert list(frag_df.columns) == ["mz", "intensity", "charge", "number"] @pytest.mark.requires_numba -def test_flatten_fragments_empty_precursor_df(): +def test_flatten_fragments_empty_precursor_df(library): """An empty library gives an empty fragment dataframe.""" - _, mz_df, intensity_df = _dense_library() + # Given + _, mz_df, intensity_df = library + # When precursor_df, frag_df = flatten_fragments( pd.DataFrame({"frag_start_idx": [], "frag_stop_idx": []}), mz_df, intensity_df ) + # Then assert len(precursor_df) == 0 assert len(frag_df) == 0 From 2b61b97b0b7cecfcaf6f849ad199d9d5f3a37da3 Mon Sep 17 00:00:00 2001 From: GeorgWa Date: Tue, 15 Sep 2026 10:20:18 +0200 Subject: [PATCH 08/14] Assert the flat fragment frames with assert_frame_equal Comparing column by column left the column set, the column order and the dtypes unchecked. Assert the whole annotated frame against an expected frame instead, and fold the row count of the filter cases into a frame comparison of mz and intensity. type and charge come out as int8, so the expected frame spells that out. Co-Authored-By: Claude Opus 5 (1M context) --- tests/unit/peptide/test_fragment.py | 36 ++++++++++++++++++----------- 1 file changed, 22 insertions(+), 14 deletions(-) diff --git a/tests/unit/peptide/test_fragment.py b/tests/unit/peptide/test_fragment.py index 0332ab3d..154ffa19 100644 --- a/tests/unit/peptide/test_fragment.py +++ b/tests/unit/peptide/test_fragment.py @@ -161,11 +161,13 @@ def test_flatten_fragments_retains_expected_fragments( ) # Then - assert len(frag_df) == keep.sum() - np.testing.assert_array_equal(frag_df["mz"].values, _flat(MZ, PEAK_MZ_DTYPE)[keep]) - np.testing.assert_array_equal( - frag_df["intensity"].values, _flat(INTENSITY, PEAK_INTENSITY_DTYPE)[keep] + expected_df = pd.DataFrame( + { + "mz": _flat(MZ, PEAK_MZ_DTYPE)[keep], + "intensity": _flat(INTENSITY, PEAK_INTENSITY_DTYPE)[keep], + } ) + pd.testing.assert_frame_equal(frag_df[["mz", "intensity"]], expected_df) @pytest.mark.requires_numba @@ -173,21 +175,27 @@ def test_flatten_fragments_annotates_retained_fragments(library): """Each annotation column describes the dense slot of its fragment.""" # Given precursor_df, mz_df, intensity_df = library + # position is the dense row within the precursor, number counts the ion series: + # b_modloss forward from position 0, y backward over ROWS_PER_PRECURSOR rows + expected_df = pd.DataFrame( + { + "mz": np.array([100.0, 110.0, 120.0, 121.0, 131.0], dtype=PEAK_MZ_DTYPE), + "intensity": np.array( + [0.11, 0.82, 0.93, 0.48, 0.61], dtype=PEAK_INTENSITY_DTYPE + ), + "type": np.array([121, 121, 121, 98, 98], dtype=np.int8), + "loss_type": np.array([0, 0, 0, 98, 98], dtype=np.int16), + "charge": np.array([1, 1, 1, 1, 1], dtype=np.int8), + "number": np.array([2, 1, 2, 1, 2], dtype=np.uint32), + "position": np.array([0, 1, 0, 0, 1], dtype=np.uint32), + } + ) # When _, frag_df = flatten_fragments(precursor_df, mz_df, intensity_df) # Then - np.testing.assert_array_equal( - frag_df["mz"].values, _flat(MZ, PEAK_MZ_DTYPE)[_flat(KEEP_UNFILTERED, bool)] - ) - # position is the dense row of the fragment within its precursor - np.testing.assert_array_equal(frag_df["position"].values, [0, 1, 0, 0, 1]) - # b_modloss counts forward, y counts backward over ROWS_PER_PRECURSOR rows - np.testing.assert_array_equal(frag_df["number"].values, [2, 1, 2, 1, 2]) - np.testing.assert_array_equal(frag_df["type"].values, [121, 121, 121, 98, 98]) - np.testing.assert_array_equal(frag_df["loss_type"].values, [0, 0, 0, 98, 98]) - np.testing.assert_array_equal(frag_df["charge"].values, [1, 1, 1, 1, 1]) + pd.testing.assert_frame_equal(frag_df, expected_df) @pytest.mark.requires_numba From dcb112f38dc2abaa9e85bd440a98f9b91e81180a Mon Sep 17 00:00:00 2001 From: GeorgWa Date: Tue, 15 Sep 2026 10:25:19 +0200 Subject: [PATCH 09/14] Fold the repeated annotation tiling into a loop The three tile calls differed only in column name and source array, and the comment above them explained the dtype of the previous implementation rather than the code it sat on. Loop over the columns instead; the dtypes stay visible where they are set, in _annotate_charged_frag_types. Trim the two comments carried into the extracted helpers. The direction comment claimed a third value of 0 that DIRECTION_MAPPING never produces. Column order of the flat fragment dataframe is unchanged. Co-Authored-By: Claude Opus 5 (1M context) --- alphabase/peptide/fragment.py | 21 ++++++++++----------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/alphabase/peptide/fragment.py b/alphabase/peptide/fragment.py index 4a6f39e3..68d401fc 100644 --- a/alphabase/peptide/fragment.py +++ b/alphabase/peptide/fragment.py @@ -991,7 +991,7 @@ def _annotate_charged_frag_types( series_ids = [] loss_ids = [] charges = [] - directions = [] # 'abc': direction=1, 'xyz': direction=-1, otherwise 0 + directions = [] # a/b/c series count forward (1), x/y/z series backward (-1) for charged_frag_type in charged_frag_types: frag_type, charge = parse_charged_frag_type(charged_frag_type) @@ -1029,8 +1029,7 @@ def _reannotate_precursor_pointers( precursor_df["flat_frag_stop_idx"] = precursor_df.frag_stop_idx precursor_df[["flat_frag_start_idx", "flat_frag_stop_idx"]] *= n_fragment_types - # cumulative sum counts the number of fragments before the given fragment which were removed. - # This sum does not include the fragment at the index position and has therefore len N +1 + # exclusive prefix sum, so cum_sum_tresh[i] counts the slots removed before slot i cum_sum_tresh = np.zeros(shape=len(excluded) + 1, dtype=np.int64) cum_sum_tresh[1:] = np.cumsum(excluded) @@ -1126,14 +1125,14 @@ def flatten_fragments( fragment_mz_df.columns.values ) - # tiling the typed arrays keeps their dtype. A tiled list of Python ints gave - # a dense int64 array first, which cost 8 bytes per dense slot. - if "type" in custom_columns: - frag_df["type"] = np.tile(series_ids, len(fragment_mz_df)) - if "loss_type" in custom_columns: - frag_df["loss_type"] = np.tile(loss_ids, len(fragment_mz_df)) - if "charge" in custom_columns: - frag_df["charge"] = np.tile(charges, len(fragment_mz_df)) + frag_type_annotations = { + "type": series_ids, + "loss_type": loss_ids, + "charge": charges, + } + for col_name, values in frag_type_annotations.items(): + if col_name in custom_columns: + frag_df[col_name] = np.tile(values, len(fragment_mz_df)) dense_directions = np.tile(directions, (len(fragment_mz_df), 1)) From 492a42298fdb5fc820e10b69914842eeeb6bab30 Mon Sep 17 00:00:00 2001 From: GeorgWa Date: Tue, 15 Sep 2026 10:42:48 +0200 Subject: [PATCH 10/14] Drop the del statements from the dense selection They free nothing. Peak numpy allocation measured with tracemalloc is 22.00 B/slot with all three dels, with one, and with none, at both 20.8M and 62.4M dense slots. The high water mark sits inside _parse_fragment, before any del in the caller runs: mz 4 + intensity 4 + directions 1 + numbers 4 + indices 4 + max_indices 4 + excluded 1. Every del frees memory that is already past that peak. Co-Authored-By: Claude Opus 5 (1M context) --- alphabase/peptide/fragment.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/alphabase/peptide/fragment.py b/alphabase/peptide/fragment.py index 77257039..c6aadf53 100644 --- a/alphabase/peptide/fragment.py +++ b/alphabase/peptide/fragment.py @@ -1068,7 +1068,6 @@ def _select_dense_fragments( intensity, n_fragment_types, ) - del dense_directions if intensity is None: excluded = mz == 0 @@ -1079,8 +1078,6 @@ def _select_dense_fragments( excluded = intensity < min_fragment_intensity excluded |= is_padding excluded |= not_top_k - del is_padding - del not_top_k # The in-place inversion prevents one more dense array. The indices stay # ascending, so the fragments keep their dense order. From 811b18e2cbff6da2f5cc47fe986bb4964c9b8bac Mon Sep 17 00:00:00 2001 From: GeorgWa Date: Tue, 15 Sep 2026 10:46:21 +0200 Subject: [PATCH 11/14] Access precursor columns with [] notation Match the bracket notation already used for the writes in the same functions, for the four column reads this branch introduced. Co-Authored-By: Claude Opus 5 (1M context) --- alphabase/peptide/fragment.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/alphabase/peptide/fragment.py b/alphabase/peptide/fragment.py index c6aadf53..e20cb5b7 100644 --- a/alphabase/peptide/fragment.py +++ b/alphabase/peptide/fragment.py @@ -1062,8 +1062,8 @@ def _select_dense_fragments( numbers, positions, not_top_k = _parse_fragment( dense_directions, - precursor_df.frag_start_idx.values, - precursor_df.frag_stop_idx.values, + precursor_df["frag_start_idx"].values, + precursor_df["frag_stop_idx"].values, keep_top_k_fragments, intensity, n_fragment_types, @@ -1171,8 +1171,8 @@ def _reannotate_precursor_pointers( number of charged fragment types, that is the number of dense columns """ - dense_start_idx = precursor_df.frag_start_idx.values.astype(np.int64) - dense_stop_idx = precursor_df.frag_stop_idx.values.astype(np.int64) + dense_start_idx = precursor_df["frag_start_idx"].values.astype(np.int64) + dense_stop_idx = precursor_df["frag_stop_idx"].values.astype(np.int64) precursor_df["flat_frag_start_idx"] = np.searchsorted( kept_indices, dense_start_idx * n_fragment_types ) From 4def3fa6741f2eb32cf5276ebd18a53fdd2c7456 Mon Sep 17 00:00:00 2001 From: GeorgWa Date: Tue, 15 Sep 2026 10:58:27 +0200 Subject: [PATCH 12/14] Point the test at the private helper name The merged tests named fill_in_indices, which this branch renamed. Co-Authored-By: Claude Opus 5 (1M context) --- tests/unit/peptide/test_fragment.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/peptide/test_fragment.py b/tests/unit/peptide/test_fragment.py index 154ffa19..e7b645fc 100644 --- a/tests/unit/peptide/test_fragment.py +++ b/tests/unit/peptide/test_fragment.py @@ -63,7 +63,7 @@ def test_filter_valid_charged_frag_types(mock_parse): CHARGED_FRAG_TYPES = ["y_z1", "b_modloss_z1"] N_PRECURSORS = 2 ROWS_PER_PRECURSOR = 2 -# mirrors the max_frag_per_peptide default of fill_in_indices +# mirrors the max_frag_per_peptide default of _fill_in_indices MAX_FRAG_PER_PEPTIDE = 300 # each mz encodes its own slot as 100 + row * 10 + column MZ = [ @@ -315,7 +315,7 @@ def test_flatten_fragments_long_precursor(): n_rows = np.iinfo(np.uint8).max + 2 assert ( n_rows <= MAX_FRAG_PER_PEPTIDE - ), "fill_in_indices cannot index a precursor this long" + ), "_fill_in_indices cannot index a precursor this long" n_types = len(CHARGED_FRAG_TYPES) rng = np.random.default_rng(0) mz_df = pd.DataFrame( From f290826b4f18822ef4ab22a2ef4e7fea7cf83459 Mon Sep 17 00:00:00 2001 From: GeorgWa Date: Tue, 15 Sep 2026 11:24:23 +0200 Subject: [PATCH 13/14] Assert every flat frame with assert_frame_equal Replace the remaining column comparisons. Each one now also pins the column set, the column order and the dtypes: - without_intensity subsumes its separate missing-column assertion - selects_custom_columns checked only the column names - filters_custom_df_columns now pins where cardinality lands - reannotates_precursor_pointers compares hardcoded flat pointers instead of reslicing mz per precursor, which was the last place the tests recomputed an expectation - empty_precursor_df compares both frames instead of their lengths Five tests assert against the same unfiltered library, so it moves into a flat_unfiltered fixture. test_flatten_fragments_long_precursor keeps its property assertions. It builds 257 rows to push fragment numbers past 255, so a frame of 1542 generated values would bury what it tests. Co-Authored-By: Claude Opus 5 (1M context) --- tests/unit/peptide/test_fragment.py | 129 +++++++++++++++------------- 1 file changed, 70 insertions(+), 59 deletions(-) diff --git a/tests/unit/peptide/test_fragment.py b/tests/unit/peptide/test_fragment.py index 154ffa19..880a226e 100644 --- a/tests/unit/peptide/test_fragment.py +++ b/tests/unit/peptide/test_fragment.py @@ -87,9 +87,10 @@ def test_filter_valid_charged_frag_types(mock_parse): [1, 1], [0, 1], ] -# keep_top_k_fragments, min_fragment_intensity, kept slots per dense row +# keep_top_k_fragments, min_fragment_intensity, kept slots per dense row, and the +# flat [start, stop] pointer of each precursor FILTER_CASES = [ - pytest.param(1000, -1, KEEP_UNFILTERED, id="padding_only"), + pytest.param(1000, -1, KEEP_UNFILTERED, [[0, 2], [2, 5]], id="padding_only"), pytest.param( 1000, 0.3, @@ -99,6 +100,7 @@ def test_filter_valid_charged_frag_types(mock_parse): [1, 1], [0, 1], ], + [[0, 1], [1, 4]], id="min_intensity", ), pytest.param( @@ -110,6 +112,7 @@ def test_filter_valid_charged_frag_types(mock_parse): [1, 0], [0, 0], ], + [[0, 1], [1, 2]], id="top_1_per_precursor", ), ] @@ -134,6 +137,24 @@ def library(): ) +@pytest.fixture +def flat_unfiltered(): + """Give the flat library expected when only the mz == 0 padding is dropped.""" + return pd.DataFrame( + { + "mz": np.array([100.0, 110.0, 120.0, 121.0, 131.0], dtype=PEAK_MZ_DTYPE), + "intensity": np.array( + [0.11, 0.82, 0.93, 0.48, 0.61], dtype=PEAK_INTENSITY_DTYPE + ), + "type": np.array([121, 121, 121, 98, 98], dtype=np.int8), + "loss_type": np.array([0, 0, 0, 98, 98], dtype=np.int16), + "charge": np.array([1, 1, 1, 1, 1], dtype=np.int8), + "number": np.array([2, 1, 2, 1, 2], dtype=np.uint32), + "position": np.array([0, 1, 0, 0, 1], dtype=np.uint32), + } + ) + + def _flat(values, dtype): """Flatten a dense grid to the slot order of the flat fragment dataframe.""" return np.array(values, dtype=dtype).reshape(-1) @@ -141,10 +162,15 @@ def _flat(values, dtype): @pytest.mark.requires_numba @pytest.mark.parametrize( - "keep_top_k_fragments, min_fragment_intensity, expected_keep", FILTER_CASES + "keep_top_k_fragments, min_fragment_intensity, expected_keep, expected_pointers", + FILTER_CASES, ) def test_flatten_fragments_retains_expected_fragments( - library, keep_top_k_fragments, min_fragment_intensity, expected_keep + library, + keep_top_k_fragments, + min_fragment_intensity, + expected_keep, + expected_pointers, # noqa: ARG001 ): """The flat library keeps only the slots listed in FILTER_CASES.""" # Given @@ -171,46 +197,41 @@ def test_flatten_fragments_retains_expected_fragments( @pytest.mark.requires_numba -def test_flatten_fragments_annotates_retained_fragments(library): - """Each annotation column describes the dense slot of its fragment.""" +def test_flatten_fragments_annotates_retained_fragments(library, flat_unfiltered): + """Each annotation column describes the dense slot of its fragment. + + position is the dense row within the precursor, number counts the ion series: + b_modloss forward from position 0, y backward over ROWS_PER_PRECURSOR rows. + """ # Given precursor_df, mz_df, intensity_df = library - # position is the dense row within the precursor, number counts the ion series: - # b_modloss forward from position 0, y backward over ROWS_PER_PRECURSOR rows - expected_df = pd.DataFrame( - { - "mz": np.array([100.0, 110.0, 120.0, 121.0, 131.0], dtype=PEAK_MZ_DTYPE), - "intensity": np.array( - [0.11, 0.82, 0.93, 0.48, 0.61], dtype=PEAK_INTENSITY_DTYPE - ), - "type": np.array([121, 121, 121, 98, 98], dtype=np.int8), - "loss_type": np.array([0, 0, 0, 98, 98], dtype=np.int16), - "charge": np.array([1, 1, 1, 1, 1], dtype=np.int8), - "number": np.array([2, 1, 2, 1, 2], dtype=np.uint32), - "position": np.array([0, 1, 0, 0, 1], dtype=np.uint32), - } - ) # When _, frag_df = flatten_fragments(precursor_df, mz_df, intensity_df) # Then - pd.testing.assert_frame_equal(frag_df, expected_df) + pd.testing.assert_frame_equal(frag_df, flat_unfiltered) @pytest.mark.requires_numba @pytest.mark.parametrize( - "keep_top_k_fragments, min_fragment_intensity, expected_keep", FILTER_CASES + "keep_top_k_fragments, min_fragment_intensity, expected_keep, expected_pointers", + FILTER_CASES, ) def test_flatten_fragments_reannotates_precursor_pointers( - library, keep_top_k_fragments, min_fragment_intensity, expected_keep + library, + keep_top_k_fragments, + min_fragment_intensity, + expected_keep, # noqa: ARG001 + expected_pointers, ): """The flat pointers of a precursor address only its own fragments.""" # Given precursor_df, mz_df, intensity_df = library - n_types = len(CHARGED_FRAG_TYPES) - mz = _flat(MZ, PEAK_MZ_DTYPE) - keep = _flat(expected_keep, bool) + expected_df = pd.DataFrame( + np.array(expected_pointers, dtype=np.int64), + columns=["flat_frag_start_idx", "flat_frag_stop_idx"], + ) # When precursor_df, frag_df = flatten_fragments( @@ -222,66 +243,55 @@ def test_flatten_fragments_reannotates_precursor_pointers( ) # Then - for row in precursor_df.itertuples(): - block = slice(row.frag_start_idx * n_types, row.frag_stop_idx * n_types) - np.testing.assert_array_equal( - frag_df["mz"].values[row.flat_frag_start_idx : row.flat_frag_stop_idx], - mz[block][keep[block]], - ) - - # the pointers must cover the fragment dataframe with no gap and no overlap - assert precursor_df.flat_frag_start_idx.iloc[0] == 0 - assert precursor_df.flat_frag_stop_idx.iloc[-1] == len(frag_df) - np.testing.assert_array_equal( - precursor_df.flat_frag_start_idx.values[1:], - precursor_df.flat_frag_stop_idx.values[:-1], + pd.testing.assert_frame_equal( + precursor_df[["flat_frag_start_idx", "flat_frag_stop_idx"]], expected_df ) + assert expected_pointers[-1][1] == len(frag_df) @pytest.mark.requires_numba -def test_flatten_fragments_filters_custom_df_columns(library): +def test_flatten_fragments_filters_custom_df_columns(library, flat_unfiltered): """flatten_fragments filters a custom_df column like the mz column.""" # Given precursor_df, mz_df, intensity_df = library cardinality = np.arange(mz_df.size, dtype=np.uint8).reshape(mz_df.shape) - cardinality_df = pd.DataFrame(cardinality, columns=CHARGED_FRAG_TYPES) + expected_df = flat_unfiltered + expected_df.insert(2, "cardinality", np.array([0, 2, 4, 5, 7], dtype=np.uint8)) # When _, frag_df = flatten_fragments( precursor_df, mz_df, intensity_df, - custom_df={"cardinality": cardinality_df}, + custom_df={ + "cardinality": pd.DataFrame(cardinality, columns=CHARGED_FRAG_TYPES) + }, ) # Then - np.testing.assert_array_equal( - frag_df["cardinality"].values, - cardinality.reshape(-1)[_flat(KEEP_UNFILTERED, bool)], - ) + pd.testing.assert_frame_equal(frag_df, expected_df) @pytest.mark.requires_numba -def test_flatten_fragments_without_intensity(library): +def test_flatten_fragments_without_intensity(library, flat_unfiltered): """Without intensities, flatten_fragments removes only the mz == 0 padding.""" # Given precursor_df, mz_df, _ = library + expected_df = flat_unfiltered.drop(columns="intensity") # When _, frag_df = flatten_fragments(precursor_df, mz_df, pd.DataFrame()) # Then - assert "intensity" not in frag_df.columns - np.testing.assert_array_equal( - frag_df["mz"].values, _flat(MZ, PEAK_MZ_DTYPE)[_flat(KEEP_UNFILTERED, bool)] - ) + pd.testing.assert_frame_equal(frag_df, expected_df) @pytest.mark.requires_numba -def test_flatten_fragments_selects_custom_columns(library): +def test_flatten_fragments_selects_custom_columns(library, flat_unfiltered): """flatten_fragments creates only the requested annotation columns.""" # Given precursor_df, mz_df, intensity_df = library + expected_df = flat_unfiltered[["mz", "intensity", "charge", "number"]] # When _, frag_df = flatten_fragments( @@ -289,7 +299,7 @@ def test_flatten_fragments_selects_custom_columns(library): ) # Then - assert list(frag_df.columns) == ["mz", "intensity", "charge", "number"] + pd.testing.assert_frame_equal(frag_df, expected_df) @pytest.mark.requires_numba @@ -297,15 +307,16 @@ def test_flatten_fragments_empty_precursor_df(library): """An empty library gives an empty fragment dataframe.""" # Given _, mz_df, intensity_df = library + empty_precursor_df = pd.DataFrame({"frag_start_idx": [], "frag_stop_idx": []}) # When - precursor_df, frag_df = flatten_fragments( - pd.DataFrame({"frag_start_idx": [], "frag_stop_idx": []}), mz_df, intensity_df - ) + precursor_df, frag_df = flatten_fragments(empty_precursor_df, mz_df, intensity_df) # Then - assert len(precursor_df) == 0 - assert len(frag_df) == 0 + pd.testing.assert_frame_equal( + precursor_df, pd.DataFrame({"frag_start_idx": [], "frag_stop_idx": []}) + ) + pd.testing.assert_frame_equal(frag_df, pd.DataFrame()) @pytest.mark.requires_numba From 6fa53b185f1d30735366161403c5d0290964aa27 Mon Sep 17 00:00:00 2001 From: GeorgWa Date: Tue, 15 Sep 2026 11:33:10 +0200 Subject: [PATCH 14/14] Compare the long precursor maxima and dtypes as series The four scalar assertions become two equality comparisons. The dtypes need their own comparison, because .max() over columns of different dtypes upcasts to a common one and would hide a narrowed column. Also drop the tuple rebuild in the parse test, expected is already the tuple being compared. Co-Authored-By: Claude Opus 5 (1M context) --- tests/unit/peptide/test_fragment.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/tests/unit/peptide/test_fragment.py b/tests/unit/peptide/test_fragment.py index 880a226e..175d64ab 100644 --- a/tests/unit/peptide/test_fragment.py +++ b/tests/unit/peptide/test_fragment.py @@ -23,7 +23,7 @@ def test_parse_charged_frag_type_with_valid_input(input_str, expected): """Test parse_charged_frag_type with valid input.""" result = parse_charged_frag_type(input_str) - assert result == (expected[0], expected[1]) + assert result == expected @pytest.mark.parametrize( @@ -343,7 +343,12 @@ def test_flatten_fragments_long_precursor(): precursor_df, mz_df, intensity_df, keep_top_k_fragments=n_rows * n_types ) - assert frag_df["position"].max() == n_rows - 1 - assert frag_df["number"].max() == n_rows - assert frag_df["position"].dtype == np.uint32 - assert frag_df["number"].dtype == np.uint32 + # .max() over mixed dtypes upcasts, so the dtypes need their own comparison + pd.testing.assert_series_equal( + frag_df[["position", "number"]].dtypes, + pd.Series({"position": np.dtype(np.uint32), "number": np.dtype(np.uint32)}), + ) + pd.testing.assert_series_equal( + frag_df[["position", "number"]].max(), + pd.Series({"position": n_rows - 1, "number": n_rows}, dtype=np.uint32), + )