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
@@ -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"] }
|
||||
@@ -0,0 +1,145 @@
|
||||
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();
|
||||
}
|
||||
@@ -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;
|
||||
@@ -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};
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,112 @@
|
||||
//! 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>();
|
||||
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"),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
use super::*;
|
||||
use obikseq::Sequence;
|
||||
use obikseq::kmer::Kmer;
|
||||
|
||||
const K: usize = 21;
|
||||
const LEVEL_MAX: usize = 6;
|
||||
|
||||
fn kmer_from_ascii(seq: &[u8]) -> CanonicalKmer {
|
||||
obikseq::set_k(K);
|
||||
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:?}");
|
||||
}
|
||||
}
|
||||
@@ -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 }
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user