Reduce peak memory of flatten_fragments - #429
Closed
GeorgWa wants to merge 6 commits into
Closed
Conversation
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. The reannotation of the precursor pointers additionally allocated an int64 cumulative sum over all dense slots. Instead, only mz and intensity are needed at dense length to compute the keep mask. All other columns are now built directly from the indices of the retained fragments, and the pointers are reannotated with a binary search on those indices, which is equivalent to subtracting the cumulative number of removed fragments. Output is bit-identical: same fragment count, column order, dtypes, values and precursor pointers. On a 2M precursor library (480M dense slots) peak RSS drops from 23.2 GB to 12.0 GB and the call is ~40% faster. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Rewrite the comments and test docstrings in Simplified Technical English (ASD-STE100): simple present, active voice, one idea per sentence. Keep only the reasons, not restatements of the code. No code changes. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Move three concerns out of the function body: the per-column annotation of the charged fragment types, the selection of the dense slots to keep, and the reannotation of the precursor pointers. flatten_fragments now reads as its five steps, and each helper carries the reason for its own memory behaviour. `_annotate_charged_frag_types` returns typed arrays, so the direction array is tiled as int8. A tiled list gave a dense int64 array before, which cost 8 bytes per dense slot. No functional change. Output stays bit-identical. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`parse_fragment` held the fragment numbers, positions and per-precursor maxima as uint32, which cost 4 bytes per dense slot each. `fill_in_indices` bounds all three values by `max_frag_per_peptide`, so uint16 holds every value with a large margin. `calculate_fragment_numbers` takes and returns uint16 accordingly. `flatten_fragments` casts the kept `number` and `position` values back to uint32, so the flat columns keep their public dtype. `fill_in_indices` also writes the per-precursor maximum as a scalar, instead of a product with an array of ones. This removes one allocation per precursor. Output stays bit-identical against alphabase main, kernels included. Peak RSS at 1.48G dense slots: 36.26 GB -> 28.00 GB. At 480M dense slots: 11.99 GB -> 9.39 GB. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`parse_fragment` held the fragment number, the position and the precursor length of every dense slot, which cost 3 arrays of 2 bytes per slot. All three values are shared by every charged fragment type of a fragment row, so `fill_in_indices` now writes one value per fragment row instead, and `parse_fragment` no longer takes the tiled directions array. `calculate_fragment_numbers` takes a direction, a row position and a row count, and `flatten_fragments` calls it for the kept fragments only. The dense pass over every slot is gone, and so is the discarded output array it needed. `flatten_fragments` also passes the intensities without a copy, because nothing writes to them any more. Zeroing the intensity of padding slots was a no-op, as padding is always excluded. Peak RSS at 1.48G dense slots: 28.00 GB -> 16.08 GB, and the runtime halves. Output stays bit-identical against alphabase main, with one exception: fragment rows that no precursor covers now report number 0 and position 0. main reports uninitialised values there, because it allocates the arrays with `np.empty`. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`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. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Desperadus
reviewed
Aug 31, 2026
Comment on lines
+1264
to
1273
| # matches. Nothing writes to `mz` or `intensity`, so both stay views on the | ||
| # input dataframes. | ||
| intensity = ( | ||
| fragment_intensity_df.values.astype(PEAK_INTENSITY_DTYPE, copy=False).reshape( | ||
| -1 | ||
| ) | ||
| if "loss_type" in custom_columns: | ||
| frag_df["loss_type"] = np.array( | ||
| np.tile(frag_loss_types, len(fragment_mz_df)), dtype=np.int16 | ||
| ) | ||
| if "charge" in custom_columns: | ||
| frag_df["charge"] = np.array( | ||
| np.tile(frag_charges, len(fragment_mz_df)), dtype=np.int8 | ||
| ) | ||
|
|
||
| frag_directions = np.array( | ||
| np.tile(frag_directions, (len(fragment_mz_df), 1)), dtype=np.int8 | ||
| if use_intensity | ||
| else None | ||
| ) | ||
|
|
There was a problem hiding this comment.
You say: "Nothing writes to mz or intensity" - should we enforce it like:
mz.flags.writeable = False
if intensity is not None:
intensity.flags.writeable = False?
This was referenced Aug 31, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
flatten_fragmentsallocated every flat column at dense length (n_fragment_rows * n_fragment_types), assembled a dataframe, then copied the kept subset out of it. Pointer reannotation added anint64cumulative sum over all dense slots.parse_fragmentheld the number, the position and the precursor length of every dense slot. For top-k libraries,mz == 0padding fills most dense slots.mzandintensityfrom the kept indices.np.searchsortedon the kept indices, instead of a dense cumulative sum.max_frag_per_peptidebounds.Peak RSS and runtime,
keep_top_k_fragments=12: