feat: add k-mer-seeded top-k ungapped local alignment

[lgruen-vcgs]

Motivation

top_k_ungapped_local_align visits every diagonal of the n × m grid, which is O(n·m) regardless of how much of the grid contains useful HSPs. For the downstream caller that prompted this change — anchorite.chained_alignment, which runs the seed step on (PDF flat string, markdown) pairs of typically 10–30 KB each — profiling showed 99% of the alignment-stage cost was a single top-level call to top_k_ungapped_local_align. On a representative 8-page paper (≈22 KB × 20 KB, A·B ≈ 4.5 × 10⁸ cells) this one call was ~870 ms.

The classical fix is BLAST/minimap2-style k-mer seeding: hash every length-k window of A, look each up against B, and only run the per-diagonal positive-segment scan on diagonals that contain at least one exact k-mer hit. For text-like alphabets this turns the seed step from O(n·m) into roughly O(n + m + hits) in practice.

New function

def top_k_ungapped_local_align_kmer(
    seqa: bytes,
    seqb: bytes,
    score_matrix: npt.NDArray[np.int32],
    k: int,
    kmer_size: int,
    max_hits_per_kmer: int,
    filter_overlap_a: bool = True,
    filter_overlap_b: bool = True,
) -> list[Alignment]: ...

Output shape and overlap-filter semantics are identical to top_k_ungapped_local_align, so the new function is a drop-in replacement at the seed step of any existing caller.

Parameters specific to this variant:

• kmer_size — seed length in bytes; must be in [1, 8] (each k-mer is packed into a u64 for the index).
• max_hits_per_kmer — k-mers whose seqa-occurrence count exceeds this are skipped, bounding low-complexity blowup (long runs of spaces, common bigrams, etc.).

Algorithm

1. Build a rolling-hash k-mer index of seqa: HashMap<u64, Vec<u32>> mapping packed-k-mer → positions in sa.
2. Walk seqb's rolling k-mer. For each lookup, collect the set of hit diagonals (diag = sa_pos - sb_pos); skip k-mers that exceed max_hits_per_kmer.
3. Run the existing per-diagonal positive-segment scan only on diagonals that have at least one hit, into the same BinaryHeap<Candidate>.
4. Pop top-k non-overlapping HSPs via the existing overlap-filter logic.

Refactor

To share code cleanly between the all-diagonals scan and the k-mer-seeded scan, two helpers are extracted out of _top_k_ungapped_local_align_core:

• process_diagonal_into_candidates — SW-style positive-segment scan along one diagonal of (sb × sa), pushing every positive HSP peak into the candidate heap.
• select_top_k_with_overlap_filter — pop top-k from the heap, skipping HSPs that overlap an already-accepted one on the A and/or B axis.

_top_k_ungapped_local_align_core now just iterates all diagonals through process_diagonal_into_candidates and calls select_top_k_with_overlap_filter. Behaviour-preserving — all existing top-k tests pass unchanged.

Correctness caveat

The k-mer seeder finds an HSP iff at least one window of kmer_size consecutive exact-match bytes lies on its diagonal between the two sequences. HSPs whose match runs are all strictly shorter than kmer_size are missed.

For text alphabets at kmer_size <= 5 over +1/-1 scoring this only occurs at very small HSP scores (the smallest construction has score 8 with match-runs of length 4 only, which requires a precise 4+4 match layout with the right mismatch arrangement — vanishingly rare on real text). Callers that need exhaustive recall at the boundary should either keep using top_k_ungapped_local_align or fall through to a full local SW (the path anchorite.chained_alignment's recursive gap-fill already uses).

A note on max_hits_per_kmer

The first end-to-end benchmark in anchorite came in at only ~1.2× speedup with max_hits_per_kmer=128. Profiling showed the seeder was still touching ~27,000 of the ~42,000 possible diagonals on an 8-page paper, because the PDF flat string and the markdown share a lot of vocabulary — moderately-common 5-mers (e.g. " the ", "tion ") hit on hundreds of off-diagonal positions, each contributing a different (sa_pos − sb_pos) to the diagonal set. The cap was filtering individual hits but not collapsing the diagonal set.

Hit-diagonal counts at kmer_size=5 on the 8-pager:

max_hits_per_kmer	unique hit diagonals	total hits
1	304	5,319
4	3,753	18,480
16	13,016	65,932
128	27,421	247,763
∞	29,417	267,918

The downstream caller (anchorite.chained_alignment) lands on max_hits_per_kmer=1 as its default — i.e. only k-mers that are unique in A are followed up. On real text, every HSP scoring at the typical seed floor (8 over +1/-1) almost always contains at least one unique-in-A 5-mer somewhere in its match runs, so the chain step is reliably anchored. No top-K HSPs were lost at this setting on the test fixtures vs. the full-scan baseline.

The new function intentionally has no default for max_hits_per_kmer — different callers may want different tradeoffs, and there's no clearly-right value across use cases.

Measured impact

End-to-end on the anchorite caller (PdfIndex(pdf, markdown=md) construction; 10 runs after 2 warmups):

Fixture	A · B	Before (full scan)	After (kmer, cap=1)	Speedup
3-page bug-repro	6,171 × 5,894	586 ms ± 7	48 ms ± 2	12.2×
8-page PLOS ONE	22,057 × 20,512	930 ms ± 34	84 ms ± 5	11.1×

The denoise path is now within ~30 ms of the raw-chars PDF extraction baseline; the new bottleneck is pypdfium2 + glyph normalisation, not alignment.

Tests

7 new cases in tests/test_seq_smith.py:

• test_top_k_ungapped_kmer_simple — simple HSP recovery, matches the existing test_top_k_ungapped_simple shape.
• test_top_k_ungapped_kmer_matches_full_scan_random_dna — k-mer seeder vs. all-diagonals scan on random DNA; same HSPs at the high end of the score distribution.
• test_top_k_ungapped_kmer_overlap — overlap filter on B, matches test_top_k_ungapped_overlap.
• test_top_k_ungapped_kmer_limit — top-k truncation.
• test_top_k_ungapped_kmer_invalid_kmer_size — kmer_size validation rejects 0 and 9+.
• test_top_k_ungapped_kmer_seqs_shorter_than_kmer — sequences shorter than kmer_size produce [].
• test_top_k_ungapped_kmer_max_hits_skips_low_complexity — max_hits_per_kmer=0 returns [] even on poly-A identity (the knob that protects against low-complexity blowup).

Test plan

• cargo check
• pytest — 60 tests pass (53 existing + 7 new)
• Smoke-tested end-to-end via anchorite test suite (135 tests pass; resolve quality identical across max_hits_per_kmer ∈ {1, 4, 16, 128} on the bug-repro)

Not in this PR

• Version bump: deferred to a separate PR per the repo's convention (e.g. Bump version to 0.5.0 #29, fix: resolve indexing issues in traceback and bump to 0.5.1 #30).
• _many variant: didn't add top_k_ungapped_local_align_kmer_many yet — happy to follow up if you want it for parity.

[lgruen-vcgs]

Subsumed by #32.