perf: optimize k-mer pipeline with compile-time tables
This commit shifts entropy and lookup table generation to compile time via a new build script, eliminating runtime overhead. It replaces heap-allocated queues in rolling statistics with a stack-allocated, const-generic ring buffer for cache-friendly operations, and implements `size_hint` on `SuperKmerIter` for efficient iterator consumption. Additionally, it establishes the baseline profile configuration and sets global k-mer parameters.
This commit is contained in:
Eric Coissac
Binary file not shown.
@@ -3,7 +3,7 @@ use std::path::PathBuf;
|
||||
|
||||
use clap::Args;
|
||||
use obifastwrite::write_scatter;
|
||||
use obikseq::RoutableSuperKmer;
|
||||
use obikseq::{RoutableSuperKmer, set_k, set_m};
|
||||
|
||||
use crate::cli::{CommonArgs, PipelineData, open_chunks};
|
||||
|
||||
@@ -41,6 +41,9 @@ pub fn run(args: SuperkmerArgs) {
|
||||
let partition_bits = args.common.partition_bits;
|
||||
let n_workers = args.common.threads.max(1);
|
||||
|
||||
set_k(k);
|
||||
set_m(m);
|
||||
|
||||
let path_source = args.common.seqfile_paths();
|
||||
|
||||
let pipe = obipipeline::make_pipe! {
|
||||
|
||||
@@ -118,7 +118,7 @@ pub struct KmerOf<L: KmerLength>(u64, PhantomData<L>);
|
||||
impl<L: KmerLength> KmerOf<L> {
|
||||
/// Wrap a raw left-aligned u64 value as a kmer.
|
||||
#[inline]
|
||||
pub fn from_raw(raw: u64) -> Self {
|
||||
pub const fn from_raw(raw: u64) -> Self {
|
||||
KmerOf(raw, PhantomData)
|
||||
}
|
||||
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
@@ -1,173 +1,108 @@
|
||||
use obikseq::kmer::Kmer;
|
||||
use std::sync::LazyLock;
|
||||
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>();
|
||||
|
||||
pub(crate) static NORMK1: LazyLock<[u64; 4]> = LazyLock::new(|| build_normalized_kmer::<4>());
|
||||
pub(crate) static NORMK2: LazyLock<[u64; 16]> = LazyLock::new(|| build_normalized_kmer::<16>());
|
||||
pub(crate) static NORMK3: LazyLock<[u64; 64]> = LazyLock::new(|| build_normalized_kmer::<64>());
|
||||
pub(crate) static NORMK4: LazyLock<[u64; 256]> = LazyLock::new(|| build_normalized_kmer::<256>());
|
||||
pub(crate) static NORMK5: LazyLock<[u64; 1024]> = LazyLock::new(|| build_normalized_kmer::<1024>());
|
||||
pub(crate) static NORMK6: LazyLock<[u64; 4096]> = LazyLock::new(|| build_normalized_kmer::<4096>());
|
||||
include!(concat!(env!("OUT_DIR"), "/ln_class_tables.rs"));
|
||||
|
||||
pub(crate) static LN_CARD_ROT1: LazyLock<[f64; 4]> =
|
||||
LazyLock::new(|| build_log_class_size::<4>(&NORMK1));
|
||||
pub(crate) static LN_CARD_ROT2: LazyLock<[f64; 16]> =
|
||||
LazyLock::new(|| build_log_class_size::<16>(&NORMK2));
|
||||
pub(crate) static LN_CARD_ROT3: LazyLock<[f64; 64]> =
|
||||
LazyLock::new(|| build_log_class_size::<64>(&NORMK3));
|
||||
pub(crate) static LN_CARD_ROT4: LazyLock<[f64; 256]> =
|
||||
LazyLock::new(|| build_log_class_size::<256>(&NORMK4));
|
||||
pub(crate) static LN_CARD_ROT5: LazyLock<[f64; 1024]> =
|
||||
LazyLock::new(|| build_log_class_size::<1024>(&NORMK5));
|
||||
pub(crate) static LN_CARD_ROT6: LazyLock<[f64; 4096]> =
|
||||
LazyLock::new(|| build_log_class_size::<4096>(&NORMK6));
|
||||
|
||||
fn ln0(x: f64) -> f64 {
|
||||
if x == 0.0 { 0.0 } else { x.ln() }
|
||||
}
|
||||
|
||||
fn normalize_circular(kmer: u64, ws: usize) -> u64 {
|
||||
const 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 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
|
||||
}
|
||||
|
||||
fn build_normalized_kmer<const N: usize>() -> [u64; N] {
|
||||
const fn build_normalized_kmer<const N: usize>() -> [u64; N] {
|
||||
let mut result = [0u64; N];
|
||||
let k = N.ilog(4) as usize;
|
||||
let k = k_from_n::<N>();
|
||||
let shift = 64 - k * 2;
|
||||
for i in 0..N {
|
||||
let mut i = 0;
|
||||
while i < N {
|
||||
let la = (i as u64) << shift;
|
||||
let ra = i as u64;
|
||||
let rc_ra = Kmer::from_raw(la).revcomp().raw() >> shift;
|
||||
let rc_ra = revcomp_raw(la, k) >> shift;
|
||||
let circ = normalize_circular(ra, k);
|
||||
let circ_rc = normalize_circular(rc_ra, k);
|
||||
result[i] = circ.min(circ_rc);
|
||||
result[i] = if circ < circ_rc { circ } else { circ_rc };
|
||||
i += 1;
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
fn build_log_class_size<const N: usize>(norm: &[u64; N]) -> [f64; N] {
|
||||
let mut sizes = [0u32; N];
|
||||
for &c in norm {
|
||||
sizes[c as usize] += 1;
|
||||
}
|
||||
let mut result = [0.0f64; N];
|
||||
for i in 0..N {
|
||||
if sizes[i] > 0 {
|
||||
result[i] = ln0(sizes[i] as f64);
|
||||
}
|
||||
}
|
||||
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)
|
||||
}
|
||||
|
||||
pub(crate) fn entropy_norm_kmer(kmer: u64, k: usize, left: bool) -> u64 {
|
||||
let shift = 64 - k * 2;
|
||||
let ra = if left { kmer >> shift } else { kmer }; // left-aligned → right-aligned index
|
||||
let canonical_ra = match k {
|
||||
1 => NORMK1[ra as usize],
|
||||
2 => NORMK2[ra as usize],
|
||||
3 => NORMK3[ra as usize],
|
||||
4 => NORMK4[ra as usize],
|
||||
5 => NORMK5[ra as usize],
|
||||
6 => NORMK6[ra as usize],
|
||||
_ => panic!("k must be 1..=6"),
|
||||
};
|
||||
if left {
|
||||
canonical_ra << shift
|
||||
} else {
|
||||
canonical_ra
|
||||
} // right-aligned → left-aligned
|
||||
}
|
||||
|
||||
pub(crate) fn ln_class_size(kmer: u64, k: usize, left: bool) -> f64 {
|
||||
let ra = if left { kmer >> (64 - k * 2) } else { kmer }; // left-aligned → right-aligned index
|
||||
match k {
|
||||
1 => LN_CARD_ROT1[NORMK1[ra as usize] as usize],
|
||||
2 => LN_CARD_ROT2[NORMK2[ra as usize] as usize],
|
||||
3 => LN_CARD_ROT3[NORMK3[ra as usize] as usize],
|
||||
4 => LN_CARD_ROT4[NORMK4[ra as usize] as usize],
|
||||
5 => LN_CARD_ROT5[NORMK5[ra as usize] as usize],
|
||||
6 => LN_CARD_ROT6[NORMK6[ra as usize] as usize],
|
||||
_ => panic!("k must be 1..=6"),
|
||||
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"),
|
||||
}
|
||||
}
|
||||
|
||||
// ── k-dependent tables (k ≤ K_MAX, ws ≤ WS_MAX) ──────────────────────────────
|
||||
|
||||
pub(crate) const K_MAX: usize = 32;
|
||||
pub(crate) const WS_MAX: usize = 6;
|
||||
|
||||
/// n·ln(n), with n_log_n[0] = 0. Indexed by n = 0..=K_MAX.
|
||||
pub(crate) static N_LOG_N: LazyLock<[f64; K_MAX + 1]> = LazyLock::new(|| build_n_log_n());
|
||||
|
||||
/// H_max[k][ws]: maximum entropy for kmer length k and word size ws.
|
||||
pub(crate) static EMAX: LazyLock<[[f64; WS_MAX + 1]; K_MAX + 1]> = LazyLock::new(|| build_emax());
|
||||
|
||||
/// ln(k − ws + 1): log of the number of ws-words in a kmer of length k.
|
||||
pub(crate) static LOG_NWORDS: LazyLock<[[f64; WS_MAX + 1]; K_MAX + 1]> =
|
||||
LazyLock::new(|| build_log_nwords());
|
||||
|
||||
pub(crate) fn n_log_n(n: usize) -> f64 {
|
||||
#[inline(always)]
|
||||
pub(crate) const fn n_log_n(n: usize) -> f64 {
|
||||
N_LOG_N[n]
|
||||
}
|
||||
|
||||
pub(crate) fn emax(k: usize, ws: usize) -> f64 {
|
||||
#[inline(always)]
|
||||
pub(crate) const fn emax(k: usize, ws: usize) -> f64 {
|
||||
EMAX[k][ws]
|
||||
}
|
||||
|
||||
pub(crate) fn log_nwords(k: usize, ws: usize) -> f64 {
|
||||
#[inline(always)]
|
||||
pub(crate) const fn log_nwords(k: usize, ws: usize) -> f64 {
|
||||
LOG_NWORDS[k][ws]
|
||||
}
|
||||
|
||||
fn build_n_log_n() -> [f64; K_MAX + 1] {
|
||||
let mut result = [0.0f64; K_MAX + 1];
|
||||
for n in 1..=K_MAX {
|
||||
result[n] = (n as f64) * (n as f64).ln();
|
||||
#[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
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
fn build_emax() -> [[f64; WS_MAX + 1]; K_MAX + 1] {
|
||||
let mut result = [[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); // 4^ws
|
||||
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()
|
||||
};
|
||||
result[k][ws] = -(t1 + t2);
|
||||
}
|
||||
#[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"),
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
fn build_log_nwords() -> [[f64; WS_MAX + 1]; K_MAX + 1] {
|
||||
let mut result = [[0.0f64; WS_MAX + 1]; K_MAX + 1];
|
||||
for k in 2..=K_MAX {
|
||||
for ws in 1..=WS_MAX.min(k - 1) {
|
||||
result[k][ws] = ((k - ws + 1) as f64).ln();
|
||||
}
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
@@ -26,6 +26,7 @@ use crate::scratch::SuperKmerScratch;
|
||||
pub struct SuperKmerIter<'a> {
|
||||
cursor: ForwardCursor<'a>,
|
||||
k: usize,
|
||||
rope_len: usize,
|
||||
theta: f64,
|
||||
scratch: SuperKmerScratch,
|
||||
stat: RollingStat,
|
||||
@@ -44,6 +45,7 @@ impl<'a> SuperKmerIter<'a> {
|
||||
Self {
|
||||
cursor: rope.fw_cursor(),
|
||||
k,
|
||||
rope_len: rope.len(),
|
||||
theta,
|
||||
scratch: SuperKmerScratch::new(),
|
||||
stat: RollingStat::new(level_max),
|
||||
@@ -71,6 +73,10 @@ impl<'a> SuperKmerIter<'a> {
|
||||
impl Iterator for SuperKmerIter<'_> {
|
||||
type Item = RoutableSuperKmer;
|
||||
|
||||
fn size_hint(&self) -> (usize, Option<usize>) {
|
||||
(0, Some(self.rope_len / self.k))
|
||||
}
|
||||
|
||||
fn next(&mut self) -> Option<RoutableSuperKmer> {
|
||||
loop {
|
||||
let byte = match self.cursor.read_next().ok() {
|
||||
|
||||
@@ -4,36 +4,118 @@ 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};
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
// ── Stack-allocated ring buffer ───────────────────────────────────────────────
|
||||
|
||||
/// Fixed-capacity ring buffer backed by a stack array.
|
||||
/// N must be a power of two; operations are branchless via `% N`.
|
||||
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]
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
buf: [T::default(); N],
|
||||
head: 0,
|
||||
len: 0,
|
||||
}
|
||||
}
|
||||
#[inline]
|
||||
fn is_empty(&self) -> bool {
|
||||
self.len == 0
|
||||
}
|
||||
#[inline]
|
||||
fn clear(&mut self) {
|
||||
self.len = 0;
|
||||
self.head = 0;
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn push_back(&mut self, val: T) {
|
||||
self.buf[(self.head + self.len) % N] = val;
|
||||
self.len += 1;
|
||||
}
|
||||
#[inline]
|
||||
fn pop_front(&mut self) -> T {
|
||||
let val = self.buf[self.head];
|
||||
self.head = (self.head + 1) % N;
|
||||
self.len -= 1;
|
||||
val
|
||||
}
|
||||
#[inline]
|
||||
fn pop_back(&mut self) -> T {
|
||||
self.len -= 1;
|
||||
self.buf[(self.head + self.len) % N]
|
||||
}
|
||||
#[inline]
|
||||
fn front(&self) -> Option<&T> {
|
||||
if self.len > 0 {
|
||||
Some(&self.buf[self.head])
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
#[inline]
|
||||
fn back(&self) -> Option<&T> {
|
||||
if self.len > 0 {
|
||||
Some(&self.buf[(self.head + self.len - 1) % N])
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// Iterate over elements front-to-back (copies, since T: Copy).
|
||||
fn iter(&self) -> impl Iterator<Item = T> + '_ {
|
||||
(0..self.len).map(move |i| self.buf[(self.head + i) % N])
|
||||
}
|
||||
}
|
||||
|
||||
// ── MmerItem ──────────────────────────────────────────────────────────────────
|
||||
|
||||
#[derive(Clone, Copy, Default)]
|
||||
struct MmerItem {
|
||||
/// 0-based position of this m-mer's first base within the current segment.
|
||||
position: usize,
|
||||
/// Raw canonical m-mer value (right-aligned), used for partition key computation.
|
||||
/// Raw canonical m-mer value (right-aligned).
|
||||
canonical: u64,
|
||||
/// mix64 hash of the canonical m-mer, used as the random ordering key.
|
||||
hash: u64,
|
||||
}
|
||||
|
||||
// ── RollingStat ───────────────────────────────────────────────────────────────
|
||||
|
||||
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,
|
||||
k1q: std::collections::VecDeque<u64>,
|
||||
k2q: std::collections::VecDeque<u64>,
|
||||
k3q: std::collections::VecDeque<u64>,
|
||||
k4q: std::collections::VecDeque<u64>,
|
||||
k5q: std::collections::VecDeque<u64>,
|
||||
k6q: std::collections::VecDeque<u64>,
|
||||
minimier: std::collections::VecDeque<MmerItem>,
|
||||
k1c: Vec<usize>,
|
||||
k2c: Vec<usize>,
|
||||
k3c: Vec<usize>,
|
||||
k4c: Vec<usize>,
|
||||
k5c: Vec<usize>,
|
||||
k6c: Vec<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>,
|
||||
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],
|
||||
}
|
||||
@@ -44,24 +126,27 @@ impl RollingStat {
|
||||
let m = params::m();
|
||||
Self {
|
||||
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: std::collections::VecDeque::with_capacity(k),
|
||||
k2q: std::collections::VecDeque::with_capacity(k - 1),
|
||||
k3q: std::collections::VecDeque::with_capacity(k - 2),
|
||||
k4q: std::collections::VecDeque::with_capacity(k - 3),
|
||||
k5q: std::collections::VecDeque::with_capacity(k - 4),
|
||||
k6q: std::collections::VecDeque::with_capacity(k - 5),
|
||||
minimier: std::collections::VecDeque::with_capacity(k - m + 2),
|
||||
k1c: vec![0; 4_usize.pow(1)],
|
||||
k2c: vec![0; 4_usize.pow(2)],
|
||||
k3c: vec![0; 4_usize.pow(3)],
|
||||
k4c: vec![0; 4_usize.pow(4)],
|
||||
k5c: vec![0; 4_usize.pow(5)],
|
||||
k6c: vec![0; 4_usize.pow(6)],
|
||||
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],
|
||||
}
|
||||
@@ -71,24 +156,22 @@ impl RollingStat {
|
||||
self.rolling_k = 0;
|
||||
self.rolling_rck = 0;
|
||||
self.received = 0;
|
||||
for &i in &self.k1q {
|
||||
self.k1c[i as usize] = 0;
|
||||
}
|
||||
for &i in &self.k2q {
|
||||
self.k2c[i as usize] = 0;
|
||||
}
|
||||
for &i in &self.k3q {
|
||||
self.k3c[i as usize] = 0;
|
||||
}
|
||||
for &i in &self.k4q {
|
||||
self.k4c[i as usize] = 0;
|
||||
}
|
||||
for &i in &self.k5q {
|
||||
self.k5c[i as usize] = 0;
|
||||
}
|
||||
for &i in &self.k6q {
|
||||
self.k6c[i as usize] = 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();
|
||||
@@ -96,37 +179,36 @@ impl RollingStat {
|
||||
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(
|
||||
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],
|
||||
ws: usize,
|
||||
canonical: u64,
|
||||
f: usize,
|
||||
) {
|
||||
sum_f_log_f[ws] += n_log_n(f - 1) - n_log_n(f);
|
||||
sum_f_log_s[ws] -= ln_class_size(canonical, ws, false);
|
||||
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(
|
||||
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],
|
||||
ws: usize,
|
||||
canonical: u64,
|
||||
g: usize,
|
||||
) {
|
||||
sum_f_log_f[ws] += n_log_n(g + 1) - n_log_n(g);
|
||||
sum_f_log_s[ws] += ln_class_size(canonical, ws, false);
|
||||
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);
|
||||
}
|
||||
|
||||
pub fn push(&mut self, nuc: u8) {
|
||||
let k = params::k();
|
||||
let m = params::m();
|
||||
let k = self.k;
|
||||
let m = self.m;
|
||||
|
||||
let bnuc = encode_nuc(nuc);
|
||||
let cnuc = bnuc ^ 3;
|
||||
@@ -135,12 +217,12 @@ 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, 1, false);
|
||||
let canonical_k2 = entropy_norm_kmer(self.rolling_k & 15, 2, false);
|
||||
let canonical_k3 = entropy_norm_kmer(self.rolling_k & 63, 3, false);
|
||||
let canonical_k4 = entropy_norm_kmer(self.rolling_k & 255, 4, false);
|
||||
let canonical_k5 = entropy_norm_kmer(self.rolling_k & 1023, 5, false);
|
||||
let canonical_k6 = entropy_norm_kmer(self.rolling_k & 4095, 6, false);
|
||||
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);
|
||||
|
||||
self.received += 1;
|
||||
|
||||
@@ -175,107 +257,147 @@ impl RollingStat {
|
||||
}
|
||||
|
||||
if self.received > k {
|
||||
let old1 = self.k1q.pop_front().unwrap();
|
||||
let f1 = self.k1c[old1 as usize];
|
||||
Self::update_sums_decrement(&mut self.sum_f_log_f, &mut self.sum_f_log_s, 1, old1, f1);
|
||||
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().unwrap();
|
||||
let f2 = self.k2c[old2 as usize];
|
||||
Self::update_sums_decrement(&mut self.sum_f_log_f, &mut self.sum_f_log_s, 2, old2, f2);
|
||||
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().unwrap();
|
||||
let f3 = self.k3c[old3 as usize];
|
||||
Self::update_sums_decrement(&mut self.sum_f_log_f, &mut self.sum_f_log_s, 3, old3, f3);
|
||||
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().unwrap();
|
||||
let f4 = self.k4c[old4 as usize];
|
||||
Self::update_sums_decrement(&mut self.sum_f_log_f, &mut self.sum_f_log_s, 4, old4, f4);
|
||||
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().unwrap();
|
||||
let f5 = self.k5c[old5 as usize];
|
||||
Self::update_sums_decrement(&mut self.sum_f_log_f, &mut self.sum_f_log_s, 5, old5, f5);
|
||||
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().unwrap();
|
||||
let f6 = self.k6c[old6 as usize];
|
||||
Self::update_sums_decrement(&mut self.sum_f_log_f, &mut self.sum_f_log_s, 6, old6, f6);
|
||||
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;
|
||||
}
|
||||
|
||||
let g1 = self.k1c[canonical_k1 as usize];
|
||||
Self::update_sums_increment(
|
||||
&mut self.sum_f_log_f,
|
||||
&mut self.sum_f_log_s,
|
||||
1,
|
||||
canonical_k1,
|
||||
g1,
|
||||
);
|
||||
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];
|
||||
Self::update_sums_increment(
|
||||
&mut self.sum_f_log_f,
|
||||
&mut self.sum_f_log_s,
|
||||
2,
|
||||
canonical_k2,
|
||||
g2,
|
||||
);
|
||||
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];
|
||||
Self::update_sums_increment(
|
||||
&mut self.sum_f_log_f,
|
||||
&mut self.sum_f_log_s,
|
||||
3,
|
||||
canonical_k3,
|
||||
g3,
|
||||
);
|
||||
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];
|
||||
Self::update_sums_increment(
|
||||
&mut self.sum_f_log_f,
|
||||
&mut self.sum_f_log_s,
|
||||
4,
|
||||
canonical_k4,
|
||||
g4,
|
||||
);
|
||||
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];
|
||||
Self::update_sums_increment(
|
||||
&mut self.sum_f_log_f,
|
||||
&mut self.sum_f_log_s,
|
||||
5,
|
||||
canonical_k5,
|
||||
g5,
|
||||
);
|
||||
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];
|
||||
Self::update_sums_increment(
|
||||
&mut self.sum_f_log_f,
|
||||
&mut self.sum_f_log_s,
|
||||
6,
|
||||
canonical_k6,
|
||||
g6,
|
||||
);
|
||||
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;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -284,7 +406,7 @@ impl RollingStat {
|
||||
}
|
||||
|
||||
pub fn ready(&self) -> bool {
|
||||
self.received >= params::k()
|
||||
self.received >= self.k
|
||||
}
|
||||
|
||||
pub fn minimizer_position(&self) -> Option<usize> {
|
||||
@@ -305,14 +427,14 @@ impl RollingStat {
|
||||
|
||||
pub fn canonical_minimizer(&self) -> Option<Minimizer> {
|
||||
self.canonical_minimizer_raw()
|
||||
.map(|raw| Minimizer::from_raw_unchecked(raw << (64 - params::m() * 2)))
|
||||
.map(|raw| Minimizer::from_raw_unchecked(raw << (64 - self.m * 2)))
|
||||
}
|
||||
|
||||
pub fn entropy(&self, order: usize) -> Option<f64> {
|
||||
if !self.ready() {
|
||||
return None;
|
||||
}
|
||||
let k = params::k();
|
||||
let k = self.k;
|
||||
let em = emax(k, order);
|
||||
if em <= 0.0 {
|
||||
return Some(1.0);
|
||||
|
||||
Reference in New Issue
Block a user