From 98450b43d80e6fc7fdad780ac2f4d5eee2ba8337 Mon Sep 17 00:00:00 2001 From: Leonhard Gruenschloss Date: Tue, 12 May 2026 10:21:26 +1000 Subject: [PATCH 1/7] feat: add k-mer-seeded top-k ungapped local alignment MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `top_k_ungapped_local_align_kmer(seqa, seqb, score_matrix, k, kmer_size, max_hits_per_kmer, filter_overlap_a=True, filter_overlap_b=True)` is a faster alternative to `top_k_ungapped_local_align` for callers that work on long inputs over a moderate-sized alphabet. Algorithm: BLAST-style k-mer seeding plus the existing per-diagonal positive-segment scan as the extension. 1. Build a rolling-hash k-mer index of `seqa` (packed k-mer -> sa positions; `kmer_size` in `[1, 8]`, packed into a `u64`). 2. Walk `seqb`'s rolling k-mer; for each lookup, collect the set of hit diagonals (`diag = sa_pos - sb_pos`). Skip k-mers whose `seqa`-occurrence count exceeds `max_hits_per_kmer` to bound low-complexity blowup (long runs of spaces, common bigrams, etc.). 3. Run the per-diagonal positive-segment scan *only* on diagonals with at least one hit, into the same candidate heap. 4. Pop top-`k` non-overlapping via the existing overlap-filter logic. The output `Alignment` shape and overlap semantics are identical to `top_k_ungapped_local_align`, so the new function is a drop-in replacement at the seed step of any caller. Implementation reuses the existing diagonal scan by extracting two helpers out of `_top_k_ungapped_local_align_core`: - `process_diagonal_into_candidates`: SW-style positive-segment scan along one diagonal of (sb x sa), pushing every positive HSP peak into a `BinaryHeap`. - `select_top_k_with_overlap_filter`: pop top-`k` from the heap, skipping HSPs that overlap an already-accepted one on `A` and/or `B`. Both `_top_k_ungapped_local_align_core` and the new k-mer variant share these helpers. Behaviour-preserving refactor — all existing top-k tests pass unchanged. Correctness caveat (documented on the new function): an HSP is found 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. Over text-like alphabets at `kmer_size <= 5` this only happens at very small HSP scores; callers that need exhaustive coverage at the boundary should keep using `top_k_ungapped_local_align` or fall through to a full local SW (as `anchorite.chained_alignment`'s recursive gap-fill does). New tests in `tests/test_seq_smith.py`: - simple HSP recovery (matches the existing `test_top_k_ungapped_simple` case), - near-identity equivalence with the full-scan top-k on random DNA at the high end of the score distribution, - overlap filter on `B` (matches `test_top_k_ungapped_overlap`), - top-`k` limit (matches `test_top_k_ungapped_limit`), - `kmer_size` validation rejects 0 and 9+ with `ValueError`, - sequences shorter than `kmer_size` return `[]`, - `max_hits_per_kmer=0` returns `[]` even on a poly-A identity pair (the cap that protects against low-complexity blowup). --- seq_smith/__init__.py | 2 + seq_smith/_seq_smith.pyi | 10 ++ src/lib.rs | 341 ++++++++++++++++++++++++++++++--------- tests/test_seq_smith.py | 140 ++++++++++++++++ 4 files changed, 417 insertions(+), 76 deletions(-) diff --git a/seq_smith/__init__.py b/seq_smith/__init__.py index 1013c98..e4f6cf0 100644 --- a/seq_smith/__init__.py +++ b/seq_smith/__init__.py @@ -11,6 +11,7 @@ overlap_align, overlap_align_many, top_k_ungapped_local_align, + top_k_ungapped_local_align_kmer, top_k_ungapped_local_align_many, ) from .python_utils import decode, encode, format_alignment_ascii, generate_cigar, make_score_matrix @@ -33,5 +34,6 @@ "overlap_align", "overlap_align_many", "top_k_ungapped_local_align", + "top_k_ungapped_local_align_kmer", "top_k_ungapped_local_align_many", ] diff --git a/seq_smith/_seq_smith.pyi b/seq_smith/_seq_smith.pyi index 64cd598..818098e 100644 --- a/seq_smith/_seq_smith.pyi +++ b/seq_smith/_seq_smith.pyi @@ -125,6 +125,16 @@ def top_k_ungapped_local_align( filter_overlap_a: bool = True, filter_overlap_b: bool = True, ) -> list[Alignment]: ... +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]: ... def top_k_ungapped_local_align_many( seqa: bytes, seqbs: Sequence[bytes], diff --git a/src/lib.rs b/src/lib.rs index 81a1627..f769912 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -6,7 +6,7 @@ use pyo3::wrap_pyfunction; use pyo3_stub_gen::{define_stub_info_gatherer, derive::*}; use rayon::prelude::*; use std::cmp::Ordering; -use std::collections::BinaryHeap; +use std::collections::{BinaryHeap, HashMap, HashSet}; /// Represents the type of an alignment fragment. #[gen_stub_pyclass_enum] @@ -1008,85 +1008,88 @@ impl PartialOrd for Candidate { } } -fn _top_k_ungapped_local_align_core( - params: UngappedAlignmentParams, - k: usize, - filter_overlap_a: bool, - filter_overlap_b: bool, -) -> PyResult> { - let sa_len = params.sa.len(); - let sb_len = params.sb.len(); - - let mut candidates: BinaryHeap = BinaryHeap::new(); +#[inline] +fn push_candidate( + candidates: &mut BinaryHeap, + score: i32, + sa_start: usize, + sb_start: usize, + len: usize, +) { + if score > 0 { + candidates.push(Candidate { + score, + sa_start, + sb_start, + len, + }); + } +} - let mut add_candidate = |score: i32, sa_start: usize, sb_start: usize, len: usize| { - if score > 0 { - candidates.push(Candidate { - score, - sa_start, - sb_start, - len, - }); +// Smith-Waterman-style positive-segment scan along a single diagonal of the +// (sb x sa) grid, pushing every positive HSP peak it finds into `candidates`. +// The diagonal starts at (start_row, start_col) and runs for `max_len` cells. +#[inline] +fn process_diagonal_into_candidates( + params: &UngappedAlignmentParams, + candidates: &mut BinaryHeap, + start_row: usize, + start_col: usize, + max_len: usize, +) { + let mut curr_score: i32 = 0; + let mut segment_start_idx: usize = 0; // index along diagonal where current positive segment started + let mut peak_score: i32 = 0; + let mut peak_idx: usize = 0; // index along diagonal where peak occurred + + for i in 0..max_len { + let row = start_row + i; + let col = start_col + i; + let val = params.match_score(row, col); + + if curr_score == 0 && val <= 0 { + continue; } - }; - - let mut process_diagonal = |start_row: usize, start_col: usize, max_len: usize| { - let mut curr_score = 0; - let mut segment_start_idx = 0; // index along diagonal where current positive segment started - let mut peak_score = 0; - let mut peak_idx = 0; // index along diagonal where peak occurred - - for i in 0..max_len { - let row = start_row + i; - let col = start_col + i; - let val = params.match_score(row, col); - - if curr_score == 0 && val <= 0 { - continue; - } - if curr_score == 0 { - segment_start_idx = i; - } - - curr_score += val; - - if curr_score <= 0 { - add_candidate( - peak_score, - start_col + segment_start_idx, - start_row + segment_start_idx, - peak_idx - segment_start_idx + 1, - ); - curr_score = 0; - peak_score = 0; - } else { - if curr_score > peak_score { - peak_score = curr_score; - peak_idx = i; - } - } + if curr_score == 0 { + segment_start_idx = i; } - add_candidate( - peak_score, - start_col + segment_start_idx, - start_row + segment_start_idx, - peak_idx - segment_start_idx + 1, - ); - }; - - // Diagonals starting at first row (row=0, col=0..sa_len) - for start_col in 0..sa_len { - let max_len = std::cmp::min(sa_len - start_col, sb_len); - process_diagonal(0, start_col, max_len); - } - // Diagonals starting at first column (row=1..sb_len, col=0) - for start_row in 1..sb_len { - let max_len = std::cmp::min(sa_len, sb_len - start_row); - process_diagonal(start_row, 0, max_len); + curr_score += val; + + if curr_score <= 0 { + push_candidate( + candidates, + peak_score, + start_col + segment_start_idx, + start_row + segment_start_idx, + peak_idx - segment_start_idx + 1, + ); + curr_score = 0; + peak_score = 0; + } else if curr_score > peak_score { + peak_score = curr_score; + peak_idx = i; + } } + push_candidate( + candidates, + peak_score, + start_col + segment_start_idx, + start_row + segment_start_idx, + peak_idx - segment_start_idx + 1, + ); +} - // Select top k non-overlapping +// Pop top-k highest-scoring candidates from the heap, dropping any that overlap +// an already-accepted candidate on the A and/or B axis. Returns Alignments in +// score-descending order. +fn select_top_k_with_overlap_filter( + params: &UngappedAlignmentParams, + mut candidates: BinaryHeap, + k: usize, + filter_overlap_a: bool, + filter_overlap_b: bool, +) -> Vec { let mut alignments: Vec = Vec::with_capacity(k); while alignments.len() < k { @@ -1134,7 +1137,7 @@ fn _top_k_ungapped_local_align_core( len: candidate.len as i32, }], score: candidate.score, - stats: stats, + stats, }); } } else { @@ -1142,7 +1145,136 @@ fn _top_k_ungapped_local_align_core( } } - Ok(alignments) + alignments +} + +fn _top_k_ungapped_local_align_core( + params: UngappedAlignmentParams, + k: usize, + filter_overlap_a: bool, + filter_overlap_b: bool, +) -> PyResult> { + let sa_len = params.sa.len(); + let sb_len = params.sb.len(); + + let mut candidates: BinaryHeap = BinaryHeap::new(); + + // Diagonals starting at first row (row=0, col=0..sa_len) + for start_col in 0..sa_len { + let max_len = std::cmp::min(sa_len - start_col, sb_len); + process_diagonal_into_candidates(¶ms, &mut candidates, 0, start_col, max_len); + } + + // Diagonals starting at first column (row=1..sb_len, col=0) + for start_row in 1..sb_len { + let max_len = std::cmp::min(sa_len, sb_len - start_row); + process_diagonal_into_candidates(¶ms, &mut candidates, start_row, 0, max_len); + } + + Ok(select_top_k_with_overlap_filter( + ¶ms, + candidates, + k, + filter_overlap_a, + filter_overlap_b, + )) +} + +// K-mer-seeded variant of `_top_k_ungapped_local_align_core`. Builds an exact- +// match k-mer index of `seqa`, finds every k-mer hit between `seqa` and `seqb`, +// and runs the diagonal positive-segment scan only on diagonals that contain +// at least one such hit. Skips diagonals where no k-mer of length `kmer_size` +// matches between the two sequences, which is the source of the speedup over +// the all-diagonals scan. +// +// Correctness: an HSP on a diagonal is found iff at least one window of +// `kmer_size` consecutive exact-match bytes lies on that diagonal between the +// two sequences. HSPs whose match runs are all strictly shorter than +// `kmer_size` are missed. For text alphabets and `kmer_size <= 5`, this only +// occurs at very small HSP scores. +fn _top_k_ungapped_local_align_kmer_core( + params: UngappedAlignmentParams, + k: usize, + kmer_size: usize, + max_hits_per_kmer: usize, + filter_overlap_a: bool, + filter_overlap_b: bool, +) -> PyResult> { + if kmer_size == 0 || kmer_size > 8 { + return Err(PyErr::new::( + "kmer_size must be in [1, 8]", + )); + } + + let sa_len = params.sa.len(); + let sb_len = params.sb.len(); + + if sa_len < kmer_size || sb_len < kmer_size { + return Ok(Vec::new()); + } + + // Pack a k-mer of size <= 8 into a u64. The mask isolates the active bytes. + let kmer_mask: u64 = if kmer_size == 8 { + u64::MAX + } else { + (1u64 << (8 * kmer_size)) - 1 + }; + + // K-mer index of seqa: rolling-hashed k-mer -> positions in sa. + let mut kmer_index: HashMap> = HashMap::new(); + { + let mut rolling: u64 = 0; + for (i, &byte) in params.sa.iter().enumerate() { + rolling = ((rolling << 8) | (byte as u64)) & kmer_mask; + if i + 1 >= kmer_size { + let pos = (i + 1 - kmer_size) as u32; + kmer_index.entry(rolling).or_default().push(pos); + } + } + } + + // Collect the set of diagonals that have at least one k-mer hit. + // diag = sa_pos - sb_pos, range [-(sb_len-1), sa_len-1]. + let mut hit_diagonals: HashSet = HashSet::new(); + { + let mut rolling: u64 = 0; + for (i, &byte) in params.sb.iter().enumerate() { + rolling = ((rolling << 8) | (byte as u64)) & kmer_mask; + if i + 1 >= kmer_size { + let sb_pos = (i + 1 - kmer_size) as i64; + if let Some(positions) = kmer_index.get(&rolling) { + // Cap on hits-per-kmer protects against quadratic blowup + // on low-complexity stretches (e.g. long runs of spaces). + if positions.len() > max_hits_per_kmer { + continue; + } + for &sa_pos in positions { + let diag = (sa_pos as i64) - sb_pos; + hit_diagonals.insert(diag); + } + } + } + } + } + + let mut candidates: BinaryHeap = BinaryHeap::new(); + for &diag in &hit_diagonals { + let (start_row, start_col) = if diag >= 0 { + (0usize, diag as usize) + } else { + ((-diag) as usize, 0usize) + }; + let max_len = std::cmp::min(sa_len - start_col, sb_len - start_row); + process_diagonal_into_candidates(¶ms, &mut candidates, start_row, start_col, max_len); + } + + Ok(select_top_k_with_overlap_filter( + ¶ms, + candidates, + k, + filter_overlap_a, + filter_overlap_b, + )) } /// Finds the top-k non-overlapping ungapped local alignments (HSPs). @@ -1236,6 +1368,62 @@ fn top_k_ungapped_local_align_many<'py>( }) } +/// Finds the top-k non-overlapping ungapped local alignments (HSPs) using k-mer seeding. +/// +/// Functionally similar to `top_k_ungapped_local_align` but much faster on long inputs +/// over moderate-sized alphabets: HSPs are sought only on diagonals of the n*m grid that +/// contain at least one exact k-mer match between `seqa` and `seqb`, instead of scanning +/// every diagonal. +/// +/// Correctness: an HSP is found 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 and +/// `kmer_size <= 5`, this only occurs at very small HSP scores. +/// +/// Args: +/// seqa (bytes): The first sequence. +/// seqb (bytes): The second sequence. +/// score_matrix (numpy.ndarray): Scoring matrix. +/// k (int): Maximum number of alignments to return. +/// kmer_size (int): Seed length in bytes; must be in [1, 8]. +/// max_hits_per_kmer (int): Skip k-mers that appear more than this many times in +/// `seqa`; protects against quadratic blowup on low-complexity stretches +/// (long runs of spaces, common short fragments, etc.). +/// filter_overlap_a (bool): Drop later HSPs that overlap an accepted HSP on A. +/// filter_overlap_b (bool): Drop later HSPs that overlap an accepted HSP on B. +/// +/// Returns: +/// list[Alignment]: Up to `k` non-overlapping alignments, in descending score order. +#[gen_stub_pyfunction] +#[pyfunction] +#[pyo3(signature = (seqa, seqb, score_matrix, k, kmer_size, max_hits_per_kmer, filter_overlap_a=true, filter_overlap_b=true))] +fn top_k_ungapped_local_align_kmer<'py>( + py: Python<'py>, + seqa: &Bound<'py, PyBytes>, + seqb: &Bound<'py, PyBytes>, + score_matrix: PyReadonlyArray2, + k: usize, + kmer_size: usize, + max_hits_per_kmer: usize, + filter_overlap_a: bool, + filter_overlap_b: bool, +) -> PyResult> { + let seqa = seqa.as_bytes().to_vec(); + let seqb = seqb.as_bytes().to_vec(); + let score_matrix = score_matrix.as_array().into_owned(); + + py.detach(move || { + _top_k_ungapped_local_align_kmer_core( + UngappedAlignmentParams::new(&seqa, &seqb, &score_matrix)?, + k, + kmer_size, + max_hits_per_kmer, + filter_overlap_a, + filter_overlap_b, + ) + }) +} + #[pymodule] fn _seq_smith(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(local_align))?; @@ -1248,6 +1436,7 @@ fn _seq_smith(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(overlap_align_many))?; m.add_wrapped(wrap_pyfunction!(top_k_ungapped_local_align))?; m.add_wrapped(wrap_pyfunction!(top_k_ungapped_local_align_many))?; + m.add_wrapped(wrap_pyfunction!(top_k_ungapped_local_align_kmer))?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/tests/test_seq_smith.py b/tests/test_seq_smith.py index 4911f2e..d3e6e29 100644 --- a/tests/test_seq_smith.py +++ b/tests/test_seq_smith.py @@ -14,6 +14,7 @@ make_score_matrix, overlap_align, top_k_ungapped_local_align, + top_k_ungapped_local_align_kmer, top_k_ungapped_local_align_many, ) @@ -589,3 +590,142 @@ def test_top_k_ungapped_many_simple() -> None: # So should be empty if score <= 0. # Our implementation returns empty if no positive peaks. assert len(alignments_list[1]) == 0 + + +# ------------------------------------------------------------------------- +# top_k_ungapped_local_align_kmer +# ------------------------------------------------------------------------- + + +def test_top_k_ungapped_kmer_simple() -> None: + """K-mer seeding finds the same two HSPs as the all-diagonals scan.""" + alphabet = "ACGT" + seqa = encode("AAAATTTTCCCC", alphabet) + seqb = encode("AAAAGGGGCCCC", alphabet) + score_matrix = make_score_matrix(alphabet, match_score=2, mismatch_score=-5) + + alignments = top_k_ungapped_local_align_kmer( + seqa, seqb, score_matrix, k=5, kmer_size=3, max_hits_per_kmer=100, + ) + + assert len(alignments) == 2 + assert alignments[0].score == 8 + assert alignments[1].score == 8 + starts = sorted([(a.fragments[0].sa_start, a.fragments[0].sb_start) for a in alignments]) + assert starts == [(0, 0), (8, 8)] + + +def test_top_k_ungapped_kmer_matches_full_scan_random_dna() -> None: + """The k-mer seeder must agree with the all-diagonals scan on any input + where HSP score >= 8 (well above the floor where seed misses can occur).""" + rng = np.random.default_rng(42) + alphabet = "ACGT" + seqa = bytes(rng.integers(0, 4, size=400, dtype=np.uint8)) + seqb = bytes(rng.integers(0, 4, size=350, dtype=np.uint8)) + score_matrix = make_score_matrix(alphabet, match_score=1, mismatch_score=-1) + + full = top_k_ungapped_local_align(seqa, seqb, score_matrix, k=20) + kmer = top_k_ungapped_local_align_kmer( + seqa, seqb, score_matrix, k=20, kmer_size=3, max_hits_per_kmer=10_000, + ) + + # On a tiny 4-letter alphabet at k=3 the seeder is dense enough that we + # expect every HSP scoring >= 5 to be caught. Compare the top-scoring + # HSPs above that floor. + def key(a) -> tuple[int, int, int]: + return (-a.score, a.fragments[0].sa_start, a.fragments[0].sb_start) + + full_top = sorted([a for a in full if a.score >= 5], key=key) + kmer_top = sorted([a for a in kmer if a.score >= 5], key=key) + assert [(a.score, a.fragments[0].sa_start, a.fragments[0].sb_start, a.fragments[0].len) for a in full_top] == [ + (a.score, a.fragments[0].sa_start, a.fragments[0].sb_start, a.fragments[0].len) for a in kmer_top + ] + + +def test_top_k_ungapped_kmer_overlap() -> None: + """Overlap filtering on B works the same as in the all-diagonals scan.""" + alphabet = "ACGT" + seqa = encode("AAAATTTTCCCCAAAATTTTCCCCAAAATTTTCCCC", alphabet) + seqb = encode("AAAAGGGGCCCC", alphabet) + score_matrix = make_score_matrix(alphabet, match_score=2, mismatch_score=-5) + + alignments = top_k_ungapped_local_align_kmer( + seqa, seqb, score_matrix, k=5, + kmer_size=3, max_hits_per_kmer=100, filter_overlap_b=False, + ) + + assert len(alignments) == 5 + assert all(a.score == 8 for a in alignments) + starts = sorted([(a.fragments[0].sa_start, a.fragments[0].sb_start) for a in alignments]) + for c, r in starts: + assert seqa[c : c + 4] == seqb[r : r + 4] + + +def test_top_k_ungapped_kmer_limit() -> None: + """Top-k truncation works the same as in the all-diagonals scan.""" + alphabet = "ACGT" + seqa = encode("AATTCCTTGG", alphabet) + seqb = encode("AAGGCCGGGG", alphabet) + score_matrix = make_score_matrix(alphabet, match_score=2, mismatch_score=-5) + + alignments = top_k_ungapped_local_align_kmer( + seqa, seqb, score_matrix, k=2, kmer_size=2, max_hits_per_kmer=100, + ) + + assert len(alignments) == 2 + assert alignments[0].score == 4 + assert alignments[1].score == 4 + + +def test_top_k_ungapped_kmer_invalid_kmer_size() -> None: + """kmer_size must be in [1, 8]; outside that range raises ValueError.""" + alphabet = "ACGT" + seqa = encode("ACGT", alphabet) + seqb = encode("ACGT", alphabet) + score_matrix = make_score_matrix(alphabet, 1, -1) + + with pytest.raises(ValueError, match="kmer_size"): + top_k_ungapped_local_align_kmer( + seqa, seqb, score_matrix, k=1, kmer_size=0, max_hits_per_kmer=10, + ) + with pytest.raises(ValueError, match="kmer_size"): + top_k_ungapped_local_align_kmer( + seqa, seqb, score_matrix, k=1, kmer_size=9, max_hits_per_kmer=10, + ) + + +def test_top_k_ungapped_kmer_seqs_shorter_than_kmer() -> None: + """Sequences shorter than `kmer_size` produce no HSPs (no k-mer index).""" + alphabet = "ACGT" + seqa = encode("AC", alphabet) + seqb = encode("AC", alphabet) + score_matrix = make_score_matrix(alphabet, 1, -1) + + alignments = top_k_ungapped_local_align_kmer( + seqa, seqb, score_matrix, k=5, kmer_size=4, max_hits_per_kmer=10, + ) + assert alignments == [] + + +def test_top_k_ungapped_kmer_max_hits_skips_low_complexity() -> None: + """A k-mer occurring more than `max_hits_per_kmer` times in seqa is skipped. + + With ``max_hits_per_kmer=0`` no k-mer hit is ever followed up, so even an + identity pair returns no HSPs. This is the knob that protects against + quadratic blowup on long runs of the same character. + """ + alphabet = "ACGT" + seqa = encode("AAAAAAAAAA", alphabet) + seqb = encode("AAAAAAAAAA", alphabet) + score_matrix = make_score_matrix(alphabet, 1, -1) + + full_alignments = top_k_ungapped_local_align_kmer( + seqa, seqb, score_matrix, k=5, kmer_size=3, max_hits_per_kmer=100, + ) + assert len(full_alignments) >= 1 + assert full_alignments[0].score == 10 + + capped_alignments = top_k_ungapped_local_align_kmer( + seqa, seqb, score_matrix, k=5, kmer_size=3, max_hits_per_kmer=0, + ) + assert capped_alignments == [] From c6d44111945c80d0657cc0343202b835225f7db0 Mon Sep 17 00:00:00 2001 From: Tobias Sargeant Date: Tue, 12 May 2026 12:53:01 +1000 Subject: [PATCH 2/7] feat: add `min_kmer_hits_per_diagonal` threshold to k-mer seeder `top_k_ungapped_local_align_kmer(..., min_kmer_hits_per_diagonal=1)` now exposes a per-diagonal hit-count floor that runs the extension scan only on diagonals whose accumulated k-mer hit count is at least the given threshold. Defaults to `1` (preserves the previous "any hit triggers extension" behaviour exactly). Why: profiling on text-like inputs (markdown vs PDF-extracted text over a ~37-character alphabet) showed `max_hits_per_kmer` alone is a weak filter. At cap=16 on a 17K x 17K pair, ~34% of grid diagonals still see a hit; most of them carry only 1-2 random k-mer collisions and contribute no real HSP, but every one still pays a full positive-segment scan. Per-diagonal hit count, in contrast, is a principled signal: under +1/-1-style scoring a match-run of length L deposits `L - kmer_size + 1` hits on its diagonal, so thresholding "hits >= T" filters out diagonals whose best possible ungapped HSP score cannot exceed a small constant. This is a per-diagonal score-floor pre-filter that costs ~O(1) per diagonal versus running the extension. End-to-end on a 39-pair markdown/PDF corpus (kmer_size=5, max_hits_per_kmer=16, k=500), full-scan vs. k-mer-seeded with the new threshold: pair (n x m) full T=1 T=5 T=8 17K x 17K 2364 ms 896 ms 206 ms 91 ms (2.6x) (11x) (25x) 73K x 72K 13038 ms 4362 ms 1209 ms 599 ms (2.7x) (10x) (20x) 50K x 55K 4379 ms 2479 ms 778 ms 413 ms (1.8x) (5.6x) (11x) Score-band recall vs. the full-scan baseline: - HSPs with score >= 50: **100% recall on every pair at every T tested (1..21)**. This is the band that drives chain anchoring in seed-and-extend callers like `anchorite.chained_alignment`. - HSPs with 10 <= score < 50: > 0.92 recall up to T=5, degrades gradually thereafter. Setting `min_kmer_hits_per_diagonal=0` is the escape hatch: the function then bypasses k-mer seeding entirely and delegates to `_top_k_ungapped_local_align_core` (the exhaustive scan). Callers who need exhaustive coverage at the low-score boundary can use this without switching APIs. Internal: the per-diagonal hit-count table is a dense `Vec` of length `sa_len + sb_len - 1`, indexed by `diag + (sb_len - 1)`. At text-alphabet densities most diagonals see >= 1 hit so a HashMap is mostly populated and pays allocator overhead for no benefit; the Vec is ~580 KB worst case on the corpus and is ~10% faster end-to-end at low T. New tests in `tests/test_seq_smith.py`: - `test_top_k_ungapped_kmer_min_hits_per_diagonal_filters`: a perfect HSP of length L over `kmer_size` k deposits `L - k + 1` hits; threshold above that drops it, threshold at/below keeps it. - `test_top_k_ungapped_kmer_min_hits_zero_falls_back_to_exhaustive`: on random DNA, `min_kmer_hits_per_diagonal=0` returns exactly the same HSPs as `top_k_ungapped_local_align`, including HSPs whose match runs are shorter than `kmer_size`. All 45 tests pass. --- seq_smith/_seq_smith.pyi | 1 + src/lib.rs | 53 +++++++++++++++++++++++++++++++------ tests/test_seq_smith.py | 57 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 103 insertions(+), 8 deletions(-) diff --git a/seq_smith/_seq_smith.pyi b/seq_smith/_seq_smith.pyi index 818098e..4ffafb8 100644 --- a/seq_smith/_seq_smith.pyi +++ b/seq_smith/_seq_smith.pyi @@ -132,6 +132,7 @@ def top_k_ungapped_local_align_kmer( k: int, kmer_size: int, max_hits_per_kmer: int, + min_kmer_hits_per_diagonal: int = 1, filter_overlap_a: bool = True, filter_overlap_b: bool = True, ) -> list[Alignment]: ... diff --git a/src/lib.rs b/src/lib.rs index f769912..a4e0ed2 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -6,7 +6,7 @@ use pyo3::wrap_pyfunction; use pyo3_stub_gen::{define_stub_info_gatherer, derive::*}; use rayon::prelude::*; use std::cmp::Ordering; -use std::collections::{BinaryHeap, HashMap, HashSet}; +use std::collections::{BinaryHeap, HashMap}; /// Represents the type of an alignment fragment. #[gen_stub_pyclass_enum] @@ -1197,6 +1197,7 @@ fn _top_k_ungapped_local_align_kmer_core( k: usize, kmer_size: usize, max_hits_per_kmer: usize, + min_kmer_hits_per_diagonal: usize, filter_overlap_a: bool, filter_overlap_b: bool, ) -> PyResult> { @@ -1206,6 +1207,16 @@ fn _top_k_ungapped_local_align_kmer_core( )); } + // `min_kmer_hits_per_diagonal == 0` is the escape hatch: skip the k-mer + // pass entirely and fall back to the exhaustive all-diagonals scan. This + // gives callers a single entry point that can degrade gracefully when the + // k-mer seeding is unsafe (very short sequences, very small alphabets, or + // when the caller cares about HSPs whose match runs are all shorter than + // `kmer_size`). + if min_kmer_hits_per_diagonal == 0 { + return _top_k_ungapped_local_align_core(params, k, filter_overlap_a, filter_overlap_b); + } + let sa_len = params.sa.len(); let sb_len = params.sb.len(); @@ -1233,9 +1244,17 @@ fn _top_k_ungapped_local_align_kmer_core( } } - // Collect the set of diagonals that have at least one k-mer hit. - // diag = sa_pos - sb_pos, range [-(sb_len-1), sa_len-1]. - let mut hit_diagonals: HashSet = HashSet::new(); + // Count k-mer hits per diagonal. Diagonals are indexed as + // `diag_idx = sa_pos - sb_pos + (sb_len - 1)`, mapping the + // [-(sb_len-1), sa_len-1] range to [0, sa_len + sb_len - 2]. + // + // We use a dense Vec instead of a HashMap: at typical text-alphabet + // densities most diagonals see >=1 hit, so the table is mostly populated + // and HashMap overhead dominates. Memory is 4 * (n + m - 1) bytes, + // ~580 KB even at n = m = 73K. + let num_diagonals = sa_len + sb_len - 1; + let diag_offset = (sb_len - 1) as i64; + let mut hits_per_diagonal: Vec = vec![0; num_diagonals]; { let mut rolling: u64 = 0; for (i, &byte) in params.sb.iter().enumerate() { @@ -1249,16 +1268,26 @@ fn _top_k_ungapped_local_align_kmer_core( continue; } for &sa_pos in positions { - let diag = (sa_pos as i64) - sb_pos; - hit_diagonals.insert(diag); + let diag_idx = ((sa_pos as i64) - sb_pos + diag_offset) as usize; + hits_per_diagonal[diag_idx] += 1; } } } } } + // Per-diagonal hit count is a lower bound on the SW score reachable on + // that diagonal: a match-run of length L deposits L - kmer_size + 1 hits, + // so thresholding "hits >= T" filters out diagonals whose best possible + // ungapped HSP score cannot exceed a small constant. See the + // `min_kmer_hits_per_diagonal` arg on the public Python function. + let min_hits = min_kmer_hits_per_diagonal as u32; let mut candidates: BinaryHeap = BinaryHeap::new(); - for &diag in &hit_diagonals { + for (diag_idx, &count) in hits_per_diagonal.iter().enumerate() { + if count < min_hits { + continue; + } + let diag = (diag_idx as i64) - diag_offset; let (start_row, start_col) = if diag >= 0 { (0usize, diag as usize) } else { @@ -1389,6 +1418,12 @@ fn top_k_ungapped_local_align_many<'py>( /// max_hits_per_kmer (int): Skip k-mers that appear more than this many times in /// `seqa`; protects against quadratic blowup on low-complexity stretches /// (long runs of spaces, common short fragments, etc.). +/// min_kmer_hits_per_diagonal (int): Require a diagonal to accumulate at least +/// this many k-mer hits before running the extension scan on it. A +/// match-run of length `L` deposits `L - kmer_size + 1` hits on its +/// diagonal, so this is effectively a per-diagonal score-floor pre-filter. +/// Defaults to 1 (any diagonal with a hit is extended). Set to 0 to +/// disable k-mer seeding entirely and fall back to the exhaustive scan. /// filter_overlap_a (bool): Drop later HSPs that overlap an accepted HSP on A. /// filter_overlap_b (bool): Drop later HSPs that overlap an accepted HSP on B. /// @@ -1396,7 +1431,7 @@ fn top_k_ungapped_local_align_many<'py>( /// list[Alignment]: Up to `k` non-overlapping alignments, in descending score order. #[gen_stub_pyfunction] #[pyfunction] -#[pyo3(signature = (seqa, seqb, score_matrix, k, kmer_size, max_hits_per_kmer, filter_overlap_a=true, filter_overlap_b=true))] +#[pyo3(signature = (seqa, seqb, score_matrix, k, kmer_size, max_hits_per_kmer, min_kmer_hits_per_diagonal=1, filter_overlap_a=true, filter_overlap_b=true))] fn top_k_ungapped_local_align_kmer<'py>( py: Python<'py>, seqa: &Bound<'py, PyBytes>, @@ -1405,6 +1440,7 @@ fn top_k_ungapped_local_align_kmer<'py>( k: usize, kmer_size: usize, max_hits_per_kmer: usize, + min_kmer_hits_per_diagonal: usize, filter_overlap_a: bool, filter_overlap_b: bool, ) -> PyResult> { @@ -1418,6 +1454,7 @@ fn top_k_ungapped_local_align_kmer<'py>( k, kmer_size, max_hits_per_kmer, + min_kmer_hits_per_diagonal, filter_overlap_a, filter_overlap_b, ) diff --git a/tests/test_seq_smith.py b/tests/test_seq_smith.py index d3e6e29..3436868 100644 --- a/tests/test_seq_smith.py +++ b/tests/test_seq_smith.py @@ -729,3 +729,60 @@ def test_top_k_ungapped_kmer_max_hits_skips_low_complexity() -> None: seqa, seqb, score_matrix, k=5, kmer_size=3, max_hits_per_kmer=0, ) assert capped_alignments == [] + + +def test_top_k_ungapped_kmer_min_hits_per_diagonal_filters() -> None: + """`min_kmer_hits_per_diagonal` thresholds out low-hit-count diagonals. + + Within an HSP of length L over `kmer_size` k, the diagonal accumulates + L - k + 1 hits. Choosing a threshold above this drops the HSP, while a + threshold at or below it keeps it. + """ + alphabet = "ACGT" + # Single perfect HSP of length 6 on the main diagonal: 6 - 3 + 1 = 4 hits. + seqa = encode("AAAAAA", alphabet) + seqb = encode("AAAAAA", alphabet) + score_matrix = make_score_matrix(alphabet, 1, -1) + + kept = top_k_ungapped_local_align_kmer( + seqa, seqb, score_matrix, k=5, kmer_size=3, max_hits_per_kmer=100, + min_kmer_hits_per_diagonal=4, + ) + assert len(kept) == 1 + assert kept[0].score == 6 + + dropped = top_k_ungapped_local_align_kmer( + seqa, seqb, score_matrix, k=5, kmer_size=3, max_hits_per_kmer=100, + min_kmer_hits_per_diagonal=5, + ) + assert dropped == [] + + +def test_top_k_ungapped_kmer_min_hits_zero_falls_back_to_exhaustive() -> None: + """`min_kmer_hits_per_diagonal=0` bypasses k-mer seeding entirely. + + The function must then return exactly what `top_k_ungapped_local_align` + returns -- including HSPs whose match runs are too short to seed under + the given `kmer_size`. + """ + alphabet = "ACGT" + rng = np.random.default_rng(7) + seqa = bytes(rng.integers(0, 4, size=200, dtype=np.uint8)) + seqb = bytes(rng.integers(0, 4, size=180, dtype=np.uint8)) + score_matrix = make_score_matrix(alphabet, 1, -1) + + full = top_k_ungapped_local_align(seqa, seqb, score_matrix, k=10) + fallback = top_k_ungapped_local_align_kmer( + seqa, seqb, score_matrix, k=10, kmer_size=5, max_hits_per_kmer=100, + min_kmer_hits_per_diagonal=0, + ) + + full_sig = [ + (a.score, a.fragments[0].sa_start, a.fragments[0].sb_start, a.fragments[0].len) + for a in full + ] + fallback_sig = [ + (a.score, a.fragments[0].sa_start, a.fragments[0].sb_start, a.fragments[0].len) + for a in fallback + ] + assert fallback_sig == full_sig From 54f8aa0cb68d0939d602cf27661ae35a149fbc9b Mon Sep 17 00:00:00 2001 From: Tobias Sargeant Date: Tue, 12 May 2026 15:40:15 +1000 Subject: [PATCH 3/7] feat: span-based extension with X-drop backward for k-mer seeder Replace the per-diagonal hit-count + full-diagonal scan with BLAST-style span emission and X-drop-bounded extension. The forward extension is the same Kadane positive-segment scan as `top_k_ungapped_local_align`, so within the seeded regime the kmer fn now reports the exact same HSP scores as the all-diagonals scan -- not just an approximation. Algorithm change. Per-diagonal seeding becomes per-span: 1. Walk seqb's k-mers; for each hit at diagonal `d` and position `p`, extend an open `(start, last, count)` span if `p - last <= max_hit_gap`, otherwise close it and start a new one. Per-diagonal open-span state plus a global `Vec`. 2. Sort spans by `(diag, start_pos)`, filter by `min_kmer_hits_per_span`. 3. For each span, walk backward along the diagonal until `back_sum < back_max - max_hit_gap` (X-drop with `max_hit_gap` as the X budget). The position with maximum back_sum is `new_start` -- the position from which a forward Kadane scan reproduces the same peak as a full-from-zero scan. 4. Forward Kadane from `new_start`, terminating at the first reset that occurs at-or-past the span's last hit. Resume position becomes the dedup floor for subsequent spans on the same diagonal. The strict zero-crossing rule on the backward walk is too tight: a lead-in of the form `(k-1) matches, 1 mismatch, k matches, ...` dips `back_sum` to -1 transiently before climbing higher, and we'd terminate prematurely with an undercounted score. Reusing `max_hit_gap` as the X-drop budget keeps a single user-facing knob ("how many missing bytes am I willing to tunnel through?") and lets `back_sum` recover from short dips. API changes. - Rename `min_kmer_hits_per_diagonal` -> `min_kmer_hits_per_span` (semantic shift: each emitted span is now the unit of filtering, not the diagonal). `=0` still falls back to the exhaustive scan. - New `max_hit_gap` parameter (default 20) -- governs both span emission (close span when next hit is > max_hit_gap past last) and backward X-drop tolerance. Tune to scoring scheme: 20 is suitable for +1 / -1 over text-like alphabets. Measured end-to-end on the 39-pair markdown/PDF corpus (kmer_size=5, max_hits_per_kmer=16, max_hit_gap=20, min_kmer_hits_per_span=1, k=500), full-scan vs new kmer-seeded: pair (n x m) full new speedup recall(s>=50) 17K x 17K 2410 ms 4.2 ms 580x 1.000 11K x 11K 955 ms 2.5 ms 381x 1.000 73K x 72K 12712 ms 17.5 ms 728x 1.000 50K x 55K 4367 ms 12.0 ms 365x 1.000 Top-band (score >= 50) recall is 100% on every pair: HSPs that anchor chained alignment are all preserved. Mid-band (10 <= s < 50) recall is 0.88-0.96 at T=1 and degrades with T. Test changes. - `test_top_k_ungapped_kmer_matches_full_scan_random_dna` now passes again exactly (was failing under min/max bounding due to lead-in truncation; X-drop backward fixes it). - Rename `min_kmer_hits_per_diagonal_filters` test to `min_kmer_hits_per_span_filters`, update assertion semantics. - `min_hits_zero_falls_back_to_exhaustive` updated to the new param name; behaviour unchanged. All 45 tests pass. --- seq_smith/_seq_smith.pyi | 3 +- src/lib.rs | 289 ++++++++++++++++++++++++++++++++------- tests/test_seq_smith.py | 19 +-- 3 files changed, 255 insertions(+), 56 deletions(-) diff --git a/seq_smith/_seq_smith.pyi b/seq_smith/_seq_smith.pyi index 4ffafb8..d0e8f83 100644 --- a/seq_smith/_seq_smith.pyi +++ b/seq_smith/_seq_smith.pyi @@ -132,7 +132,8 @@ def top_k_ungapped_local_align_kmer( k: int, kmer_size: int, max_hits_per_kmer: int, - min_kmer_hits_per_diagonal: int = 1, + max_hit_gap: int = 20, + min_kmer_hits_per_span: int = 1, filter_overlap_a: bool = True, filter_overlap_b: bool = True, ) -> list[Alignment]: ... diff --git a/src/lib.rs b/src/lib.rs index a4e0ed2..ed38d48 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1030,6 +1030,107 @@ fn push_candidate( // (sb x sa) grid, pushing every positive HSP peak it finds into `candidates`. // The diagonal starts at (start_row, start_col) and runs for `max_len` cells. #[inline] +// Backward walk from `span_start` along the diagonal, returning the position +// at which a forward Kadane scan -- started fresh with curr_score = 0 -- will +// reproduce the same peak as a full forward scan from position 0. This is +// the offset where the backward-cumulative sum reached its maximum. +// +// Termination rule (X-drop): stop when `back_sum < back_max - x_drop`. A +// pure zero-crossing rule (`back_sum < 0`) is too tight: a lead-in of the +// form "(k-1) matches, 1 mismatch, k matches, ..." dips back_sum to -1 +// transiently before climbing higher. Reusing `max_hit_gap` as the X-drop +// budget keeps a single user-facing knob: "how many missing bytes am I +// willing to tunnel through to extend a seeded region?" +fn backward_to_kadane_reset( + params: &UngappedAlignmentParams, + origin_row: usize, + origin_col: usize, + span_start: usize, + x_drop: i32, +) -> usize { + let mut back_sum: i32 = 0; + let mut back_max: i32 = 0; + let mut new_start: usize = span_start; + let mut p: usize = span_start; + while p > 0 { + p -= 1; + let val = params.match_score(origin_row + p, origin_col + p); + back_sum += val; + if back_sum > back_max { + back_max = back_sum; + new_start = p; + } + if back_sum < back_max - x_drop { + break; + } + } + new_start +} + +// Per-span variant of `process_diagonal_into_candidates`. Extends the span +// backward to the nearest Kadane reset point (cf. `backward_to_kadane_reset`), +// then runs the same positive-segment scan forward, terminating at the first +// Kadane reset that occurs at-or-past `span_end`. Returns the position +// along the diagonal at or past which the next span on the same diagonal +// may safely be processed (= one past the position where Kadane reset, or +// `max_len` if no reset occurred after `span_end`). +fn process_span_into_candidates( + params: &UngappedAlignmentParams, + candidates: &mut BinaryHeap, + origin_row: usize, + origin_col: usize, + max_len: usize, + span_start: usize, + span_end: usize, + x_drop: i32, +) -> usize { + let new_start = backward_to_kadane_reset(params, origin_row, origin_col, span_start, x_drop); + + let mut curr_score: i32 = 0; + let mut segment_start_idx: usize = new_start; + let mut peak_score: i32 = 0; + let mut peak_idx: usize = new_start; + for i in new_start..max_len { + let val = params.match_score(origin_row + i, origin_col + i); + if curr_score == 0 && val <= 0 { + continue; + } + if curr_score == 0 { + segment_start_idx = i; + } + curr_score += val; + if curr_score <= 0 { + push_candidate( + candidates, + peak_score, + origin_col + segment_start_idx, + origin_row + segment_start_idx, + peak_idx - segment_start_idx + 1, + ); + curr_score = 0; + peak_score = 0; + if i >= span_end { + // Past the span's last hit and Kadane has just reset -- + // safe to stop and resume from the next position. + return i + 1; + } + } else if curr_score > peak_score { + peak_score = curr_score; + peak_idx = i; + } + } + if peak_score > 0 { + push_candidate( + candidates, + peak_score, + origin_col + segment_start_idx, + origin_row + segment_start_idx, + peak_idx - segment_start_idx + 1, + ); + } + max_len +} + fn process_diagonal_into_candidates( params: &UngappedAlignmentParams, candidates: &mut BinaryHeap, @@ -1197,7 +1298,8 @@ fn _top_k_ungapped_local_align_kmer_core( k: usize, kmer_size: usize, max_hits_per_kmer: usize, - min_kmer_hits_per_diagonal: usize, + max_hit_gap: usize, + min_kmer_hits_per_span: usize, filter_overlap_a: bool, filter_overlap_b: bool, ) -> PyResult> { @@ -1207,13 +1309,13 @@ fn _top_k_ungapped_local_align_kmer_core( )); } - // `min_kmer_hits_per_diagonal == 0` is the escape hatch: skip the k-mer - // pass entirely and fall back to the exhaustive all-diagonals scan. This - // gives callers a single entry point that can degrade gracefully when the - // k-mer seeding is unsafe (very short sequences, very small alphabets, or - // when the caller cares about HSPs whose match runs are all shorter than - // `kmer_size`). - if min_kmer_hits_per_diagonal == 0 { + // `min_kmer_hits_per_span == 0` is the escape hatch: skip the k-mer pass + // entirely and fall back to the exhaustive all-diagonals scan. Gives + // callers a single entry point that can degrade gracefully when the k-mer + // seeding is unsafe (very short sequences, very small alphabets, or when + // the caller needs exhaustive coverage of HSPs whose match runs are all + // shorter than `kmer_size`). + if min_kmer_hits_per_span == 0 { return _top_k_ungapped_local_align_core(params, k, filter_overlap_a, filter_overlap_b); } @@ -1244,17 +1346,37 @@ fn _top_k_ungapped_local_align_kmer_core( } } - // Count k-mer hits per diagonal. Diagonals are indexed as - // `diag_idx = sa_pos - sb_pos + (sb_len - 1)`, mapping the - // [-(sb_len-1), sa_len-1] range to [0, sa_len + sb_len - 2]. + // Walk seqb's k-mers and emit closed spans per diagonal. A "span" is a + // cluster of consecutive k-mer hits on the same diagonal in which no two + // adjacent hits are more than `max_hit_gap` positions apart. Diagonals + // are indexed as `diag_idx = sa_pos - sb_pos + (sb_len - 1)`, mapping + // the [-(sb_len-1), sa_len-1] range to [0, sa_len + sb_len - 2]. // - // We use a dense Vec instead of a HashMap: at typical text-alphabet - // densities most diagonals see >=1 hit, so the table is mostly populated - // and HashMap overhead dominates. Memory is 4 * (n + m - 1) bytes, - // ~580 KB even at n = m = 73K. + // Per-diagonal open span state (`count == 0` is the "no open span" + // sentinel) plus a global Vec for closed spans. Open-span state is + // ~1.7 MB at n = m = 73K. + #[derive(Clone, Copy)] + struct OpenSpan { + count: u32, + start_pos: u32, + last_pos: u32, + } + #[derive(Clone, Copy)] + struct ClosedSpan { + diag_idx: u32, + start_pos: u32, + end_pos: u32, + count: u32, + } let num_diagonals = sa_len + sb_len - 1; let diag_offset = (sb_len - 1) as i64; - let mut hits_per_diagonal: Vec = vec![0; num_diagonals]; + let mut open_spans: Vec = vec![ + OpenSpan { count: 0, start_pos: 0, last_pos: 0 }; + num_diagonals + ]; + let mut closed_spans: Vec = Vec::new(); + let max_hit_gap_u32 = max_hit_gap as u32; + { let mut rolling: u64 = 0; for (i, &byte) in params.sb.iter().enumerate() { @@ -1269,32 +1391,93 @@ fn _top_k_ungapped_local_align_kmer_core( } for &sa_pos in positions { let diag_idx = ((sa_pos as i64) - sb_pos + diag_offset) as usize; - hits_per_diagonal[diag_idx] += 1; + // Position along the diagonal of the k-mer's start cell. + // For diag >= 0 the diagonal origin is (0, diag), so + // this is sb_pos; for diag < 0 the origin is (-diag, 0) + // and it is sa_pos. Both cases reduce to min(sa_pos, sb_pos). + let pos = std::cmp::min(sa_pos, sb_pos as u32); + let open = &mut open_spans[diag_idx]; + if open.count == 0 { + open.count = 1; + open.start_pos = pos; + open.last_pos = pos; + } else if pos - open.last_pos > max_hit_gap_u32 { + closed_spans.push(ClosedSpan { + diag_idx: diag_idx as u32, + start_pos: open.start_pos, + end_pos: open.last_pos, + count: open.count, + }); + open.count = 1; + open.start_pos = pos; + open.last_pos = pos; + } else { + open.count += 1; + open.last_pos = pos; + } } } } } } - // Per-diagonal hit count is a lower bound on the SW score reachable on - // that diagonal: a match-run of length L deposits L - kmer_size + 1 hits, - // so thresholding "hits >= T" filters out diagonals whose best possible - // ungapped HSP score cannot exceed a small constant. See the - // `min_kmer_hits_per_diagonal` arg on the public Python function. - let min_hits = min_kmer_hits_per_diagonal as u32; + // Flush remaining open spans. + for (diag_idx, open) in open_spans.iter().enumerate() { + if open.count > 0 { + closed_spans.push(ClosedSpan { + diag_idx: diag_idx as u32, + start_pos: open.start_pos, + end_pos: open.last_pos, + count: open.count, + }); + } + } + drop(open_spans); + + // Process spans per diagonal in order. Hit count on a span is a lower + // bound on the SW score reachable from its seeds: a match-run of length + // L deposits L - kmer_size + 1 hits, so the `min_kmer_hits_per_span` + // threshold is effectively a per-span score-floor pre-filter. + // + // Dedup: forward Kadane from a span's extension can run past subsequent + // spans' starts on the same diagonal. We track `resume_pos` per diagonal + // (= position past which the next span may safely be processed) and skip + // spans whose `start_pos` is below it. + closed_spans.sort_unstable_by_key(|s| (s.diag_idx, s.start_pos)); + let min_count = min_kmer_hits_per_span as u32; let mut candidates: BinaryHeap = BinaryHeap::new(); - for (diag_idx, &count) in hits_per_diagonal.iter().enumerate() { - if count < min_hits { + let mut current_diag: u32 = u32::MAX; + let mut resume_pos: usize = 0; + for span in &closed_spans { + if span.count < min_count { + continue; + } + if span.diag_idx != current_diag { + current_diag = span.diag_idx; + resume_pos = 0; + } + let span_start = span.start_pos as usize; + let span_end = span.end_pos as usize; + if span_start < resume_pos { continue; } - let diag = (diag_idx as i64) - diag_offset; - let (start_row, start_col) = if diag >= 0 { + let diag = (span.diag_idx as i64) - diag_offset; + let (origin_row, origin_col) = if diag >= 0 { (0usize, diag as usize) } else { ((-diag) as usize, 0usize) }; - let max_len = std::cmp::min(sa_len - start_col, sb_len - start_row); - process_diagonal_into_candidates(¶ms, &mut candidates, start_row, start_col, max_len); + let max_len = std::cmp::min(sa_len - origin_col, sb_len - origin_row); + resume_pos = process_span_into_candidates( + ¶ms, + &mut candidates, + origin_row, + origin_col, + max_len, + span_start, + span_end, + max_hit_gap as i32, + ); } Ok(select_top_k_with_overlap_filter( @@ -1399,15 +1582,22 @@ fn top_k_ungapped_local_align_many<'py>( /// Finds the top-k non-overlapping ungapped local alignments (HSPs) using k-mer seeding. /// -/// Functionally similar to `top_k_ungapped_local_align` but much faster on long inputs -/// over moderate-sized alphabets: HSPs are sought only on diagonals of the n*m grid that -/// contain at least one exact k-mer match between `seqa` and `seqb`, instead of scanning -/// every diagonal. +/// Functionally equivalent to `top_k_ungapped_local_align` but much faster on long +/// inputs over moderate-sized alphabets. The algorithm is BLAST-style: +/// +/// 1. Build a rolling-hash k-mer index of `seqa`. +/// 2. Walk `seqb`, collecting k-mer hits into per-diagonal spans -- consecutive +/// hits on the same diagonal whose positions are within `max_hit_gap`. +/// 3. For each span with at least `min_kmer_hits_per_span` hits, extend +/// backward along the diagonal to the nearest Kadane reset point and forward +/// with the same positive-segment scan as `top_k_ungapped_local_align`, +/// terminating at the first Kadane reset past the span's last hit. +/// 4. Pop top-`k` non-overlapping from the candidate heap. /// -/// Correctness: an HSP is found 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 and -/// `kmer_size <= 5`, this only occurs at very small HSP scores. +/// Correctness: HSPs whose match runs are all strictly shorter than `kmer_size` +/// are not seeded and so are not found. For text alphabets and `kmer_size <= 5`, +/// this only occurs at very small HSP scores. Within the seeded regime, the +/// extension reports the exact same scores as the all-diagonals scan. /// /// Args: /// seqa (bytes): The first sequence. @@ -1418,12 +1608,17 @@ fn top_k_ungapped_local_align_many<'py>( /// max_hits_per_kmer (int): Skip k-mers that appear more than this many times in /// `seqa`; protects against quadratic blowup on low-complexity stretches /// (long runs of spaces, common short fragments, etc.). -/// min_kmer_hits_per_diagonal (int): Require a diagonal to accumulate at least -/// this many k-mer hits before running the extension scan on it. A -/// match-run of length `L` deposits `L - kmer_size + 1` hits on its -/// diagonal, so this is effectively a per-diagonal score-floor pre-filter. -/// Defaults to 1 (any diagonal with a hit is extended). Set to 0 to -/// disable k-mer seeding entirely and fall back to the exhaustive scan. +/// max_hit_gap (int): Maximum allowed gap between consecutive k-mer hits on the +/// same diagonal for them to be merged into one span. Larger values merge +/// more aggressively (fewer, larger spans -> more work per span); smaller +/// values fragment more (more spans -> more book-keeping). Tune to your +/// scoring scheme. Defaults to 20, suitable for +1 / -1 over text-like +/// alphabets. +/// min_kmer_hits_per_span (int): Require a span to accumulate at least this many +/// k-mer hits before extending it. A match-run of length `L` deposits +/// `L - kmer_size + 1` hits, so this is a per-span score-floor pre-filter. +/// Defaults to 1 (any span is extended). Set to 0 to disable k-mer seeding +/// entirely and fall back to the exhaustive all-diagonals scan. /// filter_overlap_a (bool): Drop later HSPs that overlap an accepted HSP on A. /// filter_overlap_b (bool): Drop later HSPs that overlap an accepted HSP on B. /// @@ -1431,7 +1626,7 @@ fn top_k_ungapped_local_align_many<'py>( /// list[Alignment]: Up to `k` non-overlapping alignments, in descending score order. #[gen_stub_pyfunction] #[pyfunction] -#[pyo3(signature = (seqa, seqb, score_matrix, k, kmer_size, max_hits_per_kmer, min_kmer_hits_per_diagonal=1, filter_overlap_a=true, filter_overlap_b=true))] +#[pyo3(signature = (seqa, seqb, score_matrix, k, kmer_size, max_hits_per_kmer, max_hit_gap=20, min_kmer_hits_per_span=1, filter_overlap_a=true, filter_overlap_b=true))] fn top_k_ungapped_local_align_kmer<'py>( py: Python<'py>, seqa: &Bound<'py, PyBytes>, @@ -1440,7 +1635,8 @@ fn top_k_ungapped_local_align_kmer<'py>( k: usize, kmer_size: usize, max_hits_per_kmer: usize, - min_kmer_hits_per_diagonal: usize, + max_hit_gap: usize, + min_kmer_hits_per_span: usize, filter_overlap_a: bool, filter_overlap_b: bool, ) -> PyResult> { @@ -1454,7 +1650,8 @@ fn top_k_ungapped_local_align_kmer<'py>( k, kmer_size, max_hits_per_kmer, - min_kmer_hits_per_diagonal, + max_hit_gap, + min_kmer_hits_per_span, filter_overlap_a, filter_overlap_b, ) diff --git a/tests/test_seq_smith.py b/tests/test_seq_smith.py index 3436868..687a81a 100644 --- a/tests/test_seq_smith.py +++ b/tests/test_seq_smith.py @@ -731,35 +731,36 @@ def test_top_k_ungapped_kmer_max_hits_skips_low_complexity() -> None: assert capped_alignments == [] -def test_top_k_ungapped_kmer_min_hits_per_diagonal_filters() -> None: - """`min_kmer_hits_per_diagonal` thresholds out low-hit-count diagonals. +def test_top_k_ungapped_kmer_min_hits_per_span_filters() -> None: + """`min_kmer_hits_per_span` thresholds out low-hit-count spans. - Within an HSP of length L over `kmer_size` k, the diagonal accumulates - L - k + 1 hits. Choosing a threshold above this drops the HSP, while a + Within an HSP of length L over `kmer_size` k, the span accumulates + L - k + 1 hits. Choosing a threshold above this drops the span, while a threshold at or below it keeps it. """ alphabet = "ACGT" - # Single perfect HSP of length 6 on the main diagonal: 6 - 3 + 1 = 4 hits. + # Single perfect HSP of length 6 on the main diagonal: 6 - 3 + 1 = 4 hits, + # all in one span (adjacent positions, so well under max_hit_gap). seqa = encode("AAAAAA", alphabet) seqb = encode("AAAAAA", alphabet) score_matrix = make_score_matrix(alphabet, 1, -1) kept = top_k_ungapped_local_align_kmer( seqa, seqb, score_matrix, k=5, kmer_size=3, max_hits_per_kmer=100, - min_kmer_hits_per_diagonal=4, + min_kmer_hits_per_span=4, ) assert len(kept) == 1 assert kept[0].score == 6 dropped = top_k_ungapped_local_align_kmer( seqa, seqb, score_matrix, k=5, kmer_size=3, max_hits_per_kmer=100, - min_kmer_hits_per_diagonal=5, + min_kmer_hits_per_span=5, ) assert dropped == [] def test_top_k_ungapped_kmer_min_hits_zero_falls_back_to_exhaustive() -> None: - """`min_kmer_hits_per_diagonal=0` bypasses k-mer seeding entirely. + """`min_kmer_hits_per_span=0` bypasses k-mer seeding entirely. The function must then return exactly what `top_k_ungapped_local_align` returns -- including HSPs whose match runs are too short to seed under @@ -774,7 +775,7 @@ def test_top_k_ungapped_kmer_min_hits_zero_falls_back_to_exhaustive() -> None: full = top_k_ungapped_local_align(seqa, seqb, score_matrix, k=10) fallback = top_k_ungapped_local_align_kmer( seqa, seqb, score_matrix, k=10, kmer_size=5, max_hits_per_kmer=100, - min_kmer_hits_per_diagonal=0, + min_kmer_hits_per_span=0, ) full_sig = [ From ae19c16a4f4b802a8bc2c168427b9e3f80a7c588 Mon Sep 17 00:00:00 2001 From: Tobias Sargeant Date: Tue, 12 May 2026 17:28:29 +1000 Subject: [PATCH 4/7] Bump version to 0.6.0 Minor release covering the k-mer-seeded top-k ungapped local alignment work merged on this branch: - 3be4d56 feat: add k-mer-seeded top-k ungapped local alignment (top_k_ungapped_local_align_kmer; BLAST-style k-mer index of seqa + per-diagonal positive-segment scan; behaviour- preserving extraction of process_diagonal_into_candidates and select_top_k_with_overlap_filter as shared helpers.) - 51207a0 feat: add min_kmer_hits_per_diagonal threshold (Per-diagonal hit-count floor as a score-floor pre-filter; min_kmer_hits_per_diagonal=0 falls back to the exhaustive all-diagonals scan; dense Vec instead of HashMap.) - 7b8e7d2 feat: span-based extension with X-drop backward (Replaced per-diagonal hit-count + full-diagonal scan with BLAST-style span emission and X-drop-bounded extension; renamed min_kmer_hits_per_diagonal -> min_kmer_hits_per_span; added max_hit_gap parameter; 365-728x end-to-end speedup vs. top_k_ungapped_local_align on the markdown/PDF corpus with 100% recall of HSPs at score >= 50.) No breaking changes for existing callers of `top_k_ungapped_local_align` or other public APIs. `top_k_ungapped_local_align_kmer` is a new entry point introduced in this release. --- Cargo.toml | 2 +- pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 15c3f4a..5a6f146 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "seq_smith" -version = "0.5.1" +version = "0.6.0" edition = "2021" [dependencies] diff --git a/pyproject.toml b/pyproject.toml index e7663a7..0cd68a5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "maturin" [project] name = "seq-smith" -version = "0.5.1" +version = "0.6.0" authors = [ { name = "Tobias Sargeant", email = "tobias.sargeant@gmail.com" }, ] From 9e2974a1d873289b685b9ae0cd49f36087bfcc23 Mon Sep 17 00:00:00 2001 From: Tobias Sargeant Date: Tue, 12 May 2026 18:45:10 +1000 Subject: [PATCH 5/7] Fix pre-commit: ruff-format + ANN001 on test key fn - ruff-format reformatted multi-arg call sites to one-arg-per-line in the new k-mer tests (matches the file's prevailing style). - Add `Alignment` type annotation on the inner `key` function in `test_top_k_ungapped_kmer_matches_full_scan_random_dna` (ANN001). --- tests/test_seq_smith.py | 99 ++++++++++++++++++++++++++++++++--------- 1 file changed, 77 insertions(+), 22 deletions(-) diff --git a/tests/test_seq_smith.py b/tests/test_seq_smith.py index 687a81a..4dfd011 100644 --- a/tests/test_seq_smith.py +++ b/tests/test_seq_smith.py @@ -3,6 +3,7 @@ from conftest import AlignmentData from seq_smith import ( + Alignment, AlignmentFragment, FragmentType, encode, @@ -605,7 +606,12 @@ def test_top_k_ungapped_kmer_simple() -> None: score_matrix = make_score_matrix(alphabet, match_score=2, mismatch_score=-5) alignments = top_k_ungapped_local_align_kmer( - seqa, seqb, score_matrix, k=5, kmer_size=3, max_hits_per_kmer=100, + seqa, + seqb, + score_matrix, + k=5, + kmer_size=3, + max_hits_per_kmer=100, ) assert len(alignments) == 2 @@ -626,13 +632,18 @@ def test_top_k_ungapped_kmer_matches_full_scan_random_dna() -> None: full = top_k_ungapped_local_align(seqa, seqb, score_matrix, k=20) kmer = top_k_ungapped_local_align_kmer( - seqa, seqb, score_matrix, k=20, kmer_size=3, max_hits_per_kmer=10_000, + seqa, + seqb, + score_matrix, + k=20, + kmer_size=3, + max_hits_per_kmer=10_000, ) # On a tiny 4-letter alphabet at k=3 the seeder is dense enough that we # expect every HSP scoring >= 5 to be caught. Compare the top-scoring # HSPs above that floor. - def key(a) -> tuple[int, int, int]: + def key(a: Alignment) -> tuple[int, int, int]: return (-a.score, a.fragments[0].sa_start, a.fragments[0].sb_start) full_top = sorted([a for a in full if a.score >= 5], key=key) @@ -650,8 +661,13 @@ def test_top_k_ungapped_kmer_overlap() -> None: score_matrix = make_score_matrix(alphabet, match_score=2, mismatch_score=-5) alignments = top_k_ungapped_local_align_kmer( - seqa, seqb, score_matrix, k=5, - kmer_size=3, max_hits_per_kmer=100, filter_overlap_b=False, + seqa, + seqb, + score_matrix, + k=5, + kmer_size=3, + max_hits_per_kmer=100, + filter_overlap_b=False, ) assert len(alignments) == 5 @@ -669,7 +685,12 @@ def test_top_k_ungapped_kmer_limit() -> None: score_matrix = make_score_matrix(alphabet, match_score=2, mismatch_score=-5) alignments = top_k_ungapped_local_align_kmer( - seqa, seqb, score_matrix, k=2, kmer_size=2, max_hits_per_kmer=100, + seqa, + seqb, + score_matrix, + k=2, + kmer_size=2, + max_hits_per_kmer=100, ) assert len(alignments) == 2 @@ -686,11 +707,21 @@ def test_top_k_ungapped_kmer_invalid_kmer_size() -> None: with pytest.raises(ValueError, match="kmer_size"): top_k_ungapped_local_align_kmer( - seqa, seqb, score_matrix, k=1, kmer_size=0, max_hits_per_kmer=10, + seqa, + seqb, + score_matrix, + k=1, + kmer_size=0, + max_hits_per_kmer=10, ) with pytest.raises(ValueError, match="kmer_size"): top_k_ungapped_local_align_kmer( - seqa, seqb, score_matrix, k=1, kmer_size=9, max_hits_per_kmer=10, + seqa, + seqb, + score_matrix, + k=1, + kmer_size=9, + max_hits_per_kmer=10, ) @@ -702,7 +733,12 @@ def test_top_k_ungapped_kmer_seqs_shorter_than_kmer() -> None: score_matrix = make_score_matrix(alphabet, 1, -1) alignments = top_k_ungapped_local_align_kmer( - seqa, seqb, score_matrix, k=5, kmer_size=4, max_hits_per_kmer=10, + seqa, + seqb, + score_matrix, + k=5, + kmer_size=4, + max_hits_per_kmer=10, ) assert alignments == [] @@ -720,13 +756,23 @@ def test_top_k_ungapped_kmer_max_hits_skips_low_complexity() -> None: score_matrix = make_score_matrix(alphabet, 1, -1) full_alignments = top_k_ungapped_local_align_kmer( - seqa, seqb, score_matrix, k=5, kmer_size=3, max_hits_per_kmer=100, + seqa, + seqb, + score_matrix, + k=5, + kmer_size=3, + max_hits_per_kmer=100, ) assert len(full_alignments) >= 1 assert full_alignments[0].score == 10 capped_alignments = top_k_ungapped_local_align_kmer( - seqa, seqb, score_matrix, k=5, kmer_size=3, max_hits_per_kmer=0, + seqa, + seqb, + score_matrix, + k=5, + kmer_size=3, + max_hits_per_kmer=0, ) assert capped_alignments == [] @@ -746,14 +792,24 @@ def test_top_k_ungapped_kmer_min_hits_per_span_filters() -> None: score_matrix = make_score_matrix(alphabet, 1, -1) kept = top_k_ungapped_local_align_kmer( - seqa, seqb, score_matrix, k=5, kmer_size=3, max_hits_per_kmer=100, + seqa, + seqb, + score_matrix, + k=5, + kmer_size=3, + max_hits_per_kmer=100, min_kmer_hits_per_span=4, ) assert len(kept) == 1 assert kept[0].score == 6 dropped = top_k_ungapped_local_align_kmer( - seqa, seqb, score_matrix, k=5, kmer_size=3, max_hits_per_kmer=100, + seqa, + seqb, + score_matrix, + k=5, + kmer_size=3, + max_hits_per_kmer=100, min_kmer_hits_per_span=5, ) assert dropped == [] @@ -774,16 +830,15 @@ def test_top_k_ungapped_kmer_min_hits_zero_falls_back_to_exhaustive() -> None: full = top_k_ungapped_local_align(seqa, seqb, score_matrix, k=10) fallback = top_k_ungapped_local_align_kmer( - seqa, seqb, score_matrix, k=10, kmer_size=5, max_hits_per_kmer=100, + seqa, + seqb, + score_matrix, + k=10, + kmer_size=5, + max_hits_per_kmer=100, min_kmer_hits_per_span=0, ) - full_sig = [ - (a.score, a.fragments[0].sa_start, a.fragments[0].sb_start, a.fragments[0].len) - for a in full - ] - fallback_sig = [ - (a.score, a.fragments[0].sa_start, a.fragments[0].sb_start, a.fragments[0].len) - for a in fallback - ] + full_sig = [(a.score, a.fragments[0].sa_start, a.fragments[0].sb_start, a.fragments[0].len) for a in full] + fallback_sig = [(a.score, a.fragments[0].sa_start, a.fragments[0].sb_start, a.fragments[0].len) for a in fallback] assert fallback_sig == full_sig From 094d66705563458a214b0811ad445f7c9ac0fdea Mon Sep 17 00:00:00 2001 From: Tobias Sargeant Date: Tue, 12 May 2026 18:47:23 +1000 Subject: [PATCH 6/7] Apply cargo fmt --- build.rs | 2 +- src/lib.rs | 6 +++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/build.rs b/build.rs index a781ce1..dace4a9 100644 --- a/build.rs +++ b/build.rs @@ -1,3 +1,3 @@ fn main() { - pyo3_build_config::add_extension_module_link_args(); + pyo3_build_config::add_extension_module_link_args(); } diff --git a/src/lib.rs b/src/lib.rs index ed38d48..4cbb049 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1371,7 +1371,11 @@ fn _top_k_ungapped_local_align_kmer_core( let num_diagonals = sa_len + sb_len - 1; let diag_offset = (sb_len - 1) as i64; let mut open_spans: Vec = vec![ - OpenSpan { count: 0, start_pos: 0, last_pos: 0 }; + OpenSpan { + count: 0, + start_pos: 0, + last_pos: 0 + }; num_diagonals ]; let mut closed_spans: Vec = Vec::new(); From 4b41f02de3093006df10a793a5c8aff4c897ec5c Mon Sep 17 00:00:00 2001 From: Tobias Sargeant Date: Tue, 12 May 2026 18:55:42 +1000 Subject: [PATCH 7/7] Add rustfmt pre-commit hook and Rust toolchain to lint CI `rustfmt --check --edition 2021` runs as a local pre-commit hook on every staged `.rs` file. The CI lint job now installs the stable Rust toolchain with the `rustfmt` component before running `pre-commit run --all-files`, so the new hook is exercised on PR. Mechanical formatting drift caught by this hook was committed in the previous commit ('Apply cargo fmt'). --- .github/workflows/lint.yaml | 7 +++++++ .pre-commit-config.yaml | 8 ++++++++ 2 files changed, 15 insertions(+) diff --git a/.github/workflows/lint.yaml b/.github/workflows/lint.yaml index 419830e..511cb44 100644 --- a/.github/workflows/lint.yaml +++ b/.github/workflows/lint.yaml @@ -15,6 +15,13 @@ jobs: with: python-version: '3.11' + - name: Install Rust toolchain + uses: actions-rs/toolchain@v1 + with: + toolchain: stable + override: true + components: rustfmt + - name: Install packages run: pip install -r requirements-dev.txt diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index b76566c..dbf919f 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -33,6 +33,14 @@ repos: args: ["--fix"] - id: ruff-format + - repo: local + hooks: + - id: rustfmt + name: rustfmt + entry: rustfmt --check --edition 2021 + language: system + types: [rust] + - repo: https://github.com/pre-commit/mirrors-mypy rev: v1.18.2 hooks: