data-beans 0.11.121

//! Per-gene negative binomial marginal fits + PIT helpers for the
//! scDesign-style copula simulator.
//!
//! NB parametrization is `(size = r, mean = μ)` with PMF
//! `P(X=k) = C(k+r-1, k) p^r (1-p)^k`, `p = r / (r + μ)`. The log PMF is
//! computed iteratively via
//! `ln P(k+1) = ln P(k) + ln((k+r)/(k+1)) + ln(μ/(r+μ))`,
//! which avoids `lgamma` calls and stays stable for the count ranges seen
//! in single-cell data. Reference: Sun, Song, Li & Li 2021, *scDesign2*,
//! Genome Biology 22:163.

#[derive(Debug, Clone, Copy)]
pub struct NbFit {
    /// NB mean μ. For the Poisson fallback this is the rate λ.
    pub mu: f32,
    /// NB size r. `f32::INFINITY` signals the Poisson collapse (σ² ≤ μ).
    pub r: f32,
}

impl NbFit {
    pub fn is_poisson(&self) -> bool {
        !self.r.is_finite()
    }
}

/// Build the cumulative distribution table `F(0), F(1), ..., F(k_max)` by
/// iterating the log-PMF recurrence and exponentiating. Returns a vector of
/// length `k_max + 1`.
pub fn nb_cdf_table(fit: NbFit, k_max: usize) -> Vec<f64> {
    let mu = fit.mu as f64;
    let r = fit.r as f64;
    let mut cdf = Vec::with_capacity(k_max + 1);
    if mu <= 0.0 {
        cdf.resize(k_max + 1, 1.0);
        return cdf;
    }
    if !r.is_finite() {
        // Poisson(μ): ln P(0) = -μ; ln P(k+1) = ln P(k) + ln(μ) - ln(k+1).
        let mut ln_pk = -mu;
        let mut acc = ln_pk.exp();
        cdf.push(acc.min(1.0));
        for k in 0..k_max {
            ln_pk += mu.ln() - ((k + 1) as f64).ln();
            acc += ln_pk.exp();
            cdf.push(acc.min(1.0));
        }
        return cdf;
    }
    let p = r / (r + mu);
    let ln_p = p.ln();
    let ln_q = (1.0 - p).ln();
    let mut ln_pk = r * ln_p;
    let mut acc = ln_pk.exp();
    cdf.push(acc.min(1.0));
    for k in 0..k_max {
        let kf = k as f64;
        ln_pk += ((kf + r) / (kf + 1.0)).ln() + ln_q;
        acc += ln_pk.exp();
        cdf.push(acc.min(1.0));
    }
    cdf
}

/// PIT with continuity correction: `u = 0.5·F(x-1) + 0.5·F(x)`. Clamped to
/// `(EPS, 1-EPS)` so the subsequent `Φ⁻¹` is finite.
pub fn pit_continuity(table: &[f64], x: u32) -> f64 {
    const EPS: f64 = 1e-7;
    let xi = x as usize;
    let cdf_hi = if xi >= table.len() {
        *table.last().unwrap_or(&1.0)
    } else {
        table[xi]
    };
    let cdf_lo = if xi == 0 {
        0.0
    } else if xi > table.len() {
        *table.last().unwrap_or(&1.0)
    } else {
        table[xi - 1]
    };
    let u = 0.5 * (cdf_lo + cdf_hi);
    u.clamp(EPS, 1.0 - EPS)
}

/// Generous cap on the PMF table size: `μ + 20·σ + 100` past the effective
/// tail for any realistic single-cell mean. Reflects only the NB / Poisson
/// tail; zero-inflation `π` doesn't extend the support upward (it just
/// dumps mass at 0), so it's not part of the cap.
pub fn nb_table_cap(fit: NbFit) -> usize {
    let mu = fit.mu as f64;
    let r = fit.r as f64;
    let var = if r.is_finite() { mu + mu * mu / r } else { mu };
    ((mu + 20.0 * var.sqrt()).ceil() as usize).saturating_add(100)
}

/// `min { k : F(k) ≥ u }` from a pre-built CDF table. Pair with one
/// `nb_cdf_table` per fit when many quantiles share that fit; for a fit that
/// changes on every draw use [`nb_inverse_cdf`].
pub fn nb_inverse_cdf_from_table(u: f64, table: &[f64]) -> u32 {
    if u <= 0.0 {
        return 0;
    }
    for (k, &f) in table.iter().enumerate() {
        if f >= u {
            return k as u32;
        }
    }
    table.len() as u32
}

/// `min { k : F(k) ≥ u }` for one `(fit, u)` pair, capped at `k_max + 1`.
///
/// Walks the same log-PMF recurrence as [`nb_cdf_table`] and stops at the
/// first `k` whose cumulative mass reaches `u`, so it returns exactly what
/// `nb_inverse_cdf_from_table(u, &nb_cdf_table(fit, k_max))` does without
/// building the table. Use it when the fit changes on every draw (a per-cell
/// mean); keep the table when many `u` share one fit.
pub fn nb_inverse_cdf(u: f64, fit: NbFit, k_max: usize) -> u32 {
    let mu = fit.mu as f64;
    let r = fit.r as f64;
    if u <= 0.0 || mu <= 0.0 {
        return 0;
    }
    // Per-step log-PMF increment: ln P(k+1) − ln P(k).
    let poisson = !r.is_finite();
    let (mut ln_pk, ln_step) = if poisson {
        (-mu, mu.ln())
    } else {
        let p = r / (r + mu);
        (r * p.ln(), (1.0 - p).ln())
    };
    let mut acc = ln_pk.exp();
    if acc.min(1.0) >= u {
        return 0;
    }
    for k in 0..k_max {
        let kf = k as f64;
        ln_pk += if poisson {
            ln_step - (kf + 1.0).ln()
        } else {
            ((kf + r) / (kf + 1.0)).ln() + ln_step
        };
        acc += ln_pk.exp();
        if acc.min(1.0) >= u {
            return (k + 1) as u32;
        }
    }
    (k_max + 1) as u32
}

/// Standard normal CDF Φ; thin wrapper over `statrs::distribution::Normal`.
pub fn phi(z: f64) -> f64 {
    use statrs::distribution::{ContinuousCDF, Normal};
    Normal::standard().cdf(z)
}

/// Standard normal inverse CDF Φ⁻¹; thin wrapper over `statrs::distribution::Normal`.
pub fn inv_phi(u: f64) -> f64 {
    use statrs::distribution::{ContinuousCDF, Normal};
    Normal::standard().inverse_cdf(u)
}

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

    #[test]
    fn cdf_monotone_and_unit() {
        let fit = NbFit { mu: 5.0, r: 2.0 };
        let table = nb_cdf_table(fit, 200);
        for w in table.windows(2) {
            assert!(w[1] >= w[0] - 1e-12);
        }
        assert!((table.last().unwrap() - 1.0).abs() < 1e-3);
    }

    /// The table-free inverse CDF agrees with table + scan for NB and
    /// Poisson fits, across the quantile range and when the cap truncates.
    #[test]
    fn inverse_cdf_matches_table_lookup() {
        let fits = [
            NbFit { mu: 0.0, r: 2.0 },
            NbFit { mu: 1e-4, r: 0.01 },
            NbFit { mu: 0.3, r: 0.5 },
            NbFit { mu: 5.0, r: 2.0 },
            NbFit { mu: 40.0, r: 0.01 },
            NbFit { mu: 250.0, r: 8.0 },
            NbFit {
                mu: 0.7,
                r: f32::INFINITY,
            },
            NbFit {
                mu: 60.0,
                r: f32::INFINITY,
            },
        ];
        let us = [
            0.0,
            1e-7,
            0.01,
            0.2,
            0.5,
            0.731,
            0.9,
            0.99,
            0.9999,
            1.0 - 1e-7,
        ];
        for fit in fits {
            for k_max in [0usize, 3, nb_table_cap(fit)] {
                let table = nb_cdf_table(fit, k_max);
                for u in us {
                    assert_eq!(
                        nb_inverse_cdf(u, fit, k_max),
                        nb_inverse_cdf_from_table(u, &table),
                        "fit {fit:?}, k_max {k_max}, u {u}"
                    );
                }
            }
        }
    }

    #[test]
    fn phi_inv_phi_round_trip() {
        for &z in &[-3.0, -1.5, -0.5, 0.0, 0.5, 1.5, 3.0] {
            let u = phi(z);
            let z_back = inv_phi(u);
            assert!((z_back - z).abs() < 1e-9, "z={} back={}", z, z_back);
        }
    }

    #[test]
    fn phi_inv_phi_extremes() {
        // Tail behavior shouldn't blow up.
        for &u in &[1e-6, 1e-3, 0.5, 1.0 - 1e-3, 1.0 - 1e-6] {
            let z = inv_phi(u);
            assert!(z.is_finite(), "inv_phi({})={} not finite", u, z);
            let _ = rand_distr::Normal::new(0.0, 1.0)
                .unwrap()
                .sample(&mut rand::rng());
        }
    }
}