From 912f788f7f1d8d288d49843234dc2dc3cbd5616e Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Wed, 8 Jul 2026 18:16:08 +0200 Subject: [PATCH] feat: extract k-mer entropy computation into new obikentropy crate Extracts streaming entropy logic and sliding-window frequency tracking from obiskbuilder into a dedicated obikentropy crate. Introduces an EntropyTracker accumulator for O(1) per-base normalized Shannon entropy, replaces inline rolling statistics with delegated state management, and updates workspace dependencies across obikindex, obikpartitionner, and obiskbuilder. Adds criterion benchmarks to validate the refactored pipeline throughput. --- src/Cargo.lock | 10 + src/Cargo.toml | 2 +- src/obikentropy/Cargo.toml | 10 + src/{obiskbuilder => obikentropy}/build.rs | 0 src/obikentropy/src/kmer_entropy.rs | 41 +++ src/obikentropy/src/lib.rs | 17 ++ src/obikentropy/src/ring.rs | 40 +++ .../src/table.rs} | 4 + .../src/tests/kmer_entropy.rs | 2 - src/obikentropy/src/tracker.rs | 280 ++++++++++++++++++ src/obikpartitionner/Cargo.toml | 1 + src/obikpartitionner/src/filter.rs | 4 +- src/obiskbuilder/Cargo.toml | 6 + src/obiskbuilder/benches/superkmer_stream.rs | 58 ++++ src/obiskbuilder/src/kmer_entropy.rs | 49 --- src/obiskbuilder/src/lib.rs | 3 - src/obiskbuilder/src/rolling_stat.rs | 260 +--------------- 17 files changed, 482 insertions(+), 305 deletions(-) create mode 100644 src/obikentropy/Cargo.toml rename src/{obiskbuilder => obikentropy}/build.rs (100%) create mode 100644 src/obikentropy/src/kmer_entropy.rs create mode 100644 src/obikentropy/src/lib.rs create mode 100644 src/obikentropy/src/ring.rs rename src/{obiskbuilder/src/entropy_table.rs => obikentropy/src/table.rs} (93%) rename src/{obiskbuilder => obikentropy}/src/tests/kmer_entropy.rs (94%) create mode 100644 src/obikentropy/src/tracker.rs create mode 100644 src/obiskbuilder/benches/superkmer_stream.rs delete mode 100644 src/obiskbuilder/src/kmer_entropy.rs diff --git a/src/Cargo.lock b/src/Cargo.lock index 3e31a1e..f7c63e3 100644 --- a/src/Cargo.lock +++ b/src/Cargo.lock @@ -1682,6 +1682,13 @@ dependencies = [ "xxhash-rust", ] +[[package]] +name = "obikentropy" +version = "0.1.0" +dependencies = [ + "obikseq", +] + [[package]] name = "obikindex" version = "0.1.0" @@ -1742,6 +1749,7 @@ dependencies = [ "niffler 3.0.0", "obicompactvec", "obidebruinj", + "obikentropy", "obikrope", "obikseq", "obilayeredmap", @@ -1824,7 +1832,9 @@ dependencies = [ name = "obiskbuilder" version = "0.1.0" dependencies = [ + "criterion2", "lazy_static", + "obikentropy", "obikrope", "obikseq", "obiread", diff --git a/src/Cargo.toml b/src/Cargo.toml index 141df02..5684237 100644 --- a/src/Cargo.toml +++ b/src/Cargo.toml @@ -1,5 +1,5 @@ [workspace] resolver = "3" -members = ["obikseq", "obiread", "obiskbuilder", "obifastwrite", "obikmer","obikrope","obipipeline", "obikpartitionner","obiskio","obidebruinj","obilayeredmap", "obicompactvec", "obisys", "obikindex", "obitaxonomy"] +members = ["obikseq", "obiread", "obiskbuilder", "obifastwrite", "obikmer","obikrope","obipipeline", "obikpartitionner","obiskio","obidebruinj","obilayeredmap", "obicompactvec", "obisys", "obikindex", "obitaxonomy", "obikentropy"] [profile.release] debug = 1 diff --git a/src/obikentropy/Cargo.toml b/src/obikentropy/Cargo.toml new file mode 100644 index 0000000..3d1cd69 --- /dev/null +++ b/src/obikentropy/Cargo.toml @@ -0,0 +1,10 @@ +[package] +name = "obikentropy" +version = "0.1.0" +edition = "2024" + +[dependencies] +obikseq = { path = "../obikseq" } + +[dev-dependencies] +obikseq = { path = "../obikseq", features = ["test-utils"] } diff --git a/src/obiskbuilder/build.rs b/src/obikentropy/build.rs similarity index 100% rename from src/obiskbuilder/build.rs rename to src/obikentropy/build.rs diff --git a/src/obikentropy/src/kmer_entropy.rs b/src/obikentropy/src/kmer_entropy.rs new file mode 100644 index 0000000..1989972 --- /dev/null +++ b/src/obikentropy/src/kmer_entropy.rs @@ -0,0 +1,41 @@ +//! Normalized entropy of an isolated, already-built k-mer (e.g. one +//! reconstructed from an index's `unitigs.bin`, with no surrounding +//! sequence) — drives the window through [`EntropyTracker`] one base at a +//! time, exactly like the streaming path, so a `theta` threshold means the +//! same thing whether applied during index construction or after the fact +//! (e.g. `obikmer filter`). + +use obikseq::CanonicalKmer; + +use crate::tracker::{EntropyTracker, sub_word_canon}; + +/// Extension trait: compute the normalized entropy of a single canonical +/// k-mer, independent of any surrounding sequence. +pub trait KmerEntropy { + /// Normalized entropy across sub-word orders `1..=level_max` (the + /// minimum is taken across orders). Lower means less complex; `theta` + /// in `index`/`filter` rejects k-mers with a score `< theta`. + fn entropy(&self, level_max: usize) -> f64; +} + +impl KmerEntropy for CanonicalKmer { + fn entropy(&self, level_max: usize) -> f64 { + let raw = self.raw(); // left-aligned, 2 bits/base, MSB-first + let k = obikseq::params::k(); + let mask = (!0u64) >> (64 - k * 2); + + let mut tracker = EntropyTracker::new(k); + let mut rolling: u64 = 0; + for i in 0..k { + let shift = 64 - 2 * (i + 1); + let base = (raw >> shift) & 3; + rolling = ((rolling << 2) | base) & mask; + tracker.push(i + 1, sub_word_canon(rolling)); + } + tracker.normalized_entropy(level_max) + } +} + +#[cfg(test)] +#[path = "tests/kmer_entropy.rs"] +mod tests; diff --git a/src/obikentropy/src/lib.rs b/src/obikentropy/src/lib.rs new file mode 100644 index 0000000..113bdd3 --- /dev/null +++ b/src/obikentropy/src/lib.rs @@ -0,0 +1,17 @@ +//! Normalized k-mer entropy: formulas, tables, and a streaming tracker. +//! +//! This crate holds every piece of the entropy computation described in +//! `docmd/theory/entropy.md`: the compile-time tables ([`table`], private), +//! the incremental accumulator ([`EntropyTracker`]) that callers compose +//! into their own streaming state, and the [`KmerEntropy`] convenience trait +//! for scoring a single, already-built k-mer. + +#![deny(missing_docs)] + +mod kmer_entropy; +mod ring; +mod table; +mod tracker; + +pub use kmer_entropy::KmerEntropy; +pub use tracker::{EntropyTracker, SubWordCanon, sub_word_canon}; diff --git a/src/obikentropy/src/ring.rs b/src/obikentropy/src/ring.rs new file mode 100644 index 0000000..b729beb --- /dev/null +++ b/src/obikentropy/src/ring.rs @@ -0,0 +1,40 @@ +//! Stack-allocated ring buffer backing the sliding sub-word windows. + +/// Fixed-capacity ring buffer backed by a stack array. +/// N must be a power of two; operations are branchless via `% N`. +pub(crate) struct Ring { + buf: [T; N], + head: usize, + len: usize, +} + +impl Ring { + #[inline] + pub(crate) fn new() -> Self { + Self { + buf: [T::default(); N], + head: 0, + len: 0, + } + } + + #[inline] + pub(crate) fn clear(&mut self) { + self.len = 0; + self.head = 0; + } + + #[inline] + pub(crate) fn push_back(&mut self, val: T) { + self.buf[(self.head + self.len) % N] = val; + self.len += 1; + } + + #[inline] + pub(crate) fn pop_front(&mut self) -> T { + let val = self.buf[self.head]; + self.head = (self.head + 1) % N; + self.len -= 1; + val + } +} diff --git a/src/obiskbuilder/src/entropy_table.rs b/src/obikentropy/src/table.rs similarity index 93% rename from src/obiskbuilder/src/entropy_table.rs rename to src/obikentropy/src/table.rs index 9be69c3..289aa94 100644 --- a/src/obiskbuilder/src/entropy_table.rs +++ b/src/obikentropy/src/table.rs @@ -1,3 +1,7 @@ +//! Compile-time tables backing the normalized k-mer entropy formula: +//! circular-canonical sub-word classes, their log-sizes, and the max-entropy +//! correction for small samples. See `docmd/theory/entropy.md`. + pub(crate) const NORMK1: [u64; 4] = build_normalized_kmer::<4>(); pub(crate) const NORMK2: [u64; 16] = build_normalized_kmer::<16>(); pub(crate) const NORMK3: [u64; 64] = build_normalized_kmer::<64>(); diff --git a/src/obiskbuilder/src/tests/kmer_entropy.rs b/src/obikentropy/src/tests/kmer_entropy.rs similarity index 94% rename from src/obiskbuilder/src/tests/kmer_entropy.rs rename to src/obikentropy/src/tests/kmer_entropy.rs index 1456c3a..c75522e 100644 --- a/src/obiskbuilder/src/tests/kmer_entropy.rs +++ b/src/obikentropy/src/tests/kmer_entropy.rs @@ -3,12 +3,10 @@ use obikseq::Sequence; use obikseq::kmer::Kmer; const K: usize = 21; -const M: usize = 9; // RollingStat also tracks the minimizer window internally const LEVEL_MAX: usize = 6; fn kmer_from_ascii(seq: &[u8]) -> CanonicalKmer { obikseq::set_k(K); - obikseq::set_m(M); Kmer::from_ascii(seq).expect("valid k-mer sequence").canonical() } diff --git a/src/obikentropy/src/tracker.rs b/src/obikentropy/src/tracker.rs new file mode 100644 index 0000000..98686fa --- /dev/null +++ b/src/obikentropy/src/tracker.rs @@ -0,0 +1,280 @@ +//! Incremental (streaming) normalized k-mer entropy. +//! +//! [`EntropyTracker`] maintains, over a sliding window of the last `k` bases, +//! the per-sub-word-size frequency statistics needed to evaluate the +//! corrected Shannon entropy described in `docmd/theory/entropy.md`, updated +//! in O(1) per base rather than recomputed from scratch. +//! +//! It carries no notion of minimizers or superkmer segmentation — callers +//! that need both (e.g. `obiskbuilder::RollingStat`) compose an +//! `EntropyTracker` as a plain field alongside their own state, so the two +//! concerns update in the same streaming pass without being conflated in one +//! struct. + +use crate::ring::Ring; +use crate::table::{WS_MAX, emax, entropy_norm_kmer, ln_class_size, log_nwords, n_log_n}; + +/// Canonical sub-word values for word sizes 1..=6, computed by the caller +/// from its own rolling k-mer bits (see [`EntropyTracker::push`]). +pub type SubWordCanon = [u64; WS_MAX]; + +/// Compute the six canonical sub-word values (word sizes 1..=6) for the +/// current rolling right-aligned k-mer window. +#[inline] +pub fn sub_word_canon(rolling_kmer: u64) -> SubWordCanon { + [ + entropy_norm_kmer::(rolling_kmer & 3), + entropy_norm_kmer::(rolling_kmer & 15), + entropy_norm_kmer::(rolling_kmer & 63), + entropy_norm_kmer::(rolling_kmer & 255), + entropy_norm_kmer::(rolling_kmer & 1023), + entropy_norm_kmer::(rolling_kmer & 4095), + ] +} + +/// Incremental normalized-entropy accumulator over a sliding window of `k` +/// bases. Composed as a plain field by callers that also need other +/// per-base state (e.g. minimizer selection) in the same streaming pass. +pub struct EntropyTracker { + k: usize, + steady: bool, + + // Sliding-window queues over the last `k` canonical sub-words, one per + // word size — stack-allocated, capacity ≤ k ≤ 31. + k1q: Ring, + k2q: Ring, + k3q: Ring, + k4q: Ring, + k5q: Ring, + k6q: Ring, + + // Frequency count arrays. Max count per cell ≤ k ≤ 31 → u8 is sufficient. + k1c: [u8; 4], + k2c: [u8; 16], + k3c: [u8; 64], + k4c: [u8; 256], + k5c: [u8; 1024], + k6c: [u8; 4096], + + sum_f_log_f: [f64; WS_MAX + 1], + sum_f_log_s: [f64; WS_MAX + 1], +} + +impl EntropyTracker { + /// New tracker for a window of `k` bases (1..=31). + pub fn new(k: usize) -> Self { + Self { + k, + steady: false, + k1q: Ring::new(), + k2q: Ring::new(), + k3q: Ring::new(), + k4q: Ring::new(), + k5q: Ring::new(), + k6q: Ring::new(), + k1c: [0; 4], + k2c: [0; 16], + k3c: [0; 64], + k4c: [0; 256], + k5c: [0; 1024], + k6c: [0; 4096], + sum_f_log_f: [0.0; WS_MAX + 1], + sum_f_log_s: [0.0; WS_MAX + 1], + } + } + + /// Clear all accumulated state, ready to track a new window from + /// scratch (`k` is unchanged). + pub fn reset(&mut self) { + self.steady = false; + + self.k1c.fill(0); + self.k2c.fill(0); + self.k3c.fill(0); + self.k4c.fill(0); + self.k5c.fill(0); + self.k6c.fill(0); + + self.k1q.clear(); + self.k2q.clear(); + self.k3q.clear(); + self.k4q.clear(); + self.k5q.clear(); + self.k6q.clear(); + + self.sum_f_log_f = [0.0; WS_MAX + 1]; + self.sum_f_log_s = [0.0; WS_MAX + 1]; + } + + #[inline] + fn update_sums_decrement( + sum_f_log_f: &mut [f64; WS_MAX + 1], + sum_f_log_s: &mut [f64; WS_MAX + 1], + canonical: u64, + f: usize, + ) { + sum_f_log_f[K] += n_log_n(f - 1) - n_log_n(f); + sum_f_log_s[K] -= ln_class_size::(canonical); + } + + #[inline] + fn update_sums_increment( + sum_f_log_f: &mut [f64; WS_MAX + 1], + sum_f_log_s: &mut [f64; WS_MAX + 1], + canonical: u64, + g: usize, + ) { + sum_f_log_f[K] += n_log_n(g + 1) - n_log_n(g); + sum_f_log_s[K] += ln_class_size::(canonical); + } + + /// Advance the window by one base. `received` is the caller's running + /// count of bases pushed so far (1-based, i.e. after this base); + /// `canon` are the six canonical sub-word values for the current + /// rolling k-mer, from [`sub_word_canon`]. + pub fn push(&mut self, received: usize, canon: SubWordCanon) { + let [canonical_k1, canonical_k2, canonical_k3, canonical_k4, canonical_k5, canonical_k6] = + canon; + + if received > self.k { + let old1 = self.k1q.pop_front(); + let f1 = self.k1c[old1 as usize] as usize; + Self::update_sums_decrement::<1>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, old1, f1); + self.k1c[old1 as usize] -= 1; + + let old2 = self.k2q.pop_front(); + let f2 = self.k2c[old2 as usize] as usize; + Self::update_sums_decrement::<2>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, old2, f2); + self.k2c[old2 as usize] -= 1; + + let old3 = self.k3q.pop_front(); + let f3 = self.k3c[old3 as usize] as usize; + Self::update_sums_decrement::<3>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, old3, f3); + self.k3c[old3 as usize] -= 1; + + let old4 = self.k4q.pop_front(); + let f4 = self.k4c[old4 as usize] as usize; + Self::update_sums_decrement::<4>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, old4, f4); + self.k4c[old4 as usize] -= 1; + + let old5 = self.k5q.pop_front(); + let f5 = self.k5c[old5 as usize] as usize; + Self::update_sums_decrement::<5>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, old5, f5); + self.k5c[old5 as usize] -= 1; + + let old6 = self.k6q.pop_front(); + let f6 = self.k6c[old6 as usize] as usize; + Self::update_sums_decrement::<6>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, old6, f6); + self.k6c[old6 as usize] -= 1; + } + + if self.steady { + let g1 = self.k1c[canonical_k1 as usize] as usize; + Self::update_sums_increment::<1>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k1, g1); + self.k1c[canonical_k1 as usize] += 1; + self.k1q.push_back(canonical_k1); + + let g2 = self.k2c[canonical_k2 as usize] as usize; + Self::update_sums_increment::<2>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k2, g2); + self.k2c[canonical_k2 as usize] += 1; + self.k2q.push_back(canonical_k2); + + let g3 = self.k3c[canonical_k3 as usize] as usize; + Self::update_sums_increment::<3>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k3, g3); + self.k3c[canonical_k3 as usize] += 1; + self.k3q.push_back(canonical_k3); + + let g4 = self.k4c[canonical_k4 as usize] as usize; + Self::update_sums_increment::<4>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k4, g4); + self.k4c[canonical_k4 as usize] += 1; + self.k4q.push_back(canonical_k4); + + let g5 = self.k5c[canonical_k5 as usize] as usize; + Self::update_sums_increment::<5>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k5, g5); + self.k5c[canonical_k5 as usize] += 1; + self.k5q.push_back(canonical_k5); + + let g6 = self.k6c[canonical_k6 as usize] as usize; + Self::update_sums_increment::<6>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k6, g6); + self.k6c[canonical_k6 as usize] += 1; + self.k6q.push_back(canonical_k6); + } else { + self.push_warmup_increments(received, canonical_k1, canonical_k2, canonical_k3, canonical_k4, canonical_k5, canonical_k6); + } + } + + #[cold] + #[inline(never)] + fn push_warmup_increments( + &mut self, + received: usize, + canonical_k1: u64, canonical_k2: u64, canonical_k3: u64, + canonical_k4: u64, canonical_k5: u64, canonical_k6: u64, + ) { + let g1 = self.k1c[canonical_k1 as usize] as usize; + Self::update_sums_increment::<1>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k1, g1); + self.k1c[canonical_k1 as usize] += 1; + self.k1q.push_back(canonical_k1); + + if received >= 2 { + let g2 = self.k2c[canonical_k2 as usize] as usize; + Self::update_sums_increment::<2>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k2, g2); + self.k2c[canonical_k2 as usize] += 1; + self.k2q.push_back(canonical_k2); + + if received >= 3 { + let g3 = self.k3c[canonical_k3 as usize] as usize; + Self::update_sums_increment::<3>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k3, g3); + self.k3c[canonical_k3 as usize] += 1; + self.k3q.push_back(canonical_k3); + + if received >= 4 { + let g4 = self.k4c[canonical_k4 as usize] as usize; + Self::update_sums_increment::<4>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k4, g4); + self.k4c[canonical_k4 as usize] += 1; + self.k4q.push_back(canonical_k4); + + if received >= 5 { + let g5 = self.k5c[canonical_k5 as usize] as usize; + Self::update_sums_increment::<5>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k5, g5); + self.k5c[canonical_k5 as usize] += 1; + self.k5q.push_back(canonical_k5); + + if received >= 6 { + let g6 = self.k6c[canonical_k6 as usize] as usize; + Self::update_sums_increment::<6>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k6, g6); + self.k6c[canonical_k6 as usize] += 1; + self.k6q.push_back(canonical_k6); + self.steady = true; + } + } + } + } + } + } + + /// Normalized entropy at sub-word size `order` (1..=6). The caller is + /// responsible for not calling this before the window is full (`k` + /// bases pushed) — an empty/partial window yields a meaningless value. + pub fn entropy(&self, order: usize) -> f64 { + let k = self.k; + let em = emax(k, order); + if em <= 0.0 { + return 1.0; + } + let nwords = k - order + 1; + let log_nw = log_nwords(k, order); + let nw_f = nwords as f64; + let h_corr = log_nw + (self.sum_f_log_s[order] - self.sum_f_log_f[order]) / nw_f; + (h_corr / em).max(0.0) + } + + /// Minimum of [`Self::entropy`] over sub-word sizes `1..=order_max`, same + /// caller responsibility re: window readiness as `entropy`. + pub fn normalized_entropy(&self, order_max: usize) -> f64 { + let min_e = (1..=order_max) + .map(|ws| self.entropy(ws)) + .fold(f64::MAX, f64::min); + if min_e == f64::MAX { 1.0 } else { min_e } + } +} diff --git a/src/obikpartitionner/Cargo.toml b/src/obikpartitionner/Cargo.toml index a0c810e..918606f 100644 --- a/src/obikpartitionner/Cargo.toml +++ b/src/obikpartitionner/Cargo.toml @@ -13,6 +13,7 @@ obikrope = { path = "../obikrope" } niffler = "3.0.0" remove_dir_all = "0.8" obikseq = { path = "../obikseq" } +obikentropy = { path = "../obikentropy" } obiskbuilder = { path = "../obiskbuilder" } obiskio = { path = "../obiskio" } obidebruinj = { path = "../obidebruinj" } diff --git a/src/obikpartitionner/src/filter.rs b/src/obikpartitionner/src/filter.rs index db69aba..9f54a69 100644 --- a/src/obikpartitionner/src/filter.rs +++ b/src/obikpartitionner/src/filter.rs @@ -277,7 +277,7 @@ impl KmerFilter for GroupQuorumFilter { /// Reject k-mers with normalized entropy below `theta` — the same complexity /// metric `obikmer index`'s `--theta`/`--level-max` apply *during* superkmer -/// construction (see [`obiskbuilder::KmerEntropy`]), applied here after the +/// construction (see [`obikentropy::KmerEntropy`]), applied here after the /// fact, to k-mers already committed to a built index. /// /// Unlike every other filter in this module, this one needs the k-mer's own @@ -292,7 +292,7 @@ pub struct MinComplexity { impl KmerFilter for MinComplexity { fn passes(&self, kmer: CanonicalKmer, _row: &[u32], _n_genomes: usize) -> bool { - use obiskbuilder::KmerEntropy; + use obikentropy::KmerEntropy; kmer.entropy(self.level_max) >= self.theta } } diff --git a/src/obiskbuilder/Cargo.toml b/src/obiskbuilder/Cargo.toml index 5e19801..1e5c0af 100644 --- a/src/obiskbuilder/Cargo.toml +++ b/src/obiskbuilder/Cargo.toml @@ -7,7 +7,13 @@ edition = "2024" obikseq = { path = "../obikseq" } obikrope = { path = "../obikrope" } obiread = { path = "../obiread" } +obikentropy = { path = "../obikentropy" } lazy_static = "1.5.0" [dev-dependencies] obikseq = { path = "../obikseq", features = ["test-utils"] } +criterion2 = { version = "3", features = ["cargo_bench_support"] } + +[[bench]] +name = "superkmer_stream" +harness = false diff --git a/src/obiskbuilder/benches/superkmer_stream.rs b/src/obiskbuilder/benches/superkmer_stream.rs new file mode 100644 index 0000000..776e6d3 --- /dev/null +++ b/src/obiskbuilder/benches/superkmer_stream.rs @@ -0,0 +1,58 @@ +//! Throughput of the streaming superkmer pipeline (`RollingStat`'s hot path: +//! minimizer selection + entropy tracking fused in a single pass). +//! +//! Reference point for the `obikentropy` extraction: the entropy bookkeeping +//! that used to live inline in `RollingStat` was pulled out into a composed +//! `EntropyTracker`. This benchmark is run before and after that change to +//! confirm no regression. + +use criterion::{Criterion, Throughput, criterion_group, criterion_main}; +use obikrope::Rope; +use obiskbuilder::SuperKmerIter; + +const K: usize = 21; +const M: usize = 9; +const LEVEL_MAX: usize = 6; +const THETA: f64 = 0.7; +const SEQ_LEN: usize = 200_000; + +/// Deterministic pseudo-random ACGT sequence — high enough complexity that +/// the entropy filter rarely rejects, so the bench stays on the steady-state +/// path rather than repeatedly resetting. +fn make_sequence(len: usize) -> Vec { + let mut state: u64 = 0x9E3779B97F4A7C15; + (0..len) + .map(|_| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + b"ACGT"[(state % 4) as usize] + }) + .collect() +} + +fn make_rope(seq: &[u8]) -> Rope { + let mut rope = Rope::new(None); + rope.push(seq.to_vec()); + rope +} + +fn bench_build_superkmers(c: &mut Criterion) { + obikseq::set_k(K); + obikseq::set_m(M); + + let seq = make_sequence(SEQ_LEN); + let rope = make_rope(&seq); + + let mut group = c.benchmark_group("build_superkmers"); + group.throughput(Throughput::Bytes(SEQ_LEN as u64)); + group.bench_function("stream", |b| { + b.iter(|| { + SuperKmerIter::new(std::hint::black_box(&rope), K, LEVEL_MAX, THETA).count() + }); + }); + group.finish(); +} + +criterion_group!(benches, bench_build_superkmers); +criterion_main!(benches); diff --git a/src/obiskbuilder/src/kmer_entropy.rs b/src/obiskbuilder/src/kmer_entropy.rs deleted file mode 100644 index f8fcd7f..0000000 --- a/src/obiskbuilder/src/kmer_entropy.rs +++ /dev/null @@ -1,49 +0,0 @@ -//! Normalized entropy of an isolated, already-built k-mer. -//! -//! [`SuperKmerIter`](crate::SuperKmerIter) uses [`RollingStat`] to reject -//! low-complexity k-mers *during* superkmer construction, incrementally, over -//! a streaming window. [`KmerEntropy`] exposes the same metric for a single -//! k-mer taken in isolation (e.g. one already reconstructed from an index's -//! `unitigs.bin`, with no surrounding sequence) — built on the identical -//! `RollingStat` code path, not a re-derived formula, so a `theta` threshold -//! chosen for index-build-time filtering means the same thing when applied -//! after the fact (e.g. `obikmer filter`). - -use obikseq::CanonicalKmer; - -use crate::rolling_stat::RollingStat; - -/// Extension trait: compute the normalized entropy of a single canonical -/// k-mer, independent of any surrounding sequence. -pub trait KmerEntropy { - /// Normalized entropy across sub-word orders `1..=level_max` (the - /// minimum is taken across orders, same as - /// [`RollingStat::normalized_entropy`]). Lower means less complex; - /// `theta` in `index`/`filter` rejects k-mers with a score `< theta`. - /// - /// # Panics - /// - /// Requires both `obikseq::params::k()` *and* `params::m()` to already be - /// set (via `set_k`/`set_m`), even though the minimizer length plays no - /// conceptual role in an entropy score: `RollingStat` is a shared, - /// general-purpose struct (also used for minimizer selection during - /// superkmer decomposition) and unconditionally sizes an `m`-dependent - /// mask in its constructor. Callers must call `set_m` with a valid value - /// (e.g. the index's own stored `minimizer_size`) even when only the - /// entropy score is needed. - fn entropy(&self, level_max: usize) -> f64; -} - -impl KmerEntropy for CanonicalKmer { - fn entropy(&self, level_max: usize) -> f64 { - let mut stat = RollingStat::new(level_max); - for &base in self.to_ascii().iter() { - stat.push(base); - } - stat.normalized_entropy().unwrap_or(1.0) - } -} - -#[cfg(test)] -#[path = "tests/kmer_entropy.rs"] -mod tests; diff --git a/src/obiskbuilder/src/lib.rs b/src/obiskbuilder/src/lib.rs index 61d8cf1..2dd9c2f 100644 --- a/src/obiskbuilder/src/lib.rs +++ b/src/obiskbuilder/src/lib.rs @@ -6,16 +6,13 @@ #![deny(missing_docs)] pub mod iter; -pub mod kmer_entropy; pub mod stream_iter; mod scratch; pub(crate) mod encoding; -pub(crate) mod entropy_table; pub(crate) mod rolling_stat; pub use iter::SuperKmerIter; -pub use kmer_entropy::KmerEntropy; pub use scratch::SuperKmerScratch; pub use stream_iter::SuperKmerStreamIter; diff --git a/src/obiskbuilder/src/rolling_stat.rs b/src/obiskbuilder/src/rolling_stat.rs index 28d49a6..c824c99 100644 --- a/src/obiskbuilder/src/rolling_stat.rs +++ b/src/obiskbuilder/src/rolling_stat.rs @@ -1,8 +1,8 @@ +use obikentropy::{EntropyTracker, sub_word_canon}; use obikseq::kmer::{Minimizer, hash_kmer}; use obikseq::params; use crate::encoding::encode_nuc; -use crate::entropy_table::{WS_MAX, emax, entropy_norm_kmer, ln_class_size, log_nwords, n_log_n}; // ── Stack-allocated ring buffer ─────────────────────────────────────────────── @@ -83,33 +83,19 @@ pub struct RollingStat { entropy_max_k: usize, k: usize, m: usize, - steady: bool, rolling_k: u64, rolling_rck: u64, k_mask: u64, m_mask: u64, received: usize, - // Sliding-window queues — stack-allocated, capacity ≤ k ≤ 31. - k1q: Ring, - k2q: Ring, - k3q: Ring, - k4q: Ring, - k5q: Ring, - k6q: Ring, + // Minimizer selection state. minimier: Ring, - // Frequency count arrays. - // Max count per cell ≤ k ≤ 31 → u8 is sufficient. - k1c: [u8; 4], - k2c: [u8; 16], - k3c: [u8; 64], - k4c: [u8; 256], - k5c: [u8; 1024], - k6c: [u8; 4096], - - sum_f_log_f: [f64; WS_MAX + 1], - sum_f_log_s: [f64; WS_MAX + 1], + // Entropy tracking, composed as a plain inline field so both concerns + // update in the same streaming pass without being conflated in one + // struct — see `obikentropy::EntropyTracker`. + entropy: EntropyTracker, } impl RollingStat { @@ -120,27 +106,13 @@ impl RollingStat { entropy_max_k, k, m, - steady: false, rolling_k: 0, rolling_rck: 0, k_mask: (!0u64) >> (64 - k * 2), m_mask: (!0u64) >> (64 - m * 2), received: 0, - k1q: Ring::new(), - k2q: Ring::new(), - k3q: Ring::new(), - k4q: Ring::new(), - k5q: Ring::new(), - k6q: Ring::new(), minimier: Ring::new(), - k1c: [0; 4], - k2c: [0; 16], - k3c: [0; 64], - k4c: [0; 256], - k5c: [0; 1024], - k6c: [0; 4096], - sum_f_log_f: [0.0; WS_MAX + 1], - sum_f_log_s: [0.0; WS_MAX + 1], + entropy: EntropyTracker::new(k), } } @@ -148,54 +120,9 @@ impl RollingStat { self.rolling_k = 0; self.rolling_rck = 0; self.received = 0; - self.steady = false; - // for i in self.k1q.iter() { self.k1c[i as usize] = 0; } - // for i in self.k2q.iter() { self.k2c[i as usize] = 0; } - // for i in self.k3q.iter() { self.k3c[i as usize] = 0; } - // for i in self.k4q.iter() { self.k4c[i as usize] = 0; } - // for i in self.k5q.iter() { self.k5c[i as usize] = 0; } - // for i in self.k6q.iter() { self.k6c[i as usize] = 0; } - - self.k1c.fill(0); - self.k2c.fill(0); - self.k3c.fill(0); - self.k4c.fill(0); - self.k5c.fill(0); - self.k6c.fill(0); - - self.k1q.clear(); - self.k2q.clear(); - self.k3q.clear(); - self.k4q.clear(); - self.k5q.clear(); - self.k6q.clear(); self.minimier.clear(); - - self.sum_f_log_f = [0.0; WS_MAX + 1]; - self.sum_f_log_s = [0.0; WS_MAX + 1]; - } - - #[inline] - fn update_sums_decrement( - sum_f_log_f: &mut [f64; WS_MAX + 1], - sum_f_log_s: &mut [f64; WS_MAX + 1], - canonical: u64, - f: usize, - ) { - sum_f_log_f[K] += n_log_n(f - 1) - n_log_n(f); - sum_f_log_s[K] -= ln_class_size::(canonical); - } - - #[inline] - fn update_sums_increment( - sum_f_log_f: &mut [f64; WS_MAX + 1], - sum_f_log_s: &mut [f64; WS_MAX + 1], - canonical: u64, - g: usize, - ) { - sum_f_log_f[K] += n_log_n(g + 1) - n_log_n(g); - sum_f_log_s[K] += ln_class_size::(canonical); + self.entropy.reset(); } pub fn push(&mut self, nuc: u8) { @@ -209,12 +136,7 @@ impl RollingStat { self.rolling_rck = ((self.rolling_rck >> 2) | ((cnuc as u64) << ((k - 1) * 2))) & self.k_mask; - let canonical_k1 = entropy_norm_kmer::(self.rolling_k & 3); - let canonical_k2 = entropy_norm_kmer::(self.rolling_k & 15); - let canonical_k3 = entropy_norm_kmer::(self.rolling_k & 63); - let canonical_k4 = entropy_norm_kmer::(self.rolling_k & 255); - let canonical_k5 = entropy_norm_kmer::(self.rolling_k & 1023); - let canonical_k6 = entropy_norm_kmer::(self.rolling_k & 4095); + let canon = sub_word_canon(self.rolling_k); self.received += 1; @@ -248,153 +170,7 @@ impl RollingStat { } } - if self.received > k { - let old1 = self.k1q.pop_front(); - let f1 = self.k1c[old1 as usize] as usize; - Self::update_sums_decrement::<1>( - &mut self.sum_f_log_f, - &mut self.sum_f_log_s, - old1, - f1, - ); - self.k1c[old1 as usize] -= 1; - - let old2 = self.k2q.pop_front(); - let f2 = self.k2c[old2 as usize] as usize; - Self::update_sums_decrement::<2>( - &mut self.sum_f_log_f, - &mut self.sum_f_log_s, - old2, - f2, - ); - self.k2c[old2 as usize] -= 1; - - let old3 = self.k3q.pop_front(); - let f3 = self.k3c[old3 as usize] as usize; - Self::update_sums_decrement::<3>( - &mut self.sum_f_log_f, - &mut self.sum_f_log_s, - old3, - f3, - ); - self.k3c[old3 as usize] -= 1; - - let old4 = self.k4q.pop_front(); - let f4 = self.k4c[old4 as usize] as usize; - Self::update_sums_decrement::<4>( - &mut self.sum_f_log_f, - &mut self.sum_f_log_s, - old4, - f4, - ); - self.k4c[old4 as usize] -= 1; - - let old5 = self.k5q.pop_front(); - let f5 = self.k5c[old5 as usize] as usize; - Self::update_sums_decrement::<5>( - &mut self.sum_f_log_f, - &mut self.sum_f_log_s, - old5, - f5, - ); - self.k5c[old5 as usize] -= 1; - - let old6 = self.k6q.pop_front(); - let f6 = self.k6c[old6 as usize] as usize; - Self::update_sums_decrement::<6>( - &mut self.sum_f_log_f, - &mut self.sum_f_log_s, - old6, - f6, - ); - self.k6c[old6 as usize] -= 1; - } - - if self.steady { - let g1 = self.k1c[canonical_k1 as usize] as usize; - Self::update_sums_increment::<1>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k1, g1); - self.k1c[canonical_k1 as usize] += 1; - self.k1q.push_back(canonical_k1); - - let g2 = self.k2c[canonical_k2 as usize] as usize; - Self::update_sums_increment::<2>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k2, g2); - self.k2c[canonical_k2 as usize] += 1; - self.k2q.push_back(canonical_k2); - - let g3 = self.k3c[canonical_k3 as usize] as usize; - Self::update_sums_increment::<3>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k3, g3); - self.k3c[canonical_k3 as usize] += 1; - self.k3q.push_back(canonical_k3); - - let g4 = self.k4c[canonical_k4 as usize] as usize; - Self::update_sums_increment::<4>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k4, g4); - self.k4c[canonical_k4 as usize] += 1; - self.k4q.push_back(canonical_k4); - - let g5 = self.k5c[canonical_k5 as usize] as usize; - Self::update_sums_increment::<5>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k5, g5); - self.k5c[canonical_k5 as usize] += 1; - self.k5q.push_back(canonical_k5); - - let g6 = self.k6c[canonical_k6 as usize] as usize; - Self::update_sums_increment::<6>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k6, g6); - self.k6c[canonical_k6 as usize] += 1; - self.k6q.push_back(canonical_k6); - } else { - self.push_warmup_increments( - canonical_k1, canonical_k2, canonical_k3, - canonical_k4, canonical_k5, canonical_k6, - ); - } - } - - #[cold] - #[inline(never)] - fn push_warmup_increments( - &mut self, - canonical_k1: u64, canonical_k2: u64, canonical_k3: u64, - canonical_k4: u64, canonical_k5: u64, canonical_k6: u64, - ) { - let g1 = self.k1c[canonical_k1 as usize] as usize; - Self::update_sums_increment::<1>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k1, g1); - self.k1c[canonical_k1 as usize] += 1; - self.k1q.push_back(canonical_k1); - - if self.received >= 2 { - let g2 = self.k2c[canonical_k2 as usize] as usize; - Self::update_sums_increment::<2>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k2, g2); - self.k2c[canonical_k2 as usize] += 1; - self.k2q.push_back(canonical_k2); - - if self.received >= 3 { - let g3 = self.k3c[canonical_k3 as usize] as usize; - Self::update_sums_increment::<3>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k3, g3); - self.k3c[canonical_k3 as usize] += 1; - self.k3q.push_back(canonical_k3); - - if self.received >= 4 { - let g4 = self.k4c[canonical_k4 as usize] as usize; - Self::update_sums_increment::<4>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k4, g4); - self.k4c[canonical_k4 as usize] += 1; - self.k4q.push_back(canonical_k4); - - if self.received >= 5 { - let g5 = self.k5c[canonical_k5 as usize] as usize; - Self::update_sums_increment::<5>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k5, g5); - self.k5c[canonical_k5 as usize] += 1; - self.k5q.push_back(canonical_k5); - - if self.received >= 6 { - let g6 = self.k6c[canonical_k6 as usize] as usize; - Self::update_sums_increment::<6>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k6, g6); - self.k6c[canonical_k6 as usize] += 1; - self.k6q.push_back(canonical_k6); - self.steady = true; - } - } - } - } - } + self.entropy.push(self.received, canon); } pub fn ready(&self) -> bool { @@ -426,25 +202,13 @@ impl RollingStat { if !self.ready() { return None; } - let k = self.k; - let em = emax(k, order); - if em <= 0.0 { - return Some(1.0); - } - let nwords = k - order + 1; - let log_nw = log_nwords(k, order); - let nw_f = nwords as f64; - let h_corr = log_nw + (self.sum_f_log_s[order] - self.sum_f_log_f[order]) / nw_f; - Some((h_corr / em).max(0.0)) + Some(self.entropy.entropy(order)) } pub fn normalized_entropy(&self) -> Option { if !self.ready() { return None; } - let min_e = (1..=self.entropy_max_k) - .filter_map(|ws| self.entropy(ws)) - .fold(f64::MAX, f64::min); - Some(if min_e == f64::MAX { 1.0 } else { min_e }) + Some(self.entropy.normalized_entropy(self.entropy_max_k)) } }