Skip to content

Add sort-free Rust FDR kernels (q-values, keep-best, fused finalize) - #130

Draft
GeorgWa wants to merge 1 commit into
mainfrom
feature/rust-qval-calc
Draft

Add sort-free Rust FDR kernels (q-values, keep-best, fused finalize)#130
GeorgWa wants to merge 1 commit into
mainfrom
feature/rust-qval-calc

Conversation

@GeorgWa

@GeorgWa GeorgWa commented Aug 24, 2026

Copy link
Copy Markdown
Collaborator

Adds src/fdr.rs with array kernels for the FDR hot spots, mirroring the pandas implementations in alphadia/fdr/fdr.py:

  • fdr_q_values: lexicographic argsort by (proba, decoy, tiebreak) plus the running target/decoy FDR and reverse running minimum. Parallel sort via rayon.
  • fdr_keep_best: per-group arg-min over the score, ties broken by lowest original index.
  • fdr_finalize: fused and sort-free — a per-group arg-min for keep-best, then a fixed-bin histogram of target/decoy counts over proba whose cumulative sums give the q-value per bin. Replaces a chain of up to four O(N log N) sorts.

Motivation is the pandas multi-key sort / cumsum / groupby overhead, which dominates FDR at ~10M+ candidates.

Python side: MannLabs/alphadia#841.

🤖 Generated with Claude Code

Adds an fdr module with three PyO3 functions used to accelerate the
pandas-heavy FDR pipeline in alphadia/fdr/fdr.py:

- fdr_q_values: parallel multi-key argsort + cumsum + reverse-cummin,
  an exact drop-in for get_q_values returning (order, qvalues).
- fdr_keep_best: per-group arg-min over proba (sort-free).
- fdr_finalize: fused sort-free finalization. Keeps the best PSM per
  group and assigns q-values by histogramming target/decoy counts over
  proba bins (cumsum + reverse running minimum), avoiding any O(N) sort
  or ordering array. q-values are quantized to bin resolution
  (1/n_bins), far below any FDR threshold.

Unit tests cover argsort ordering, exact q-values, histogram q-values,
and keep-best selection.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@GeorgWa

GeorgWa commented Aug 24, 2026

Copy link
Copy Markdown
Collaborator Author

Context: the tie-break argument, and why the signature is not frozen

Draft for now — this wants to be reviewed together with MannLabs/alphadia#841, which is the only consumer.

Why fdr_q_values takes a tie-break

q_values_in_order accumulates target/decoy counts along the sort order, so rows with an equal proba have no defined position relative to each other and the resulting q-values depend on which one comes first. argsort_fdr therefore sorts by (proba, decoy, tiebreak) to make the order total, mirroring df.sort_values([score, decoy, extra_sort_columns]) in the Python reference.

The tie-break arrives as &[i64] because Rust cannot compare arbitrary Python objects. On the Python side that is fine for precursor FDR (precursor_idx) but not for protein FDR, which breaks ties on a string accession; #841 maps non-integer columns to order-preserving integer codes rather than branching on dtype.

The constraint is that the pandas implementation stays live — always for fragment competition, and everywhere while these kernels are unreleased — so the two must agree numerically, not merely each be deterministic.

The signature may lose its third argument

There is an open proposal to assign one q-value per exact equal-proba block instead of ranking within ties. If that lands:

  • fdr_q_values drops tiebreak entirely;
  • argsort_fdr becomes a single-key sort, and par_sort_unstable_by becomes safe because intra-block order stops being load-bearing — currently the total order is what makes the unstable sort acceptable only in combination with the tie-break key;
  • the two extra comparison levels and one i64 array come out of the hot path.

So this is a small perf win on top of a correctness one, and it is worth not treating the current signature as settled while reviewing.

Grouping would be on exact f64 equality, deliberately. The ties in question come from bit-identical classifier outputs; scores differing by one ULP are genuinely different and belong in separate blocks. That is a different mechanism from q_values_histogram here, which bins by value ranges — lumping distinct scores together and quantizing q-values to bin width. That approximation is defensible for fdr_finalize at 10M+ candidates but should not be confused with exact tie grouping, and it is why fdr_finalize needs no tie-break argument today.

Status

Draft until reviewed alongside #841. Note that until these kernels ship in a released alphadia-search-rs, the Python side has _RUST_FDR_AVAILABLE = False, so nothing exercises them end to end and the parity tests there skip at module level.

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.

1 participant