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
+145
View File
@@ -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();
}