legume-numeric 0.8.11

#![allow(dead_code)]

extern crate special;

use crate::param::io::*;
use crate::param::traits::*;
use nalgebra::{DMatrix, DVector};
use rayon::prelude::*;

#[derive(Debug, Clone)]
pub struct GammaMatrix {
    num_rows: usize,
    num_columns: usize,
    //////////////////////
    // hyper parameters //
    //////////////////////
    a0: f32,
    b0: f32,
    /// Per-row prior `(a0, b0)`, overriding the scalar pair when set.
    /// Lengths equal `num_rows`. See [`Self::with_row_prior`].
    row_prior: Option<(DVector<f32>, DVector<f32>)>,
    ///////////////////////////
    // sufficient statistics //
    ///////////////////////////
    a_stat: DMatrix<f32>,
    b_stat: DMatrix<f32>,
    //////////////////////////
    // estimated parameters //
    //////////////////////////
    estimated_mean: DMatrix<f32>,
    estimated_sd: DMatrix<f32>,
    estimated_log_mean: DMatrix<f32>,
    estimated_log_sd: DMatrix<f32>,
}

impl ParamIo for GammaMatrix {
    type Mat = DMatrix<f32>;
}

impl TwoStatParam for GammaMatrix {
    type Mat = DMatrix<f32>;
    type Scalar = f32;

    fn new(dims: (usize, usize), a: Self::Scalar, b: Self::Scalar) -> Self {
        Self {
            num_rows: dims.0,
            num_columns: dims.1,
            a0: a,
            b0: b,
            row_prior: None,
            a_stat: DMatrix::from_element(dims.0, dims.1, a),
            b_stat: DMatrix::from_element(dims.0, dims.1, b),
            // `estimated_mean` is eager: the coordinate descent reads
            // `posterior_mean()` before the first calibration (relying on a
            // zero start), and it gets allocated immediately anyway.
            estimated_mean: DMatrix::zeros(dims.0, dims.1),
            // The sd / log_mean / log_sd planes are lazily allocated by
            // `map_calibrate_*` (via `calibrate_with`). An iterative fit that
            // only reads `posterior_mean()` (calibrating `MeanOnly`) never
            // pays for them — they're materialized only when output needs
            // them (a calibrate with `All` / `MeanAndLogMean`).
            estimated_sd: DMatrix::zeros(0, 0),
            estimated_log_mean: DMatrix::zeros(0, 0),
            estimated_log_sd: DMatrix::zeros(0, 0),
        }
    }

    fn add_stat(&mut self, add_a: &Self::Mat, add_b: &Self::Mat) {
        self.a_stat += add_a;
        self.b_stat += add_b;
    }
    fn update_stat(&mut self, update_a: &Self::Mat, update_b: &Self::Mat) {
        self.reset_stat();
        self.add_stat(update_a, update_b);
    }
    fn reset_stat(&mut self) {
        match &self.row_prior {
            None => {
                self.a_stat.fill(self.a0);
                self.b_stat.fill(self.b0);
            }
            Some((a0, b0)) => {
                // Column-wise copies are contiguous in column-major storage and
                // are a no-op on a released (0 x 0) plane.
                for mut col in self.a_stat.column_iter_mut() {
                    col.copy_from(a0);
                }
                for mut col in self.b_stat.column_iter_mut() {
                    col.copy_from(b0);
                }
            }
        }
    }
    fn update_stat_col(&mut self, update_a: &Self::Mat, update_b: &Self::Mat, k: usize) {
        match &self.row_prior {
            None => {
                self.a_stat
                    .column_mut(k)
                    .copy_from(&update_a.map(|x| x + self.a0));
                self.b_stat
                    .column_mut(k)
                    .copy_from(&update_b.map(|x| x + self.b0));
            }
            Some((a0, b0)) => {
                let mut a = self.a_stat.column_mut(k);
                a.copy_from(update_a);
                a += a0;
                let mut b = self.b_stat.column_mut(k);
                b.copy_from(update_b);
                b += b0;
            }
        }
    }

    // fn nrows(&self) -> usize {
    //     self.num_rows
    // }

    // fn ncols(&self) -> usize {
    //     self.num_columns
    // }

    // fn len(&self) -> usize {
    //     self.num_rows * self.num_columns
    // }
    fn map_calibrate_mean(&mut self) {
        self.estimated_mean = self.a_stat.zip_map(&self.b_stat, |a, b| a / b);
    }
    fn map_calibrate_sd(&mut self) {
        self.estimated_sd = self.a_stat.zip_map(&self.b_stat, |a, b| a.sqrt() / b);
    }
    fn map_calibrate_log_mean(&mut self) {
        use special::Gamma;
        self.estimated_log_mean = self
            .a_stat
            .zip_map(&self.b_stat, |a, b| a.digamma() - b.ln());
    }
    fn map_calibrate_log_sd(&mut self) {
        // `sd[ln X] = sqrt(trigamma(a))` exactly, for `X ~ Gamma(a, b)` — note it
        // does not depend on the rate, which is why `b_stat` plays no part here.
        //
        // This replaced `1/sqrt(a - 1)`, the large-`a` asymptote, which was
        // wrong in the regime that dominates sparse count data: 46% high at
        // `a = 1.5`, 21% high at `a = 2.2` (a typical detected feature), and
        // agreeing only past `a ~ 100`. Below `a = 1` it has no real value at
        // all, and the old code returned 0 there — i.e. it reported PERFECT
        // certainty for a feature with no counts, whose posterior is the prior
        // and whose true `sd` is the largest in the matrix (1.283 at `a = 1`).
        // Anything reading `log_sd` as a precision was being handed the
        // inversion of the truth.
        use special::Gamma;
        self.estimated_log_sd = self.a_stat.map(|a| a.trigamma().sqrt());
    }
}

impl Inference for GammaMatrix {
    type Mat = DMatrix<f32>;
    type Scalar = f32;

    fn posterior_mean(&self) -> &Self::Mat {
        &self.estimated_mean
    }

    fn posterior_sd(&self) -> &Self::Mat {
        &self.estimated_sd
    }

    fn posterior_log_mean(&self) -> &Self::Mat {
        &self.estimated_log_mean
    }

    fn posterior_log_sd(&self) -> &Self::Mat {
        &self.estimated_log_sd
    }

    fn posterior_sample(&self) -> anyhow::Result<Self::Mat> {
        use rand_distr::{Distribution, Gamma};
        let eps = 1e-8;

        let sampled = self
            .a_stat
            .as_slice()
            .par_iter()
            .zip(self.b_stat.as_slice().par_iter())
            .map_init(rand::rng, |rng, (&a, &b)| -> anyhow::Result<f32> {
                let shape = a + eps;
                let scale = (b + eps).recip();
                let pdf = Gamma::new(shape, scale)?;
                Ok(pdf.sample(rng))
            })
            .collect::<anyhow::Result<Vec<_>>>()?;

        Ok(Self::Mat::from_vec(self.nrows(), self.ncols(), sampled))
    }

    fn posterior_log_sample(&self, seed: u64) -> anyhow::Result<Self::Mat> {
        use rand::rngs::SmallRng;
        use rand::SeedableRng;
        use rand_distr::{Distribution, StandardNormal};

        // Fixed chunk width, so which elements share an RNG is a property of
        // the data shape and not of how rayon happened to split the work.
        const CHUNK: usize = 1024;
        let m_slice = self.estimated_log_mean.as_slice();
        let s_slice = self.estimated_log_sd.as_slice();
        let mut sampled = vec![0.0f32; m_slice.len()];
        sampled
            .par_chunks_mut(CHUNK)
            .enumerate()
            .for_each(|(ci, out)| {
                let mut rng =
                    SmallRng::seed_from_u64(seed ^ (ci as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15));
                let base = ci * CHUNK;
                for (k, o) in out.iter_mut().enumerate() {
                    let z: f32 = StandardNormal.sample(&mut rng);
                    *o = m_slice[base + k] + s_slice[base + k] * z;
                }
            });

        Ok(Self::Mat::from_vec(self.nrows(), self.ncols(), sampled))
    }

    fn nrows(&self) -> usize {
        self.num_rows
    }

    fn ncols(&self) -> usize {
        self.num_columns
    }
}

/// Row-stack one plane across blocks. Returns an empty matrix (and skips
/// the work) when `enabled` is false or the plane is lazily-unallocated in
/// the first block, so empty/dropped planes never get materialized.
fn stack_field<F>(
    blocks: &[GammaMatrix],
    nrows: usize,
    ncols: usize,
    enabled: bool,
    sel: F,
) -> DMatrix<f32>
where
    F: Fn(&GammaMatrix) -> &DMatrix<f32>,
{
    if !enabled || blocks.is_empty() || sel(&blocks[0]).nrows() == 0 {
        return DMatrix::zeros(0, 0);
    }
    let mut out = DMatrix::zeros(nrows, ncols);
    let mut r0 = 0;
    for b in blocks {
        let src = sel(b);
        out.rows_mut(r0, src.nrows()).copy_from(src);
        r0 += src.nrows();
    }
    out
}

impl GammaMatrix {
    /// Drop the sufficient-stat planes (`a_stat` / `b_stat`) after
    /// calibration, keeping only the posterior estimates. Use when the
    /// consumer reads posterior means / log-means but never
    /// `posterior_sample` (which is the only reader of `a_stat`/`b_stat`).
    /// Halves the resident footprint of a calibrated parameter.
    pub fn release_stats(&mut self) {
        self.a_stat = DMatrix::zeros(0, 0);
        self.b_stat = DMatrix::zeros(0, 0);
    }

    /// Zero every `estimated_mean` entry whose corresponding `numerator` is
    /// zero, collapsing the per-column Gamma prior baseline (`a0/denom`,
    /// present at *every* unobserved cell) to exact zero. This lets a
    /// downstream triplet-ization of the mean be **sparse** — only the
    /// observed support survives. It's the lossy-but-correct choice for
    /// count-based consumers (the baseline is a regularization floor, not
    /// signal). `numerator` must match the mean's shape; only meaningful
    /// after a mean calibration.
    pub fn sparsify_mean_to_support(&mut self, numerator: &DMatrix<f32>) {
        debug_assert_eq!(self.estimated_mean.shape(), numerator.shape());
        self.estimated_mean
            .iter_mut()
            .zip(numerator.iter())
            .for_each(|(m, &n)| {
                if n == 0.0 {
                    *m = 0.0;
                }
            });
    }

    /// Whether the posterior at `(row, col)` carries any data beyond the
    /// prior: `a_stat > a0`. The read-only counterpart of
    /// [`Self::sparsify_mean_to_support`], for consumers that serialize the
    /// mean without owning the numerator — an unsupported entry's mean is the
    /// prior floor `a0 / (b0 + denom)`, which is regularization, not signal.
    /// Writing those floors out turns a sparse posterior dense: a carried
    /// pseudobulk reference measured 100.0% dense (34M of 34M entries) before
    /// its writer checked this.
    #[must_use]
    pub fn has_data_support(&self, row: usize, col: usize) -> bool {
        self.a_stat[(row, col)] > self.a0_at(row)
    }

    /// The unregularized rate at `(row, col)`: `(a_stat − a0) / (b_stat − b0)`
    /// — data sum over data denominator, no prior in either. Zero when the
    /// entry has no data support.
    ///
    /// This is what a *serialized* posterior should usually store: paired with
    /// its denominator, it is a bijection of the sufficient statistics, so a
    /// consumer reconstructs `a_stat`/`b_stat` exactly. The posterior mean
    /// `(a0 + sum)/(b0 + n)` is the right *estimate* but the wrong *carrier* —
    /// its prior shrinkage (1.85× at `sum = 1, n = 12`) gets re-ingested as if
    /// it were data, and a second posterior forms around an already-shrunk
    /// value.
    #[must_use]
    pub fn evidence_mean(&self, row: usize, col: usize) -> f32 {
        let a = self.a_stat[(row, col)] - self.a0_at(row);
        let b = self.b_stat[(row, col)] - self.b0_at(row);
        if a > 0.0 && b > 0.0 {
            a / b
        } else {
            0.0
        }
    }

    /// Prior shape for `row`: its row prior when set, else the scalar `a0`.
    #[inline]
    fn a0_at(&self, row: usize) -> f32 {
        self.row_prior.as_ref().map_or(self.a0, |(a, _)| a[row])
    }

    /// Prior rate for `row`: its row prior when set, else the scalar `b0`.
    #[inline]
    fn b0_at(&self, row: usize) -> f32 {
        self.row_prior.as_ref().map_or(self.b0, |(_, b)| b[row])
    }

    /// A matrix whose row `d` has prior `Gamma(a0[d], b0[d])`. The statistics
    /// start at the prior, as with [`TwoStatParam::new`].
    pub fn with_row_prior(dims: (usize, usize), a0: &DVector<f32>, b0: &DVector<f32>) -> Self {
        let mut out = Self::new(dims, 0.0, 0.0);
        out.set_row_prior(a0, b0);
        out.reset_stat();
        out
    }

    /// Replace the per-row prior. Accumulated statistics are left as they
    /// are; the new prior applies at the next `reset_stat` / `update_stat`.
    pub fn set_row_prior(&mut self, a0: &DVector<f32>, b0: &DVector<f32>) {
        assert_eq!(a0.len(), self.num_rows, "row prior a0 length != rows");
        assert_eq!(b0.len(), self.num_rows, "row prior b0 length != rows");
        self.row_prior = Some((a0.clone(), b0.clone()));
    }

    /// The per-row prior, if one is set.
    #[must_use]
    pub fn row_prior(&self) -> Option<(&DVector<f32>, &DVector<f32>)> {
        self.row_prior.as_ref().map(|(a, b)| (a, b))
    }

    /// Row-stack per-feature-block parameters (from a gene-blocked fit)
    /// into one `[Σrowsᵢ × K]` parameter. All blocks must share the column
    /// count and either share the scalar hyper-params or all carry a row
    /// prior, in which case the row priors are concatenated. Calibrated
    /// planes present in the first block
    /// are stacked; lazily-empty planes stay empty. `stack_stats` controls
    /// whether `a_stat`/`b_stat` are carried through — pass `false` when the
    /// output only needs posterior estimates, so the heavy sufficient-stat
    /// planes are never assembled at full width.
    pub fn vconcat(blocks: Vec<GammaMatrix>, stack_stats: bool) -> Self {
        assert!(!blocks.is_empty(), "vconcat of empty block list");
        let ncols = blocks[0].num_columns;
        let a0 = blocks[0].a0;
        let b0 = blocks[0].b0;
        let nrows: usize = blocks.iter().map(|b| b.num_rows).sum();
        let row_prior = if blocks[0].row_prior.is_some() {
            assert!(
                blocks.iter().all(|b| b.row_prior.is_some()),
                "vconcat: blocks mix scalar and row priors"
            );
            let a = DVector::from_iterator(
                nrows,
                blocks
                    .iter()
                    .flat_map(|b| b.row_prior.as_ref().expect("checked").0.iter().copied()),
            );
            let b = DVector::from_iterator(
                nrows,
                blocks
                    .iter()
                    .flat_map(|b| b.row_prior.as_ref().expect("checked").1.iter().copied()),
            );
            Some((a, b))
        } else {
            assert!(
                blocks
                    .iter()
                    .all(|b| b.row_prior.is_none() && b.a0 == a0 && b.b0 == b0),
                "vconcat: blocks must share hyper-params"
            );
            None
        };
        let a_stat = stack_field(&blocks, nrows, ncols, stack_stats, |g| &g.a_stat);
        let b_stat = stack_field(&blocks, nrows, ncols, stack_stats, |g| &g.b_stat);
        let estimated_mean = stack_field(&blocks, nrows, ncols, true, |g| &g.estimated_mean);
        let estimated_sd = stack_field(&blocks, nrows, ncols, true, |g| &g.estimated_sd);
        let estimated_log_mean =
            stack_field(&blocks, nrows, ncols, true, |g| &g.estimated_log_mean);
        let estimated_log_sd = stack_field(&blocks, nrows, ncols, true, |g| &g.estimated_log_sd);
        Self {
            num_rows: nrows,
            num_columns: ncols,
            a0,
            b0,
            row_prior,
            a_stat,
            b_stat,
            estimated_mean,
            estimated_sd,
            estimated_log_mean,
            estimated_log_sd,
        }
    }
}