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.
This commit is contained in:
Eric Coissac
2026-07-08 18:36:16 +02:00
parent e725523898
commit 912f788f7f
17 changed files with 482 additions and 305 deletions
+10
View File
@@ -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",
+1 -1
View File
@@ -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
+10
View File
@@ -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"] }
+41
View File
@@ -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;
+17
View File
@@ -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};
+40
View File
@@ -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<T: Copy + Default, const N: usize> {
buf: [T; N],
head: usize,
len: usize,
}
impl<T: Copy + Default, const N: usize> Ring<T, N> {
#[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
}
}
@@ -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>();
@@ -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()
}
+280
View File
@@ -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::<false, 1>(rolling_kmer & 3),
entropy_norm_kmer::<false, 2>(rolling_kmer & 15),
entropy_norm_kmer::<false, 3>(rolling_kmer & 63),
entropy_norm_kmer::<false, 4>(rolling_kmer & 255),
entropy_norm_kmer::<false, 5>(rolling_kmer & 1023),
entropy_norm_kmer::<false, 6>(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<u64, 32>,
k2q: Ring<u64, 32>,
k3q: Ring<u64, 32>,
k4q: Ring<u64, 32>,
k5q: Ring<u64, 32>,
k6q: Ring<u64, 32>,
// 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<const K: usize>(
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::<false, K>(canonical);
}
#[inline]
fn update_sums_increment<const K: usize>(
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::<false, K>(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 }
}
}
+1
View File
@@ -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" }
+2 -2
View File
@@ -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
}
}
+6
View File
@@ -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
@@ -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<u8> {
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);
-49
View File
@@ -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;
-3
View File
@@ -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;
+12 -248
View File
@@ -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<u64, 32>,
k2q: Ring<u64, 32>,
k3q: Ring<u64, 32>,
k4q: Ring<u64, 32>,
k5q: Ring<u64, 32>,
k6q: Ring<u64, 32>,
// Minimizer selection state.
minimier: Ring<MmerItem, 32>,
// 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<const K: usize>(
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::<false, K>(canonical);
}
#[inline]
fn update_sums_increment<const K: usize>(
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::<false, K>(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::<false, 1>(self.rolling_k & 3);
let canonical_k2 = entropy_norm_kmer::<false, 2>(self.rolling_k & 15);
let canonical_k3 = entropy_norm_kmer::<false, 3>(self.rolling_k & 63);
let canonical_k4 = entropy_norm_kmer::<false, 4>(self.rolling_k & 255);
let canonical_k5 = entropy_norm_kmer::<false, 5>(self.rolling_k & 1023);
let canonical_k6 = entropy_norm_kmer::<false, 6>(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<f64> {
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))
}
}