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
@@ -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);
|
||||
@@ -1,145 +0,0 @@
|
||||
use std::fs;
|
||||
use std::path::PathBuf;
|
||||
|
||||
const K_MAX: usize = 32;
|
||||
const WS_MAX: usize = 6;
|
||||
|
||||
fn normalize_circular(kmer: u64, ws: usize) -> u64 {
|
||||
let mask = (1u64 << (ws * 2)) - 1;
|
||||
let mut canonical = kmer & mask;
|
||||
let mut current = canonical;
|
||||
for _ in 0..ws - 1 {
|
||||
let top = (current >> ((ws - 1) * 2)) & 3;
|
||||
current = ((current << 2) | top) & mask;
|
||||
if current < canonical {
|
||||
canonical = current;
|
||||
}
|
||||
}
|
||||
canonical
|
||||
}
|
||||
|
||||
fn revcomp_raw(x: u64, k: usize) -> u64 {
|
||||
let x = !x;
|
||||
let x = x.swap_bytes();
|
||||
let x = ((x >> 4) & 0x0F0F0F0F0F0F0F0F) | ((x & 0x0F0F0F0F0F0F0F0F) << 4);
|
||||
let x = ((x >> 2) & 0x3333333333333333) | ((x & 0x3333333333333333) << 2);
|
||||
x << (64 - 2 * k)
|
||||
}
|
||||
|
||||
fn build_normalized_kmer(k: usize) -> Vec<u64> {
|
||||
let n = 1usize << (k * 2);
|
||||
let shift = 64 - k * 2;
|
||||
let mut result = vec![0u64; n];
|
||||
for i in 0..n {
|
||||
let la = (i as u64) << shift;
|
||||
let ra = i as u64;
|
||||
let rc_ra = revcomp_raw(la, k) >> shift;
|
||||
let circ = normalize_circular(ra, k);
|
||||
let circ_rc = normalize_circular(rc_ra, k);
|
||||
result[i] = if circ < circ_rc { circ } else { circ_rc };
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
fn build_ln_class(norm: &[u64]) -> Vec<f64> {
|
||||
let n = norm.len();
|
||||
let mut sizes = vec![0u32; n];
|
||||
for &c in norm {
|
||||
sizes[c as usize] += 1;
|
||||
}
|
||||
norm.iter()
|
||||
.map(|&c| {
|
||||
let s = sizes[c as usize];
|
||||
if s > 0 { (s as f64).ln() } else { 0.0 }
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn build_n_log_n() -> [f64; K_MAX + 1] {
|
||||
let mut t = [0.0f64; K_MAX + 1];
|
||||
for n in 1..=K_MAX {
|
||||
t[n] = (n as f64) * (n as f64).ln();
|
||||
}
|
||||
t
|
||||
}
|
||||
|
||||
fn build_emax() -> [[f64; WS_MAX + 1]; K_MAX + 1] {
|
||||
let mut t = [[0.0f64; WS_MAX + 1]; K_MAX + 1];
|
||||
for k in 2..=K_MAX {
|
||||
for ws in 1..=WS_MAX.min(k - 1) {
|
||||
let n_raw = 1usize << (ws * 2);
|
||||
let nwords = k - ws + 1;
|
||||
let c = nwords / n_raw;
|
||||
let r = nwords % n_raw;
|
||||
let nf = nwords as f64;
|
||||
let t1 = if c == 0 || n_raw == r {
|
||||
0.0
|
||||
} else {
|
||||
let f1 = c as f64 / nf;
|
||||
(n_raw - r) as f64 * f1 * f1.ln()
|
||||
};
|
||||
let t2 = if r == 0 {
|
||||
0.0
|
||||
} else {
|
||||
let f2 = (c + 1) as f64 / nf;
|
||||
r as f64 * f2 * f2.ln()
|
||||
};
|
||||
t[k][ws] = -(t1 + t2);
|
||||
}
|
||||
}
|
||||
t
|
||||
}
|
||||
|
||||
fn build_log_nwords() -> [[f64; WS_MAX + 1]; K_MAX + 1] {
|
||||
let mut t = [[0.0f64; WS_MAX + 1]; K_MAX + 1];
|
||||
for k in 2..=K_MAX {
|
||||
for ws in 1..=WS_MAX.min(k - 1) {
|
||||
t[k][ws] = ((k - ws + 1) as f64).ln();
|
||||
}
|
||||
}
|
||||
t
|
||||
}
|
||||
|
||||
fn emit_f64_1d(out: &mut String, name: &str, n: usize, values: &[f64]) {
|
||||
out.push_str(&format!("pub(crate) const {name}: [f64; {n}] = [\n"));
|
||||
for v in values {
|
||||
out.push_str(&format!(" {v:?},\n"));
|
||||
}
|
||||
out.push_str("];\n");
|
||||
}
|
||||
|
||||
fn emit_f64_2d(out: &mut String, name: &str, rows: usize, cols: usize, values: &[[f64; WS_MAX + 1]]) {
|
||||
out.push_str(&format!("pub(crate) const {name}: [[f64; {cols}]; {rows}] = [\n"));
|
||||
for row in values {
|
||||
out.push_str(" [");
|
||||
for (i, v) in row.iter().enumerate() {
|
||||
if i > 0 { out.push_str(", "); }
|
||||
out.push_str(&format!("{v:?}"));
|
||||
}
|
||||
out.push_str("],\n");
|
||||
}
|
||||
out.push_str("];\n");
|
||||
}
|
||||
|
||||
fn main() {
|
||||
let out_dir = PathBuf::from(std::env::var("OUT_DIR").unwrap());
|
||||
let mut out = String::new();
|
||||
|
||||
for k in 1..=6usize {
|
||||
let n = 1usize << (k * 2);
|
||||
let norm = build_normalized_kmer(k);
|
||||
let ln_class = build_ln_class(&norm);
|
||||
emit_f64_1d(&mut out, &format!("LN_CLASS{k}"), n, &ln_class);
|
||||
}
|
||||
|
||||
let n_log_n = build_n_log_n();
|
||||
emit_f64_1d(&mut out, "N_LOG_N", K_MAX + 1, &n_log_n);
|
||||
|
||||
let emax = build_emax();
|
||||
emit_f64_2d(&mut out, "EMAX", K_MAX + 1, WS_MAX + 1, &emax);
|
||||
|
||||
let log_nwords = build_log_nwords();
|
||||
emit_f64_2d(&mut out, "LOG_NWORDS", K_MAX + 1, WS_MAX + 1, &log_nwords);
|
||||
|
||||
fs::write(out_dir.join("ln_class_tables.rs"), out).unwrap();
|
||||
}
|
||||
@@ -1,108 +0,0 @@
|
||||
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>();
|
||||
pub(crate) const NORMK4: [u64; 256] = build_normalized_kmer::<256>();
|
||||
pub(crate) const NORMK5: [u64; 1024] = build_normalized_kmer::<1024>();
|
||||
pub(crate) const NORMK6: [u64; 4096] = build_normalized_kmer::<4096>();
|
||||
|
||||
include!(concat!(env!("OUT_DIR"), "/ln_class_tables.rs"));
|
||||
|
||||
const fn normalize_circular(kmer: u64, ws: usize) -> u64 {
|
||||
let mask = (1u64 << (ws * 2)) - 1;
|
||||
let mut canonical = kmer & mask;
|
||||
let mut current = canonical;
|
||||
let mut i = 0;
|
||||
while i < (ws - 1) {
|
||||
let top = (current >> ((ws - 1) * 2)) & 3;
|
||||
current = ((current << 2) | top) & mask;
|
||||
if current < canonical {
|
||||
canonical = current;
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
canonical
|
||||
}
|
||||
|
||||
const fn build_normalized_kmer<const N: usize>() -> [u64; N] {
|
||||
let mut result = [0u64; N];
|
||||
let k = k_from_n::<N>();
|
||||
let shift = 64 - k * 2;
|
||||
let mut i = 0;
|
||||
while i < N {
|
||||
let la = (i as u64) << shift;
|
||||
let ra = i as u64;
|
||||
let rc_ra = revcomp_raw(la, k) >> shift;
|
||||
let circ = normalize_circular(ra, k);
|
||||
let circ_rc = normalize_circular(rc_ra, k);
|
||||
result[i] = if circ < circ_rc { circ } else { circ_rc };
|
||||
i += 1;
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
const fn revcomp_raw(x: u64, k: usize) -> u64 {
|
||||
let x = !x;
|
||||
let x = x.swap_bytes();
|
||||
let x = ((x >> 4) & 0x0F0F0F0F0F0F0F0F) | ((x & 0x0F0F0F0F0F0F0F0F) << 4);
|
||||
let x = ((x >> 2) & 0x3333333333333333) | ((x & 0x3333333333333333) << 2);
|
||||
x << (64 - 2 * k)
|
||||
}
|
||||
|
||||
const fn k_from_n<const N: usize>() -> usize {
|
||||
match N {
|
||||
4 => 1,
|
||||
16 => 2,
|
||||
64 => 3,
|
||||
256 => 4,
|
||||
1024 => 5,
|
||||
4096 => 6,
|
||||
_ => panic!("N must be a power of 4"),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) const WS_MAX: usize = 6;
|
||||
|
||||
#[inline(always)]
|
||||
pub(crate) const fn n_log_n(n: usize) -> f64 {
|
||||
N_LOG_N[n]
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
pub(crate) const fn emax(k: usize, ws: usize) -> f64 {
|
||||
EMAX[k][ws]
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
pub(crate) const fn log_nwords(k: usize, ws: usize) -> f64 {
|
||||
LOG_NWORDS[k][ws]
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
pub(crate) const fn entropy_norm_kmer<const LEFT: bool, const K: usize>(kmer: u64) -> u64 {
|
||||
const SHIFT: [usize; 7] = [0, 62, 60, 58, 56, 54, 52];
|
||||
const NORM: [&[u64]; 7] = [&[], &NORMK1, &NORMK2, &NORMK3, &NORMK4, &NORMK5, &NORMK6];
|
||||
|
||||
let shift = SHIFT[K];
|
||||
let ra = if LEFT { kmer >> shift } else { kmer };
|
||||
let canonical_ra = NORM[K][ra as usize];
|
||||
if LEFT {
|
||||
canonical_ra << shift
|
||||
} else {
|
||||
canonical_ra
|
||||
}
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
pub(crate) const fn ln_class_size<const LEFT: bool, const K: usize>(kmer: u64) -> f64 {
|
||||
const SHIFT: [usize; 7] = [0, 62, 60, 58, 56, 54, 52];
|
||||
let ra = if LEFT { kmer >> SHIFT[K] } else { kmer };
|
||||
match K {
|
||||
1 => LN_CLASS1[ra as usize],
|
||||
2 => LN_CLASS2[ra as usize],
|
||||
3 => LN_CLASS3[ra as usize],
|
||||
4 => LN_CLASS4[ra as usize],
|
||||
5 => LN_CLASS5[ra as usize],
|
||||
6 => LN_CLASS6[ra as usize],
|
||||
_ => panic!("k must be 1..=6"),
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
@@ -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;
|
||||
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,54 +0,0 @@
|
||||
use super::*;
|
||||
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()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn homopolymer_scores_lower_than_diverse_sequence() {
|
||||
let homopolymer = kmer_from_ascii(b"AAAAAAAAAAAAAAAAAAAAA"); // 21 bases
|
||||
let diverse = kmer_from_ascii(b"CATTAGCGTACCTGATCAGGT"); // 21 bases, same as used elsewhere in this workspace's tests
|
||||
|
||||
let e_homopolymer = homopolymer.entropy(LEVEL_MAX);
|
||||
let e_diverse = diverse.entropy(LEVEL_MAX);
|
||||
|
||||
assert!(
|
||||
e_homopolymer < e_diverse,
|
||||
"homopolymer ({e_homopolymer}) should score lower than a diverse sequence ({e_diverse})"
|
||||
);
|
||||
// A pure homopolymer is the most degenerate case representable — its
|
||||
// score should sit near the bottom of the range, not just "somewhat lower".
|
||||
assert!(e_homopolymer < 0.3, "homopolymer entropy unexpectedly high: {e_homopolymer}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn entropy_is_deterministic_for_the_same_kmer() {
|
||||
let a = kmer_from_ascii(b"CATTAGCGTACCTGATCAGGT");
|
||||
let b = kmer_from_ascii(b"CATTAGCGTACCTGATCAGGT");
|
||||
assert_eq!(a.entropy(LEVEL_MAX), b.entropy(LEVEL_MAX));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn entropy_is_within_zero_one_range() {
|
||||
let mut repeat = "AT".repeat(K / 2 + 1);
|
||||
repeat.truncate(K);
|
||||
|
||||
for seq in [
|
||||
"AAAAAAAAAAAAAAAAAAAAA".to_string(),
|
||||
repeat,
|
||||
"CATTAGCGTACCTGATCAGGT".to_string(),
|
||||
] {
|
||||
assert_eq!(seq.len(), K, "test sequence must be exactly K bases: {seq:?}");
|
||||
let kmer = kmer_from_ascii(seq.as_bytes());
|
||||
let e = kmer.entropy(LEVEL_MAX);
|
||||
assert!((0.0..=1.0).contains(&e), "entropy {e} out of [0,1] for {seq:?}");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user