refactor: remove equivalence class folding from entropy pipeline

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.
This commit is contained in:
Eric Coissac
2026-07-08 19:36:30 +02:00
parent 912f788f7f
commit 3e026ee80f
10 changed files with 142 additions and 286 deletions
+32 -16
View File
@@ -1,6 +1,6 @@
# Kmer entropy filter # 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 ## 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$$ $$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 ## 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: This is a plain Shannon entropy over the observed raw-word frequencies.
$$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.
## Maximum entropy correction for small samples ## 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. 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.330.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 ## 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}}$$ $$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 ## 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: 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. - **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.
+7 -3
View File
@@ -3,10 +3,14 @@
## Code couvert ## Code couvert
- `obiskbuilder/src/entropy_table.rs` — filtre Shannon sur les kmers à basse complexité - `obikentropy/src/table.rs`, `obikentropy/src/tracker.rs` — formule d'entropie et tables de correction petits effectifs
- `obiskbuilder/src/lib.rs` — application du filtre lors du scatter (phase 1) - `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 ## 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. (`obikmer/src/cli.rs``CommonArgs`) correspondent bien à ce qui est décrit.
+1 -1
View File
@@ -1711,7 +1711,7 @@ dependencies = [
[[package]] [[package]]
name = "obikmer" name = "obikmer"
version = "1.1.38" version = "1.1.39"
dependencies = [ dependencies = [
"clap", "clap",
"csv", "csv",
+4 -59
View File
@@ -4,57 +4,6 @@ use std::path::PathBuf;
const K_MAX: usize = 32; const K_MAX: usize = 32;
const WS_MAX: usize = 6; const WS_MAX: usize = 6;
fn normalize_circular(kmer: u64, ws: usize) -> u64 {
let mask = (1u64 << (ws * 2)) - 1;
let mut canonical = kmer & mask;
let mut current = canonical;
for _ in 0..ws - 1 {
let top = (current >> ((ws - 1) * 2)) & 3;
current = ((current << 2) | top) & mask;
if current < canonical {
canonical = current;
}
}
canonical
}
fn revcomp_raw(x: u64, k: usize) -> u64 {
let x = !x;
let x = x.swap_bytes();
let x = ((x >> 4) & 0x0F0F0F0F0F0F0F0F) | ((x & 0x0F0F0F0F0F0F0F0F) << 4);
let x = ((x >> 2) & 0x3333333333333333) | ((x & 0x3333333333333333) << 2);
x << (64 - 2 * k)
}
fn build_normalized_kmer(k: usize) -> Vec<u64> {
let n = 1usize << (k * 2);
let shift = 64 - k * 2;
let mut result = vec![0u64; n];
for i in 0..n {
let la = (i as u64) << shift;
let ra = i as u64;
let rc_ra = revcomp_raw(la, k) >> shift;
let circ = normalize_circular(ra, k);
let circ_rc = normalize_circular(rc_ra, k);
result[i] = if circ < circ_rc { circ } else { circ_rc };
}
result
}
fn build_ln_class(norm: &[u64]) -> Vec<f64> {
let n = norm.len();
let mut sizes = vec![0u32; n];
for &c in norm {
sizes[c as usize] += 1;
}
norm.iter()
.map(|&c| {
let s = sizes[c as usize];
if s > 0 { (s as f64).ln() } else { 0.0 }
})
.collect()
}
fn build_n_log_n() -> [f64; K_MAX + 1] { fn build_n_log_n() -> [f64; K_MAX + 1] {
let mut t = [0.0f64; K_MAX + 1]; let mut t = [0.0f64; K_MAX + 1];
for n in 1..=K_MAX { for n in 1..=K_MAX {
@@ -63,6 +12,9 @@ fn build_n_log_n() -> [f64; K_MAX + 1] {
t 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] { fn build_emax() -> [[f64; WS_MAX + 1]; K_MAX + 1] {
let mut t = [[0.0f64; WS_MAX + 1]; K_MAX + 1]; let mut t = [[0.0f64; WS_MAX + 1]; K_MAX + 1];
for k in 2..=K_MAX { for k in 2..=K_MAX {
@@ -125,13 +77,6 @@ fn main() {
let out_dir = PathBuf::from(std::env::var("OUT_DIR").unwrap()); let out_dir = PathBuf::from(std::env::var("OUT_DIR").unwrap());
let mut out = String::new(); 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(); let n_log_n = build_n_log_n();
emit_f64_1d(&mut out, "N_LOG_N", K_MAX + 1, &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(); let log_nwords = build_log_nwords();
emit_f64_2d(&mut out, "LOG_NWORDS", K_MAX + 1, WS_MAX + 1, &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();
} }
+2 -2
View File
@@ -7,7 +7,7 @@
use obikseq::CanonicalKmer; use obikseq::CanonicalKmer;
use crate::tracker::{EntropyTracker, sub_word_canon}; use crate::tracker::EntropyTracker;
/// Extension trait: compute the normalized entropy of a single canonical /// Extension trait: compute the normalized entropy of a single canonical
/// k-mer, independent of any surrounding sequence. /// k-mer, independent of any surrounding sequence.
@@ -30,7 +30,7 @@ impl KmerEntropy for CanonicalKmer {
let shift = 64 - 2 * (i + 1); let shift = 64 - 2 * (i + 1);
let base = (raw >> shift) & 3; let base = (raw >> shift) & 3;
rolling = ((rolling << 2) | base) & mask; rolling = ((rolling << 2) | base) & mask;
tracker.push(i + 1, sub_word_canon(rolling)); tracker.push(i + 1, rolling);
} }
tracker.normalized_entropy(level_max) tracker.normalized_entropy(level_max)
} }
+1 -1
View File
@@ -14,4 +14,4 @@ mod table;
mod tracker; mod tracker;
pub use kmer_entropy::KmerEntropy; pub use kmer_entropy::KmerEntropy;
pub use tracker::{EntropyTracker, SubWordCanon, sub_word_canon}; pub use tracker::EntropyTracker;
+12 -94
View File
@@ -1,68 +1,16 @@
//! Compile-time tables backing the normalized k-mer entropy formula: //! Compile-time tables backing the normalized k-mer entropy formula: the
//! circular-canonical sub-word classes, their log-sizes, and the max-entropy //! max-entropy correction for small samples. See `docmd/theory/entropy.md`.
//! 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>(); include!(concat!(env!("OUT_DIR"), "/entropy_tables.rs"));
pub(crate) const NORMK2: [u64; 16] = build_normalized_kmer::<16>();
pub(crate) const NORMK3: [u64; 64] = build_normalized_kmer::<64>();
pub(crate) const NORMK4: [u64; 256] = build_normalized_kmer::<256>();
pub(crate) const NORMK5: [u64; 1024] = build_normalized_kmer::<1024>();
pub(crate) const NORMK6: [u64; 4096] = build_normalized_kmer::<4096>();
include!(concat!(env!("OUT_DIR"), "/ln_class_tables.rs"));
const fn normalize_circular(kmer: u64, ws: usize) -> u64 {
let mask = (1u64 << (ws * 2)) - 1;
let mut canonical = kmer & mask;
let mut current = canonical;
let mut i = 0;
while i < (ws - 1) {
let top = (current >> ((ws - 1) * 2)) & 3;
current = ((current << 2) | top) & mask;
if current < canonical {
canonical = current;
}
i += 1;
}
canonical
}
const fn build_normalized_kmer<const N: usize>() -> [u64; N] {
let mut result = [0u64; N];
let k = k_from_n::<N>();
let shift = 64 - k * 2;
let mut i = 0;
while i < N {
let la = (i as u64) << shift;
let ra = i as u64;
let rc_ra = revcomp_raw(la, k) >> shift;
let circ = normalize_circular(ra, k);
let circ_rc = normalize_circular(rc_ra, k);
result[i] = if circ < circ_rc { circ } else { circ_rc };
i += 1;
}
result
}
const fn revcomp_raw(x: u64, k: usize) -> u64 {
let x = !x;
let x = x.swap_bytes();
let x = ((x >> 4) & 0x0F0F0F0F0F0F0F0F) | ((x & 0x0F0F0F0F0F0F0F0F) << 4);
let x = ((x >> 2) & 0x3333333333333333) | ((x & 0x3333333333333333) << 2);
x << (64 - 2 * k)
}
const fn k_from_n<const N: usize>() -> usize {
match N {
4 => 1,
16 => 2,
64 => 3,
256 => 4,
1024 => 5,
4096 => 6,
_ => panic!("N must be a power of 4"),
}
}
pub(crate) const WS_MAX: usize = 6; 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 { pub(crate) const fn log_nwords(k: usize, ws: usize) -> f64 {
LOG_NWORDS[k][ws] LOG_NWORDS[k][ws]
} }
#[inline(always)]
pub(crate) const fn entropy_norm_kmer<const LEFT: bool, const K: usize>(kmer: u64) -> u64 {
const SHIFT: [usize; 7] = [0, 62, 60, 58, 56, 54, 52];
const NORM: [&[u64]; 7] = [&[], &NORMK1, &NORMK2, &NORMK3, &NORMK4, &NORMK5, &NORMK6];
let shift = SHIFT[K];
let ra = if LEFT { kmer >> shift } else { kmer };
let canonical_ra = NORM[K][ra as usize];
if LEFT {
canonical_ra << shift
} else {
canonical_ra
}
}
#[inline(always)]
pub(crate) const fn ln_class_size<const LEFT: bool, const K: usize>(kmer: u64) -> f64 {
const SHIFT: [usize; 7] = [0, 62, 60, 58, 56, 54, 52];
let ra = if LEFT { kmer >> SHIFT[K] } else { kmer };
match K {
1 => LN_CLASS1[ra as usize],
2 => LN_CLASS2[ra as usize],
3 => LN_CLASS3[ra as usize],
4 => LN_CLASS4[ra as usize],
5 => LN_CLASS5[ra as usize],
6 => LN_CLASS6[ra as usize],
_ => panic!("k must be 1..=6"),
}
}
+80 -105
View File
@@ -1,9 +1,12 @@
//! Incremental (streaming) normalized k-mer entropy. //! Incremental (streaming) normalized k-mer entropy.
//! //!
//! [`EntropyTracker`] maintains, over a sliding window of the last `k` bases, //! [`EntropyTracker`] maintains, over a sliding window of the last `k` bases,
//! the per-sub-word-size frequency statistics needed to evaluate the //! the per-sub-word-size raw-word frequency statistics needed to evaluate
//! corrected Shannon entropy described in `docmd/theory/entropy.md`, updated //! the corrected Shannon entropy described in `docmd/theory/entropy.md`,
//! in O(1) per base rather than recomputed from scratch. //! 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 //! It carries no notion of minimizers or superkmer segmentation — callers
//! that need both (e.g. `obiskbuilder::RollingStat`) compose an //! that need both (e.g. `obiskbuilder::RollingStat`) compose an
@@ -12,25 +15,7 @@
//! struct. //! struct.
use crate::ring::Ring; use crate::ring::Ring;
use crate::table::{WS_MAX, emax, entropy_norm_kmer, ln_class_size, log_nwords, n_log_n}; use crate::table::{WS_MAX, emax, log_nwords, n_log_n};
/// Canonical sub-word values for word sizes 1..=6, computed by the caller
/// from its own rolling k-mer bits (see [`EntropyTracker::push`]).
pub type SubWordCanon = [u64; WS_MAX];
/// Compute the six canonical sub-word values (word sizes 1..=6) for the
/// current rolling right-aligned k-mer window.
#[inline]
pub fn sub_word_canon(rolling_kmer: u64) -> SubWordCanon {
[
entropy_norm_kmer::<false, 1>(rolling_kmer & 3),
entropy_norm_kmer::<false, 2>(rolling_kmer & 15),
entropy_norm_kmer::<false, 3>(rolling_kmer & 63),
entropy_norm_kmer::<false, 4>(rolling_kmer & 255),
entropy_norm_kmer::<false, 5>(rolling_kmer & 1023),
entropy_norm_kmer::<false, 6>(rolling_kmer & 4095),
]
}
/// Incremental normalized-entropy accumulator over a sliding window of `k` /// Incremental normalized-entropy accumulator over a sliding window of `k`
/// bases. Composed as a plain field by callers that also need other /// bases. Composed as a plain field by callers that also need other
@@ -39,8 +24,8 @@ pub struct EntropyTracker {
k: usize, k: usize,
steady: bool, steady: bool,
// Sliding-window queues over the last `k` canonical sub-words, one per // Sliding-window queues over the last `k` raw sub-words, one per word
// word size — stack-allocated, capacity ≤ k ≤ 31. // size — stack-allocated, capacity ≤ k ≤ 31.
k1q: Ring<u64, 32>, k1q: Ring<u64, 32>,
k2q: Ring<u64, 32>, k2q: Ring<u64, 32>,
k3q: Ring<u64, 32>, k3q: Ring<u64, 32>,
@@ -48,7 +33,8 @@ pub struct EntropyTracker {
k5q: Ring<u64, 32>, k5q: Ring<u64, 32>,
k6q: Ring<u64, 32>, k6q: Ring<u64, 32>,
// 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], k1c: [u8; 4],
k2c: [u8; 16], k2c: [u8; 16],
k3c: [u8; 64], k3c: [u8; 64],
@@ -57,7 +43,6 @@ pub struct EntropyTracker {
k6c: [u8; 4096], k6c: [u8; 4096],
sum_f_log_f: [f64; WS_MAX + 1], sum_f_log_f: [f64; WS_MAX + 1],
sum_f_log_s: [f64; WS_MAX + 1],
} }
impl EntropyTracker { impl EntropyTracker {
@@ -79,7 +64,6 @@ impl EntropyTracker {
k5c: [0; 1024], k5c: [0; 1024],
k6c: [0; 4096], k6c: [0; 4096],
sum_f_log_f: [0.0; WS_MAX + 1], 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.k6q.clear();
self.sum_f_log_f = [0.0; WS_MAX + 1]; self.sum_f_log_f = [0.0; WS_MAX + 1];
self.sum_f_log_s = [0.0; WS_MAX + 1];
} }
#[inline] #[inline]
fn update_sums_decrement<const K: usize>( fn update_sums_decrement<const K: usize>(sum_f_log_f: &mut [f64; WS_MAX + 1], f: usize) {
sum_f_log_f: &mut [f64; WS_MAX + 1],
sum_f_log_s: &mut [f64; WS_MAX + 1],
canonical: u64,
f: usize,
) {
sum_f_log_f[K] += n_log_n(f - 1) - n_log_n(f); sum_f_log_f[K] += n_log_n(f - 1) - n_log_n(f);
sum_f_log_s[K] -= ln_class_size::<false, K>(canonical);
} }
#[inline] #[inline]
fn update_sums_increment<const K: usize>( fn update_sums_increment<const K: usize>(sum_f_log_f: &mut [f64; WS_MAX + 1], g: usize) {
sum_f_log_f: &mut [f64; WS_MAX + 1],
sum_f_log_s: &mut [f64; WS_MAX + 1],
canonical: u64,
g: usize,
) {
sum_f_log_f[K] += n_log_n(g + 1) - n_log_n(g); sum_f_log_f[K] += n_log_n(g + 1) - n_log_n(g);
sum_f_log_s[K] += ln_class_size::<false, K>(canonical);
} }
/// Advance the window by one base. `received` is the caller's running /// 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); /// 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_kmer` is the current right-aligned, 2-bit-packed k-mer
/// rolling k-mer, from [`sub_word_canon`]. /// window (same convention as `obiskbuilder::RollingStat::rolling_k`).
pub fn push(&mut self, received: usize, canon: SubWordCanon) { pub fn push(&mut self, received: usize, rolling_kmer: u64) {
let [canonical_k1, canonical_k2, canonical_k3, canonical_k4, canonical_k5, canonical_k6] = let raw1 = rolling_kmer & 3;
canon; 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 { if received > self.k {
let old1 = self.k1q.pop_front(); let old1 = self.k1q.pop_front();
let f1 = self.k1c[old1 as usize] as usize; 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; self.k1c[old1 as usize] -= 1;
let old2 = self.k2q.pop_front(); let old2 = self.k2q.pop_front();
let f2 = self.k2c[old2 as usize] as usize; 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; self.k2c[old2 as usize] -= 1;
let old3 = self.k3q.pop_front(); let old3 = self.k3q.pop_front();
let f3 = self.k3c[old3 as usize] as usize; 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; self.k3c[old3 as usize] -= 1;
let old4 = self.k4q.pop_front(); let old4 = self.k4q.pop_front();
let f4 = self.k4c[old4 as usize] as usize; 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; self.k4c[old4 as usize] -= 1;
let old5 = self.k5q.pop_front(); let old5 = self.k5q.pop_front();
let f5 = self.k5c[old5 as usize] as usize; 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; self.k5c[old5 as usize] -= 1;
let old6 = self.k6q.pop_front(); let old6 = self.k6q.pop_front();
let f6 = self.k6c[old6 as usize] as usize; 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; self.k6c[old6 as usize] -= 1;
} }
if self.steady { if self.steady {
let g1 = self.k1c[canonical_k1 as usize] as usize; let g1 = self.k1c[raw1 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::update_sums_increment::<1>(&mut self.sum_f_log_f, g1);
self.k1c[canonical_k1 as usize] += 1; self.k1c[raw1 as usize] += 1;
self.k1q.push_back(canonical_k1); self.k1q.push_back(raw1);
let g2 = self.k2c[canonical_k2 as usize] as usize; let g2 = self.k2c[raw2 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::update_sums_increment::<2>(&mut self.sum_f_log_f, g2);
self.k2c[canonical_k2 as usize] += 1; self.k2c[raw2 as usize] += 1;
self.k2q.push_back(canonical_k2); self.k2q.push_back(raw2);
let g3 = self.k3c[canonical_k3 as usize] as usize; let g3 = self.k3c[raw3 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::update_sums_increment::<3>(&mut self.sum_f_log_f, g3);
self.k3c[canonical_k3 as usize] += 1; self.k3c[raw3 as usize] += 1;
self.k3q.push_back(canonical_k3); self.k3q.push_back(raw3);
let g4 = self.k4c[canonical_k4 as usize] as usize; let g4 = self.k4c[raw4 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::update_sums_increment::<4>(&mut self.sum_f_log_f, g4);
self.k4c[canonical_k4 as usize] += 1; self.k4c[raw4 as usize] += 1;
self.k4q.push_back(canonical_k4); self.k4q.push_back(raw4);
let g5 = self.k5c[canonical_k5 as usize] as usize; let g5 = self.k5c[raw5 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::update_sums_increment::<5>(&mut self.sum_f_log_f, g5);
self.k5c[canonical_k5 as usize] += 1; self.k5c[raw5 as usize] += 1;
self.k5q.push_back(canonical_k5); self.k5q.push_back(raw5);
let g6 = self.k6c[canonical_k6 as usize] as usize; let g6 = self.k6c[raw6 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::update_sums_increment::<6>(&mut self.sum_f_log_f, g6);
self.k6c[canonical_k6 as usize] += 1; self.k6c[raw6 as usize] += 1;
self.k6q.push_back(canonical_k6); self.k6q.push_back(raw6);
} else { } 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( fn push_warmup_increments(
&mut self, &mut self,
received: usize, received: usize,
canonical_k1: u64, canonical_k2: u64, canonical_k3: u64, raw1: u64, raw2: u64, raw3: u64,
canonical_k4: u64, canonical_k5: u64, canonical_k6: u64, raw4: u64, raw5: u64, raw6: u64,
) { ) {
let g1 = self.k1c[canonical_k1 as usize] as usize; let g1 = self.k1c[raw1 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::update_sums_increment::<1>(&mut self.sum_f_log_f, g1);
self.k1c[canonical_k1 as usize] += 1; self.k1c[raw1 as usize] += 1;
self.k1q.push_back(canonical_k1); self.k1q.push_back(raw1);
if received >= 2 { if received >= 2 {
let g2 = self.k2c[canonical_k2 as usize] as usize; let g2 = self.k2c[raw2 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::update_sums_increment::<2>(&mut self.sum_f_log_f, g2);
self.k2c[canonical_k2 as usize] += 1; self.k2c[raw2 as usize] += 1;
self.k2q.push_back(canonical_k2); self.k2q.push_back(raw2);
if received >= 3 { if received >= 3 {
let g3 = self.k3c[canonical_k3 as usize] as usize; let g3 = self.k3c[raw3 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::update_sums_increment::<3>(&mut self.sum_f_log_f, g3);
self.k3c[canonical_k3 as usize] += 1; self.k3c[raw3 as usize] += 1;
self.k3q.push_back(canonical_k3); self.k3q.push_back(raw3);
if received >= 4 { if received >= 4 {
let g4 = self.k4c[canonical_k4 as usize] as usize; let g4 = self.k4c[raw4 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::update_sums_increment::<4>(&mut self.sum_f_log_f, g4);
self.k4c[canonical_k4 as usize] += 1; self.k4c[raw4 as usize] += 1;
self.k4q.push_back(canonical_k4); self.k4q.push_back(raw4);
if received >= 5 { if received >= 5 {
let g5 = self.k5c[canonical_k5 as usize] as usize; let g5 = self.k5c[raw5 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::update_sums_increment::<5>(&mut self.sum_f_log_f, g5);
self.k5c[canonical_k5 as usize] += 1; self.k5c[raw5 as usize] += 1;
self.k5q.push_back(canonical_k5); self.k5q.push_back(raw5);
if received >= 6 { if received >= 6 {
let g6 = self.k6c[canonical_k6 as usize] as usize; let g6 = self.k6c[raw6 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::update_sums_increment::<6>(&mut self.sum_f_log_f, g6);
self.k6c[canonical_k6 as usize] += 1; self.k6c[raw6 as usize] += 1;
self.k6q.push_back(canonical_k6); self.k6q.push_back(raw6);
self.steady = true; self.steady = true;
} }
} }
@@ -265,7 +240,7 @@ impl EntropyTracker {
let nwords = k - order + 1; let nwords = k - order + 1;
let log_nw = log_nwords(k, order); let log_nw = log_nwords(k, order);
let nw_f = nwords as f64; 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) (h_corr / em).max(0.0)
} }
+1 -1
View File
@@ -1,6 +1,6 @@
[package] [package]
name = "obikmer" name = "obikmer"
version = "1.1.38" version = "1.1.39"
edition = "2024" edition = "2024"
[[bin]] [[bin]]
+2 -4
View File
@@ -1,4 +1,4 @@
use obikentropy::{EntropyTracker, sub_word_canon}; use obikentropy::EntropyTracker;
use obikseq::kmer::{Minimizer, hash_kmer}; use obikseq::kmer::{Minimizer, hash_kmer};
use obikseq::params; use obikseq::params;
@@ -136,8 +136,6 @@ impl RollingStat {
self.rolling_rck = self.rolling_rck =
((self.rolling_rck >> 2) | ((cnuc as u64) << ((k - 1) * 2))) & self.k_mask; ((self.rolling_rck >> 2) | ((cnuc as u64) << ((k - 1) * 2))) & self.k_mask;
let canon = sub_word_canon(self.rolling_k);
self.received += 1; self.received += 1;
if self.received >= m { 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 { pub fn ready(&self) -> bool {