obikmer/src/obikentropy/src/tracker.rs

//! Incremental (streaming) normalized k-mer entropy.
//!
//! [`EntropyTracker`] maintains, over a sliding window of the last `k` bases,
//! 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
//! `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, 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
/// 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` raw sub-words, one per word
    // size — stack-allocated, capacity ≤ k ≤ 31.
    k1q: Ring<u64, 32>,
    k2q: Ring<u64, 32>,
    k3q: Ring<u64, 32>,
    k4q: Ring<u64, 32>,
    k5q: Ring<u64, 32>,
    k6q: Ring<u64, 32>,

    // 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],
    k4c: [u8; 256],
    k5c: [u8; 1024],
    k6c: [u8; 4096],

    sum_f_log_f: [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],
        }
    }

    /// 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];
    }

    #[inline]
    fn update_sums_decrement<const K: usize>(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);
    }

    #[inline]
    fn update_sums_increment<const K: usize>(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);
    }

    /// 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);
    /// `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, 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, 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, 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, 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, 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, f6);
            self.k6c[old6 as usize] -= 1;
        }

        if self.steady {
            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[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[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[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[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[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, raw1, raw2, raw3, raw4, raw5, raw6);
        }
    }

    #[cold]
    #[inline(never)]
    fn push_warmup_increments(
        &mut self,
        received: usize,
        raw1: u64, raw2: u64, raw3: u64,
        raw4: u64, raw5: u64, raw6: u64,
    ) {
        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[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[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[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[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[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;
                        }
                    }
                }
            }
        }
    }

    /// 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_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 }
    }
}