From e725523898b1b6e310b38fb069b4890878d6eb50 Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Wed, 8 Jul 2026 12:39:21 +0200 Subject: [PATCH 1/3] feat: add entropy-driven k-mer complexity filtering Introduces a MinComplexity filter driven by new CLI arguments, enabling sequence-aware threshold checks during index reconstruction and partitioning. Adds the kmer_entropy module for normalized complexity scoring, updates the KmerFilter trait to evaluate per-kmer context, and refactors test modules for better organization. --- src/obikmer/src/cmd/filter.rs | 21 ++- src/obikmer/src/main.rs | 2 +- src/obikpartitionner/Cargo.toml | 2 +- src/obikpartitionner/src/dump_layer.rs | 38 ++--- src/obikpartitionner/src/filter.rs | 62 ++++++-- src/obikpartitionner/src/rebuild_layer.rs | 4 +- src/obiskbuilder/src/iter.rs | 156 +-------------------- src/obiskbuilder/src/kmer_entropy.rs | 49 +++++++ src/obiskbuilder/src/lib.rs | 2 + src/obiskbuilder/src/tests/iter.rs | 152 ++++++++++++++++++++ src/obiskbuilder/src/tests/kmer_entropy.rs | 54 +++++++ 11 files changed, 350 insertions(+), 192 deletions(-) create mode 100644 src/obiskbuilder/src/kmer_entropy.rs create mode 100644 src/obiskbuilder/src/tests/iter.rs create mode 100644 src/obiskbuilder/src/tests/kmer_entropy.rs diff --git a/src/obikmer/src/cmd/filter.rs b/src/obikmer/src/cmd/filter.rs index 993f80a..ae3f45f 100644 --- a/src/obikmer/src/cmd/filter.rs +++ b/src/obikmer/src/cmd/filter.rs @@ -2,14 +2,14 @@ use std::path::PathBuf; use clap::Args; use obikindex::{KmerIndex, MergeMode}; -use obikpartitionner::filter::{MaxTotalCount, MinTotalCount}; +use obikpartitionner::filter::{MaxTotalCount, MinComplexity, MinTotalCount}; use obisys::Reporter; use tracing::info; use super::predicate::FilterArgs as KmerFilterArgs; #[derive(Args)] -pub struct FilterArgs { +pub struct FilterCmdArgs { /// Source index directory pub source: PathBuf, @@ -28,6 +28,18 @@ pub struct FilterArgs { #[arg(long)] pub max_total_count: Option, + /// Minimum normalized entropy (complexity) to keep a k-mer — same metric + /// as `obikmer index`'s --theta, applied here to k-mers already committed + /// to the source index (reconstructed from unitigs.bin). K-mers scoring + /// below this are removed. + #[arg(long)] + pub min_complexity: Option, + + /// Maximum sub-word size for the complexity computation (see `obikmer + /// index`'s --level-max). Only used when --min-complexity is set. + #[arg(long, default_value_t = 6)] + pub complexity_level_max: usize, + /// Output as presence/absence instead of counts #[arg(long)] pub presence: bool, @@ -37,7 +49,7 @@ pub struct FilterArgs { pub force: bool, } -pub fn run(args: FilterArgs) { +pub fn run(args: FilterCmdArgs) { let src = KmerIndex::open(&args.source).unwrap_or_else(|e| { eprintln!("error opening source index: {e}"); std::process::exit(1); @@ -62,6 +74,9 @@ pub fn run(args: FilterArgs) { if let Some(v) = args.max_total_count { filters.push(Box::new(MaxTotalCount { total: v })); } + if let Some(theta) = args.min_complexity { + filters.push(Box::new(MinComplexity { level_max: args.complexity_level_max, theta })); + } let mut rep = Reporter::new(); KmerIndex::rebuild(&args.output, &src, &filters, mode, args.force, &mut rep) diff --git a/src/obikmer/src/main.rs b/src/obikmer/src/main.rs index a0b270b..cb701e6 100644 --- a/src/obikmer/src/main.rs +++ b/src/obikmer/src/main.rs @@ -21,7 +21,7 @@ enum Commands { /// Merge multiple built indexes into one Merge(cmd::merge::MergeArgs), /// Apply row-level selection (σ) to an index: retain only k-mers matching the predicates - Filter(cmd::filter::FilterArgs), + Filter(cmd::filter::FilterCmdArgs), /// Project and/or aggregate genome columns into a new or in-place index Select(cmd::select::SelectArgs), /// Query an index with sequences and annotate matches diff --git a/src/obikpartitionner/Cargo.toml b/src/obikpartitionner/Cargo.toml index 1c42cab..a0c810e 100644 --- a/src/obikpartitionner/Cargo.toml +++ b/src/obikpartitionner/Cargo.toml @@ -6,7 +6,6 @@ edition = "2024" [dev-dependencies] tempfile = "3" obikseq = { path = "../obikseq", features = ["test-utils"] } -obiskbuilder = { path = "../obiskbuilder" } obiread = { path = "../obiread" } obikrope = { path = "../obikrope" } @@ -14,6 +13,7 @@ obikrope = { path = "../obikrope" } niffler = "3.0.0" remove_dir_all = "0.8" obikseq = { path = "../obikseq" } +obiskbuilder = { path = "../obiskbuilder" } obiskio = { path = "../obiskio" } obidebruinj = { path = "../obidebruinj" } obilayeredmap = { path = "../obilayeredmap" } diff --git a/src/obikpartitionner/src/dump_layer.rs b/src/obikpartitionner/src/dump_layer.rs index e1c8226..00a8507 100644 --- a/src/obikpartitionner/src/dump_layer.rs +++ b/src/obikpartitionner/src/dump_layer.rs @@ -62,7 +62,7 @@ impl KmerPartition { for (kmer, _, _) in reader.iter_indexed_canonical_kmers() { if let Some(slot) = mphf.find(kmer) { let row = mat.row(slot); - if passes_all(filters, &row, n_genomes) { + if passes_all(filters, kmer, &row, n_genomes) { cont = cb(kmer, row); if !cont { break; } } @@ -75,7 +75,7 @@ impl KmerPartition { for (kmer, _, _) in reader.iter_indexed_canonical_kmers() { if let Some(slot) = mphf.find(kmer) { let row: Box<[u32]> = mat.row(slot).iter().map(|&b| b as u32).collect(); - if passes_all(filters, &row, n_genomes) { + if passes_all(filters, kmer, &row, n_genomes) { cont = cb(kmer, row); if !cont { break; } } @@ -83,16 +83,17 @@ impl KmerPartition { } cont } else { - // No data matrix: implicit presence — all values are 1. - // The filter result is identical for every kmer, so evaluate once. + // No data matrix: implicit presence — all values are 1. `row` + // is identical for every kmer, but a filter can still depend + // on the kmer's own sequence (e.g. MinComplexity), so this + // cannot be evaluated once for the whole layer — filters must + // still be tested per kmer. let all_present: Box<[u32]> = vec![1u32; n_genomes].into(); let mut cont = true; - if passes_all(filters, &all_present, n_genomes) { - for (kmer, _, _) in reader.iter_indexed_canonical_kmers() { - if mphf.find(kmer).is_some() { - cont = cb(kmer, all_present.clone()); - if !cont { break; } - } + for (kmer, _, _) in reader.iter_indexed_canonical_kmers() { + if mphf.find(kmer).is_some() && passes_all(filters, kmer, &all_present, n_genomes) { + cont = cb(kmer, all_present.clone()); + if !cont { break; } } } cont @@ -140,7 +141,7 @@ impl KmerPartition { for (kmer, _, _) in reader.iter_indexed_canonical_kmers() { if let Some(slot) = mphf.find(kmer) { let row = mat.row(slot); - if passes_all(filters, &row, n_genomes) { + if passes_all(filters, kmer, &row, n_genomes) { cont = cb(part, layer, kmer, row); if !cont { break; } } @@ -153,7 +154,7 @@ impl KmerPartition { for (kmer, _, _) in reader.iter_indexed_canonical_kmers() { if let Some(slot) = mphf.find(kmer) { let row: Box<[u32]> = mat.row(slot).iter().map(|&b| b as u32).collect(); - if passes_all(filters, &row, n_genomes) { + if passes_all(filters, kmer, &row, n_genomes) { cont = cb(part, layer, kmer, row); if !cont { break; } } @@ -161,14 +162,15 @@ impl KmerPartition { } cont } else { + // Same as iter_partition_kmers: row is constant but a filter + // may still depend on the kmer's own sequence, so this must + // be tested per kmer, not once for the whole layer. let all_present: Box<[u32]> = vec![1u32; n_genomes].into(); let mut cont = true; - if passes_all(filters, &all_present, n_genomes) { - for (kmer, _, _) in reader.iter_indexed_canonical_kmers() { - if mphf.find(kmer).is_some() { - cont = cb(part, layer, kmer, all_present.clone()); - if !cont { break; } - } + for (kmer, _, _) in reader.iter_indexed_canonical_kmers() { + if mphf.find(kmer).is_some() && passes_all(filters, kmer, &all_present, n_genomes) { + cont = cb(part, layer, kmer, all_present.clone()); + if !cont { break; } } } cont diff --git a/src/obikpartitionner/src/filter.rs b/src/obikpartitionner/src/filter.rs index 00f3b03..db69aba 100644 --- a/src/obikpartitionner/src/filter.rs +++ b/src/obikpartitionner/src/filter.rs @@ -1,17 +1,24 @@ use obicompactvec::FilterMask; +use obikseq::CanonicalKmer; -/// Trait for kmer row filters. +/// Trait for kmer filters. /// +/// `kmer` is the k-mer's own canonical sequence, reconstructed from the +/// source index's `unitigs.bin` (always present — see `rebuild_layer.rs`); /// `row` contains raw per-genome counts (or 0/1 for presence/absence data). -/// `n_genomes` equals `row.len()`. +/// `n_genomes` equals `row.len()`. Most filters only need `row` — `kmer` is +/// there for filters that reason about the k-mer's sequence itself (e.g. +/// [`MinComplexity`]). pub trait KmerFilter: Send + Sync { - fn passes(&self, row: &[u32], n_genomes: usize) -> bool; + fn passes(&self, kmer: CanonicalKmer, row: &[u32], n_genomes: usize) -> bool; /// Express this filter as a [`FilterMask`] column-operation expression. /// /// Returns `Some(expr)` if the filter can be evaluated solely from matrix /// column aggregates (no per-kmer row scan needed). Returns `None` if the - /// filter requires row-level inspection. + /// filter requires row-level inspection — always the case for a filter + /// that needs the k-mer's sequence, since a `FilterMask` only expresses + /// per-genome column aggregates, never per-slot sequence data. /// /// `threshold` semantics in the returned mask use `>= threshold`, matching /// [`obicompactvec::MatrixGroupOps`]. Implementations must add 1 to any @@ -23,8 +30,13 @@ pub trait KmerFilter: Send + Sync { /// True when `row` passes every filter in `filters`. /// Returns `true` if `filters` is empty. -pub fn passes_all(filters: &[Box], row: &[u32], n_genomes: usize) -> bool { - filters.iter().all(|f| f.passes(row, n_genomes)) +pub fn passes_all( + filters: &[Box], + kmer: CanonicalKmer, + row: &[u32], + n_genomes: usize, +) -> bool { + filters.iter().all(|f| f.passes(kmer, row, n_genomes)) } // ── Quorum filters ───────────────────────────────────────────────────────────── @@ -40,7 +52,7 @@ pub struct MinGenomeFraction { } impl KmerFilter for MinGenomeFraction { - fn passes(&self, row: &[u32], n_genomes: usize) -> bool { + fn passes(&self, _kmer: CanonicalKmer, row: &[u32], n_genomes: usize) -> bool { let p = present_count(row, self.threshold); p as f64 / n_genomes as f64 >= self.frac } @@ -63,7 +75,7 @@ pub struct MaxGenomeFraction { } impl KmerFilter for MaxGenomeFraction { - fn passes(&self, row: &[u32], n_genomes: usize) -> bool { + fn passes(&self, _kmer: CanonicalKmer, row: &[u32], n_genomes: usize) -> bool { let p = present_count(row, self.threshold); p as f64 / n_genomes as f64 <= self.frac } @@ -86,7 +98,7 @@ pub struct MinGenomeCount { } impl KmerFilter for MinGenomeCount { - fn passes(&self, row: &[u32], _n_genomes: usize) -> bool { + fn passes(&self, _kmer: CanonicalKmer, row: &[u32], _n_genomes: usize) -> bool { present_count(row, self.threshold) >= self.count } @@ -107,7 +119,7 @@ pub struct MaxGenomeCount { } impl KmerFilter for MaxGenomeCount { - fn passes(&self, row: &[u32], _n_genomes: usize) -> bool { + fn passes(&self, _kmer: CanonicalKmer, row: &[u32], _n_genomes: usize) -> bool { present_count(row, self.threshold) <= self.count } @@ -129,7 +141,7 @@ pub struct MinTotalCount { } impl KmerFilter for MinTotalCount { - fn passes(&self, row: &[u32], _n_genomes: usize) -> bool { + fn passes(&self, _kmer: CanonicalKmer, row: &[u32], _n_genomes: usize) -> bool { row.iter().sum::() >= self.total } @@ -147,7 +159,7 @@ pub struct MaxTotalCount { } impl KmerFilter for MaxTotalCount { - fn passes(&self, row: &[u32], _n_genomes: usize) -> bool { + fn passes(&self, _kmer: CanonicalKmer, row: &[u32], _n_genomes: usize) -> bool { row.iter().sum::() <= self.total } @@ -212,7 +224,7 @@ impl GroupQuorumFilter { } impl KmerFilter for GroupQuorumFilter { - fn passes(&self, row: &[u32], _n_genomes: usize) -> bool { + fn passes(&self, _kmer: CanonicalKmer, row: &[u32], _n_genomes: usize) -> bool { if !self.ingroup_idx.is_empty() { let n = self.ingroup_idx.iter() .filter(|&&i| row.get(i).copied().unwrap_or(0) > self.threshold) @@ -260,3 +272,27 @@ impl KmerFilter for GroupQuorumFilter { Some(FilterMask::And(parts)) } } + +// ── Complexity filter (post-hoc, sequence-based) ────────────────────────────── + +/// Reject k-mers with normalized entropy below `theta` — the same complexity +/// metric `obikmer index`'s `--theta`/`--level-max` apply *during* superkmer +/// construction (see [`obiskbuilder::KmerEntropy`]), applied here after the +/// fact, to k-mers already committed to a built index. +/// +/// Unlike every other filter in this module, this one needs the k-mer's own +/// sequence, not its per-genome row — `column_mask_expr` is never overridden +/// (stays `None`), so this filter always forces the row-level scan path in +/// `rebuild_layer.rs` (which reconstructs the sequence from `unitigs.bin` +/// regardless, so no extra I/O beyond what filtering already requires). +pub struct MinComplexity { + pub level_max: usize, + pub theta: f64, +} + +impl KmerFilter for MinComplexity { + fn passes(&self, kmer: CanonicalKmer, _row: &[u32], _n_genomes: usize) -> bool { + use obiskbuilder::KmerEntropy; + kmer.entropy(self.level_max) >= self.theta + } +} diff --git a/src/obikpartitionner/src/rebuild_layer.rs b/src/obikpartitionner/src/rebuild_layer.rs index b8893ef..e608c49 100644 --- a/src/obikpartitionner/src/rebuild_layer.rs +++ b/src/obikpartitionner/src/rebuild_layer.rs @@ -126,7 +126,7 @@ fn iter_src_kmers_masked( Some(m) => m.get(slot), None => { let row = src_data.fill_row_by_slot(slot, n_genomes); - filters.iter().all(|f| f.passes(&row, n_genomes)) + filters.iter().all(|f| f.passes(kmer, &row, n_genomes)) } }; if passes { cb(kmer); } @@ -165,7 +165,7 @@ fn iter_src_layers( cb(kmer, row.into_boxed_slice()); } else { let row = src_data.fill_row_by_slot(slot, n_genomes); - if filters.iter().all(|f| f.passes(&row, n_genomes)) { + if filters.iter().all(|f| f.passes(kmer, &row, n_genomes)) { cb(kmer, row.into_boxed_slice()); } } diff --git a/src/obiskbuilder/src/iter.rs b/src/obiskbuilder/src/iter.rs index 0839b97..8a33123 100644 --- a/src/obiskbuilder/src/iter.rs +++ b/src/obiskbuilder/src/iter.rs @@ -149,157 +149,5 @@ impl Iterator for SuperKmerIter<'_> { // ── tests ───────────────────────────────────────────────────────────────────── #[cfg(test)] -mod tests { - use super::*; - use obikrope::Rope; - - fn setup() { - obikseq::params::set_k(K); - obikseq::params::set_m(5); - } - - fn make_rope(data: &[u8]) -> Rope { - let mut r = Rope::new(None); - r.push(data.to_vec()); - r - } - - fn run_nofilter(data: &[u8], k: usize) -> Vec> { - let rope = make_rope(data); - SuperKmerIter::new(&rope, k, 1, 0.0) - .map(|rsk| rsk.superkmer().to_ascii()) - .collect() - } - - // k=11, m=5 — valeurs minimales du projet (k ∈ [11,31]) - const K: usize = 11; - - /// Collect the set of canonical k-mers from a raw ASCII sequence (no NUL). - fn direct_canonical_kmers(seq: &[u8]) -> std::collections::HashSet> { - (0..seq.len().saturating_sub(K - 1)) - .map(|i| obikseq::SuperKmer::from_ascii(&seq[i..i + K]).to_ascii()) - .collect() - } - - /// Collect the set of canonical k-mers emitted by SuperKmerIter over a rope. - fn iter_canonical_kmers(rope: &Rope) -> std::collections::HashSet> { - SuperKmerIter::new(rope, K, 1, 0.0) - .flat_map(|rsk| { - rsk.superkmer() - .iter_canonical_kmers() - .map(|km| km.to_ascii()) - .collect::>() - }) - .collect() - } - - #[test] - fn coverage_single_segment() { - setup(); - let seq = b"ACGTACGTACGTACGTACGT"; - let rope = make_rope(&[seq.as_ref(), b"\x00"].concat()); - let direct = direct_canonical_kmers(seq); - let from_iter = iter_canonical_kmers(&rope); - let missing: Vec<_> = direct.difference(&from_iter).collect(); - assert!( - missing.is_empty(), - "k-mers perdus dans segment unique : {missing:?}" - ); - } - - #[test] - fn coverage_two_segments() { - setup(); - let seg1 = b"ACGTACGTACGTACGTACGT"; - let seg2 = b"TGCATGCATGCATGCATGCA"; - let rope = make_rope(&[seg1.as_ref(), b"\x00", seg2.as_ref(), b"\x00"].concat()); - let mut direct = direct_canonical_kmers(seg1); - direct.extend(direct_canonical_kmers(seg2)); - let from_iter = iter_canonical_kmers(&rope); - let missing: Vec<_> = direct.difference(&from_iter).collect(); - assert!( - missing.is_empty(), - "k-mers perdus dans deux segments : {missing:?}" - ); - } - - #[test] - fn coverage_minimizer_boundary() { - setup(); - // sequence assez longue pour forcer plusieurs changements de minimiseur - let seq: Vec = (0..80).map(|i| b"ACGT"[i % 4]).collect(); - let rope = make_rope(&[seq.as_slice(), b"\x00"].concat()); - let direct = direct_canonical_kmers(&seq); - let from_iter = iter_canonical_kmers(&rope); - let missing: Vec<_> = direct.difference(&from_iter).collect(); - assert!( - missing.is_empty(), - "k-mers perdus à la frontière de minimiseur : {missing:?}" - ); - } - - #[test] - fn single_segment_one_superkmer() { - setup(); - let out = run_nofilter(b"ACGTACGTACGTACGTACGT\x00", K); - assert!(!out.is_empty()); - let total: Vec = out.into_iter().flatten().collect(); - assert!(total.len() >= K); - } - - #[test] - fn segment_shorter_than_k_emits_nothing() { - setup(); - let out = run_nofilter(b"ACGTACGT\x00", K); - assert_eq!(out, Vec::>::new()); - } - - #[test] - fn empty_input_emits_nothing() { - setup(); - let out = run_nofilter(b"", K); - assert_eq!(out, Vec::>::new()); - } - - #[test] - fn two_segments_both_emitted() { - setup(); - let out = run_nofilter(b"ACGTACGTACGTACGT\x00TGCATGCATGCATGCA\x00", K); - assert!(!out.is_empty()); - } - - #[test] - fn low_complexity_kmer_is_rejected() { - setup(); - let out_pass = run_nofilter(b"AAAAAAAAAAAACGTACGTACGT\x00", K); - assert!(!out_pass.is_empty()); - - let rope = make_rope(b"AAAAAAAAAAAAAAAAAAAA\x00"); - let out_reject: Vec> = SuperKmerIter::new(&rope, K, 6, 0.9) - .map(|rsk| rsk.superkmer().to_ascii()) - .collect(); - assert!(out_reject.is_empty()); - } - - #[test] - fn multi_slice_rope() { - setup(); - let data = b"ACGTACGTACGTACGTACGT\x00"; - let mid = data.len() / 2; - let mut rope = Rope::new(None); - rope.push(data[..mid].to_vec()); - rope.push(data[mid..].to_vec()); - let out: Vec> = SuperKmerIter::new(&rope, K, 1, 0.0) - .map(|rsk| rsk.superkmer().to_ascii()) - .collect(); - assert!(!out.is_empty()); - } - - #[test] - fn yields_minimizer_value() { - setup(); - let rope = make_rope(b"ACGTACGTACGTACGTACGT\x00"); - let results: Vec = SuperKmerIter::new(&rope, K, 1, 0.0).collect(); - assert!(!results.is_empty()); - } -} +#[path = "tests/iter.rs"] +mod tests; diff --git a/src/obiskbuilder/src/kmer_entropy.rs b/src/obiskbuilder/src/kmer_entropy.rs new file mode 100644 index 0000000..f8fcd7f --- /dev/null +++ b/src/obiskbuilder/src/kmer_entropy.rs @@ -0,0 +1,49 @@ +//! 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; diff --git a/src/obiskbuilder/src/lib.rs b/src/obiskbuilder/src/lib.rs index 6dc2de1..61d8cf1 100644 --- a/src/obiskbuilder/src/lib.rs +++ b/src/obiskbuilder/src/lib.rs @@ -6,6 +6,7 @@ #![deny(missing_docs)] pub mod iter; +pub mod kmer_entropy; pub mod stream_iter; mod scratch; @@ -14,6 +15,7 @@ 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; diff --git a/src/obiskbuilder/src/tests/iter.rs b/src/obiskbuilder/src/tests/iter.rs new file mode 100644 index 0000000..4e03036 --- /dev/null +++ b/src/obiskbuilder/src/tests/iter.rs @@ -0,0 +1,152 @@ +use super::*; +use obikrope::Rope; + +fn setup() { + obikseq::params::set_k(K); + obikseq::params::set_m(5); +} + +fn make_rope(data: &[u8]) -> Rope { + let mut r = Rope::new(None); + r.push(data.to_vec()); + r +} + +fn run_nofilter(data: &[u8], k: usize) -> Vec> { + let rope = make_rope(data); + SuperKmerIter::new(&rope, k, 1, 0.0) + .map(|rsk| rsk.superkmer().to_ascii()) + .collect() +} + +// k=11, m=5 — valeurs minimales du projet (k ∈ [11,31]) +const K: usize = 11; + +/// Collect the set of canonical k-mers from a raw ASCII sequence (no NUL). +fn direct_canonical_kmers(seq: &[u8]) -> std::collections::HashSet> { + (0..seq.len().saturating_sub(K - 1)) + .map(|i| obikseq::SuperKmer::from_ascii(&seq[i..i + K]).to_ascii()) + .collect() +} + +/// Collect the set of canonical k-mers emitted by SuperKmerIter over a rope. +fn iter_canonical_kmers(rope: &Rope) -> std::collections::HashSet> { + SuperKmerIter::new(rope, K, 1, 0.0) + .flat_map(|rsk| { + rsk.superkmer() + .iter_canonical_kmers() + .map(|km| km.to_ascii()) + .collect::>() + }) + .collect() +} + +#[test] +fn coverage_single_segment() { + setup(); + let seq = b"ACGTACGTACGTACGTACGT"; + let rope = make_rope(&[seq.as_ref(), b"\x00"].concat()); + let direct = direct_canonical_kmers(seq); + let from_iter = iter_canonical_kmers(&rope); + let missing: Vec<_> = direct.difference(&from_iter).collect(); + assert!( + missing.is_empty(), + "k-mers perdus dans segment unique : {missing:?}" + ); +} + +#[test] +fn coverage_two_segments() { + setup(); + let seg1 = b"ACGTACGTACGTACGTACGT"; + let seg2 = b"TGCATGCATGCATGCATGCA"; + let rope = make_rope(&[seg1.as_ref(), b"\x00", seg2.as_ref(), b"\x00"].concat()); + let mut direct = direct_canonical_kmers(seg1); + direct.extend(direct_canonical_kmers(seg2)); + let from_iter = iter_canonical_kmers(&rope); + let missing: Vec<_> = direct.difference(&from_iter).collect(); + assert!( + missing.is_empty(), + "k-mers perdus dans deux segments : {missing:?}" + ); +} + +#[test] +fn coverage_minimizer_boundary() { + setup(); + // sequence assez longue pour forcer plusieurs changements de minimiseur + let seq: Vec = (0..80).map(|i| b"ACGT"[i % 4]).collect(); + let rope = make_rope(&[seq.as_slice(), b"\x00"].concat()); + let direct = direct_canonical_kmers(&seq); + let from_iter = iter_canonical_kmers(&rope); + let missing: Vec<_> = direct.difference(&from_iter).collect(); + assert!( + missing.is_empty(), + "k-mers perdus à la frontière de minimiseur : {missing:?}" + ); +} + +#[test] +fn single_segment_one_superkmer() { + setup(); + let out = run_nofilter(b"ACGTACGTACGTACGTACGT\x00", K); + assert!(!out.is_empty()); + let total: Vec = out.into_iter().flatten().collect(); + assert!(total.len() >= K); +} + +#[test] +fn segment_shorter_than_k_emits_nothing() { + setup(); + let out = run_nofilter(b"ACGTACGT\x00", K); + assert_eq!(out, Vec::>::new()); +} + +#[test] +fn empty_input_emits_nothing() { + setup(); + let out = run_nofilter(b"", K); + assert_eq!(out, Vec::>::new()); +} + +#[test] +fn two_segments_both_emitted() { + setup(); + let out = run_nofilter(b"ACGTACGTACGTACGT\x00TGCATGCATGCATGCA\x00", K); + assert!(!out.is_empty()); +} + +#[test] +fn low_complexity_kmer_is_rejected() { + setup(); + let out_pass = run_nofilter(b"AAAAAAAAAAAACGTACGTACGT\x00", K); + assert!(!out_pass.is_empty()); + + let rope = make_rope(b"AAAAAAAAAAAAAAAAAAAA\x00"); + let out_reject: Vec> = SuperKmerIter::new(&rope, K, 6, 0.9) + .map(|rsk| rsk.superkmer().to_ascii()) + .collect(); + assert!(out_reject.is_empty()); +} + +#[test] +fn multi_slice_rope() { + setup(); + let data = b"ACGTACGTACGTACGTACGT\x00"; + let mid = data.len() / 2; + let mut rope = Rope::new(None); + rope.push(data[..mid].to_vec()); + rope.push(data[mid..].to_vec()); + let out: Vec> = SuperKmerIter::new(&rope, K, 1, 0.0) + .map(|rsk| rsk.superkmer().to_ascii()) + .collect(); + assert!(!out.is_empty()); +} + +#[test] +fn yields_minimizer_value() { + setup(); + let rope = make_rope(b"ACGTACGTACGTACGTACGT\x00"); + let results: Vec = SuperKmerIter::new(&rope, K, 1, 0.0).collect(); + assert!(!results.is_empty()); +} diff --git a/src/obiskbuilder/src/tests/kmer_entropy.rs b/src/obiskbuilder/src/tests/kmer_entropy.rs new file mode 100644 index 0000000..1456c3a --- /dev/null +++ b/src/obiskbuilder/src/tests/kmer_entropy.rs @@ -0,0 +1,54 @@ +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:?}"); + } +} -- 2.54.0 From 912f788f7f1d8d288d49843234dc2dc3cbd5616e Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Wed, 8 Jul 2026 18:16:08 +0200 Subject: [PATCH 2/3] 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. --- src/Cargo.lock | 10 + src/Cargo.toml | 2 +- src/obikentropy/Cargo.toml | 10 + src/{obiskbuilder => obikentropy}/build.rs | 0 src/obikentropy/src/kmer_entropy.rs | 41 +++ src/obikentropy/src/lib.rs | 17 ++ src/obikentropy/src/ring.rs | 40 +++ .../src/table.rs} | 4 + .../src/tests/kmer_entropy.rs | 2 - src/obikentropy/src/tracker.rs | 280 ++++++++++++++++++ src/obikpartitionner/Cargo.toml | 1 + src/obikpartitionner/src/filter.rs | 4 +- src/obiskbuilder/Cargo.toml | 6 + src/obiskbuilder/benches/superkmer_stream.rs | 58 ++++ src/obiskbuilder/src/kmer_entropy.rs | 49 --- src/obiskbuilder/src/lib.rs | 3 - src/obiskbuilder/src/rolling_stat.rs | 260 +--------------- 17 files changed, 482 insertions(+), 305 deletions(-) create mode 100644 src/obikentropy/Cargo.toml rename src/{obiskbuilder => obikentropy}/build.rs (100%) create mode 100644 src/obikentropy/src/kmer_entropy.rs create mode 100644 src/obikentropy/src/lib.rs create mode 100644 src/obikentropy/src/ring.rs rename src/{obiskbuilder/src/entropy_table.rs => obikentropy/src/table.rs} (93%) rename src/{obiskbuilder => obikentropy}/src/tests/kmer_entropy.rs (94%) create mode 100644 src/obikentropy/src/tracker.rs create mode 100644 src/obiskbuilder/benches/superkmer_stream.rs delete mode 100644 src/obiskbuilder/src/kmer_entropy.rs diff --git a/src/Cargo.lock b/src/Cargo.lock index 3e31a1e..f7c63e3 100644 --- a/src/Cargo.lock +++ b/src/Cargo.lock @@ -1682,6 +1682,13 @@ dependencies = [ "xxhash-rust", ] +[[package]] +name = "obikentropy" +version = "0.1.0" +dependencies = [ + "obikseq", +] + [[package]] name = "obikindex" version = "0.1.0" @@ -1742,6 +1749,7 @@ dependencies = [ "niffler 3.0.0", "obicompactvec", "obidebruinj", + "obikentropy", "obikrope", "obikseq", "obilayeredmap", @@ -1824,7 +1832,9 @@ dependencies = [ name = "obiskbuilder" version = "0.1.0" dependencies = [ + "criterion2", "lazy_static", + "obikentropy", "obikrope", "obikseq", "obiread", diff --git a/src/Cargo.toml b/src/Cargo.toml index 141df02..5684237 100644 --- a/src/Cargo.toml +++ b/src/Cargo.toml @@ -1,5 +1,5 @@ [workspace] resolver = "3" -members = ["obikseq", "obiread", "obiskbuilder", "obifastwrite", "obikmer","obikrope","obipipeline", "obikpartitionner","obiskio","obidebruinj","obilayeredmap", "obicompactvec", "obisys", "obikindex", "obitaxonomy"] +members = ["obikseq", "obiread", "obiskbuilder", "obifastwrite", "obikmer","obikrope","obipipeline", "obikpartitionner","obiskio","obidebruinj","obilayeredmap", "obicompactvec", "obisys", "obikindex", "obitaxonomy", "obikentropy"] [profile.release] debug = 1 diff --git a/src/obikentropy/Cargo.toml b/src/obikentropy/Cargo.toml new file mode 100644 index 0000000..3d1cd69 --- /dev/null +++ b/src/obikentropy/Cargo.toml @@ -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"] } diff --git a/src/obiskbuilder/build.rs b/src/obikentropy/build.rs similarity index 100% rename from src/obiskbuilder/build.rs rename to src/obikentropy/build.rs diff --git a/src/obikentropy/src/kmer_entropy.rs b/src/obikentropy/src/kmer_entropy.rs new file mode 100644 index 0000000..1989972 --- /dev/null +++ b/src/obikentropy/src/kmer_entropy.rs @@ -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; diff --git a/src/obikentropy/src/lib.rs b/src/obikentropy/src/lib.rs new file mode 100644 index 0000000..113bdd3 --- /dev/null +++ b/src/obikentropy/src/lib.rs @@ -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}; diff --git a/src/obikentropy/src/ring.rs b/src/obikentropy/src/ring.rs new file mode 100644 index 0000000..b729beb --- /dev/null +++ b/src/obikentropy/src/ring.rs @@ -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 { + buf: [T; N], + head: usize, + len: usize, +} + +impl Ring { + #[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 + } +} diff --git a/src/obiskbuilder/src/entropy_table.rs b/src/obikentropy/src/table.rs similarity index 93% rename from src/obiskbuilder/src/entropy_table.rs rename to src/obikentropy/src/table.rs index 9be69c3..289aa94 100644 --- a/src/obiskbuilder/src/entropy_table.rs +++ b/src/obikentropy/src/table.rs @@ -1,3 +1,7 @@ +//! 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>(); diff --git a/src/obiskbuilder/src/tests/kmer_entropy.rs b/src/obikentropy/src/tests/kmer_entropy.rs similarity index 94% rename from src/obiskbuilder/src/tests/kmer_entropy.rs rename to src/obikentropy/src/tests/kmer_entropy.rs index 1456c3a..c75522e 100644 --- a/src/obiskbuilder/src/tests/kmer_entropy.rs +++ b/src/obikentropy/src/tests/kmer_entropy.rs @@ -3,12 +3,10 @@ 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() } diff --git a/src/obikentropy/src/tracker.rs b/src/obikentropy/src/tracker.rs new file mode 100644 index 0000000..98686fa --- /dev/null +++ b/src/obikentropy/src/tracker.rs @@ -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::(rolling_kmer & 3), + entropy_norm_kmer::(rolling_kmer & 15), + entropy_norm_kmer::(rolling_kmer & 63), + entropy_norm_kmer::(rolling_kmer & 255), + entropy_norm_kmer::(rolling_kmer & 1023), + entropy_norm_kmer::(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, + k2q: Ring, + k3q: Ring, + k4q: Ring, + k5q: Ring, + k6q: 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], +} + +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( + 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::(canonical); + } + + #[inline] + fn update_sums_increment( + 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::(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 } + } +} diff --git a/src/obikpartitionner/Cargo.toml b/src/obikpartitionner/Cargo.toml index a0c810e..918606f 100644 --- a/src/obikpartitionner/Cargo.toml +++ b/src/obikpartitionner/Cargo.toml @@ -13,6 +13,7 @@ obikrope = { path = "../obikrope" } niffler = "3.0.0" remove_dir_all = "0.8" obikseq = { path = "../obikseq" } +obikentropy = { path = "../obikentropy" } obiskbuilder = { path = "../obiskbuilder" } obiskio = { path = "../obiskio" } obidebruinj = { path = "../obidebruinj" } diff --git a/src/obikpartitionner/src/filter.rs b/src/obikpartitionner/src/filter.rs index db69aba..9f54a69 100644 --- a/src/obikpartitionner/src/filter.rs +++ b/src/obikpartitionner/src/filter.rs @@ -277,7 +277,7 @@ impl KmerFilter for GroupQuorumFilter { /// Reject k-mers with normalized entropy below `theta` — the same complexity /// metric `obikmer index`'s `--theta`/`--level-max` apply *during* superkmer -/// construction (see [`obiskbuilder::KmerEntropy`]), applied here after the +/// construction (see [`obikentropy::KmerEntropy`]), applied here after the /// fact, to k-mers already committed to a built index. /// /// Unlike every other filter in this module, this one needs the k-mer's own @@ -292,7 +292,7 @@ pub struct MinComplexity { impl KmerFilter for MinComplexity { fn passes(&self, kmer: CanonicalKmer, _row: &[u32], _n_genomes: usize) -> bool { - use obiskbuilder::KmerEntropy; + use obikentropy::KmerEntropy; kmer.entropy(self.level_max) >= self.theta } } diff --git a/src/obiskbuilder/Cargo.toml b/src/obiskbuilder/Cargo.toml index 5e19801..1e5c0af 100644 --- a/src/obiskbuilder/Cargo.toml +++ b/src/obiskbuilder/Cargo.toml @@ -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 diff --git a/src/obiskbuilder/benches/superkmer_stream.rs b/src/obiskbuilder/benches/superkmer_stream.rs new file mode 100644 index 0000000..776e6d3 --- /dev/null +++ b/src/obiskbuilder/benches/superkmer_stream.rs @@ -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 { + 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); diff --git a/src/obiskbuilder/src/kmer_entropy.rs b/src/obiskbuilder/src/kmer_entropy.rs deleted file mode 100644 index f8fcd7f..0000000 --- a/src/obiskbuilder/src/kmer_entropy.rs +++ /dev/null @@ -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; diff --git a/src/obiskbuilder/src/lib.rs b/src/obiskbuilder/src/lib.rs index 61d8cf1..2dd9c2f 100644 --- a/src/obiskbuilder/src/lib.rs +++ b/src/obiskbuilder/src/lib.rs @@ -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; diff --git a/src/obiskbuilder/src/rolling_stat.rs b/src/obiskbuilder/src/rolling_stat.rs index 28d49a6..c824c99 100644 --- a/src/obiskbuilder/src/rolling_stat.rs +++ b/src/obiskbuilder/src/rolling_stat.rs @@ -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, - k2q: Ring, - k3q: Ring, - k4q: Ring, - k5q: Ring, - k6q: Ring, + // Minimizer selection state. 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], + // 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( - 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::(canonical); - } - - #[inline] - fn update_sums_increment( - 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::(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::(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); + 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 { 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)) } } -- 2.54.0 From fd2c23e7df12e71d18685be967f32c1210f347d9 Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Wed, 8 Jul 2026 19:26:33 +0200 Subject: [PATCH 3/3] refactor: remove equivalence class folding from entropy pipeline MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Removes circular-reverse complement machinery and explicit k-mer canonicalization across the entropy pipeline. Frequency tallying and Shannon entropy computation now operate directly on raw k-mer values, eliminating prior score inflation and alignment-dependent artifacts while preserving orientation invariance. Updates build scripts to generate normalized lookup tables for k-mer lengths 1–6, restricts the public API to `EntropyTracker`, and bumps crate versions. Documentation is updated to reflect the simplified raw-value approach and revised module structure. --- docmd/theory/entropy.md | 48 ++++--- docmd/theory/entropy.refs.md | 10 +- src/Cargo.lock | 2 +- src/obikentropy/build.rs | 63 +-------- src/obikentropy/src/kmer_entropy.rs | 4 +- src/obikentropy/src/lib.rs | 2 +- src/obikentropy/src/table.rs | 106 ++------------- src/obikentropy/src/tracker.rs | 185 ++++++++++++--------------- src/obikmer/Cargo.toml | 2 +- src/obiskbuilder/src/rolling_stat.rs | 6 +- 10 files changed, 142 insertions(+), 286 deletions(-) diff --git a/docmd/theory/entropy.md b/docmd/theory/entropy.md index 4cd159e..4811240 100644 --- a/docmd/theory/entropy.md +++ b/docmd/theory/entropy.md @@ -1,6 +1,6 @@ # Kmer entropy filter -Low-complexity kmers (polyA, polyT, tandem repeats) are detected and excluded during phase 1. The filter computes a **normalized Shannon entropy** over sub-words of multiple sizes, corrected for two sources of bias: the small number of observations within a single kmer, and the unequal sizes of circular equivalence classes. +Low-complexity kmers (polyA, polyT, tandem repeats) are detected and excluded during phase 1. The filter computes a **normalized Shannon entropy** over sub-words of multiple sizes, corrected for one source of bias: the small number of observations within a single kmer relative to the number of possible sub-words. ## Sub-word frequencies @@ -8,17 +8,15 @@ For a kmer of length k and a sub-word size ws (1 ≤ ws ≤ ws_max, typically ws $$w_i = \text{kmer}[i \mathinner{..} i+ws-1], \quad i = 0, \ldots, n_{\text{words}}-1$$ -Each sub-word is mapped to its **circular canonical form**: the lexicographic minimum among all cyclic rotations of the word **and all cyclic rotations of its reverse complement**. This extended equivalence relation ensures that entropy(K) = entropy(revcomp(K)) — the filter is strand-symmetric. Let $s_j$ be the size of equivalence class $j$ (number of distinct raw words mapping to canonical form $j$), and $f_j$ the count of canonical form $j$ among the $n_{\text{words}}$ sub-words ($\sum_j f_j = n_{\text{words}}$). +Each sub-word is tallied under its own raw 2-bit-packed value — **no canonicalization**. Let $f_j$ be the count of raw word $j$ among the $n_{\text{words}}$ sub-words ($\sum_j f_j = n_{\text{words}}$), over the $4^{ws}$ possible raw words. + +An earlier version of this filter first folded each sub-word into a circular+reverse-complement equivalence class, then "unfolded" the observed class frequency back onto its members to correct for unequal class sizes. That machinery bought nothing it was claimed for — see *Why no equivalence classes* below — while measurably weakening detection of the very sequences the filter exists to catch, so it was removed. ## Corrected Shannon entropy -The circular equivalence classes have unequal sizes: under a uniform distribution over all $4^{ws}$ raw words, class $j$ is visited with probability $s_j / 4^{ws}$, not $1/n_a$. Computing entropy directly over canonical classes therefore underestimates the entropy of a random sequence. +$$H_{\text{corr}} = \log(n_{\text{words}}) - \frac{1}{n_{\text{words}}} \sum_j f_j \log f_j$$ -The correction "unfolds" each canonical class back to its member raw words, redistributing each observation of class $j$ equally among its $s_j$ members: - -$$H_{\text{corr}} = \log(n_{\text{words}}) - \frac{1}{n_{\text{words}}} \sum_j f_j \log f_j + \frac{1}{n_{\text{words}}} \sum_j f_j \log s_j$$ - -The last term is the correction for unequal class sizes. For a uniformly random sequence ($f_j \approx n_{\text{words}} \cdot s_j / 4^{ws}$), this gives $H_{\text{corr}} \approx \log(4^{ws}) = 2 \cdot ws \cdot \log 2$, the maximum entropy over raw words. +This is a plain Shannon entropy over the observed raw-word frequencies. ## Maximum entropy correction for small samples @@ -42,27 +40,45 @@ $$\text{entropy}(kmer) = \min_{ws=1}^{ws_{\max}} \hat{H}(ws)$$ A value near 0 indicates low complexity (e.g. AAAA…); near 1 indicates high complexity. A kmer is rejected if $\text{entropy}(kmer) < \theta$, where $\theta$ is a collection parameter (default 0.7). The minimum across word sizes ensures that any scale of repetition is detected independently: polyA is caught at ws=1, dinucleotide repeats at ws=2, etc. +## Why no equivalence classes + +A prior design folded each sub-word into the canonical form of its circular-rotation + reverse-complement equivalence class before tallying, on the reasoning that (a) it guarantees $\text{entropy}(K) = \text{entropy}(\text{revcomp}(K))$, and (b) collapsing phase-shifted repeats (e.g. `ATG` ≡ `TGA` ≡ `GAT`) into one class better reflects that they are "the same" low-complexity pattern. + +Both properties already hold for the raw, unfolded entropy above, without any class machinery: + +- **Reverse complement**: for any K of length n, window $j$ of $\text{revcomp}(K)$ equals $\text{revcomp}$ of window $(n{-}ws{-}j)$ of K. This is a bijection between the window sets under which each window maps to its own revcomp — and revcomp is itself a bijection (involution) on the space of raw ws-mers. So the multiset of raw-word frequencies for $\text{revcomp}(K)$ is exactly a relabeling of the multiset for K, and Shannon entropy — a function of the frequency multiset alone — is exactly invariant. No folding required, for any K. +- **Tandem repeats**: a period-p repeat sampled by a stride-1 sliding window naturally cycles through its own rotations as raw tokens (e.g. `ATGATGATG…` yields the raw words `ATG`, `TGA`, `GAT` in rotation as the window slides). The low diversity this represents (few distinct raw words out of $4^{ws}$ possible) is already visible in the raw frequency distribution — no folding needed to detect it. + +What the fold-then-unfold step actually did was credit each observed class with the frequency of equivalence-class members that were **never observed on the read strand**, inflating $H_{\text{corr}}$ for genuine repeats. Worked example: k=31, ws=3, kmer = `ATG` repeated ($n_{\text{words}}=29$, all 29 windows fall into one class of size 6 under the old scheme — 3 rotations × forward/revcomp): + +| | $H_{\text{corr}}$ | normalized | +|---|---|---| +| old (folded, class size 6) | $\log 6 \approx 1.79$ | $\approx 0.53$ | +| current (raw, unfolded) | $\log 3 \approx 1.10$ | $\approx 0.33$ | + +The gap is not a rounding artifact: per sub-word order, the folded score for this same repeat swings from 0.53 (ws=3, aligned with the period) up to **1.03** (ws=5, misaligned with the period) — i.e. a period-3 repeat could score *above* the theoretical maximum for a random sequence, depending on which ws happens to divide the repeat's period. The raw formula stays flat at ≈0.33–0.40 across ws=2..6 regardless of alignment, which is the robustness the "minimum across ws" design was meant to provide in the first place. + ## Interpretation as an effective number of classes -$H_{\text{corr}}$ is a standard Shannon entropy over raw words (after unfolding the equivalence classes), so the classical perplexity interpretation holds directly: $N_{\text{eff}} = e^{H_{\text{corr}}}$ is the number of equiprobable classes that would yield the same entropy. +$H_{\text{corr}}$ is a standard Shannon entropy over raw words, so the classical perplexity interpretation holds directly: $N_{\text{eff}} = e^{H_{\text{corr}}}$ is the number of equiprobable raw words that would yield the same entropy. -For the normalised score $\hat{H}$, dividing by $H_{\text{max}}$ changes the logarithm base: +For the normalised score $\hat{H}$, dividing by $H_{\max}$ changes the logarithm base: -$$\hat{H} = \frac{\log N_{\text{eff}}}{\log N_{\text{max}}} = \log_{N_{\text{max}}} N_{\text{eff}} \quad \Longleftrightarrow \quad N_{\text{eff}} = N_{\text{max}}^{\,\hat{H}}$$ +$$\hat{H} = \frac{\log N_{\text{eff}}}{\log N_{\max}} = \log_{N_{\max}} N_{\text{eff}} \quad \Longleftrightarrow \quad N_{\text{eff}} = N_{\max}^{\,\hat{H}}$$ -The property is preserved: $\hat{H}$ is the logarithm (in base $N_{\text{max}}$) of the effective number of equi-represented classes. +The property is preserved: $\hat{H}$ is the logarithm (in base $N_{\max}$) of the effective number of equi-represented raw words. -In the large-sample limit ($n_{\text{words}} \gg 4^{ws}$), $N_{\text{max}} \approx 4^{ws}$, giving: +In the large-sample limit ($n_{\text{words}} \gg 4^{ws}$), $N_{\max} \approx 4^{ws}$, giving: $$N_{\text{eff}} \approx 4^{ws \cdot \hat{H}}$$ -This has a clean interpretation: $ws \cdot \hat{H}$ is the **effective word length** (in bases) of a perfectly uniform distribution that would produce the same entropy. At $\hat{H} = 1$ the full space of $4^{ws}$ words is used; at $\hat{H} = 0.5$ with ws=2, only $4^1 = 4$ effective classes out of 16 are occupied. +This has a clean interpretation: $ws \cdot \hat{H}$ is the **effective word length** (in bases) of a perfectly uniform distribution that would produce the same entropy. At $\hat{H} = 1$ the full space of $4^{ws}$ words is used; at $\hat{H} = 0.5$ with ws=2, only $4^1 = 4$ effective words out of 16 are occupied. -In our actual regime, $n_{\text{words}}$ is small and $4^{ws}$ can exceed $n_{\text{words}}$, so $H_{\text{max}} < \log(4^{ws})$ due to the small-sample correction. The exact effective count is $N_{\text{max}}^{\hat{H}}$, not $4^{ws \cdot \hat{H}}$. +In our actual regime, $n_{\text{words}}$ is small and $4^{ws}$ can exceed $n_{\text{words}}$, so $H_{\max} < \log(4^{ws})$ due to the small-sample correction. The exact effective count is $N_{\max}^{\hat{H}}$, not $4^{ws \cdot \hat{H}}$. ## Properties The entropy score is a function of the kmer sequence alone — it does not depend on the surrounding context or on the position within any genome. Two consequences: -- **Orientation invariance**: $\text{entropy}(K) = \text{entropy}(\text{revcomp}(K))$, guaranteed by the strand-symmetric canonical form. +- **Orientation invariance**: $\text{entropy}(K) = \text{entropy}(\text{revcomp}(K))$ — see *Why no equivalence classes* above for why this holds without any explicit strand-folding step. - **Context independence**: the same kmer is always rejected or always kept, regardless of which genome it occurs in, where in that genome it appears, or which strand is considered. The filter defines a fixed partition of the kmer space into low-complexity and valid kmers. diff --git a/docmd/theory/entropy.refs.md b/docmd/theory/entropy.refs.md index 79ca55c..7701468 100644 --- a/docmd/theory/entropy.refs.md +++ b/docmd/theory/entropy.refs.md @@ -3,10 +3,14 @@ ## Code couvert -- `obiskbuilder/src/entropy_table.rs` — filtre Shannon sur les kmers à basse complexité -- `obiskbuilder/src/lib.rs` — application du filtre lors du scatter (phase 1) +- `obikentropy/src/table.rs`, `obikentropy/src/tracker.rs` — formule d'entropie et tables de correction petits effectifs +- `obikentropy/src/kmer_entropy.rs` — entropie d'un kmer isolé (`KmerEntropy`) +- `obiskbuilder/src/rolling_stat.rs` — composition de `obikentropy::EntropyTracker` dans le suivi streaming (sélection de minimiseur + entropie) +- `obiskbuilder/src/iter.rs`, `obiskbuilder/src/stream_iter.rs` — application du filtre lors du scatter (phase 1) ## Notes -Document théorique stable. Vérifier que les paramètres `theta` et `level_max` dans le CLI +Le repli en classes d'équivalence circulaires + brin inverse (décrit dans une version antérieure de ce document) a été supprimé : voir la section « Why no equivalence classes » de `entropy.md` pour la justification théorique et numérique. + +Vérifier que les paramètres `theta` et `level_max` dans le CLI (`obikmer/src/cli.rs` → `CommonArgs`) correspondent bien à ce qui est décrit. diff --git a/src/Cargo.lock b/src/Cargo.lock index f7c63e3..0dc087d 100644 --- a/src/Cargo.lock +++ b/src/Cargo.lock @@ -1711,7 +1711,7 @@ dependencies = [ [[package]] name = "obikmer" -version = "1.1.38" +version = "1.1.39" dependencies = [ "clap", "csv", diff --git a/src/obikentropy/build.rs b/src/obikentropy/build.rs index f422279..d2c33b9 100644 --- a/src/obikentropy/build.rs +++ b/src/obikentropy/build.rs @@ -4,57 +4,6 @@ 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 { @@ -63,6 +12,9 @@ fn build_n_log_n() -> [f64; K_MAX + 1] { t } +/// Max achievable entropy over `4^ws` raw sub-words given only `nwords` +/// observations (most-uniform integer partition), per +/// `docmd/theory/entropy.md`. 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 { @@ -125,13 +77,6 @@ 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); @@ -141,5 +86,5 @@ fn main() { 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(); + fs::write(out_dir.join("entropy_tables.rs"), out).unwrap(); } diff --git a/src/obikentropy/src/kmer_entropy.rs b/src/obikentropy/src/kmer_entropy.rs index 1989972..1efd64d 100644 --- a/src/obikentropy/src/kmer_entropy.rs +++ b/src/obikentropy/src/kmer_entropy.rs @@ -7,7 +7,7 @@ use obikseq::CanonicalKmer; -use crate::tracker::{EntropyTracker, sub_word_canon}; +use crate::tracker::EntropyTracker; /// Extension trait: compute the normalized entropy of a single canonical /// k-mer, independent of any surrounding sequence. @@ -30,7 +30,7 @@ impl KmerEntropy for CanonicalKmer { 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.push(i + 1, rolling); } tracker.normalized_entropy(level_max) } diff --git a/src/obikentropy/src/lib.rs b/src/obikentropy/src/lib.rs index 113bdd3..dc18360 100644 --- a/src/obikentropy/src/lib.rs +++ b/src/obikentropy/src/lib.rs @@ -14,4 +14,4 @@ mod table; mod tracker; pub use kmer_entropy::KmerEntropy; -pub use tracker::{EntropyTracker, SubWordCanon, sub_word_canon}; +pub use tracker::EntropyTracker; diff --git a/src/obikentropy/src/table.rs b/src/obikentropy/src/table.rs index 289aa94..37ef4df 100644 --- a/src/obikentropy/src/table.rs +++ b/src/obikentropy/src/table.rs @@ -1,68 +1,16 @@ -//! 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`. +//! Compile-time tables backing the normalized k-mer entropy formula: the +//! max-entropy correction for small samples. See `docmd/theory/entropy.md`. +//! +//! Entropy is computed directly on raw (non-canonicalized) sub-words — no +//! equivalence-class folding. Empirically (see the discussion that produced +//! this crate's history), folding sub-words into circular/revcomp classes +//! before unfolding them back buys nothing for the invariances it was meant +//! to guarantee (both hold for raw sub-word entropy already, by a direct +//! bijection argument for revcomp and by the sliding window's own dynamics +//! for tandem repeats), while it measurably *weakens* detection of the +//! low-complexity sequences the filter exists to catch. -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() -> [u64; N] { - let mut result = [0u64; N]; - let k = k_from_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() -> usize { - match N { - 4 => 1, - 16 => 2, - 64 => 3, - 256 => 4, - 1024 => 5, - 4096 => 6, - _ => panic!("N must be a power of 4"), - } -} +include!(concat!(env!("OUT_DIR"), "/entropy_tables.rs")); pub(crate) const WS_MAX: usize = 6; @@ -80,33 +28,3 @@ pub(crate) const fn emax(k: usize, ws: usize) -> f64 { pub(crate) const fn log_nwords(k: usize, ws: usize) -> f64 { LOG_NWORDS[k][ws] } - -#[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 - } -} - -#[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"), - } -} diff --git a/src/obikentropy/src/tracker.rs b/src/obikentropy/src/tracker.rs index 98686fa..eb71bab 100644 --- a/src/obikentropy/src/tracker.rs +++ b/src/obikentropy/src/tracker.rs @@ -1,9 +1,12 @@ //! 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. +//! the per-sub-word-size raw-word 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. No +//! canonicalization is applied — each sub-word is tallied under its own raw +//! 2-bit-packed value; only the small-sample max-entropy correction departs +//! from a textbook Shannon entropy. //! //! It carries no notion of minimizers or superkmer segmentation — callers //! that need both (e.g. `obiskbuilder::RollingStat`) compose an @@ -12,25 +15,7 @@ //! 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::(rolling_kmer & 3), - entropy_norm_kmer::(rolling_kmer & 15), - entropy_norm_kmer::(rolling_kmer & 63), - entropy_norm_kmer::(rolling_kmer & 255), - entropy_norm_kmer::(rolling_kmer & 1023), - entropy_norm_kmer::(rolling_kmer & 4095), - ] -} +use crate::table::{WS_MAX, emax, log_nwords, n_log_n}; /// Incremental normalized-entropy accumulator over a sliding window of `k` /// bases. Composed as a plain field by callers that also need other @@ -39,8 +24,8 @@ 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. + // Sliding-window queues over the last `k` raw sub-words, one per word + // size — stack-allocated, capacity ≤ k ≤ 31. k1q: Ring, k2q: Ring, k3q: Ring, @@ -48,7 +33,8 @@ pub struct EntropyTracker { k5q: Ring, k6q: Ring, - // Frequency count arrays. Max count per cell ≤ k ≤ 31 → u8 is sufficient. + // Frequency count arrays, indexed by the raw sub-word value (2 bits per + // base). Max count per cell ≤ k ≤ 31 → u8 is sufficient. k1c: [u8; 4], k2c: [u8; 16], k3c: [u8; 64], @@ -57,7 +43,6 @@ pub struct EntropyTracker { k6c: [u8; 4096], sum_f_log_f: [f64; WS_MAX + 1], - sum_f_log_s: [f64; WS_MAX + 1], } impl EntropyTracker { @@ -79,7 +64,6 @@ impl EntropyTracker { k5c: [0; 1024], k6c: [0; 4096], sum_f_log_f: [0.0; WS_MAX + 1], - sum_f_log_s: [0.0; WS_MAX + 1], } } @@ -103,103 +87,94 @@ impl EntropyTracker { 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( - sum_f_log_f: &mut [f64; WS_MAX + 1], - sum_f_log_s: &mut [f64; WS_MAX + 1], - canonical: u64, - f: usize, - ) { + fn update_sums_decrement(sum_f_log_f: &mut [f64; WS_MAX + 1], f: usize) { 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( - sum_f_log_f: &mut [f64; WS_MAX + 1], - sum_f_log_s: &mut [f64; WS_MAX + 1], - canonical: u64, - g: usize, - ) { + fn update_sums_increment(sum_f_log_f: &mut [f64; WS_MAX + 1], g: usize) { sum_f_log_f[K] += n_log_n(g + 1) - n_log_n(g); - sum_f_log_s[K] += ln_class_size::(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; + /// `rolling_kmer` is the current right-aligned, 2-bit-packed k-mer + /// window (same convention as `obiskbuilder::RollingStat::rolling_k`). + pub fn push(&mut self, received: usize, rolling_kmer: u64) { + let raw1 = rolling_kmer & 3; + let raw2 = rolling_kmer & 15; + let raw3 = rolling_kmer & 63; + let raw4 = rolling_kmer & 255; + let raw5 = rolling_kmer & 1023; + let raw6 = rolling_kmer & 4095; 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::update_sums_decrement::<1>(&mut self.sum_f_log_f, 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::update_sums_decrement::<2>(&mut self.sum_f_log_f, 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::update_sums_decrement::<3>(&mut self.sum_f_log_f, 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::update_sums_decrement::<4>(&mut self.sum_f_log_f, 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::update_sums_decrement::<5>(&mut self.sum_f_log_f, 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::update_sums_decrement::<6>(&mut self.sum_f_log_f, 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 g1 = self.k1c[raw1 as usize] as usize; + Self::update_sums_increment::<1>(&mut self.sum_f_log_f, g1); + self.k1c[raw1 as usize] += 1; + self.k1q.push_back(raw1); - 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 g2 = self.k2c[raw2 as usize] as usize; + Self::update_sums_increment::<2>(&mut self.sum_f_log_f, g2); + self.k2c[raw2 as usize] += 1; + self.k2q.push_back(raw2); - 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 g3 = self.k3c[raw3 as usize] as usize; + Self::update_sums_increment::<3>(&mut self.sum_f_log_f, g3); + self.k3c[raw3 as usize] += 1; + self.k3q.push_back(raw3); - 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 g4 = self.k4c[raw4 as usize] as usize; + Self::update_sums_increment::<4>(&mut self.sum_f_log_f, g4); + self.k4c[raw4 as usize] += 1; + self.k4q.push_back(raw4); - 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 g5 = self.k5c[raw5 as usize] as usize; + Self::update_sums_increment::<5>(&mut self.sum_f_log_f, g5); + self.k5c[raw5 as usize] += 1; + self.k5q.push_back(raw5); - 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); + let g6 = self.k6c[raw6 as usize] as usize; + Self::update_sums_increment::<6>(&mut self.sum_f_log_f, g6); + self.k6c[raw6 as usize] += 1; + self.k6q.push_back(raw6); } else { - self.push_warmup_increments(received, canonical_k1, canonical_k2, canonical_k3, canonical_k4, canonical_k5, canonical_k6); + self.push_warmup_increments(received, raw1, raw2, raw3, raw4, raw5, raw6); } } @@ -208,43 +183,43 @@ impl EntropyTracker { 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, + raw1: u64, raw2: u64, raw3: u64, + raw4: u64, raw5: u64, raw6: 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); + let g1 = self.k1c[raw1 as usize] as usize; + Self::update_sums_increment::<1>(&mut self.sum_f_log_f, g1); + self.k1c[raw1 as usize] += 1; + self.k1q.push_back(raw1); 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); + let g2 = self.k2c[raw2 as usize] as usize; + Self::update_sums_increment::<2>(&mut self.sum_f_log_f, g2); + self.k2c[raw2 as usize] += 1; + self.k2q.push_back(raw2); 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); + let g3 = self.k3c[raw3 as usize] as usize; + Self::update_sums_increment::<3>(&mut self.sum_f_log_f, g3); + self.k3c[raw3 as usize] += 1; + self.k3q.push_back(raw3); 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); + let g4 = self.k4c[raw4 as usize] as usize; + Self::update_sums_increment::<4>(&mut self.sum_f_log_f, g4); + self.k4c[raw4 as usize] += 1; + self.k4q.push_back(raw4); 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); + let g5 = self.k5c[raw5 as usize] as usize; + Self::update_sums_increment::<5>(&mut self.sum_f_log_f, g5); + self.k5c[raw5 as usize] += 1; + self.k5q.push_back(raw5); 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); + let g6 = self.k6c[raw6 as usize] as usize; + Self::update_sums_increment::<6>(&mut self.sum_f_log_f, g6); + self.k6c[raw6 as usize] += 1; + self.k6q.push_back(raw6); self.steady = true; } } @@ -265,7 +240,7 @@ impl EntropyTracker { 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; + let h_corr = log_nw - self.sum_f_log_f[order] / nw_f; (h_corr / em).max(0.0) } diff --git a/src/obikmer/Cargo.toml b/src/obikmer/Cargo.toml index d9ea7a6..9d00ffd 100644 --- a/src/obikmer/Cargo.toml +++ b/src/obikmer/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "obikmer" -version = "1.1.38" +version = "1.1.39" edition = "2024" [[bin]] diff --git a/src/obiskbuilder/src/rolling_stat.rs b/src/obiskbuilder/src/rolling_stat.rs index c824c99..09b3f8d 100644 --- a/src/obiskbuilder/src/rolling_stat.rs +++ b/src/obiskbuilder/src/rolling_stat.rs @@ -1,4 +1,4 @@ -use obikentropy::{EntropyTracker, sub_word_canon}; +use obikentropy::EntropyTracker; use obikseq::kmer::{Minimizer, hash_kmer}; use obikseq::params; @@ -136,8 +136,6 @@ impl RollingStat { self.rolling_rck = ((self.rolling_rck >> 2) | ((cnuc as u64) << ((k - 1) * 2))) & self.k_mask; - let canon = sub_word_canon(self.rolling_k); - self.received += 1; if self.received >= m { @@ -170,7 +168,7 @@ impl RollingStat { } } - self.entropy.push(self.received, canon); + self.entropy.push(self.received, self.rolling_k); } pub fn ready(&self) -> bool { -- 2.54.0