Skip to content

[4/5] Perf: build the flat fragment columns from the kept indices - #437

Merged
GeorgWa merged 24 commits into
mainfrom
perf/build-flat-columns-from-kept-indices
Sep 15, 2026
Merged

GeorgWa merged 24 commits into
mainfrom
perf/build-flat-columns-from-kept-indices

Conversation

@GeorgWa

@GeorgWa GeorgWa commented Aug 31, 2026

Copy link
Copy Markdown
Collaborator

Performance. Output bit-identical to main.

flatten_fragments allocated 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 an int64 cumulative sum over all dense slots. For top-k libraries, mz == 0 padding fills most dense slots.

  • Build all columns except mz and intensity from the kept indices.
  • Reannotate the pointers with np.searchsorted on the kept indices, instead of a dense cumulative sum.
  • Accumulate the masks in place.

48M dense slots: peak RSS 3.49 GB -> 1.77 GB.

Stacked on #436. Splitting #429.

🤖 Generated with Claude Code

GeorgWa and others added 4 commits August 31, 2026 15:56
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) <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 for the refactoring
that follows.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…tion

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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
@GeorgWa GeorgWa changed the title Perf: build the flat fragment columns from the kept indices [4/5] Perf: build the flat fragment columns from the kept indices Aug 31, 2026
if "position" in custom_columns:
frag_df["position"] = positions.reshape(-1)

frag_df = {}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

frag_dict?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a bit tricky. we use this pattern a lot as this is how we create a pandas df. calling it frag dict would suggest somethign else.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

well, _df suggests it's a dataframe .. but if this follows current code convention, fair enough

Comment thread alphabase/peptide/fragment.py Outdated
intensity,
n_fragment_types,
)
del dense_directions

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

did you check if this is garbage collected immediately? or does this need an explicit call to gc?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I only checked if this improved memory, not if it triggered gc. but assume so.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

isn't gc only triggered automatically at the exit of a method?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I removed the del statements 👍🏻

del dense_directions

if intensity is None:
excluded = mz == 0

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

mz == 0 is shorthand for all array elements are zero? is there a more numpythonic way of doing that?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes exactly, this is the pythonic way tbh

Comment thread alphabase/peptide/fragment.py Outdated
Comment on lines +1082 to +1083
del is_padding
del not_top_k

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

how large can the mz array be? is it worth the del noise?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

While I see the point. I'm not sure if it's worth the 30min investigation.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

removed them

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you.

Note: What would take you 30 minutes of investigation now, will take other people (including future you) 90 minutes, nice "every LoC is a liability" example :-)

Comment thread alphabase/peptide/fragment.py Outdated
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)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

please use [] notation to access df columns consistently

Comment thread alphabase/peptide/fragment.py Outdated
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)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

so frag_stop_idx is int32 (I assume), but needs to be int64 for the dense representation? worth a comment..

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes, and we want to be explicit with the dtype before constructing the pandas dataframes. Otherwise weird parsing can happen.

GeorgWa and others added 19 commits September 15, 2026 09:57
_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) <noreply@anthropic.com>
"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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
The merged tests named fill_in_indices, which this branch renamed.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…ns-from-kept-indices

Two conflicts, both where this branch had already deleted the code the
other branch edited: the cumsum in _reannotate_precursor_pointers, which
searchsorted replaced, and the dense annotation tiling in
flatten_fragments, which _annotate_kept_fragments replaced. Kept this
branch's version of both.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
@GeorgWa
GeorgWa requested a review from mschwoer September 15, 2026 11:48
Base automatically changed from refactor/annotate-charged-frag-types to main September 15, 2026 11:50
Main carries PRs #434 to #436, so the three branches below this one
arrive from main and only fragment.py still differs.

Six conflicts, all where this branch had rewritten what main still has:
the dense annotation tiling, the cumsum in _reannotate_precursor_pointers
that searchsorted replaces, and the helper signatures around them. Kept
this branch's version of each, and main's tuple annotation, since main
dropped Python 3.8 and the typing imports with it.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@GeorgWa
GeorgWa merged commit bf628cd into main Sep 15, 2026
5 checks passed
@GeorgWa
GeorgWa deleted the perf/build-flat-columns-from-kept-indices branch September 15, 2026 15:19
GeorgWa pushed a commit that referenced this pull request Sep 15, 2026
Main carries PR #437, so only the per-row rework still differs. Sixteen
conflicts, all where this branch replaced the per-dense-slot arrays that
main now has. Kept the per-row version of each.

Reword the two comments mschwoer asked about:

- the position cast said what the line already says. It now names the
  reason for the cast, that row_positions is the narrower uint16.
- the copy=False comment said a copy is kept "out of memory". It now
  says a view is returned instead of a second dense array being
  allocated.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants