perf: optimize k-mer queries with sparse index and run-based aggregation
Replaces the dense `KmerResults` matrix with a sparse `SmerIndex` (`Vec<bool>` + offsets) that tracks k-mer presence independently of per-genome counts. Introduces a new aggregation pass, `sparse_findere_for_genome`, which sorts hits, detects contiguous runs, and applies monotone-deque scans to compute sliding-window minimums. This reduces query complexity from O(n_smers) to O(hits log hits), significantly lowering memory overhead and computational cost for low-density queries. Adds a deterministic PRNG and dense reference oracle in tests to validate correctness against randomized inputs without external property-testing crates.
This commit is contained in:
Eric Coissac
+67
-39
@@ -16,37 +16,43 @@ Given a set of query sequences, determine for each sequence how many of its k-me
|
|||||||
|
|
||||||
## Algorithm
|
## Algorithm
|
||||||
|
|
||||||
The query follows the same superkmer-based partitioning strategy used at indexing time. Everything below happens inside `process_chunk` (`query.rs`); there is no separate per-stage function — dereplication, MPHF/matrix lookup, and the Findere window are three phases of one function, operating on a single flat allocation per chunk.
|
The query follows the same superkmer-based partitioning strategy used at indexing time. Everything below happens inside `process_chunk` (`query.rs`); there is no separate per-stage function, but the internal data flow is staged: k-mer-level dereplication, a two-part MPHF/column-major matrix lookup (`obikpartitionner::query_partition_with`), and a sparse Findere pass, each producing sparse intermediate structures rather than one dense allocation for the whole chunk.
|
||||||
|
|
||||||
```
|
```
|
||||||
for each chunk of sequences (parallel workers via obipipeline, one call to process_chunk):
|
for each chunk of sequences (parallel workers via obipipeline, one call to process_chunk):
|
||||||
build QueryBatch (QueryBatch::from_records):
|
build QueryBatch (QueryBatch::from_records):
|
||||||
decompose all sequences into superkmers (SuperKmerIter)
|
decompose all sequences into superkmers (SuperKmerIter) — construction only,
|
||||||
deduplicate identical superkmers into
|
not the dedup key
|
||||||
map: HashMap<RoutableSuperKmer, Vec<SKDesc>> ← SKDesc = (seq_idx, kmer_offset)
|
deduplicate at k-mer granularity, split by partition in the same pass:
|
||||||
allocate KmerResults (KmerResults::new): one flat allocation for the whole chunk —
|
by_partition: Vec<HashMap<CanonicalKmer, Vec<KmerDesc>>> ← KmerDesc = (seq_idx, pos)
|
||||||
data: Vec<u32> sized total_smers × n_genomes (row-major, zero-initialised)
|
allocate SmerIndex (SmerIndex::new): in_index: Vec<bool>, sized total_smers —
|
||||||
in_index: Vec<bool> sized total_smers
|
NOT multiplied by n_genomes
|
||||||
split the batch's unique superkmers by partition, via minimiser hash
|
allocate by_genome: Vec<Vec<(seq_idx, pos, value)>>, one empty Vec per genome —
|
||||||
|
stays empty (zero cost) for every genome this chunk never matches
|
||||||
for each partition p:
|
for each partition p:
|
||||||
query_partition_with(p, superkmers_routed_to_p, on_hit):
|
query_partition_with(p, kmers_for_p, on_event):
|
||||||
→ load QueryLayer(s) for p
|
stage 1 (MPHF-only): for each unique k-mer, try each layer's MphfLayer::find
|
||||||
→ for each s-mer of each superkmer: try each layer's MphfLayer::find in turn,
|
in turn, stop at the first hit; bucket confirmed hits by (layer, slot);
|
||||||
stop at the first hit (a k-mer belongs to at most one layer)
|
emit QueryHit::Found(descs) once per hit k-mer
|
||||||
→ on_hit(sk_idx, kmer_idx, row): broadcast `row` into KmerResults, once per
|
stage 2 (column-major fetch): for each layer with ≥1 hit, for each genome
|
||||||
SKDesc referencing this superkmer (all (seq_idx, position) occurrences)
|
column g in 0..layer.n_cols(): scan that layer's bucketed slots, look up
|
||||||
for each sequence, for each genome (inlined sliding window, no separate function):
|
col_value(g, slot); emit QueryHit::Value(descs, g, value) on nonzero
|
||||||
monotone-deque scan over s-mer positions → win_min[pos][g] = min over the
|
on_event dispatches: Found → SmerIndex::mark_found for every desc;
|
||||||
z-window [pos, pos+z), "not in index" treated as 0
|
Value → push (seq_idx, pos, value) into by_genome[g]
|
||||||
for each sequence, for each output position (0..n_kmers_out):
|
for each genome g with ≥1 hit (sparse_findere_for_genome):
|
||||||
accumulate confirmed k_user-mer results into acc (kmer_count / kmer_missing /
|
sort by_genome[g] by (seq_idx, pos); detect maximal runs of consecutive pos
|
||||||
per-genome totals) and, if --detail, into cov
|
within one seq_idx; monotone-deque window-minimum scoped to each run →
|
||||||
|
confirmed_by_genome[g]: Vec<(seq_idx, pos_out, value)>
|
||||||
|
accumulate genome_totals per sequence from confirmed_by_genome (per genome, direct)
|
||||||
|
accumulate kmer_count / kmer_missing per (sequence, output position), O(1) each,
|
||||||
|
using only the confirmed-any bitmap and SmerIndex — independent of n_genomes
|
||||||
|
if --detail: densify confirmed_by_genome into per-(seq, genome) coverage arrays
|
||||||
emit annotated sequences (emit_batch)
|
emit annotated sequences (emit_batch)
|
||||||
```
|
```
|
||||||
|
|
||||||
Superkmers that appear more than once in the batch (same sequence or across sequences) are deduplicated: each unique `RoutableSuperKmer` is queried once per partition, and the result is broadcast to every `SKDesc` entry that references it. See [Future work](#throughput--parallelism--identified-potential-not-yet-implemented) for a discussion of pushing this dereplication down to k-mer granularity.
|
Superkmers that appear more than once in the batch (same sequence or across sequences), or different superkmers that happen to share a k-mer (read overlaps, repeats, a SNP splitting an otherwise-identical run), are deduplicated at k-mer granularity: each unique `CanonicalKmer` triggers at most one MPHF lookup and, on hit, one matrix fetch, broadcast to every `KmerDesc` occurrence referencing it.
|
||||||
|
|
||||||
**Findere requires full-sequence aggregation.** The sliding window is run once per `(sequence, genome)` pair on the complete s-mer result range, after all partitions have contributed to `KmerResults`. Applying it per superkmer would produce false negatives at superkmer boundaries, where the z-window spans two superkmers.
|
**Findere requires full-sequence aggregation.** The sliding window (now per-run, not per-sequence — see [Findere z-window filter](#findere-z-window-filter)) only ever runs after all partitions have contributed their hits to `by_genome`. Applying it per superkmer would produce false negatives at superkmer boundaries, where the z-window spans two superkmers.
|
||||||
|
|
||||||
Batches are processed in parallel via `obipipeline` workers; the `--threads` flag controls the number of worker threads.
|
Batches are processed in parallel via `obipipeline` workers; the `--threads` flag controls the number of worker threads.
|
||||||
|
|
||||||
@@ -56,31 +62,44 @@ Batches are processed in parallel via `obipipeline` workers; the `--threads` fla
|
|||||||
|
|
||||||
For approximate index modes, the index physically stores s-mers of size `s = k_user − z + 1`; `idx.kmer_size()` (bound to `k` in `process_chunk`) is this physically-indexed s-mer size, so decomposing the query at `k` naturally produces s-mer results.
|
For approximate index modes, the index physically stores s-mers of size `s = k_user − z + 1`; `idx.kmer_size()` (bound to `k` in `process_chunk`) is this physically-indexed s-mer size, so decomposing the query at `k` naturally produces s-mer results.
|
||||||
|
|
||||||
The z-window aggregation is inlined per `(sequence, genome)` in `process_chunk` — a monotone-deque sliding-window minimum, not a separate function operating on an intermediate `Option<Box<[u32]>>` vector:
|
The z-window aggregation is **sparse**, per genome, implemented in `sparse_findere_for_genome` (`query.rs`) — a run-detection pass followed by a monotone-deque sliding-window minimum scoped to each run, not a dense scan over every s-mer position of every sequence:
|
||||||
|
|
||||||
```
|
```
|
||||||
for each sequence, for each genome g independently:
|
sparse_findere_for_genome(hits, z, presence, threshold):
|
||||||
dq: VecDeque<(pos, value)> // reused across all (sequence, genome) pairs
|
// hits: raw (seq_idx, pos_smer, value) triples for this genome, as delivered
|
||||||
for i in 0..n_smers:
|
// by query_partition_with's QueryHit::Value — only ever nonzero entries;
|
||||||
v = if results.is_in_index(seq, i) { results.val(seq, i, g) } else { 0 }
|
// a position with no hit for this genome simply has no entry at all.
|
||||||
evict dq entries that left the window [i-z+1, i]
|
sort hits by (seq_idx, pos_smer)
|
||||||
maintain dq monotone non-decreasing (pop back while back.value >= v), push (i, v)
|
|
||||||
if i + 1 >= z:
|
for each maximal run of consecutive pos_smer values within the same seq_idx:
|
||||||
pos = i + 1 - z
|
dq: VecDeque<(run-relative index, value)>
|
||||||
win_min[pos][g] = dq.front().value // minimum over the z-window
|
for k, (_, pos, value) in enumerate(run):
|
||||||
|
maintain dq monotone non-decreasing (pop back while back.value >= value)
|
||||||
|
push (k, value)
|
||||||
|
evict dq entries with run-relative index <= k - z
|
||||||
|
if k + 1 >= z:
|
||||||
|
win_min = dq.front().value
|
||||||
|
if win_min > 0:
|
||||||
|
pos_out = pos + 1 - z
|
||||||
|
confirmed.push((seq_idx, pos_out, adjust(win_min)))
|
||||||
|
return confirmed
|
||||||
```
|
```
|
||||||
|
|
||||||
Input length `n_smers`, output length `n_kmers_out = n_smers − z + 1` (k_user-mer positions). The scan is `O(n_smers)` per genome (each position enters and leaves the deque once).
|
A window can only be confirmed (`win_min > 0`) when all `z` s-mers in it are present *and* nonzero for this genome — which, by construction, can only happen strictly inside one contiguous run of hits (any gap — an absent or zero-valued s-mer — forces `win_min = 0` for every window spanning it, exactly matching the old dense scan's "not in index counts as 0" rule, just never materialising the zero). The deque logic is otherwise identical to the pre-sparsification version; it's scoped to run-relative indices instead of the whole sequence.
|
||||||
|
|
||||||
Output position `pos` is confirmed for genome `g` iff `win_min[pos][g] > 0` — equivalent to "all z consecutive s-mer values in the window are nonzero for `g`", since a monotone-deque minimum is `0` as soon as any element in the window is `0`.
|
This runs once per genome that has at least one hit in the chunk (`process_chunk` iterates `by_genome`, one `Vec<(seq_idx, pos_smer, value)>` per genome, built from `QueryHit::Value` during the partition loop — genomes with zero hits in this chunk have an empty `Vec` and cost nothing beyond the iteration itself). Total work is `O(hits log hits)` per genome (the sort) rather than `O(n_smers)` per genome regardless of hit count — a genuine complexity win on top of the memory one, for the common case where most `(chunk, genome)` pairs have no or few hits.
|
||||||
|
|
||||||
**The value reported per confirmed position is the window minimum, not the leftmost s-mer's raw value.** For presence indexes (0/1 values) this is equivalent to a logical AND either way. For count indexes it is not: the accumulated count for genome `g` at position `pos` is `min(results[pos..pos+z][g])`, the weakest link across the window — not `results[pos][g]`.
|
Output position `pos_out` is confirmed for genome `g` iff its run produced a nonzero `win_min` — equivalent to "all `z` consecutive s-mer values in the window are nonzero for `g`", same semantics as before.
|
||||||
|
|
||||||
**`kmer_missing` bookkeeping is separate from the window value** and does look specifically at the leftmost s-mer: a position with no genome confirmed (`win_min[pos][*] == 0` for all `g`) counts as `kmer_missing` iff the leftmost s-mer of that window, `results.is_in_index(seq, pos)`, is absent from the index entirely (see [`kmer_missing` semantics](#kmer_missing-semantics)).
|
**The value reported per confirmed position is the window minimum, not the leftmost s-mer's raw value** — unchanged from the dense version. For presence indexes (0/1 values) this is equivalent to a logical AND either way. For count indexes it is not: the accumulated count for genome `g` at position `pos_out` is the minimum across the window, the weakest link — not the leftmost s-mer's own count. The presence/count adjustment (`u32::from(win_min >= threshold)` vs. raw `win_min`) is applied once, inside `sparse_findere_for_genome`, rather than later during accumulation.
|
||||||
|
|
||||||
**Short sequences**: when the s-mer count is less than `z`, `n_kmers_out` is `0` (`process_chunk` skips the sequence: `if out_n == 0 { continue; }`) — no k_user-mer is emitted for that sequence.
|
**`kmer_missing` bookkeeping is independent of the per-genome sparse structures**, by design (see roadmap point 9): a lightweight dense `SmerIndex` (`in_index: Vec<bool>`, sized `total_smers` — **not** multiplied by `n_genomes`) is populated from `QueryHit::Found` during the partition loop, one entry per hit k-mer regardless of which genome(s) it matched. A position with no genome confirmed counts as `kmer_missing` iff the leftmost s-mer of that window is absent from `SmerIndex` entirely (see [`kmer_missing` semantics](#kmer_missing-semantics)).
|
||||||
|
|
||||||
**Exact indexes**: `z = 1`, the window degenerates to a single element — the "minimum" is just that element's value, and `pos == i`, i.e. a passthrough.
|
**Coverage (`--detail`)** is built by re-scanning each genome's confirmed-hit list (already computed, no extra pass over raw data) and densifying into the `[u32; n_kmers_out]` arrays the JSON output format requires — but only when `--detail` is actually requested; the sparse structures cost nothing extra when it isn't.
|
||||||
|
|
||||||
|
**Short sequences**: when a sequence's s-mer count is less than `z`, its run(s) — if any hits exist at all — can never reach length `z`, so no window is ever confirmed for it; no k_user-mer is emitted, same outcome as the dense version's `n_kmers_out == 0` early-skip, reached here as a natural consequence rather than a separate check.
|
||||||
|
|
||||||
|
**Exact indexes**: `z = 1`, every single-hit "run" of length 1 immediately satisfies `k + 1 >= z`, so every hit is its own confirmed window with `win_min` equal to its own value — a passthrough, as before.
|
||||||
|
|
||||||
### Effective z at query time
|
### Effective z at query time
|
||||||
|
|
||||||
@@ -144,7 +163,7 @@ Coverage reflects confirmed k_user-mers only. The vectors are emitted in the JSO
|
|||||||
|
|
||||||
## `kmer_missing` semantics
|
## `kmer_missing` semantics
|
||||||
|
|
||||||
`kmer_missing` counts k_user-mer positions where the leftmost s-mer of the window (`results.is_in_index(seq_idx, pos)`, `KmerResults`) is `false` — i.e. absent from the index entirely. K-mers where the z-window fails because a later s-mer is absent or zero (but the leftmost one is present) are not counted as missing — the leftmost s-mer being present is used as proxy for index membership.
|
`kmer_missing` counts k_user-mer positions where the leftmost s-mer of the window (`smer_index.is_in_index(seq_idx, pos)`, `SmerIndex`) is `false` — i.e. absent from the index entirely. K-mers where the z-window fails because a later s-mer is absent or zero (but the leftmost one is present) are not counted as missing — the leftmost s-mer being present is used as proxy for index membership.
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
@@ -367,6 +386,15 @@ The column-major loop in stage 2 is therefore a **plain sequential loop** for no
|
|||||||
- **Validation — this is the correctness-critical phase**: property-test comparing old (dense, pre-phase-3) and new (sparse) implementations on the same randomized input/index fixtures, asserting identical `kmer_count`, `kmer_missing`, `kmer_strict_matches`, and (with `--detail`) `coverage` for every sequence. Keep both implementations compiled side by side (behind a debug-only flag or a temporary parallel code path) only for the duration of this validation; delete the dense path once parity is confirmed — per this project's own convention, superseded code is not kept "just in case."
|
- **Validation — this is the correctness-critical phase**: property-test comparing old (dense, pre-phase-3) and new (sparse) implementations on the same randomized input/index fixtures, asserting identical `kmer_count`, `kmer_missing`, `kmer_strict_matches`, and (with `--detail`) `coverage` for every sequence. Keep both implementations compiled side by side (behind a debug-only flag or a temporary parallel code path) only for the duration of this validation; delete the dense path once parity is confirmed — per this project's own convention, superseded code is not kept "just in case."
|
||||||
- **Update `docmd/architecture/query.md` itself**: once this phase lands, the "Findere z-window filter" section (which currently — correctly — describes the dense deque-over-`0..n_smers` scan) needs another pass to describe the sparse run-detection algorithm instead, as already flagged when this phase was discussed.
|
- **Update `docmd/architecture/query.md` itself**: once this phase lands, the "Findere z-window filter" section (which currently — correctly — describes the dense deque-over-`0..n_smers` scan) needs another pass to describe the sparse run-detection algorithm instead, as already flagged when this phase was discussed.
|
||||||
|
|
||||||
|
**Implemented as planned, no deviations discovered this time.** What shipped:
|
||||||
|
- `KmerResults` removed entirely, replaced by `SmerIndex` (`in_index: Vec<bool>` + `offsets`, unchanged size/purpose, renamed since it's no longer "results" — just the O(1)-per-position "was this k-mer found at all" bookkeeping) and `by_genome: Vec<Vec<(seq_idx, pos, value)>>` (one empty `Vec` per genome until a hit arrives — genomes with zero hits in a chunk cost nothing beyond the outer `Vec`'s own allocation).
|
||||||
|
- New `sparse_findere_for_genome(hits, z, presence, threshold) -> (Vec<ConfirmedHit>, n_runs, total_run_len)` (`query.rs`): sorts one genome's raw hits by `(seq_idx, pos)`, detects maximal runs of consecutive `pos` within one sequence, runs the same monotone-deque window-minimum as before but scoped to each run (run-relative indices for eviction, absolute `pos` for computing `pos_out`). Presence/count adjustment (`u32::from(win_min >= threshold)` vs. raw) is applied inside this function, once per confirmed hit, rather than later during accumulation.
|
||||||
|
- `process_chunk` restructured into three passes after the partition loop: (1) run `sparse_findere_for_genome` per genome, collecting `confirmed_by_genome` and run-detection stats; (2) accumulate `genome_totals` directly from `confirmed_by_genome` and mark a `confirmed_any: Vec<bool>` (sized `total_kmers_out`, not `× n_genomes`); (3) a position-only pass (`O(total_kmers_out)`, no genome factor) computing `kmer_count`/`kmer_missing` from `confirmed_any` + `SmerIndex`. `cov` (`--detail`) is populated by re-scanning `confirmed_by_genome` — only when `--detail` is actually set, otherwise skipped entirely.
|
||||||
|
- Debug log added (`"sparse Findere"`): `n_dense_would_be` (`n_occurrences × n_genomes` — what the deleted dense path would have allocated), `n_sparse_entries` (what's actually retained), `n_runs`/`avg_run_len` (per the plan's ask, to see whether hits mostly fail to form complete windows).
|
||||||
|
- **Unit tests**: `sparse_findere_matches_dense_reference_on_random_inputs` (`obikmer/src/cmd/tests/query.rs`) — 200 randomized cases (sequence count/length, `z`, presence/count mode, threshold, hit density from sparse to fully-dense) comparing `sparse_findere_for_genome` against `dense_reference_findere`, a faithful reimplementation of the deleted dense algorithm kept only as a test-local correctness oracle (no property-testing crate added — checked first, none was a workspace dependency; a small `std`-only xorshift64 PRNG stands in for one, deterministic and dependency-free). All 200 cases pass.
|
||||||
|
- **Functional validation performed**: full workspace build + `cargo test --workspace`, zero failures. End-to-end against real indexes: baseline output (no flags) unchanged from pre-phase-5 recorded values on the same fixtures; `--count-missing` correct (`kmer_missing: 0` on a self-match); `--detail` correct — coverage array length matches `kmer_count`, and critically, re-ran the two-genome cross-contamination check from phase 4 with `--detail --count-missing`: `genomeA` reads show coverage sum `106` for `genomeA` and `0` for `genomeB` (and vice versa) — confirms the sparse-to-dense `cov` reconstruction doesn't leak across genomes either, not just the scalar `kmer_strict_matches` path.
|
||||||
|
- This phase's roadmap item ("update the Findere z-window filter section") — done, see above; the "Algorithm" section's pseudocode was also updated, since it still named `KmerResults`/`SKDesc` from before phases 3–4.
|
||||||
|
|
||||||
### Phase 6 — Parallel gzip decompression (independent, optional)
|
### Phase 6 — Parallel gzip decompression (independent, optional)
|
||||||
|
|
||||||
Tracked separately in [chunkreader.md](../implementation/chunkreader.md#future-work--parallel-gzip-decompression-in-xopen); parked pending validation of `rapidgzip-rs` on real data. Not a dependency of, or a dependency for, phases 0–5 — `xopen` is shared infrastructure (`obiread`), phase 1 benefits from it but doesn't require it (phase 1 parallelises *across* files; this phase would additionally parallelise *within* one large file).
|
Tracked separately in [chunkreader.md](../implementation/chunkreader.md#future-work--parallel-gzip-decompression-in-xopen); parked pending validation of `rapidgzip-rs` on real data. Not a dependency of, or a dependency for, phases 0–5 — `xopen` is shared infrastructure (`obiread`), phase 1 benefits from it but doesn't require it (phase 1 parallelises *across* files; this phase would additionally parallelise *within* one large file).
|
||||||
|
|||||||
+167
-107
@@ -165,21 +165,22 @@ impl QueryBatch {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// ── KmerResults — allocation-free ragged result matrix ───────────────────────
|
// ── SmerIndex — sparse "was this k-mer found at all" bookkeeping ─────────────
|
||||||
|
|
||||||
/// Flat storage for per-kmer query results across all sequences in a chunk.
|
/// Tracks, per (sequence, s-mer position), whether the k-mer was found in the
|
||||||
///
|
/// index at all — independent of *which* genome(s) matched. Sized
|
||||||
/// Replaces `Vec<Vec<Option<Box<[u32]>>>>` — a single allocation for the whole
|
/// `total_smers` (one `bool` per s-mer occurrence in the chunk), **not**
|
||||||
/// chunk instead of one `Box<[u32]>` per found k-mer.
|
/// multiplied by `n_genomes`: this is the O(1)-per-position bookkeeping that
|
||||||
struct KmerResults {
|
/// `kmer_missing` needs (the leftmost-s-mer-of-window membership test), kept
|
||||||
data: Vec<u32>, // total_kmers × n_genomes, row-major
|
/// dense because it's already cheap — the `n_genomes`-scaled data lives in
|
||||||
in_index: Vec<bool>, // total_kmers — true if the kmer was found in the index
|
/// the sparse per-genome hit lists built alongside it (see `process_chunk`).
|
||||||
offsets: Vec<usize>, // offsets[i]..offsets[i+1] = kmer range for sequence i
|
struct SmerIndex {
|
||||||
n_genomes: usize,
|
in_index: Vec<bool>, // total_smers
|
||||||
|
offsets: Vec<usize>, // offsets[i]..offsets[i+1] = s-mer range for sequence i
|
||||||
}
|
}
|
||||||
|
|
||||||
impl KmerResults {
|
impl SmerIndex {
|
||||||
fn new(n_kmers_per_seq: &[u32], n_genomes: usize) -> Self {
|
fn new(n_kmers_per_seq: &[u32]) -> Self {
|
||||||
let mut offsets = Vec::with_capacity(n_kmers_per_seq.len() + 1);
|
let mut offsets = Vec::with_capacity(n_kmers_per_seq.len() + 1);
|
||||||
let mut total = 0usize;
|
let mut total = 0usize;
|
||||||
offsets.push(0);
|
offsets.push(0);
|
||||||
@@ -188,43 +189,96 @@ impl KmerResults {
|
|||||||
offsets.push(total);
|
offsets.push(total);
|
||||||
}
|
}
|
||||||
Self {
|
Self {
|
||||||
data: vec![0u32; total * n_genomes],
|
|
||||||
in_index: vec![false; total],
|
in_index: vec![false; total],
|
||||||
offsets,
|
offsets,
|
||||||
n_genomes,
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn n_kmers_for(&self, seq: usize) -> usize {
|
|
||||||
self.offsets[seq + 1] - self.offsets[seq]
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Mark the k-mer at (seq, kmer) as found in the index — independent of
|
/// Mark the k-mer at (seq, kmer) as found in the index — independent of
|
||||||
/// any particular genome's value. Called once per hit k-mer (stage 1 of
|
/// any particular genome's value. Called once per hit k-mer (stage 1 of
|
||||||
/// `query_partition_with`), regardless of how the column-major fetch
|
/// `query_partition_with`), regardless of how the column-major fetch
|
||||||
/// (stage 2) later fills in per-genome values.
|
/// (stage 2) later reports per-genome values.
|
||||||
fn mark_found(&mut self, seq: usize, kmer: usize) {
|
fn mark_found(&mut self, seq: usize, kmer: usize) {
|
||||||
let abs = self.offsets[seq] + kmer;
|
let abs = self.offsets[seq] + kmer;
|
||||||
self.in_index[abs] = true;
|
self.in_index[abs] = true;
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Set the value for one genome at (seq, kmer). Called once per nonzero
|
|
||||||
/// `(k-mer, genome)` pair delivered by the column-major fetch (stage 2).
|
|
||||||
fn set_one(&mut self, seq: usize, kmer: usize, g: usize, value: u32) {
|
|
||||||
let abs = self.offsets[seq] + kmer;
|
|
||||||
self.data[abs * self.n_genomes + g] = value;
|
|
||||||
}
|
|
||||||
|
|
||||||
#[inline]
|
#[inline]
|
||||||
fn is_in_index(&self, seq: usize, kmer: usize) -> bool {
|
fn is_in_index(&self, seq: usize, kmer: usize) -> bool {
|
||||||
self.in_index[self.offsets[seq] + kmer]
|
self.in_index[self.offsets[seq] + kmer]
|
||||||
}
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Value for genome `g` at (seq, kmer); meaningful only when `is_in_index`.
|
// ── Sparse Findere: per-genome run detection + sliding-window minimum ────────
|
||||||
#[inline]
|
|
||||||
fn val(&self, seq: usize, kmer: usize, g: usize) -> u32 {
|
/// One confirmed z-window: genome `g`'s window ending at k-mer `pos` (the
|
||||||
self.data[(self.offsets[seq] + kmer) * self.n_genomes + g]
|
/// *leftmost* s-mer of the window, i.e. the k_user-mer's output position) is
|
||||||
|
/// fully present and nonzero, with window-minimum `value`.
|
||||||
|
type ConfirmedHit = (u32, u32, u32); // (seq_idx, pos_out, value)
|
||||||
|
|
||||||
|
/// Reduce one genome's raw sparse s-mer hits — `(seq_idx, pos_smer, raw_value)`,
|
||||||
|
/// unsorted, exactly as delivered by `QueryHit::Value` — into confirmed
|
||||||
|
/// z-windows, without ever visiting a position that had no hit at all.
|
||||||
|
///
|
||||||
|
/// A z-window is confirmed only when all z s-mers in it are present *and*
|
||||||
|
/// nonzero for this genome (matching the dense sliding-window's semantics,
|
||||||
|
/// where "not in index" or a zero value both contribute 0 to the window
|
||||||
|
/// minimum) — which can only happen inside a maximal run of consecutive
|
||||||
|
/// `pos_smer` values for the same sequence. `hits` is sorted in place by
|
||||||
|
/// `(seq_idx, pos_smer)` to expose those runs; the monotone-deque
|
||||||
|
/// window-minimum then runs per run, on run-relative indices, identical in
|
||||||
|
/// spirit to the dense version's whole-sequence scan.
|
||||||
|
///
|
||||||
|
/// Returns the confirmed hits plus `(n_runs, total_run_len)` for logging —
|
||||||
|
/// a low average run length relative to `z` means most hits fail to form a
|
||||||
|
/// complete window.
|
||||||
|
fn sparse_findere_for_genome(
|
||||||
|
hits: &mut [(u32, u32, u32)],
|
||||||
|
z: usize,
|
||||||
|
presence: bool,
|
||||||
|
threshold: u32,
|
||||||
|
) -> (Vec<ConfirmedHit>, usize, usize) {
|
||||||
|
hits.sort_unstable_by_key(|&(seq, pos, _)| (seq, pos));
|
||||||
|
|
||||||
|
let mut confirmed = Vec::new();
|
||||||
|
let mut n_runs = 0usize;
|
||||||
|
let mut total_run_len = 0usize;
|
||||||
|
let mut dq: VecDeque<(usize, u32)> = VecDeque::new(); // (run-relative index, value)
|
||||||
|
|
||||||
|
let mut i = 0;
|
||||||
|
while i < hits.len() {
|
||||||
|
let seq = hits[i].0;
|
||||||
|
let mut j = i + 1;
|
||||||
|
while j < hits.len() && hits[j].0 == seq && hits[j].1 == hits[j - 1].1 + 1 {
|
||||||
|
j += 1;
|
||||||
}
|
}
|
||||||
|
let run = &hits[i..j];
|
||||||
|
n_runs += 1;
|
||||||
|
total_run_len += run.len();
|
||||||
|
|
||||||
|
dq.clear();
|
||||||
|
for (k, &(_, pos, val)) in run.iter().enumerate() {
|
||||||
|
while dq.back().map_or(false, |&(_, v)| v >= val) {
|
||||||
|
dq.pop_back();
|
||||||
|
}
|
||||||
|
dq.push_back((k, val));
|
||||||
|
while dq.front().map_or(false, |&(fk, _)| fk + z <= k) {
|
||||||
|
dq.pop_front();
|
||||||
|
}
|
||||||
|
if k + 1 >= z {
|
||||||
|
let win_min = dq.front().unwrap().1;
|
||||||
|
if win_min > 0 {
|
||||||
|
let pos_out = pos + 1 - z as u32;
|
||||||
|
let c = if presence { u32::from(win_min >= threshold) } else { win_min };
|
||||||
|
confirmed.push((seq, pos_out, c));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
i = j;
|
||||||
|
}
|
||||||
|
|
||||||
|
(confirmed, n_runs, total_run_len)
|
||||||
}
|
}
|
||||||
|
|
||||||
// ── Per-sequence accumulator ──────────────────────────────────────────────────
|
// ── Per-sequence accumulator ──────────────────────────────────────────────────
|
||||||
@@ -271,8 +325,14 @@ fn process_chunk(
|
|||||||
let batch = QueryBatch::from_records(records, k, 6, 0.7, n_partitions);
|
let batch = QueryBatch::from_records(records, k, 6, 0.7, n_partitions);
|
||||||
let n_seqs = batch.ids.len();
|
let n_seqs = batch.ids.len();
|
||||||
|
|
||||||
// Flat result matrix — one allocation for the whole chunk.
|
// Sparse bookkeeping for the whole chunk:
|
||||||
let mut results = KmerResults::new(&batch.n_kmers, n_genomes);
|
// - smer_index: O(total_smers) — is this s-mer in the index at all.
|
||||||
|
// - by_genome[g]: raw (seq_idx, pos_smer, value) hits for genome g, only
|
||||||
|
// ever containing nonzero entries (query_partition_with never emits a
|
||||||
|
// QueryHit::Value for a zero value) — empty for every genome this chunk
|
||||||
|
// never matched, which is the common case for unrelated queries.
|
||||||
|
let mut smer_index = SmerIndex::new(&batch.n_kmers);
|
||||||
|
let mut by_genome: Vec<Vec<(u32, u32, u32)>> = (0..n_genomes).map(|_| Vec::new()).collect();
|
||||||
|
|
||||||
// Dedup-ratio bookkeeping: occurrences (from batch.n_kmers, computed
|
// Dedup-ratio bookkeeping: occurrences (from batch.n_kmers, computed
|
||||||
// before dedup) vs. unique k-mers actually queried (query_stats) — the
|
// before dedup) vs. unique k-mers actually queried (query_stats) — the
|
||||||
@@ -296,12 +356,12 @@ fn process_chunk(
|
|||||||
|event| match event {
|
|event| match event {
|
||||||
QueryHit::Found(descs) => {
|
QueryHit::Found(descs) => {
|
||||||
for desc in descs {
|
for desc in descs {
|
||||||
results.mark_found(desc.seq_idx as usize, desc.pos as usize);
|
smer_index.mark_found(desc.seq_idx as usize, desc.pos as usize);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
QueryHit::Value(descs, g, v) => {
|
QueryHit::Value(descs, g, v) => {
|
||||||
for desc in descs {
|
for desc in descs {
|
||||||
results.set_one(desc.seq_idx as usize, desc.pos as usize, g, v);
|
by_genome[g].push((desc.seq_idx, desc.pos, v));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
@@ -323,94 +383,94 @@ fn process_chunk(
|
|||||||
"k-mer dedup + column-major fetch"
|
"k-mer dedup + column-major fetch"
|
||||||
);
|
);
|
||||||
|
|
||||||
// Sliding window minimum — one reusable buffer and one deque per batch.
|
// ── Sparse Findere: per-genome run detection + sliding-window minimum ────
|
||||||
//
|
//
|
||||||
// win_min[pos * n_genomes + g] = min count across the z-window [pos, pos+z)
|
// Confirmed z-windows, per genome, replace the dense win_min matrix:
|
||||||
// for genome g, where "not in index" counts as 0.
|
// total retained memory is O(actual hits), not O(total_smers × n_genomes)
|
||||||
//
|
// — the whole point of this pass. See sparse_findere_for_genome's doc for
|
||||||
// win_min > 0 ↔ all z consecutive kmers are in the index with count > 0
|
// why run detection is equivalent to the dense scan's semantics.
|
||||||
// ↔ Findere confirmation (for z=1 this degenerates to the
|
let presence = force_presence || !with_counts;
|
||||||
// simple case with no overhead).
|
let threshold = presence_threshold;
|
||||||
//
|
let z = effective_z;
|
||||||
// Works uniformly for count matrices and presence/absence (0/1) matrices.
|
|
||||||
let max_n_kmers = batch.n_kmers.iter().map(|&n| n as usize).max().unwrap_or(0);
|
|
||||||
let mut win_min = vec![0u32; max_n_kmers * n_genomes];
|
|
||||||
|
|
||||||
let mut accs: Vec<SeqAcc> = (0..n_seqs).map(|_| SeqAcc::new(n_genomes)).collect();
|
|
||||||
|
|
||||||
let n_kmers_out: Vec<usize> = batch
|
let n_kmers_out: Vec<usize> = batch
|
||||||
.n_kmers
|
.n_kmers
|
||||||
.iter()
|
.iter()
|
||||||
.map(|&n| {
|
.map(|&n| {
|
||||||
let n = n as usize;
|
let n = n as usize;
|
||||||
if n >= effective_z { n - effective_z + 1 } else { 0 }
|
if n >= z { n - z + 1 } else { 0 }
|
||||||
})
|
})
|
||||||
.collect();
|
.collect();
|
||||||
|
let mut out_offsets = Vec::with_capacity(n_seqs + 1);
|
||||||
|
{
|
||||||
|
let mut total = 0usize;
|
||||||
|
out_offsets.push(0);
|
||||||
|
for &n in &n_kmers_out {
|
||||||
|
total += n;
|
||||||
|
out_offsets.push(total);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let total_out = *out_offsets.last().unwrap_or(&0);
|
||||||
|
|
||||||
|
let n_dense_would_be = n_occurrences as u64 * n_genomes as u64;
|
||||||
|
let mut n_sparse_entries = 0u64;
|
||||||
|
let mut n_runs_total = 0usize;
|
||||||
|
let mut run_len_total = 0usize;
|
||||||
|
|
||||||
|
let mut confirmed_by_genome: Vec<Vec<ConfirmedHit>> = Vec::with_capacity(n_genomes);
|
||||||
|
for hits in &mut by_genome {
|
||||||
|
n_sparse_entries += hits.len() as u64;
|
||||||
|
let (confirmed, n_runs, run_len) = sparse_findere_for_genome(hits, z, presence, threshold);
|
||||||
|
n_runs_total += n_runs;
|
||||||
|
run_len_total += run_len;
|
||||||
|
confirmed_by_genome.push(confirmed);
|
||||||
|
}
|
||||||
|
|
||||||
|
debug!(
|
||||||
|
n_dense_would_be,
|
||||||
|
n_sparse_entries,
|
||||||
|
n_runs = n_runs_total,
|
||||||
|
avg_run_len = if n_runs_total > 0 { run_len_total as f64 / n_runs_total as f64 } else { 0.0 },
|
||||||
|
z,
|
||||||
|
"sparse Findere"
|
||||||
|
);
|
||||||
|
|
||||||
|
// ── Accumulate: genome totals (per genome, from confirmed hits) ──────────
|
||||||
|
let mut accs: Vec<SeqAcc> = (0..n_seqs).map(|_| SeqAcc::new(n_genomes)).collect();
|
||||||
|
let mut confirmed_any = vec![false; total_out];
|
||||||
|
|
||||||
|
for (g, hits) in confirmed_by_genome.iter().enumerate() {
|
||||||
|
for &(seq_idx, pos_out, c) in hits {
|
||||||
|
let abs_out = out_offsets[seq_idx as usize] + pos_out as usize;
|
||||||
|
confirmed_any[abs_out] = true;
|
||||||
|
accs[seq_idx as usize].genome_totals[g] += c;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Accumulate: kmer_count / kmer_missing (per position, genome-independent) ─
|
||||||
|
for seq_idx in 0..n_seqs {
|
||||||
|
let out_n = n_kmers_out[seq_idx];
|
||||||
|
let acc = &mut accs[seq_idx];
|
||||||
|
for pos in 0..out_n {
|
||||||
|
let abs_out = out_offsets[seq_idx] + pos;
|
||||||
|
if confirmed_any[abs_out] {
|
||||||
|
acc.kmer_count += 1;
|
||||||
|
} else if !smer_index.is_in_index(seq_idx, pos) {
|
||||||
|
acc.kmer_missing += 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Coverage (--detail): densify only when actually requested ────────────
|
||||||
let mut cov: Vec<Vec<Vec<u32>>> = if detail {
|
let mut cov: Vec<Vec<Vec<u32>>> = if detail {
|
||||||
n_kmers_out.iter().map(|&n| vec![vec![0u32; n]; n_genomes]).collect()
|
n_kmers_out.iter().map(|&n| vec![vec![0u32; n]; n_genomes]).collect()
|
||||||
} else {
|
} else {
|
||||||
Vec::new()
|
Vec::new()
|
||||||
};
|
};
|
||||||
|
if detail {
|
||||||
let presence = force_presence || !with_counts;
|
for (g, hits) in confirmed_by_genome.iter().enumerate() {
|
||||||
let threshold = presence_threshold;
|
for &(seq_idx, pos_out, c) in hits {
|
||||||
let z = effective_z;
|
cov[seq_idx as usize][g][pos_out as usize] += c;
|
||||||
|
|
||||||
// Deque reused across all (seq, genome) pairs.
|
|
||||||
let mut dq: VecDeque<(usize, u32)> = VecDeque::with_capacity(z + 1);
|
|
||||||
|
|
||||||
for seq_idx in 0..n_seqs {
|
|
||||||
let n = results.n_kmers_for(seq_idx);
|
|
||||||
let out_n = n_kmers_out[seq_idx];
|
|
||||||
if out_n == 0 { continue; }
|
|
||||||
|
|
||||||
let mins = &mut win_min[..out_n * n_genomes];
|
|
||||||
mins.fill(0);
|
|
||||||
|
|
||||||
// ── Per-genome sliding window minimum ─────────────────────────────────
|
|
||||||
for g in 0..n_genomes {
|
|
||||||
dq.clear();
|
|
||||||
for i in 0..n {
|
|
||||||
let v_i = if results.is_in_index(seq_idx, i) {
|
|
||||||
results.val(seq_idx, i, g)
|
|
||||||
} else {
|
|
||||||
0
|
|
||||||
};
|
|
||||||
// Evict elements that have left the window.
|
|
||||||
while dq.front().map_or(false, |&(f, _)| f + z <= i) {
|
|
||||||
dq.pop_front();
|
|
||||||
}
|
|
||||||
// Maintain monotone non-decreasing back→front for minimum at front.
|
|
||||||
while dq.back().map_or(false, |&(_, v)| v >= v_i) {
|
|
||||||
dq.pop_back();
|
|
||||||
}
|
|
||||||
dq.push_back((i, v_i));
|
|
||||||
// Window [pos, pos+z) is complete when i = pos + z - 1.
|
|
||||||
if i + 1 >= z {
|
|
||||||
let pos = i + 1 - z;
|
|
||||||
mins[pos * n_genomes + g] = dq.front().unwrap().1;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── Accumulate ────────────────────────────────────────────────────────
|
|
||||||
let acc = &mut accs[seq_idx];
|
|
||||||
for pos in 0..out_n {
|
|
||||||
let any = (0..n_genomes).any(|g| mins[pos * n_genomes + g] > 0);
|
|
||||||
if !any {
|
|
||||||
if !results.is_in_index(seq_idx, pos) {
|
|
||||||
acc.kmer_missing += 1;
|
|
||||||
}
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
acc.kmer_count += 1;
|
|
||||||
for g in 0..n_genomes {
|
|
||||||
let v = mins[pos * n_genomes + g];
|
|
||||||
if v == 0 { continue; }
|
|
||||||
let c = if presence { u32::from(v >= threshold) } else { v };
|
|
||||||
acc.genome_totals[g] += c;
|
|
||||||
if detail { cov[seq_idx][g][pos] += c; }
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -122,3 +122,107 @@ fn partition_routing_is_a_pure_function_of_the_kmer() {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ── sparse_findere_for_genome vs. a dense reference implementation ──────────
|
||||||
|
//
|
||||||
|
// No property-testing crate (proptest/quickcheck) is a workspace dependency
|
||||||
|
// (checked before writing this — not adding one for a single test module,
|
||||||
|
// per this project's dependency-approval rule). A tiny deterministic xorshift
|
||||||
|
// PRNG, std-only, stands in for one.
|
||||||
|
|
||||||
|
/// Faithful reimplementation of the pre-phase-5 dense sliding-window scan —
|
||||||
|
/// the algorithm `sparse_findere_for_genome` replaced — used here only as a
|
||||||
|
/// correctness oracle, not in production code. Operates on one genome's
|
||||||
|
/// hits across possibly many sequences, exactly like the sparse version.
|
||||||
|
fn dense_reference_findere(
|
||||||
|
hits: &[(u32, u32, u32)],
|
||||||
|
seq_lens: &[usize],
|
||||||
|
z: usize,
|
||||||
|
presence: bool,
|
||||||
|
threshold: u32,
|
||||||
|
) -> Vec<(u32, u32, u32)> {
|
||||||
|
let mut by_seq: Vec<Vec<u32>> = seq_lens.iter().map(|&n| vec![0u32; n]).collect();
|
||||||
|
for &(seq, pos, val) in hits {
|
||||||
|
by_seq[seq as usize][pos as usize] = val;
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut confirmed = Vec::new();
|
||||||
|
for (seq_idx, values) in by_seq.iter().enumerate() {
|
||||||
|
let n = values.len();
|
||||||
|
let mut dq: std::collections::VecDeque<(usize, u32)> = std::collections::VecDeque::new();
|
||||||
|
for i in 0..n {
|
||||||
|
let v_i = values[i];
|
||||||
|
while dq.front().map_or(false, |&(f, _)| f + z <= i) {
|
||||||
|
dq.pop_front();
|
||||||
|
}
|
||||||
|
while dq.back().map_or(false, |&(_, v)| v >= v_i) {
|
||||||
|
dq.pop_back();
|
||||||
|
}
|
||||||
|
dq.push_back((i, v_i));
|
||||||
|
if i + 1 >= z {
|
||||||
|
let win_min = dq.front().unwrap().1;
|
||||||
|
if win_min > 0 {
|
||||||
|
let pos_out = (i + 1 - z) as u32;
|
||||||
|
let c = if presence { u32::from(win_min >= threshold) } else { win_min };
|
||||||
|
confirmed.push((seq_idx as u32, pos_out, c));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
confirmed
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Minimal std-only xorshift64 PRNG — deterministic, seedable, no dependency.
|
||||||
|
struct Xorshift64(u64);
|
||||||
|
impl Xorshift64 {
|
||||||
|
fn next(&mut self) -> u64 {
|
||||||
|
self.0 ^= self.0 << 13;
|
||||||
|
self.0 ^= self.0 >> 7;
|
||||||
|
self.0 ^= self.0 << 17;
|
||||||
|
self.0
|
||||||
|
}
|
||||||
|
fn range(&mut self, n: u32) -> u32 {
|
||||||
|
(self.next() % n as u64) as u32
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn sparse_findere_matches_dense_reference_on_random_inputs() {
|
||||||
|
let mut rng = Xorshift64(0x5eed_5eed_5eed_5eedu64);
|
||||||
|
|
||||||
|
for case in 0..200 {
|
||||||
|
let n_seqs = 1 + rng.range(4) as usize;
|
||||||
|
let seq_lens: Vec<usize> = (0..n_seqs).map(|_| 1 + rng.range(30) as usize).collect();
|
||||||
|
let z = 1 + rng.range(4) as usize;
|
||||||
|
let presence = rng.range(2) == 0;
|
||||||
|
let threshold = 1 + rng.range(3);
|
||||||
|
|
||||||
|
// Sparse density varies across cases, including edge cases (empty,
|
||||||
|
// fully dense) — deliberately not uniform, to stress both few-hits
|
||||||
|
// and many-overlapping-runs scenarios.
|
||||||
|
let density = rng.range(101);
|
||||||
|
let mut hits: Vec<(u32, u32, u32)> = Vec::new();
|
||||||
|
for (seq_idx, &len) in seq_lens.iter().enumerate() {
|
||||||
|
for pos in 0..len {
|
||||||
|
if rng.range(100) < density {
|
||||||
|
let val = 1 + rng.range(5); // never 0 — matches QueryHit::Value's invariant
|
||||||
|
hits.push((seq_idx as u32, pos as u32, val));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut sparse_input = hits.clone();
|
||||||
|
let (mut sparse_result, _, _) =
|
||||||
|
sparse_findere_for_genome(&mut sparse_input, z, presence, threshold);
|
||||||
|
let mut dense_result = dense_reference_findere(&hits, &seq_lens, z, presence, threshold);
|
||||||
|
|
||||||
|
sparse_result.sort_unstable();
|
||||||
|
dense_result.sort_unstable();
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
sparse_result, dense_result,
|
||||||
|
"case {case}: n_seqs={n_seqs} seq_lens={seq_lens:?} z={z} presence={presence} \
|
||||||
|
threshold={threshold} density={density} hits={hits:?}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user