diff --git a/obikmer 2026-05-20 12.35 profile.json.gz b/obikmer 2026-05-20 12.35 profile.json.gz new file mode 100644 index 0000000..57da44b Binary files /dev/null and b/obikmer 2026-05-20 12.35 profile.json.gz differ diff --git a/src/obikmer/src/cmd/superkmer.rs b/src/obikmer/src/cmd/superkmer.rs index 8dc938a..2ea1060 100644 --- a/src/obikmer/src/cmd/superkmer.rs +++ b/src/obikmer/src/cmd/superkmer.rs @@ -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! { diff --git a/src/obikseq/src/kmer.rs b/src/obikseq/src/kmer.rs index 70e49bb..9091661 100644 --- a/src/obikseq/src/kmer.rs +++ b/src/obikseq/src/kmer.rs @@ -118,7 +118,7 @@ pub struct KmerOf(u64, PhantomData); impl KmerOf { /// 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) } diff --git a/src/obiskbuilder/build.rs b/src/obiskbuilder/build.rs new file mode 100644 index 0000000..f422279 --- /dev/null +++ b/src/obiskbuilder/build.rs @@ -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 { + 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 { + 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(); +} diff --git a/src/obiskbuilder/src/entropy_table.rs b/src/obiskbuilder/src/entropy_table.rs index 4132ef7..9be69c3 100644 --- a/src/obiskbuilder/src/entropy_table.rs +++ b/src/obiskbuilder/src/entropy_table.rs @@ -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() -> [u64; N] { +const fn build_normalized_kmer() -> [u64; N] { let mut result = [0u64; N]; - let k = N.ilog(4) as usize; + let k = k_from_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(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() -> 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(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(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 } diff --git a/src/obiskbuilder/src/iter.rs b/src/obiskbuilder/src/iter.rs index f874349..0839b97 100644 --- a/src/obiskbuilder/src/iter.rs +++ b/src/obiskbuilder/src/iter.rs @@ -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) { + (0, Some(self.rope_len / self.k)) + } + fn next(&mut self) -> Option { loop { let byte = match self.cursor.read_next().ok() { diff --git a/src/obiskbuilder/src/rolling_stat.rs b/src/obiskbuilder/src/rolling_stat.rs index fb61d20..17838a0 100644 --- a/src/obiskbuilder/src/rolling_stat.rs +++ b/src/obiskbuilder/src/rolling_stat.rs @@ -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 { + buf: [T; N], + head: usize, + len: usize, +} + +impl Ring { + #[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 + '_ { + (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, - k2q: std::collections::VecDeque, - k3q: std::collections::VecDeque, - k4q: std::collections::VecDeque, - k5q: std::collections::VecDeque, - k6q: std::collections::VecDeque, - minimier: std::collections::VecDeque, - k1c: Vec, - k2c: Vec, - k3c: Vec, - k4c: Vec, - k5c: Vec, - k6c: Vec, + + // Sliding-window queues — stack-allocated, capacity ≤ k ≤ 31. + k1q: Ring, + k2q: Ring, + k3q: Ring, + k4q: Ring, + k5q: Ring, + k6q: Ring, + minimier: Ring, + + // 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( 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::(canonical); } #[inline] - fn update_sums_increment( + fn update_sums_increment( 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::(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::(self.rolling_k & 3); + let canonical_k2 = entropy_norm_kmer::(self.rolling_k & 15); + let canonical_k3 = entropy_norm_kmer::(self.rolling_k & 63); + let canonical_k4 = entropy_norm_kmer::(self.rolling_k & 255); + let canonical_k5 = entropy_norm_kmer::(self.rolling_k & 1023); + let canonical_k6 = entropy_norm_kmer::(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 { @@ -305,14 +427,14 @@ impl RollingStat { pub fn canonical_minimizer(&self) -> Option { 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 { 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);