gam-sae 0.3.155

//! Per-point sparse atom codes for multi-manifold reconstruction.
//!
//! This module owns the storage of per-observation soft assignments over a
//! library of `K` candidate manifold-atoms (see [`crate::assignment::SaeAssignment`]
//! for the surrounding selection/gate layer). The two key types are:
//!
//! * [`BitVec`] — a minimal dependency-free bitset used to record the *active
//!   support* `S_n ⊆ {0, …, K−1}` of each observation. We avoid pulling in
//!   the external `bitvec` crate to keep this module aligned with the rest of
//!   `gam`'s "no extra deps for new primitives" policy.
//! * [`SparseAtomCode`] — the per-point pair `(active_mask, weights)` whose
//!   semantics are documented on the type. Reconstruction at point `n` is
//!
//!   ```text
//!   Ẑ_n  =  Σ_{k ∈ S_n}  w_{n,k}  ·  decoder_k(t_{n,k})
//!   ```
//!
//!   so `weights[k]` is meaningful only when `active_mask.get(k) == true`.
//!   We store `weights` densely (`Vec<f64>` of length `K`) rather than
//!   sparsely; for the typical SAE workload `K` is small (tens to low
//!   hundreds), and the dense layout lets us reuse [`ndarray`] views and
//!   simple BLAS-shaped loops downstream. The mask carries the discrete
//!   active-set information; the weights carry the soft amplitudes.
//!
//! ## Per-point block locality (arrow structure)
//!
//! Each [`SparseAtomCode`] is the per-row ext-coordinate block for observation `n`
//! restricted to the `K` atoms. Combined with the per-atom on-manifold
//! coordinate `t_{n,k} ∈ ℝ^{d_k}` (held in the per-atom latent-coordinate
//! blocks of [`crate::assignment::SaeAssignment`]), the row-local
//! ext-coordinate vector is
//!
//! ```text
//!   ext_n  =  ( a_{n,1..K}  ;  t_{n,1,·}  ;  …  ;  t_{n,K,·} )
//! ```
//!
//! whose interaction graph with the shared decoder coefficients `B_1..B_K`
//! is exactly the arrow / bordered-Hessian pattern from `latent_coord.md`
//! §2.2. The Schur complement that Piece 1 uses to eliminate β before the
//! per-row solve generalises here with one change: the row-`n` block now
//! couples to *only the active subset* `S_n` of decoder borders, not to all
//! K of them. That is the structural fact this module records.

use std::collections::BTreeMap;

/// Minimal bit-vector. Backing storage is `Vec<u64>` words.
///
/// We expose only the operations the atom-selection layer needs: construction,
/// `get`, `set`, `count_ones`, and iteration of set indices. This is
/// deliberately tiny — adding the external `bitvec` crate would be overkill
/// for a few hundred bits per observation.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BitVec {
    words: Vec<u64>,
    len: usize,
}

impl BitVec {
    /// All-zero bitset of length `len`.
    pub fn zeros(len: usize) -> Self {
        let words = vec![0u64; len.div_ceil(64)];
        Self { words, len }
    }

    /// All-ones bitset of length `len`.
    pub fn ones(len: usize) -> Self {
        let mut bv = Self::zeros(len);
        for i in 0..len {
            bv.set(i, true);
        }
        bv
    }

    pub fn len(&self) -> usize {
        self.len
    }

    pub fn is_empty(&self) -> bool {
        self.len == 0
    }

    #[inline]
    pub fn get(&self, i: usize) -> bool {
        assert!(
            i < self.len,
            "BitVec::get index {i} out of bounds {}",
            self.len
        );
        let (w, b) = (i / 64, i % 64);
        (self.words[w] >> b) & 1 == 1
    }

    #[inline]
    pub fn set(&mut self, i: usize, v: bool) {
        assert!(
            i < self.len,
            "BitVec::set index {i} out of bounds {}",
            self.len
        );
        let (w, b) = (i / 64, i % 64);
        if v {
            self.words[w] |= 1u64 << b;
        } else {
            self.words[w] &= !(1u64 << b);
        }
    }

    /// Number of set bits.
    pub fn count_ones(&self) -> usize {
        self.words.iter().map(|w| w.count_ones() as usize).sum()
    }

    /// Iterator over set indices in ascending order.
    pub fn iter_ones(&self) -> impl Iterator<Item = usize> + '_ {
        (0..self.len).filter(move |&i| self.get(i))
    }

    /// Zero all bits in place.
    pub fn clear(&mut self) {
        for w in self.words.iter_mut() {
            *w = 0;
        }
    }
}

/// Per-point sparse code over `K` candidate atoms.
///
/// Invariants (checked in debug builds):
///
/// * `active_mask.len() == weights.len() == K`.
/// * For any `k` with `active_mask.get(k) == false`, the value `weights[k]`
///   is a nuisance — it must not influence reconstruction. Selection
///   strategies that lower a weight to zero (e.g. an L¹-relaxed gate after
///   thresholding) are responsible for clearing the
///   corresponding mask bit *and* zeroing `weights[k]`.
///
/// We do not require `weights[k] >= 0`; some strategies (entropic softmax,
/// TopK projection) keep the simplex, while others (L¹-relaxed) only enforce
/// non-negativity at the active-set step. The owning
/// gate/selection layer ([`crate::assignment::SaeAssignment`]) documents which
/// invariant it maintains.
#[derive(Debug, Clone)]
pub struct SparseAtomCode {
    /// Length-`K` bitmask of active atoms for this point.
    pub active_mask: BitVec,
    /// Length-`K` dense weight vector. Only entries at active indices are
    /// semantically meaningful.
    pub weights: Vec<f64>,
}

impl SparseAtomCode {
    /// Cold-start: no atoms active, all weights zero.
    pub fn empty(k_atoms: usize) -> Self {
        Self {
            active_mask: BitVec::zeros(k_atoms),
            weights: vec![0.0; k_atoms],
        }
    }

    /// Total number of candidate atoms `K` this code is sized for.
    pub fn k_atoms(&self) -> usize {
        self.weights.len()
    }

    /// Cardinality of the active support `|S_n|`.
    pub fn n_active(&self) -> usize {
        self.active_mask.count_ones()
    }

    /// Set the weight for atom `k` and mark it active.
    pub fn assign(&mut self, k: usize, w: f64) {
        assert!(k < self.k_atoms());
        self.active_mask.set(k, true);
        self.weights[k] = w;
    }

}

/// Storage for the per-row codes of all `N` observations.
///
/// Held column-of-structs rather than struct-of-columns: each row's
/// `(active_mask, weights)` lives together because the atom-selection
/// strategies all touch a single row at a time. Cross-row vectorization
/// happens through ndarray views built on demand.
#[derive(Debug, Clone)]
pub struct SparseAtomCodes {
    codes: Vec<SparseAtomCode>,
    k_atoms: usize,
}

impl SparseAtomCodes {
    /// Allocate `n_obs` empty codes, each sized for `k_atoms`.
    pub fn empty(n_obs: usize, k_atoms: usize) -> Self {
        let codes = (0..n_obs).map(|_| SparseAtomCode::empty(k_atoms)).collect();
        Self { codes, k_atoms }
    }

    pub fn n_obs(&self) -> usize {
        self.codes.len()
    }

    pub fn k_atoms(&self) -> usize {
        self.k_atoms
    }

    pub fn row(&self, n: usize) -> &SparseAtomCode {
        &self.codes[n]
    }

    pub fn row_mut(&mut self, n: usize) -> &mut SparseAtomCode {
        &mut self.codes[n]
    }

    pub fn iter(&self) -> impl Iterator<Item = &SparseAtomCode> {
        self.codes.iter()
    }

    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut SparseAtomCode> {
        self.codes.iter_mut()
    }

    /// All atom pairs that co-fire at least once, with their support and
    /// amplitude-code statistics, computed in one sparse pass over row supports.
    ///
    /// This is the structure-search candidate index: rows contribute only their
    /// active-set pairs, so the producer cost is `Σ_row |S_row|²` and the output
    /// is bounded by observed co-firings, not by `K²`.
    pub fn coactive_pair_stats(&self) -> Vec<(usize, usize, CoactivationStats)> {
        #[derive(Clone, Copy, Debug, Default)]
        struct PairAccum {
            n_joint: usize,
            sum_a: f64,
            sum_b: f64,
            sum_a2: f64,
            sum_b2: f64,
            sum_ab: f64,
        }

        let n_obs = self.n_obs();
        let mut marg = vec![0usize; self.k_atoms];
        let mut pairs: BTreeMap<(usize, usize), PairAccum> = BTreeMap::new();
        for code in &self.codes {
            let active: Vec<usize> = code.active_mask.iter_ones().collect();
            for &atom in &active {
                marg[atom] += 1;
            }
            for (idx, &u) in active.iter().enumerate() {
                for &v in &active[idx + 1..] {
                    let (a, b) = if u < v { (u, v) } else { (v, u) };
                    let wa = code.weights[a];
                    let wb = code.weights[b];
                    let acc = pairs.entry((a, b)).or_default();
                    acc.n_joint += 1;
                    acc.sum_a += wa;
                    acc.sum_b += wb;
                    acc.sum_a2 += wa * wa;
                    acc.sum_b2 += wb * wb;
                    acc.sum_ab += wa * wb;
                }
            }
        }

        pairs
            .into_iter()
            .map(|((a, b), acc)| {
                let weight_correlation = if acc.n_joint < 2 {
                    0.0
                } else {
                    let n = acc.n_joint as f64;
                    let cov = acc.sum_ab - acc.sum_a * acc.sum_b / n;
                    let var_a = acc.sum_a2 - acc.sum_a * acc.sum_a / n;
                    let var_b = acc.sum_b2 - acc.sum_b * acc.sum_b / n;
                    if var_a > 0.0 && var_b > 0.0 {
                        (cov / (var_a.sqrt() * var_b.sqrt())).clamp(-1.0, 1.0)
                    } else {
                        0.0
                    }
                };
                let stats = CoactivationStats::from_counts(
                    n_obs,
                    marg[a],
                    marg[b],
                    acc.n_joint,
                    weight_correlation,
                );
                (a, b, stats)
            })
            .collect()
    }

    /// Universal per-token code lengths for the binary support process, in bits.
    ///
    /// The MDL *selection* price names, per token, WHICH atoms fired. The uniform
    /// (combinatorial) price `log₂ C(G, k)` assumes every `k`-subset is equally
    /// likely — it is the WORST case, and it grossly OVERPAYS a dictionary whose
    /// co-firing supports are predictable (a tiling SAE where adjacent atoms fire
    /// together): charging that worst case would let an MDL comparison argue with
    /// itself. The honest currency is the entropy of the EMPIRICAL support
    /// distribution `H(S) = −Σ_s p(s) log₂ p(s)` over the observed supports
    /// `s ⊆ {0,…,G−1}`. `H(S)` cannot be read off directly (the support space is
    /// exponential and each token is one sample), so it is priced by the
    /// achievable code length of a LOW-ORDER model of the binary support process
    /// fit to the data — no magic constants, only empirical marginals and
    /// pairwise co-occurrence counts:
    ///
    /// Both learned models use Krichevsky–Trofimov sequential probabilities, so
    /// no fitted Bernoulli parameter is transmitted for free. The Chow–Liu code
    /// additionally transmits its data-selected labeled tree using Cayley's
    /// `G^(G−2)` possibilities before conditionally KT-coding every child. The
    /// combinatorial reference transmits each row's cardinality and then its
    /// subset, so variable support sizes are charged exactly rather than at a
    /// rounded mean.
    pub fn support_entropy(&self) -> SupportEntropy {
        let n = self.n_obs();
        let g = self.k_atoms();
        if n == 0 || g == 0 {
            return SupportEntropy {
                tree_bits: 0.0,
                independent_bits: 0.0,
                combinatorial_bits: 0.0,
                mean_support: 0.0,
            };
        }

        // Marginal firing counts and pairwise co-occurrence counts, in ONE
        // streaming pass over the active-support masks. Only the (small) active
        // set of each row contributes to the pairwise counts, so the cost is
        // `Σ_row |S_row|²` — cheap for the sparse SAE codes this consumes.
        let mut marg = vec![0.0_f64; g];
        let mut co = vec![0.0_f64; g * g]; // symmetric; upper triangle used
        let mut total_active = 0.0_f64;
        for code in &self.codes {
            let active: Vec<usize> = code.active_mask.iter_ones().collect();
            total_active += active.len() as f64;
            for (idx, &u) in active.iter().enumerate() {
                marg[u] += 1.0;
                for &v in &active[idx + 1..] {
                    co[u * g + v] += 1.0;
                }
            }
        }

        let nn = n as f64;
        let independent_total: f64 = (0..g)
            .map(|atom| kt_bernoulli_bits(self.codes.iter().map(|code| code.active_mask.get(atom))))
            .sum();

        // Maximum-weight (mutual-information) spanning tree by Prim over the dense
        // pairwise-MI graph. `best_mi[v]` is the largest MI connecting an
        // out-of-tree node `v` to the current tree; grow the tree by the
        // out-of-tree node of largest such MI, accumulating its edge weight.
        let mi = |u: usize, v: usize| -> f64 {
            let (a, b) = if u < v { (u, v) } else { (v, u) };
            mutual_information_bits(nn, marg[a], marg[b], co[a * g + b])
        };
        let mut in_tree = vec![false; g];
        let mut best_mi = vec![f64::NEG_INFINITY; g];
        let mut best_parent = vec![0usize; g];
        let mut parent = vec![0usize; g];
        in_tree[0] = true;
        for v in 1..g {
            best_mi[v] = mi(0, v);
            best_parent[v] = 0;
        }
        for _ in 1..g {
            // Pick the out-of-tree node joined to the tree by the strongest edge.
            let mut pick = usize::MAX;
            let mut pick_w = f64::NEG_INFINITY;
            for v in 0..g {
                if !in_tree[v] && best_mi[v] > pick_w {
                    pick_w = best_mi[v];
                    pick = v;
                }
            }
            if pick == usize::MAX {
                break;
            }
            in_tree[pick] = true;
            parent[pick] = best_parent[pick];
            for v in 0..g {
                if !in_tree[v] {
                    let w = mi(pick, v);
                    if w > best_mi[v] {
                        best_mi[v] = w;
                        best_parent[v] = pick;
                    }
                }
            }
        }

        let tree_structure_bits = if g <= 2 {
            0.0
        } else {
            (g as f64 - 2.0) * (g as f64).log2()
        };
        let mut tree_total = tree_structure_bits
            + kt_bernoulli_bits(self.codes.iter().map(|code| code.active_mask.get(0)));
        for child in 1..g {
            tree_total += kt_conditional_bernoulli_bits(self.codes.iter().map(|code| {
                (
                    code.active_mask.get(parent[child]),
                    code.active_mask.get(child),
                )
            }));
        }

        let mean_support = total_active / nn;
        let cardinality_bits = (g as f64 + 1.0).log2();
        let combinatorial_bits = cardinality_bits
            + self
                .codes
                .iter()
                .map(|code| log2_binom(g as i64, code.active_mask.count_ones() as i64))
                .sum::<f64>()
                / nn;

        SupportEntropy {
            tree_bits: tree_total / nn,
            independent_bits: independent_total / nn,
            combinatorial_bits,
            mean_support,
        }
    }
}

/// The empirical per-token support-entropy estimate produced by
/// [`SparseAtomCodes::support_entropy`], with the two references (independent
/// model and combinatorial worst case) it is read against. All fields are BITS
/// per token except [`Self::mean_support`]. See the method for the derivation.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct SupportEntropy {
    /// Complete Chow–Liu universal code per token, including tree transmission
    /// and conditional KT parameter learning.
    pub tree_bits: f64,
    /// Independent-Bernoulli KT universal code per token.
    pub independent_bits: f64,
    /// `log₂(G+1) + mean_i log₂ C(G, |S_i|)`: a valid variable-cardinality
    /// fixed code, unlike `log₂ C(G, round(mean |S|))`.
    pub combinatorial_bits: f64,
    /// Mean support size `k̄` (mean number of active atoms per token).
    pub mean_support: f64,
}

/// Sequential Krichevsky–Trofimov code for a binary sequence. The Jeffreys
/// half-counts make the code universal and give every one-sample outcome one
/// bit instead of the zero-bit plug-in pathology.
fn kt_bernoulli_bits(values: impl IntoIterator<Item = bool>) -> f64 {
    let mut counts = [0.5_f64, 0.5_f64];
    let mut bits = 0.0;
    for value in values {
        let index = usize::from(value);
        let probability = counts[index] / (counts[0] + counts[1]);
        bits -= probability.log2();
        counts[index] += 1.0;
    }
    bits
}

/// Conditional KT code with a separate Bernoulli predictor in each binary
/// parent context.
fn kt_conditional_bernoulli_bits(values: impl IntoIterator<Item = (bool, bool)>) -> f64 {
    let mut counts = [[0.5_f64, 0.5_f64], [0.5_f64, 0.5_f64]];
    let mut bits = 0.0;
    for (context, value) in values {
        let context_index = usize::from(context);
        let value_index = usize::from(value);
        let probability = counts[context_index][value_index]
            / (counts[context_index][0] + counts[context_index][1]);
        bits -= probability.log2();
        counts[context_index][value_index] += 1.0;
    }
    bits
}

/// Pairwise mutual information `I(x_u; x_v)` in bits of two binary indicators,
/// from the `2×2` empirical joint implied by the counts `n` (rows), `n_u`, `n_v`
/// (marginals), and `n_uv` (joint). Non-negative up to floating-point noise; the
/// `0 · log 0` cells are dropped.
fn mutual_information_bits(n: f64, n_u: f64, n_v: f64, n_uv: f64) -> f64 {
    if n <= 0.0 {
        return 0.0;
    }
    let p1x = n_u / n;
    let px1 = n_v / n;
    let p11 = n_uv / n;
    let p10 = (p1x - p11).max(0.0);
    let p01 = (px1 - p11).max(0.0);
    let p00 = (1.0 - p11 - p10 - p01).max(0.0);
    let cell = |p: f64, pa: f64, pb: f64| -> f64 {
        if p > 0.0 && pa > 0.0 && pb > 0.0 {
            p * (p / (pa * pb)).log2()
        } else {
            0.0
        }
    };
    let mi = cell(p11, p1x, px1)
        + cell(p10, p1x, 1.0 - px1)
        + cell(p01, 1.0 - p1x, px1)
        + cell(p00, 1.0 - p1x, 1.0 - px1);
    mi.max(0.0)
}

/// `log₂ C(g, k)`: bits to name which `k` of `g` atoms fired under the uniform
/// support prior. Computed as `Σ_{i=1..k} log₂((g−k+i)/i)` so it never overflows
/// a binomial. Zero when `g ≤ 0` or `k ≤ 0`; `k` is capped at `g`. (The same
/// combinatorial bound [`crate::description_length::selection_bits`] reports; kept
/// local so the support-entropy estimator stays self-contained.)
fn log2_binom(g: i64, k: i64) -> f64 {
    if g <= 0 || k <= 0 {
        return 0.0;
    }
    let k = k.min(g);
    let mut bits = 0.0;
    for i in 1..=k {
        bits += ((g - k + i) as f64 / i as f64).log2();
    }
    bits
}

/// Pairwise co-activation summary for two atoms (see
/// `SparseAtomCodes::coactivation`). All probabilities are empirical
/// popcount ratios over the active-support masks.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct CoactivationStats {
    /// Total number of observations the codes cover.
    pub n_obs: usize,
    /// Rows where atom `a` is active.
    pub n_a: usize,
    /// Rows where atom `b` is active.
    pub n_b: usize,
    /// Rows where both are active.
    pub n_joint: usize,
    /// `P(a active | b active)`; `0` when `b` is never active.
    pub p_a_given_b: f64,
    /// `P(b active | a active)`; `0` when `a` is never active.
    pub p_b_given_a: f64,
    /// `P(a∧b) / (P(a)·P(b))`; `1` for independent atoms, `0` when either
    /// marginal is empty.
    pub lift: f64,
    /// Pearson correlation of the two atoms' activation WEIGHTS over the
    /// jointly-active rows (see `SparseAtomCodes::weight_codependence`) — the
    /// amplitude/interaction half of the fusion criterion. `0` when the joint
    /// support is too small or a weight is constant there.
    pub weight_correlation: f64,
}

impl CoactivationStats {
    fn from_counts(
        n_obs: usize,
        n_a: usize,
        n_b: usize,
        n_joint: usize,
        weight_correlation: f64,
    ) -> Self {
        let cond = |joint: usize, marg: usize| {
            if marg == 0 {
                0.0
            } else {
                joint as f64 / marg as f64
            }
        };
        let lift = if n_a == 0 || n_b == 0 || n_obs == 0 {
            0.0
        } else {
            (n_joint as f64 * n_obs as f64) / (n_a as f64 * n_b as f64)
        };
        Self {
            n_obs,
            n_a,
            n_b,
            n_joint,
            p_a_given_b: cond(n_joint, n_b),
            p_b_given_a: cond(n_joint, n_a),
            lift,
            weight_correlation,
        }
    }

    /// Symmetric code dependence `min(P(a|b), P(b|a))` — the canonical-order
    /// trigger for FUSION proposals (descending). Near 0 for independent or
    /// disjoint atoms; near 1 only when the two supports essentially coincide,
    /// which is the shattering signature.
    pub fn dependence(&self) -> f64 {
        self.p_a_given_b.min(self.p_b_given_a)
    }

    /// Conditional asymmetry `|P(a|b) − P(b|a)|` — large when one atom's
    /// support nests inside the other's (the A⇒B absorption signature, where
    /// `P(parent|child) ≈ 1` but not conversely). Flags the pair for a
    /// targeted within-atom substructure audit; it is never itself an
    /// acceptance criterion.
    pub fn absorption_asymmetry(&self) -> f64 {
        (self.p_a_given_b - self.p_b_given_a).abs()
    }

}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn bitvec_basic() {
        let mut bv = BitVec::zeros(70);
        assert_eq!(bv.len(), 70);
        assert!(!bv.get(5));
        bv.set(5, true);
        bv.set(64, true);
        assert!(bv.get(5));
        assert!(bv.get(64));
        assert_eq!(bv.count_ones(), 2);
        let ones: Vec<usize> = bv.iter_ones().collect();
        assert_eq!(ones, vec![5, 64]);
        bv.set(5, false);
        assert_eq!(bv.count_ones(), 1);
    }

}