edge-rs 0.3.1

//! The marginal Poisson-gamma mixed model NEBULA starts from.
//!
//! The subject random effect is integrated out analytically: each subject has
//! one gamma frailty shared by its cells, and each cell has its own gamma
//! overdispersion. The closed-form marginal likelihood in `(beta, sigma, phi)`
//! is minimised, by projected Newton on dense designs and L-BFGS-B where the
//! zero-count tables apply, to give the starting values every later NEBULA
//! stage refines. Ports `ptmg_*_eigen` from `nebula`'s C++ plus the R-level
//! gamma corrections. `ptmg_value_gradient_hessian` supplies the marginal
//! Hessian edgePython lacks; see `UPSTREAM_DEVIATIONS.md` A20.
//!
//! ### Layout
//!
//! Cells are sorted by subject and [`GeneData::fid`] holds the boundaries:
//! subject `s` owns cells `fid[s]..fid[s + 1]`. The design is row-major
//! `n_cells * n_coef`. The parameter vector is `[beta (n_coef), sigma, phi]`.
//!
//! One gene is the unit of work and kernels are sequential; the caller
//! parallelises over genes. The exception is [`cumsum_y`], which fans out over
//! genes itself. Public kernels allocate scratch per call; the optimiser holds
//! a `PtmgScratch` per gene and uses `ptmg_value_and_gradient_with`.
//!
//! ### References
//!
//! He et al., Communications Biology 4, 629, 2021

use rayon::prelude::*;

use crate::numeric::gamma::{digamma, ln_gamma, trigamma};
use crate::prelude::*;
use crate::sc::pml::positive_runs;
use crate::sc::zeros::ZeroCells;

///////////////
// Gene data //
///////////////

/// Per-gene data laid out the way the kernels want it.
///
/// Counts are sparse: non-zero entries plus their cells, as
/// [`positive_indices`] returns them. Build with [`GeneData::new`], the only
/// place the kernels' invariants are checked.
#[derive(Clone, Copy, Debug)]
pub struct GeneData<'a> {
    /// Design matrix, row-major `n_cells * n_coef`, one row per cell.
    pub design: &'a [f64],
    /// Log offset per cell, length `n_cells`.
    pub log_offset: &'a [f64],
    /// Non-zero counts for this gene, in increasing cell order.
    pub counts: &'a [f64],
    /// Cell index of each entry of `counts`, strictly increasing.
    pub cells: &'a [usize],
    /// Count total per subject, length `n_subjects`. `cumsumy` in the reference.
    pub subject_totals: &'a [f64],
    /// Subject boundaries, length `n_subjects + 1`. Subject `s` owns cells
    /// `fid[s]..fid[s + 1]`, so `fid[0] == 0` and `fid[n_subjects] == n_cells`.
    pub fid: &'a [usize],
    /// Number of design columns.
    pub n_coef: usize,
}

impl<'a> GeneData<'a> {
    /// Assembles and validates the per-gene inputs.
    ///
    /// ### Params
    ///
    /// * `design` - Row-major `n_cells * n_coef` design matrix
    /// * `log_offset` - Log offset per cell, its length fixes `n_cells`
    /// * `counts` - Non-zero counts for this gene
    /// * `cells` - Cell index of each count, strictly increasing
    /// * `subject_totals` - Count total per subject
    /// * `fid` - Subject boundaries, length `n_subjects + 1`
    /// * `n_coef` - Number of design columns
    ///
    /// ### Returns
    ///
    /// The validated view, or [`EdgeErrors`] if the pieces disagree in length,
    /// `fid` is not a partition of `0..n_cells`, or a cell index is out of
    /// range or out of order.
    pub fn new(
        design: &'a [f64],
        log_offset: &'a [f64],
        counts: &'a [f64],
        cells: &'a [usize],
        subject_totals: &'a [f64],
        fid: &'a [usize],
        n_coef: usize,
    ) -> Result<Self, EdgeErrors> {
        if n_coef == 0 {
            return Err(EdgeErrors::MustBePositive("n_coef".to_string()));
        }
        let n_cells = log_offset.len();
        if n_cells == 0 {
            return Err(EdgeErrors::MustBePositive("n_cells".to_string()));
        }
        if design.len() != n_cells * n_coef {
            return Err(EdgeErrors::LengthMismatch {
                name: "design",
                expected: n_cells * n_coef,
                got: design.len(),
            });
        }
        if fid.len() < 2 {
            return Err(EdgeErrors::TooFewSubjects {
                required: 1,
                got: fid.len().saturating_sub(1),
            });
        }
        let n_subjects = fid.len() - 1;
        if subject_totals.len() != n_subjects {
            return Err(EdgeErrors::LengthMismatch {
                name: "subject_totals",
                expected: n_subjects,
                got: subject_totals.len(),
            });
        }
        if fid[0] != 0 || fid[n_subjects] != n_cells {
            return Err(EdgeErrors::InvalidArgument(format!(
                "fid must run from 0 to {n_cells}; got {} to {}.",
                fid[0], fid[n_subjects]
            )));
        }
        for s in 0..n_subjects {
            if fid[s + 1] <= fid[s] {
                return Err(EdgeErrors::InvalidArgument(format!(
                    "Subject {s} spans no cells: fid[{s}] = {}, fid[{}] = {}.",
                    fid[s],
                    s + 1,
                    fid[s + 1]
                )));
            }
        }
        if counts.len() != cells.len() {
            return Err(EdgeErrors::LengthMismatch {
                name: "cells",
                expected: counts.len(),
                got: cells.len(),
            });
        }
        let mut previous: Option<usize> = None;
        for &c in cells {
            if c >= n_cells {
                return Err(EdgeErrors::InvalidArgument(format!(
                    "Cell index {c} is out of range for {n_cells} cells."
                )));
            }
            if previous.is_some_and(|p| c <= p) {
                return Err(EdgeErrors::InvalidArgument(
                    "Cell indices must be strictly increasing.".to_string(),
                ));
            }
            previous = Some(c);
        }

        Ok(Self {
            design,
            log_offset,
            counts,
            cells,
            subject_totals,
            fid,
            n_coef,
        })
    }

    /// Number of cells for this gene.
    ///
    /// ### Returns
    ///
    /// `log_offset.len()`.
    #[inline]
    pub fn n_cells(&self) -> usize {
        self.log_offset.len()
    }

    /// Number of subjects for this gene.
    ///
    /// ### Returns
    ///
    /// `fid.len() - 1`.
    #[inline]
    pub fn n_subjects(&self) -> usize {
        self.fid.len() - 1
    }

    /// Length the parameter vector must have.
    ///
    /// ### Returns
    ///
    /// `n_coef + 2`, for `[beta, sigma, phi]`.
    #[inline]
    pub fn n_params(&self) -> usize {
        self.n_coef + 2
    }
}

/// Widest design whose per-subject gradient sums live on the stack, where the
/// compiler keeps them in registers (236 against 258 us per evaluation on a
/// 20000-cell, three-coefficient gene). Wider designs use the heap.
const LOCAL_COEF: usize = 8;

/// Per-gene buffers for the kernels, so repeated evaluations of one gene
/// allocate nothing.
pub(crate) struct PtmgScratch<'z> {
    /// Count per cell, zero where the gene is not expressed. Constant per gene,
    /// so filled once.
    y_cell: Vec<f64>,
    /// `exp(eta)` per cell, overwritten by every evaluation. With tables, only
    /// the positive and the loose cells are written.
    extb: Vec<f64>,
    /// Sum of `extb` per subject, overwritten by every evaluation.
    cumsumxtb: Vec<f64>,
    /// The run's zero-count tables, or `None` to sweep every cell.
    zeros: Option<&'z ZeroCells>,
    /// Where each subject's run of positive counts starts, length
    /// `n_subjects + 1`. Empty without tables.
    run: Vec<usize>,
}

impl<'z> PtmgScratch<'z> {
    /// Buffers for one gene.
    ///
    /// ### Params
    ///
    /// * `data` - The gene the scratch will serve, and only that gene
    ///
    /// ### Returns
    ///
    /// The scratch with the per-cell counts scattered in.
    pub(crate) fn new(data: &GeneData<'_>) -> Self {
        let mut y_cell = vec![0.0; data.n_cells()];
        for (&c, &y) in data.cells.iter().zip(data.counts.iter()) {
            y_cell[c] = y;
        }
        Self {
            y_cell,
            extb: vec![0.0; data.n_cells()],
            cumsumxtb: vec![0.0; data.n_subjects()],
            zeros: None,
            run: Vec::new(),
        }
    }

    /// Buffers for one gene whose value and gradient sum the zero counts from
    /// the run's tables.
    ///
    /// ### Params
    ///
    /// * `data` - The gene the scratch will serve, and only that gene
    /// * `zeros` - Tables built from this gene's design, offsets and subjects
    ///
    /// ### Returns
    ///
    /// The scratch, with the positive-count runs located.
    pub(crate) fn tabled(data: &GeneData<'_>, zeros: &'z ZeroCells) -> Self {
        Self {
            zeros: Some(zeros),
            run: positive_runs(data.cells, data.fid),
            ..Self::new(data)
        }
    }
}

//////////////////////////
// Reparametrised terms //
//////////////////////////

/// The scalars every kernel derives from `(sigma, phi)`, and their derivatives.
///
/// `sigma` is the log of one plus the subject-level variance: the gamma frailty
/// has shape `alpha = 1 / (exp(sigma) - 1)` and rate
/// `lambda = alpha / sqrt(exp(sigma))`. `phi`, the cell-level gamma shape,
/// enters directly.
#[derive(Clone, Copy, Debug)]
struct Terms {
    /// `sqrt(exp(sigma))`, which is also `alpha / lambda`.
    exps_s: f64,
    /// Subject frailty shape `1 / (exp(sigma) - 1)`.
    alpha: f64,
    /// Subject frailty rate `alpha / sqrt(exp(sigma))`.
    lambda: f64,
    /// Cell-level gamma shape, the `phi` of the parameter vector.
    gamma: f64,
    /// `ln(lambda)`.
    log_lambda: f64,
    /// `ln(gamma)`.
    log_gamma: f64,
    /// `d alpha / d sigma`.
    alpha_pr: f64,
    /// `d lambda / d sigma`.
    lambda_pr: f64,
    /// `d^2 alpha / d sigma^2`.
    alpha_dpr: f64,
    /// `d^2 lambda / d sigma^2`.
    lambda_dpr: f64,
}

impl Terms {
    /// Derives the reparametrised scalars at a point.
    ///
    /// Reproduces the reference arithmetic verbatim, including
    /// `exp(sigma) - 1` in place of `expm1(sigma)`.
    ///
    /// ### Params
    ///
    /// * `sigma` - Subject-level variance parameter
    /// * `phi` - Cell-level gamma shape
    ///
    /// ### Returns
    ///
    /// The scalars and their first two derivatives in `sigma`.
    fn new(sigma: f64, phi: f64) -> Self {
        let exps = sigma.exp();
        let exps_m1 = exps - 1.0;
        let exps_s = exps.sqrt();
        let alpha = 1.0 / exps_m1;
        let lambda = alpha / exps_s;
        let exps_sq = exps_m1 * exps_m1;

        Self {
            exps_s,
            alpha,
            lambda,
            gamma: phi,
            log_lambda: lambda.ln(),
            log_gamma: phi.ln(),
            alpha_pr: -exps / exps_sq,
            lambda_pr: (1.0 - 3.0 * exps) / (2.0 * exps_s * exps_sq),
            alpha_dpr: 2.0 * exps * exps / (exps_sq * exps_m1) - exps / exps_sq,
            lambda_dpr: -3.0 * exps / (2.0 * exps_s * exps_sq)
                + (3.0 * exps - 1.0) / (4.0 * exps_s * exps_sq)
                + (3.0 * exps - 1.0) * exps_s / (exps_sq * exps_m1),
        }
    }
}

///////////////////////////
// Gamma-function pieces //
///////////////////////////

/// How far to differentiate the gamma-function part of the likelihood.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum GammaOrder {
    /// Value only.
    Value,
    /// Value and first derivative.
    Gradient,
    /// Value, first and second derivative.
    Hessian,
}

/// The gamma-function terms the C++ kernels leave to the R wrappers.
///
/// Subject block: `sum lgamma(cumsumy + alpha) - lgamma(alpha)` over subjects
/// with a non-zero total, differentiated through `alpha`. Cell block:
/// `sum lgamma(y + phi) - lgamma(phi)` over cells with a non-zero count,
/// differentiated in `phi`.
#[derive(Clone, Copy, Debug, Default)]
struct GammaTerms {
    /// Subject block, evaluated.
    subject_value: f64,
    /// `d subject_value / d alpha`.
    subject_digamma: f64,
    /// `d^2 subject_value / d alpha^2`.
    subject_trigamma: f64,
    /// Cell block, evaluated.
    cell_value: f64,
    /// `d cell_value / d phi`.
    cell_digamma: f64,
    /// `d^2 cell_value / d phi^2`.
    cell_trigamma: f64,
}

/// Evaluates the gamma-function terms and, on request, their derivatives.
///
/// Counts of one and two use closed forms, as the reference does:
/// `lgamma(1 + g) - lgamma(g) = log(g)` and
/// `lgamma(2 + g) - lgamma(g) = log(g) + log(g + 1)`. Only counts of three and
/// above call `lgamma`. Subjects with a zero total are skipped (`posind`).
///
/// ### Params
///
/// * `data` - The gene
/// * `alpha` - Subject frailty shape
/// * `gamma` - Cell-level gamma shape, the `phi` of the parameter vector
/// * `order` - How far to differentiate
///
/// ### Returns
///
/// The two blocks and the derivatives `order` asks for; the rest stay zero.
fn gamma_terms(data: &GeneData<'_>, alpha: f64, gamma: f64, order: GammaOrder) -> GammaTerms {
    let mut out = GammaTerms::default();

    // ------------- //
    // subject block //
    // ------------- //
    let mut n_positive = 0.0;
    let mut ln_sum = 0.0;
    let mut di_sum = 0.0;
    let mut tri_sum = 0.0;
    for &total in data.subject_totals {
        if total > 0.0 {
            n_positive += 1.0;
            ln_sum += ln_gamma(total + alpha);
            if order != GammaOrder::Value {
                di_sum += digamma(total + alpha);
            }
            if order == GammaOrder::Hessian {
                tri_sum += trigamma(total + alpha);
            }
        }
    }
    out.subject_value = ln_sum - n_positive * ln_gamma(alpha);
    if order != GammaOrder::Value {
        out.subject_digamma = di_sum - n_positive * digamma(alpha);
    }
    if order == GammaOrder::Hessian {
        out.subject_trigamma = tri_sum - n_positive * trigamma(alpha);
    }

    // ---------- //
    // cell block //
    // ---------- //
    let mut n_one = 0.0;
    let mut n_two = 0.0;
    let mut n_big = 0.0;
    let mut ln_big = 0.0;
    let mut di_big = 0.0;
    let mut tri_big = 0.0;
    for &y in data.counts {
        if y == 1.0 {
            n_one += 1.0;
        } else if y == 2.0 {
            n_two += 1.0;
        } else {
            n_big += 1.0;
            ln_big += ln_gamma(y + gamma);
            if order != GammaOrder::Value {
                di_big += digamma(y + gamma);
            }
            if order == GammaOrder::Hessian {
                tri_big += trigamma(y + gamma);
            }
        }
    }
    let n_small = n_one + n_two;
    out.cell_value =
        ln_big - n_big * ln_gamma(gamma) + n_small * gamma.ln() + n_two * (gamma + 1.0).ln();
    if order != GammaOrder::Value {
        out.cell_digamma =
            di_big - n_big * digamma(gamma) + n_small / gamma + n_two / (gamma + 1.0);
    }
    if order == GammaOrder::Hessian {
        out.cell_trigamma = tri_big
            - n_big * trigamma(gamma)
            - n_small / (gamma * gamma)
            - n_two / ((gamma + 1.0) * (gamma + 1.0));
    }

    out
}

/////////////
// Kernels //
/////////////

/// Marginal negative log-likelihood of one gene.
///
/// Port of `ptmg_ll_eigen` plus the gamma-function terms `R/ptmg.R` adds, so
/// this equals `nebula:::ptmg_ll`. The value-only path skips every buffer the
/// gradient needs.
///
/// ### Params
///
/// * `data` - The gene, as validated by [`GeneData::new`]
/// * `params` - `[beta (n_coef), sigma, phi]`, length [`GeneData::n_params`]
///
/// ### Returns
///
/// The negative log-likelihood. Lower is better, matching what the optimiser
/// minimises.
pub fn ptmg_neg_log_likelihood(data: &GeneData<'_>, params: &[f64]) -> f64 {
    debug_assert_eq!(params.len(), data.n_params());

    let nb = data.n_coef;
    let n_cells = data.n_cells();
    let n_subjects = data.n_subjects();
    let beta = &params[..nb];
    let terms = Terms::new(params[nb], params[nb + 1]);
    let gamma = terms.gamma;

    let mut scratch = PtmgScratch::new(data);
    let mut total = linear_predictor(data, beta, &mut scratch);
    let PtmgScratch {
        y_cell,
        extb,
        cumsumxtb,
        ..
    } = &scratch;

    for s in 0..n_subjects {
        let span = data.fid[s]..data.fid[s + 1];
        let cumsumxtb = cumsumxtb[s];
        let ystar = data.subject_totals[s] + terms.alpha;
        let mustar = cumsumxtb + terms.lambda;
        total -= ystar * mustar.ln();

        let ymustar = ystar / mustar;
        total += ymustar * cumsumxtb;
        for i in span {
            let w = ymustar * extb[i] + gamma;
            total -= (gamma + y_cell[i]) * w.ln();
        }
    }

    total += (n_subjects as f64) * terms.alpha * terms.log_lambda;
    total += (n_cells as f64) * gamma * terms.log_gamma;

    let extra = gamma_terms(data, terms.alpha, gamma, GammaOrder::Value);

    -total - extra.subject_value - extra.cell_value
}

/// Gradient of the marginal negative log-likelihood.
///
/// Port of `ptmg_der_eigen` plus the R-level corrections, equal to
/// `nebula:::ptmg_der`. Runs [`ptmg_value_and_gradient`] and drops the value;
/// prefer that when both are needed.
///
/// ### Params
///
/// * `data` - The gene, as validated by [`GeneData::new`]
/// * `params` - `[beta (n_coef), sigma, phi]`, length [`GeneData::n_params`]
///
/// ### Returns
///
/// The gradient, length `n_coef + 2`.
pub fn ptmg_gradient(data: &GeneData<'_>, params: &[f64]) -> Vec<f64> {
    evaluate(data, params, &mut PtmgScratch::new(data)).1
}

/// Value and gradient of the marginal negative log-likelihood together.
///
/// Port of `ptmg_ll_der_eigen` plus the R-level corrections, so this equals
/// `nebula:::ptmg_ll_der`. This is the optimiser's objective.
///
/// ### Params
///
/// * `data` - The gene, as validated by [`GeneData::new`]
/// * `params` - `[beta (n_coef), sigma, phi]`, length [`GeneData::n_params`]
///
/// ### Returns
///
/// The negative log-likelihood and its gradient, length `n_coef + 2`.
pub fn ptmg_value_and_gradient(data: &GeneData<'_>, params: &[f64]) -> (f64, Vec<f64>) {
    ptmg_value_and_gradient_with(data, params, &mut PtmgScratch::new(data))
}

/// [`ptmg_value_and_gradient`] on buffers the caller keeps across evaluations
/// of one gene.
///
/// ### Params
///
/// * `data` - The gene, as validated by [`GeneData::new`]
/// * `params` - `[beta (n_coef), sigma, phi]`, length [`GeneData::n_params`]
/// * `scratch` - Built by [`PtmgScratch::new`] from this same gene
///
/// ### Returns
///
/// The negative log-likelihood and its gradient: bit for bit what
/// [`ptmg_value_and_gradient`] returns without tables, to rounding with them.
pub(crate) fn ptmg_value_and_gradient_with(
    data: &GeneData<'_>,
    params: &[f64],
    scratch: &mut PtmgScratch<'_>,
) -> (f64, Vec<f64>) {
    if let Some(zeros) = scratch.zeros {
        return evaluate_tabled(data, params, zeros, scratch);
    }
    evaluate(data, params, scratch)
}

/// Value, gradient and Hessian of the marginal negative log-likelihood.
///
/// Port of `ptmg_ll_der_hes_eigen` plus the R-level corrections. The covariance
/// of `beta` is read off this observed information; edgePython has no
/// equivalent (A20).
///
/// The Hessian is in the direct parametrisation `[beta, sigma, phi]`. The R
/// package's `ptmg_ll_der_hes2` and `ptmg_ll_der_hes3` also apply
/// `sigma -> exp(sigma)` for `trust`; that chain rule is the driver's job.
///
/// ### Params
///
/// * `data` - The gene, as validated by [`GeneData::new`]
/// * `params` - `[beta (n_coef), sigma, phi]`, length [`GeneData::n_params`]
///
/// ### Returns
///
/// The negative log-likelihood, its gradient of length `n_coef + 2`, and its
/// Hessian as a row-major `(n_coef + 2) * (n_coef + 2)` matrix.
pub fn ptmg_value_gradient_hessian(
    data: &GeneData<'_>,
    params: &[f64],
) -> (f64, Vec<f64>, Vec<f64>) {
    evaluate_hessian(data, params, &mut PtmgScratch::new(data))
}

/// [`ptmg_value_gradient_hessian`] on buffers the caller keeps across
/// evaluations of one gene.
///
/// ### Params
///
/// * `data` - The gene, as validated by [`GeneData::new`]
/// * `params` - `[beta (n_coef), sigma, phi]`, length [`GeneData::n_params`]
/// * `scratch` - Built by [`PtmgScratch::new`] from this same gene
///
/// ### Returns
///
/// As [`ptmg_value_gradient_hessian`].
pub(crate) fn ptmg_value_gradient_hessian_with(
    data: &GeneData<'_>,
    params: &[f64],
    scratch: &mut PtmgScratch<'_>,
) -> (f64, Vec<f64>, Vec<f64>) {
    evaluate_hessian(data, params, scratch)
}

/// Fills the exponentiated linear predictor and its subject sums, and
/// accumulates `sum y * eta`.
///
/// ### Params
///
/// * `data` - The gene
/// * `beta` - Coefficients, length `n_coef`
/// * `scratch` - The gene's buffers; `extb` and `cumsumxtb` are overwritten
///
/// ### Returns
///
/// `sum_i y_i * eta_i`, the only part of the likelihood linear in `eta`.
fn linear_predictor(data: &GeneData<'_>, beta: &[f64], scratch: &mut PtmgScratch<'_>) -> f64 {
    let nb = data.n_coef;
    let mut total = 0.0;

    for (s, csx_s) in scratch.cumsumxtb.iter_mut().enumerate() {
        let span = data.fid[s]..data.fid[s + 1];
        let mut csx = 0.0;
        for (((row, &offset), &y), extb_i) in data.design[span.start * nb..span.end * nb]
            .chunks_exact(nb)
            .zip(&data.log_offset[span.clone()])
            .zip(&scratch.y_cell[span.clone()])
            .zip(&mut scratch.extb[span])
        {
            let mut eta = offset;
            for (&x, &b) in row.iter().zip(beta.iter()) {
                eta += x * b;
            }
            // Guarded so a huge eta paired with a zero count cannot make `0 * inf`.
            if y != 0.0 {
                total += eta * y;
            }
            let e = eta.exp();
            *extb_i = e;
            csx += e;
        }
        *csx_s = csx;
    }

    total
}

/// Subject-level quantities every kernel derives from the subject sums of
/// `exp(eta)`.
struct Subjects {
    /// `ln(mustar)`, `mustar = cumsumxtb + lambda`.
    mustar_log: Vec<f64>,
    /// `ystar / mustar`, `ystar = cumsumy + alpha`.
    ymustar: Vec<f64>,
    /// `ystar / mustar^2`.
    ymumustar: Vec<f64>,
    /// `1 / mustar`.
    imustar: Vec<f64>,
}

impl Subjects {
    /// Derives the subject scalars and folds their terms into the likelihood:
    /// `-ystar ln(mustar)` per subject, then `k alpha ln(lambda)` and
    /// `n gamma ln(gamma)`.
    ///
    /// ### Params
    ///
    /// * `data` - The gene
    /// * `terms` - The reparametrised variance components
    /// * `cumsumxtb` - Sum of `exp(eta)` per subject
    /// * `total` - The log-likelihood so far, updated
    ///
    /// ### Returns
    ///
    /// The scalars.
    #[allow(clippy::needless_range_loop)]
    fn new(data: &GeneData<'_>, terms: &Terms, cumsumxtb: &[f64], total: &mut f64) -> Self {
        let k = data.n_subjects();
        let mut out = Self {
            mustar_log: vec![0.0; k],
            ymustar: vec![0.0; k],
            ymumustar: vec![0.0; k],
            imustar: vec![0.0; k],
        };
        for s in 0..k {
            let ystar = data.subject_totals[s] + terms.alpha;
            let mu = cumsumxtb[s] + terms.lambda;
            let log_mu = mu.ln();
            out.mustar_log[s] = log_mu;
            out.ymustar[s] = ystar / mu;
            out.ymumustar[s] = ystar / mu / mu;
            out.imustar[s] = 1.0 / mu;
            *total -= ystar * log_mu;
        }
        *total += (k as f64) * terms.alpha * terms.log_lambda;
        *total += (data.n_cells() as f64) * terms.gamma * terms.log_gamma;
        out
    }
}

/// The gradient in log-likelihood sign, and the intermediates the Hessian
/// reuses.
struct GradientParts {
    /// Gradient of the log-likelihood, length `n_coef + 2`, before the gamma
    /// corrections.
    gradient: Vec<f64>,
    /// `ymumustar * d42c` per subject.
    ymm_d: Vec<f64>,
    /// `d42c / mustar` per subject.
    hbl0: Vec<f64>,
    /// `sum hbl0`.
    dbim: f64,
    /// `sum ymm_d`.
    sum_ymm_d: f64,
}

/// The gradient from the per-subject sums, as `ptmg_der_eigen` assembles it.
///
/// ### Params
///
/// * `data` - The gene
/// * `terms` - The reparametrised variance components
/// * `subjects` - The subject scalars
/// * `xexb_f` - `sum x e` per subject, row-major `k * n_coef`
/// * `dbeta_41` - `sum g x e` per subject, row-major `k * n_coef`
/// * `d42c` - `sum g e - sum e` per subject
/// * `slpey` - `sum log w` over every cell
/// * `gstar_sum` - `sum g` over every cell
///
/// ### Returns
///
/// The gradient and its intermediates.
#[allow(clippy::too_many_arguments)]
fn gradient_parts(
    data: &GeneData<'_>,
    terms: &Terms,
    subjects: &Subjects,
    xexb_f: &[f64],
    dbeta_41: &[f64],
    d42c: &[f64],
    slpey: f64,
    gstar_sum: f64,
) -> GradientParts {
    let nb = data.n_coef;
    let k = data.n_subjects();
    let n_cells = data.n_cells() as f64;
    let Subjects {
        mustar_log,
        ymustar,
        ymumustar,
        imustar,
    } = subjects;

    let ymm_d: Vec<f64> = (0..k).map(|s| ymumustar[s] * d42c[s]).collect();
    let mut gradient = vec![0.0; nb + 2];
    for j in 0..nb {
        let mut acc = 0.0;
        for s in 0..k {
            acc += xexb_f[s * nb + j] * ymm_d[s] - dbeta_41[s * nb + j] * ymustar[s];
        }
        gradient[j] = acc;
    }
    for (&c, &y) in data.cells.iter().zip(data.counts.iter()) {
        let row = &data.design[c * nb..(c + 1) * nb];
        for j in 0..nb {
            gradient[j] += row[j] * y;
        }
    }

    let ldm = terms.log_lambda * (k as f64) - mustar_log.iter().sum::<f64>();
    let adlmy = terms.exps_s * (k as f64) - ymustar.iter().sum::<f64>();
    let hbl0: Vec<f64> = (0..k).map(|s| d42c[s] * imustar[s]).collect();
    let dbim: f64 = hbl0.iter().sum();
    let sum_ymm_d: f64 = ymm_d.iter().sum();
    gradient[nb] = -terms.alpha_pr * dbim
        + terms.lambda_pr * sum_ymm_d
        + terms.alpha_pr * ldm
        + terms.lambda_pr * adlmy;
    gradient[nb + 1] = terms.log_gamma * n_cells + n_cells - slpey - gstar_sum;

    GradientParts {
        gradient,
        ymm_d,
        hbl0,
        dbim,
        sum_ymm_d,
    }
}

/// Negates the log-likelihood and its gradient and subtracts the gamma terms
/// `R/ptmg.R` adds.
///
/// ### Params
///
/// * `data` - The gene
/// * `terms` - The reparametrised variance components
/// * `total` - The log-likelihood
/// * `gradient` - Its gradient, from [`gradient_parts`]
///
/// ### Returns
///
/// The negative log-likelihood and its gradient.
fn finish_gradient(
    data: &GeneData<'_>,
    terms: &Terms,
    total: f64,
    mut gradient: Vec<f64>,
) -> (f64, Vec<f64>) {
    let nb = data.n_coef;
    for g in gradient.iter_mut() {
        *g = -*g;
    }
    let extra = gamma_terms(data, terms.alpha, terms.gamma, GammaOrder::Gradient);
    let value = -total - extra.subject_value - extra.cell_value;
    gradient[nb] -= terms.alpha_pr * extra.subject_digamma;
    gradient[nb + 1] -= extra.cell_digamma;
    (value, gradient)
}

/// Value and gradient, in one pass over the gene.
///
/// Shared body of [`ptmg_value_and_gradient`] and [`ptmg_gradient`].
/// Accumulates in log-likelihood sign and negates once at the end; the gamma
/// corrections are then subtracted from the negated quantities, the order
/// `R/ptmg.R` uses.
///
/// ### Params
///
/// * `data` - The gene, as validated by [`GeneData::new`]
/// * `params` - `[beta (n_coef), sigma, phi]`, length [`GeneData::n_params`]
/// * `scratch` - Built by [`PtmgScratch::new`] from this same gene
///
/// ### Returns
///
/// The negative log-likelihood and its gradient of length `n_coef + 2`.
fn evaluate(data: &GeneData<'_>, params: &[f64], scratch: &mut PtmgScratch<'_>) -> (f64, Vec<f64>) {
    debug_assert_eq!(params.len(), data.n_params());

    let nb = data.n_coef;
    let k = data.n_subjects();
    let fid = data.fid;
    let beta = &params[..nb];
    let terms = Terms::new(params[nb], params[nb + 1]);
    let gamma = terms.gamma;

    let mut total = linear_predictor(data, beta, scratch);
    let PtmgScratch {
        y_cell,
        extb,
        cumsumxtb,
        ..
    } = &*scratch;

    let subjects = Subjects::new(data, &terms, cumsumxtb, &mut total);
    let ymustar = &subjects.ymustar;

    // One pass per cell. Sum order matches the reference.
    let mut xexb_f = vec![0.0; k * nb];
    let mut dbeta_41 = vec![0.0; k * nb];
    let mut d42c = vec![0.0; k];
    // One subject's cells, adding into its rows of `xexb_f` and `dbeta_41`.
    // Running sums pass by value to stay in registers; returns them with the
    // subject's `sum gstar * extb`.
    let cells = |s: usize, sums: [f64; 3], xf: &mut [f64], d41: &mut [f64]| {
        let [mut total, mut slpey, mut gstar_sum] = sums;
        let ym = ymustar[s];
        let mut acc = 0.0;
        for i in fid[s]..fid[s + 1] {
            let e = extb[i];
            let we = ym * e;
            let w = we + gamma;
            let log_w = w.ln();
            let inv_w = 1.0 / w;
            total += we;
            total -= (gamma + y_cell[i]) * log_w;
            slpey += log_w;
            let g = (gamma + y_cell[i]) * inv_w;
            gstar_sum += g;

            let row = &data.design[i * nb..(i + 1) * nb];
            for j in 0..nb {
                let v = row[j] * e;
                xf[j] += v;
                d41[j] += g * v;
            }
            acc += g * e;
        }
        ([total, slpey, gstar_sum], acc)
    };
    let mut sums = [total, 0.0, 0.0];
    for s in 0..k {
        let rows = s * nb..(s + 1) * nb;
        let acc;
        if nb <= LOCAL_COEF {
            let mut xf = [0.0; LOCAL_COEF];
            let mut d41 = [0.0; LOCAL_COEF];
            (sums, acc) = cells(s, sums, &mut xf[..nb], &mut d41[..nb]);
            xexb_f[rows.clone()].copy_from_slice(&xf[..nb]);
            dbeta_41[rows].copy_from_slice(&d41[..nb]);
        } else {
            (sums, acc) = cells(s, sums, &mut xexb_f[rows.clone()], &mut dbeta_41[rows]);
        }
        d42c[s] = acc - cumsumxtb[s];
    }
    let [total, slpey, gstar_sum] = sums;

    let parts = gradient_parts(
        data, &terms, &subjects, &xexb_f, &dbeta_41, &d42c, slpey, gstar_sum,
    );
    finish_gradient(data, &terms, total, parts.gradient)
}

/// Value, gradient and Hessian in one pass over the gene.
///
/// Every Hessian entry is a sum over subjects of per-subject sums over cells,
/// so one pass per subject accumulates them all: the vectors `sum x e`,
/// `sum g x e`, `sum x e (g/w - 1/w)` and `sum x e g e/w`, three scalars, and
/// the three symmetric `x x' e` blocks weighted by `1`, `g` and `g e/w`. The
/// formulas that follow are [`evaluate`]'s, applied to those sums.
///
/// ### Params
///
/// * `data` - The gene, as validated by [`GeneData::new`]
/// * `params` - `[beta (n_coef), sigma, phi]`, length [`GeneData::n_params`]
/// * `scratch` - Built by [`PtmgScratch::new`] from this same gene
///
/// ### Returns
///
/// The negative log-likelihood, its gradient of length `n_coef + 2`, and its
/// row-major Hessian.
fn evaluate_hessian(
    data: &GeneData<'_>,
    params: &[f64],
    scratch: &mut PtmgScratch<'_>,
) -> (f64, Vec<f64>, Vec<f64>) {
    debug_assert_eq!(params.len(), data.n_params());

    let nb = data.n_coef;
    let nn = nb * nb;
    let n_cells = data.n_cells();
    let k = data.n_subjects();
    let n_params = nb + 2;
    let fid = data.fid;
    let beta = &params[..nb];
    let terms = Terms::new(params[nb], params[nb + 1]);
    let gamma = terms.gamma;

    let mut total = linear_predictor(data, beta, scratch);
    let PtmgScratch {
        y_cell,
        extb,
        cumsumxtb,
        ..
    } = &*scratch;

    let subjects = Subjects::new(data, &terms, cumsumxtb, &mut total);
    let Subjects {
        mustar_log,
        ymustar,
        ymumustar,
        imustar,
    } = &subjects;

    // Per-subject sums, row-major `k * nb` vectors and `k * nb * nb` blocks
    // (upper triangle filled).
    let mut xexb_f = vec![0.0; k * nb];
    let mut dbeta_41 = vec![0.0; k * nb];
    let mut tmp1_f = vec![0.0; k * nb];
    let mut tmp2_f = vec![0.0; k * nb];
    let mut block_plain = vec![0.0; k * nn];
    let mut block_g = vec![0.0; k * nn];
    let mut block_gte = vec![0.0; k * nn];
    let mut hb2_f = vec![0.0; k];
    let mut b2f = vec![0.0; k];
    let mut d42c = vec![0.0; k];
    let mut scalars = HessianScalars {
        total,
        slpey: 0.0,
        gstar_sum: 0.0,
        sum_tempa: 0.0,
        sum_gt: 0.0,
    };
    for s in 0..k {
        let v = s * nb..(s + 1) * nb;
        let m = s * nn..(s + 1) * nn;
        let block = SubjectBlock {
            xf: &mut xexb_f[v.clone()],
            d41: &mut dbeta_41[v.clone()],
            t1: &mut tmp1_f[v.clone()],
            t2: &mut tmp2_f[v],
            plain: &mut block_plain[m.clone()],
            g: &mut block_g[m.clone()],
            gte: &mut block_gte[m],
        };
        macro_rules! width {
            ($nb:literal) => {
                hessian_subject::<$nb>(
                    data,
                    fid[s]..fid[s + 1],
                    extb,
                    y_cell,
                    ymustar[s],
                    gamma,
                    block,
                    &mut scalars,
                )
            };
        }
        let (acc, hb2, b2) = match nb {
            1 => width!(1),
            2 => width!(2),
            3 => width!(3),
            4 => width!(4),
            5 => width!(5),
            6 => width!(6),
            7 => width!(7),
            8 => width!(8),
            _ => width!(0),
        };
        hb2_f[s] = hb2;
        b2f[s] = b2;
        d42c[s] = acc - cumsumxtb[s];
    }
    let HessianScalars {
        total,
        slpey,
        gstar_sum,
        sum_tempa,
        sum_gt,
    } = scalars;

    let GradientParts {
        mut gradient,
        ymm_d,
        hbl0,
        dbim,
        sum_ymm_d,
    } = gradient_parts(
        data, &terms, &subjects, &xexb_f, &dbeta_41, &d42c, slpey, gstar_sum,
    );
    let value = -total;
    for g in gradient.iter_mut() {
        *g = -*g;
    }

    // Hessian, in log-likelihood sign, as in `evaluate`.
    let mut hessian = vec![0.0; n_params * n_params];
    let lam_ratio = terms.lambda_pr / terms.lambda;
    let mut hes_sigma = (terms.alpha_dpr * terms.log_lambda + 2.0 * terms.alpha_pr * lam_ratio
        - terms.alpha * lam_ratio * lam_ratio
        + terms.alpha / terms.lambda * terms.lambda_dpr)
        * (k as f64);
    let sum_hbl0_imu: f64 = (0..k).map(|s| hbl0[s] * imustar[s]).sum();
    let sum_ymmd_imu: f64 = (0..k).map(|s| ymm_d[s] * imustar[s]).sum();
    hes_sigma -= dbim * terms.alpha_dpr
        - 2.0 * terms.alpha_pr * terms.lambda_pr * sum_hbl0_imu
        - terms.lambda_dpr * sum_ymm_d
        + 2.0 * terms.lambda_pr * terms.lambda_pr * sum_ymmd_imu;
    hes_sigma -= 2.0 * terms.alpha_pr * terms.lambda_pr * imustar.iter().sum::<f64>()
        + terms.alpha_dpr * mustar_log.iter().sum::<f64>()
        + terms.lambda_dpr * ymustar.iter().sum::<f64>()
        - terms.lambda_pr * terms.lambda_pr * ymumustar.iter().sum::<f64>();

    let mut hes_sigma_beta = vec![0.0; nb];
    for j in 0..nb {
        let mut d41_imu = 0.0;
        let mut d41_ymm = 0.0;
        let mut xf_hbl_imu = 0.0;
        let mut xf_hbl_ymm = 0.0;
        for s in 0..k {
            d41_imu += dbeta_41[s * nb + j] * imustar[s];
            d41_ymm += dbeta_41[s * nb + j] * ymumustar[s];
            xf_hbl_imu += xexb_f[s * nb + j] * hbl0[s] * imustar[s];
            xf_hbl_ymm += xexb_f[s * nb + j] * hbl0[s] * ymumustar[s];
        }
        hes_sigma_beta[j] = -terms.alpha_pr * d41_imu + terms.lambda_pr * d41_ymm
            - (-terms.alpha_pr * xf_hbl_imu + 2.0 * terms.lambda_pr * xf_hbl_ymm);
    }
    let apm: Vec<f64> = (0..k)
        .map(|s| terms.alpha_pr * imustar[s] - terms.lambda_pr * ymumustar[s])
        .collect();

    let hes_gamma = (n_cells as f64) / gamma - 2.0 * sum_tempa + sum_gt;
    let mut hes_gamma_beta = vec![0.0; nb];
    for j in 0..nb {
        let mut a = 0.0;
        let mut b = 0.0;
        for s in 0..k {
            a += tmp1_f[s * nb + j] * ymustar[s];
            b += xexb_f[s * nb + j] * hb2_f[s] * ymumustar[s];
        }
        hes_gamma_beta[j] = a - b;
    }
    let hes_sigma_gamma: f64 = (0..k).map(|s| apm[s] * hb2_f[s]).sum();
    for j in 0..nb {
        let mut a = 0.0;
        let mut b = 0.0;
        for s in 0..k {
            a += tmp2_f[s * nb + j] * apm[s] * ymustar[s];
            b += xexb_f[s * nb + j] * apm[s] * b2f[s] * ymumustar[s];
        }
        hes_sigma_beta[j] += a - b;
    }
    hes_sigma += (0..k).map(|s| b2f[s] * apm[s] * apm[s]).sum::<f64>();

    let ymu2: Vec<f64> = ymustar.iter().map(|v| v * v).collect();
    let ymuymuimu: Vec<f64> = (0..k).map(|s| ymu2[s] * imustar[s]).collect();
    let b2f_scaled: Vec<f64> = (0..k).map(|s| b2f[s] * ymuymuimu[s] * imustar[s]).collect();
    let ymmd_imu: Vec<f64> = (0..k).map(|s| ymm_d[s] * imustar[s]).collect();
    for p in 0..nb {
        for q in p..nb {
            let mut acc = 0.0;
            for s in 0..k {
                let at = s * nn + p * nb + q;
                acc -= block_g[at] * ymustar[s];
                acc += block_gte[at] * ymu2[s];
                acc += ymm_d[s] * block_plain[at];
            }
            for s in 0..k {
                let xf_p = xexb_f[s * nb + p];
                let xf_q = xexb_f[s * nb + q];
                acc += (dbeta_41[s * nb + p] * xf_q + dbeta_41[s * nb + q] * xf_p) * ymumustar[s];
                acc -= ymuymuimu[s] * (tmp2_f[s * nb + p] * xf_q + tmp2_f[s * nb + q] * xf_p);
                let cross = xf_p * xf_q;
                acc -= cross * ymumustar[s];
                acc += b2f_scaled[s] * cross;
                acc -= 2.0 * cross * ymmd_imu[s];
            }
            hessian[p * n_params + q] = acc;
            hessian[q * n_params + p] = acc;
        }
    }

    hessian[nb * n_params + nb] = hes_sigma;
    hessian[(nb + 1) * n_params + nb + 1] = hes_gamma;
    hessian[nb * n_params + nb + 1] = hes_sigma_gamma;
    hessian[(nb + 1) * n_params + nb] = hes_sigma_gamma;
    for j in 0..nb {
        hessian[nb * n_params + j] = hes_sigma_beta[j];
        hessian[j * n_params + nb] = hes_sigma_beta[j];
        hessian[(nb + 1) * n_params + j] = hes_gamma_beta[j];
        hessian[j * n_params + nb + 1] = hes_gamma_beta[j];
    }
    for h in hessian.iter_mut() {
        *h = -*h;
    }

    let extra = gamma_terms(data, terms.alpha, gamma, GammaOrder::Hessian);
    let value = value - extra.subject_value - extra.cell_value;
    gradient[nb] -= terms.alpha_pr * extra.subject_digamma;
    gradient[nb + 1] -= extra.cell_digamma;
    hessian[nb * n_params + nb] -= terms.alpha_dpr * extra.subject_digamma
        + terms.alpha_pr * terms.alpha_pr * extra.subject_trigamma;
    hessian[(nb + 1) * n_params + nb + 1] -= extra.cell_trigamma;

    (value, gradient, hessian)
}

/// Running scalars of [`evaluate_hessian`]'s sweep over the cells.
struct HessianScalars {
    /// The log-likelihood so far.
    total: f64,
    /// `sum log w`.
    slpey: f64,
    /// `sum g`.
    gstar_sum: f64,
    /// `sum 1 / w`.
    sum_tempa: f64,
    /// `sum g / w`.
    sum_gt: f64,
}

/// One subject's slices of [`evaluate_hessian`]'s per-subject sums.
struct SubjectBlock<'a> {
    /// `sum x e`, length `nb`.
    xf: &'a mut [f64],
    /// `sum g x e`, length `nb`.
    d41: &'a mut [f64],
    /// `sum x e (g / w - 1 / w)`, length `nb`.
    t1: &'a mut [f64],
    /// `sum x e g e / w`, length `nb`.
    t2: &'a mut [f64],
    /// `sum x x' e`, row-major `nb * nb`, upper triangle.
    plain: &'a mut [f64],
    /// `sum x x' e g`, upper triangle.
    g: &'a mut [f64],
    /// `sum x x' e g e / w`, upper triangle.
    gte: &'a mut [f64],
}

/// Widest design whose per-subject Hessian sums live on the stack.
const HESSIAN_LOCAL_COEF: usize = 8;

/// One subject's sweep for [`evaluate_hessian`], at a compile-time design
/// width (`NB == 0` reads it from `data`).
///
/// With the width a constant the loops over columns unroll and the sums sit in
/// stack arrays the compiler keeps in registers; they are written to `block`
/// once at the end.
///
/// ### Params
///
/// * `data` - The gene
/// * `cells` - The subject's cells
/// * `extb` - `exp(eta)` per cell
/// * `y_cell` - Count per cell
/// * `ym` - The subject's `ystar / mustar`
/// * `gamma` - Cell-level gamma shape
/// * `block` - The subject's sums, written
/// * `scalars` - Running scalars, updated
///
/// ### Returns
///
/// `sum g e`, `sum (g e / w - e / w)` and `sum g e^2 / w` over the subject.
#[allow(clippy::too_many_arguments)]
fn hessian_subject<const NB: usize>(
    data: &GeneData<'_>,
    cells: std::ops::Range<usize>,
    extb: &[f64],
    y_cell: &[f64],
    ym: f64,
    gamma: f64,
    block: SubjectBlock<'_>,
    scalars: &mut HessianScalars,
) -> (f64, f64, f64) {
    let nb = if NB == 0 { data.n_coef } else { NB };
    const V: usize = HESSIAN_LOCAL_COEF;
    let mut vecs = [[0.0f64; V]; 4];
    let mut blocks = [[0.0f64; V * V]; 3];
    let local = NB != 0;
    let SubjectBlock {
        xf,
        d41,
        t1,
        t2,
        plain,
        g: bg,
        gte: bgte,
    } = block;
    let mut acc = 0.0;
    let mut hb2 = 0.0;
    let mut b2 = 0.0;
    let HessianScalars {
        total,
        slpey,
        gstar_sum,
        sum_tempa,
        sum_gt,
    } = scalars;
    {
        let [lxf, ld41, lt1, lt2] = &mut vecs;
        let [lbp, lbg, lbgte] = &mut blocks;
        let SubjectBlock {
            xf,
            d41,
            t1,
            t2,
            plain: bp,
            g: bg,
            gte: bgte,
        } = if local {
            SubjectBlock {
                xf: &mut lxf[..nb],
                d41: &mut ld41[..nb],
                t1: &mut lt1[..nb],
                t2: &mut lt2[..nb],
                plain: &mut lbp[..nb * nb],
                g: &mut lbg[..nb * nb],
                gte: &mut lbgte[..nb * nb],
            }
        } else {
            SubjectBlock {
                xf: &mut *xf,
                d41: &mut *d41,
                t1: &mut *t1,
                t2: &mut *t2,
                plain: &mut *plain,
                g: &mut *bg,
                gte: &mut *bgte,
            }
        };
        for i in cells {
            let e = extb[i];
            let y = y_cell[i];
            let we = ym * e;
            let w = we + gamma;
            let log_w = w.ln();
            let inv_w = 1.0 / w;
            *total += we;
            *total -= (gamma + y) * log_w;
            *slpey += log_w;
            let g = (gamma + y) * inv_w;
            *gstar_sum += g;
            let gt = g * inv_w;
            let gte = gt * e;
            *sum_tempa += inv_w;
            *sum_gt += gt;
            acc += g * e;
            hb2 += gte - e * inv_w;
            b2 += gte * e;
            let f1 = gt - inv_w;

            let row = &data.design[i * nb..(i + 1) * nb];
            for a in 0..nb {
                let xe = row[a] * e;
                xf[a] += xe;
                d41[a] += g * xe;
                t1[a] += xe * f1;
                t2[a] += xe * gte;
                for b in a..nb {
                    let c = xe * row[b];
                    bp[a * nb + b] += c;
                    bg[a * nb + b] += c * g;
                    bgte[a * nb + b] += c * gte;
                }
            }
        }
    }
    if local {
        xf.copy_from_slice(&vecs[0][..nb]);
        d41.copy_from_slice(&vecs[1][..nb]);
        t1.copy_from_slice(&vecs[2][..nb]);
        t2.copy_from_slice(&vecs[3][..nb]);
        plain.copy_from_slice(&blocks[0][..nb * nb]);
        bg.copy_from_slice(&blocks[1][..nb * nb]);
        bgte.copy_from_slice(&blocks[2][..nb * nb]);
    }
    (acc, hb2, b2)
}

/// Value and gradient with the zero counts summed from the run's tables.
///
/// [`evaluate`] without the Hessian, split as the penalised fit splits it: the
/// terms of every cell at a count of zero, then what the positive counts
/// change. At a count of zero a cell's terms depend on `v = ym e / gamma`, so
/// per group of [`ZeroCells`] they are `L` and `A` at
/// `s = ln(ym) + x'beta - ln(gamma)`, and `sum ym e` is `ym exp(x'beta) sum O`
/// exactly. Loose cells are summed one by one.
///
/// ### Params
///
/// * `data` - The gene, as validated by [`GeneData::new`]
/// * `params` - `[beta (n_coef), sigma, phi]`, length [`GeneData::n_params`]
/// * `zeros` - Tables built from this gene's design, offsets and subjects
/// * `scratch` - Built by [`PtmgScratch::tabled`] from this same gene
///
/// ### Returns
///
/// The negative log-likelihood and its gradient, length `n_coef + 2`.
#[allow(clippy::needless_range_loop)]
fn evaluate_tabled(
    data: &GeneData<'_>,
    params: &[f64],
    zeros: &ZeroCells,
    scratch: &mut PtmgScratch<'_>,
) -> (f64, Vec<f64>) {
    let nb = data.n_coef;
    let k = data.n_subjects();
    let beta = &params[..nb];
    let terms = Terms::new(params[nb], params[nb + 1]);
    let gamma = terms.gamma;
    let ln_gamma = terms.log_gamma;
    let row_of = |r: usize| &data.design[r * nb..(r + 1) * nb];
    let dot = |row: &[f64]| row.iter().zip(beta).fold(0.0, |acc, (&x, &b)| acc + x * b);

    // Subject sums of `e` and `x e`, and the `exp` of every cell summed one by
    // one later.
    let mut total = 0.0;
    let mut group_xb = vec![0.0; zeros.groups.len()];
    let mut xexb_f = vec![0.0; k * nb];
    for s in 0..k {
        let xf = &mut xexb_f[s * nb..(s + 1) * nb];
        let mut csx = 0.0;
        for gi in zeros.subject_groups[s]..zeros.subject_groups[s + 1] {
            let g = &zeros.groups[gi];
            let row = row_of(g.row);
            let xb = dot(row);
            group_xb[gi] = xb;
            let e = xb.exp() * g.sum_o;
            csx += e;
            for j in 0..nb {
                xf[j] += row[j] * e;
            }
        }
        for &r in &zeros.loose[zeros.subject_loose[s]..zeros.subject_loose[s + 1]] {
            let row = row_of(r);
            let e = (data.log_offset[r] + dot(row)).exp();
            scratch.extb[r] = e;
            csx += e;
            for j in 0..nb {
                xf[j] += row[j] * e;
            }
        }
        for i in scratch.run[s]..scratch.run[s + 1] {
            let c = data.cells[i];
            let eta = data.log_offset[c] + dot(row_of(c));
            total += eta * data.counts[i];
            scratch.extb[c] = eta.exp();
        }
        scratch.cumsumxtb[s] = csx;
    }

    let subjects = Subjects::new(data, &terms, &scratch.cumsumxtb, &mut total);
    let ymustar = &subjects.ymustar;

    let mut slpey = 0.0;
    let mut gstar_sum = 0.0;
    let mut dbeta_41 = vec![0.0; k * nb];
    let mut d42c = vec![0.0; k];
    for s in 0..k {
        let ym = ymustar[s];
        let ln_ym = ym.ln();
        let d41 = &mut dbeta_41[s * nb..(s + 1) * nb];
        let mut acc = 0.0;
        // `sum ym e` over every cell of the subject.
        total += ym * scratch.cumsumxtb[s];
        for gi in zeros.subject_groups[s]..zeros.subject_groups[s + 1] {
            let g = &zeros.groups[gi];
            let row = row_of(g.row);
            let [l, a] = g.eval::<0, 2>(ln_ym + group_xb[gi] - ln_gamma);
            let sum_log_w = g.n * ln_gamma + l;
            total -= gamma * sum_log_w;
            slpey += sum_log_w;
            gstar_sum += g.n - a;
            // `sum gamma e / w = (gamma / ym) A`.
            let ge = gamma / ym * a;
            acc += ge;
            for j in 0..nb {
                d41[j] += row[j] * ge;
            }
        }
        let mut cell = |r: usize, y: f64, dense: bool| {
            let e = scratch.extb[r];
            let w = ym * e + gamma;
            let log_w = w.ln();
            let gw = if dense { gamma / w } else { y / w };
            if dense {
                total -= gamma * log_w;
                slpey += log_w;
            } else {
                total -= y * log_w;
            }
            gstar_sum += gw;
            let ge = gw * e;
            acc += ge;
            let row = row_of(r);
            for j in 0..nb {
                d41[j] += row[j] * ge;
            }
        };
        for &r in &zeros.loose[zeros.subject_loose[s]..zeros.subject_loose[s + 1]] {
            cell(r, 0.0, true);
        }
        for i in scratch.run[s]..scratch.run[s + 1] {
            cell(data.cells[i], data.counts[i], false);
        }
        d42c[s] = acc - scratch.cumsumxtb[s];
    }

    let parts = gradient_parts(
        data, &terms, &subjects, &xexb_f, &dbeta_41, &d42c, slpey, gstar_sum,
    );
    finish_gradient(data, &terms, total, parts.gradient)
}

/////////////
// Helpers //
/////////////

/// Per-subject count totals for every gene.
///
/// Port of `call_cumsumy`. Counts are gene-major here (the reference is
/// cell-major), so the subject is looked up per non-zero entry.
///
/// ### Params
///
/// * `counts` - Gene-major counts, [`SparseFormat::Csr`] over
///   `(n_genes, n_cells)`
/// * `fid` - Subject boundaries, length `n_subjects + 1`, running to `n_cells`
///
/// ### Returns
///
/// Row-major `n_genes * n_subjects` totals, or [`EdgeErrors`] if the matrix is
/// not gene-major or `fid` does not partition the cells.
pub fn cumsum_y(counts: &CompressedSparse<f64>, fid: &[usize]) -> Result<Vec<f64>, EdgeErrors> {
    if counts.format != SparseFormat::Csr {
        return Err(EdgeErrors::MalformedSparse(
            "cumsum_y needs gene-major counts in CSR form.".to_string(),
        ));
    }
    let n_genes = counts.nrows();
    let n_cells = counts.ncols();
    if fid.len() < 2 {
        return Err(EdgeErrors::TooFewSubjects {
            required: 1,
            got: fid.len().saturating_sub(1),
        });
    }
    let n_subjects = fid.len() - 1;
    if fid[0] != 0 || fid[n_subjects] != n_cells {
        return Err(EdgeErrors::InvalidArgument(format!(
            "fid must run from 0 to {n_cells}; got {} to {}.",
            fid[0], fid[n_subjects]
        )));
    }

    // Cell to subject, so the inner loop is a lookup rather than a search.
    let mut subject_of = vec![0usize; n_cells];
    for s in 0..n_subjects {
        if fid[s + 1] < fid[s] || fid[s + 1] > n_cells {
            return Err(EdgeErrors::InvalidArgument(
                "fid must be non-decreasing and bounded by the cell count.".to_string(),
            ));
        }
        for slot in subject_of.iter_mut().take(fid[s + 1]).skip(fid[s]) {
            *slot = s;
        }
    }

    let mut out = vec![0.0; n_genes * n_subjects];
    out.par_chunks_mut(n_subjects)
        .enumerate()
        .for_each(|(gene, row)| {
            let (indices, values) = counts.outer(gene);
            for (&cell, &value) in indices.iter().zip(values.iter()) {
                if value > 0.0 {
                    row[subject_of[cell as usize]] += value;
                }
            }
        });

    Ok(out)
}

////////////////
// GeneCounts //
////////////////

/// The non-zero counts of one gene, and the summaries NEBULA derives from them.
#[derive(Clone, Debug)]
pub struct GeneCounts {
    /// Cell index of each non-zero count, increasing.
    pub cells: Vec<usize>,
    /// The non-zero counts themselves.
    pub counts: Vec<f64>,
    /// Mean count across all cells, zeros included. `mct` in the reference, and
    /// what the intercept is initialised from.
    pub mean_count: f64,
    /// How many counts equal one.
    pub n_one: usize,
    /// How many counts equal two.
    pub n_two: usize,
}

/// Extracts one gene's non-zero counts.
///
/// Port of `call_posindy`. Counts of one and two are tallied separately for the
/// closed forms in the gamma terms.
///
/// ### Params
///
/// * `counts` - Gene-major counts, [`SparseFormat::Csr`] over
///   `(n_genes, n_cells)`
/// * `gene` - Row index of the gene
///
/// ### Returns
///
/// The gene's non-zero counts with their cells and summaries, or [`EdgeErrors`]
/// if the matrix is not gene-major or the gene index is out of range.
pub fn positive_indices(
    counts: &CompressedSparse<f64>,
    gene: usize,
) -> Result<GeneCounts, EdgeErrors> {
    if counts.format != SparseFormat::Csr {
        return Err(EdgeErrors::MalformedSparse(
            "positive_indices needs gene-major counts in CSR form.".to_string(),
        ));
    }
    let n_genes = counts.nrows();
    if gene >= n_genes {
        return Err(EdgeErrors::InvalidArgument(format!(
            "Gene index {gene} is out of range for {n_genes} genes."
        )));
    }
    let n_cells = counts.ncols();
    if n_cells == 0 {
        return Err(EdgeErrors::MustBePositive("n_cells".to_string()));
    }

    let (indices, values) = counts.outer(gene);
    let mut cells = Vec::with_capacity(indices.len());
    let mut kept = Vec::with_capacity(indices.len());
    let mut sum = 0.0;
    let mut n_one = 0;
    let mut n_two = 0;

    for (&cell, &value) in indices.iter().zip(values.iter()) {
        if value > 0.0 {
            cells.push(cell as usize);
            kept.push(value);
            sum += value;
            if value == 1.0 {
                n_one += 1;
            } else if value == 2.0 {
                n_two += 1;
            }
        }
    }

    Ok(GeneCounts {
        cells,
        counts: kept,
        mean_count: sum / (n_cells as f64),
        n_one,
        n_two,
    })
}

/// Centres and scales a design matrix.
///
/// Port of `center_m`. Each column is centred and divided by its population
/// standard deviation. A column with no spread is not scaled: it becomes ones if
/// its first entry is non-zero (the intercept), otherwise it stays zero with a
/// standard deviation of `-1` so the caller can drop it.
///
/// ### Params
///
/// * `design` - Row-major `n_rows * n_cols` design matrix
/// * `n_rows` - Number of cells
/// * `n_cols` - Number of columns
///
/// ### Returns
///
/// The centred and scaled design, row-major and the same shape, together with
/// the per-column standard deviations, or [`EdgeErrors::ShapeMismatch`] if the
/// slice length disagrees with the shape.
pub fn centre_design(
    design: &[f64],
    n_rows: usize,
    n_cols: usize,
) -> Result<(Vec<f64>, Vec<f64>), EdgeErrors> {
    if design.len() != n_rows * n_cols {
        return Err(EdgeErrors::ShapeMismatch {
            expected: (n_rows, n_cols),
            got: (design.len(), 1),
        });
    }
    if n_rows == 0 || n_cols == 0 {
        return Err(EdgeErrors::MustBePositive("design shape".to_string()));
    }

    let mut centred = vec![0.0; n_rows * n_cols];
    let mut sds = vec![0.0; n_cols];

    for j in 0..n_cols {
        let mut mean = 0.0;
        for i in 0..n_rows {
            mean += design[i * n_cols + j];
        }
        mean /= n_rows as f64;

        let mut variance = 0.0;
        for i in 0..n_rows {
            let d = design[i * n_cols + j] - mean;
            centred[i * n_cols + j] = d;
            variance += d * d;
        }
        let sd = (variance / n_rows as f64).sqrt();

        if sd > 0.0 {
            for i in 0..n_rows {
                centred[i * n_cols + j] /= sd;
            }
            sds[j] = sd;
        } else if design[j] != 0.0 {
            for i in 0..n_rows {
                centred[i * n_cols + j] = 1.0;
            }
            sds[j] = sd;
        } else {
            sds[j] = -1.0;
        }
    }

    Ok((centred, sds))
}

/// Log offsets and the summaries the NEBULA driver keys off.
#[derive(Clone, Debug)]
pub struct OffsetSummary {
    /// Log of the supplied offsets, one per cell.
    pub log_offset: Vec<f64>,
    /// Mean of `log_offset`, or zero when no offset was supplied.
    pub mean_log_offset: f64,
    /// Mean of the offsets on their original scale.
    pub mean_offset: f64,
    /// Coefficient of variation of the offsets on their original scale.
    pub cv: f64,
}

/// Prepares the offsets and measures their spread.
///
/// Port of `cv_offset`. With no offset every cell gets one (log zero, no
/// spread). The coefficient of variation decides whether NEBULA-LN's
/// approximation is safe.
///
/// ### Params
///
/// * `offset` - Offsets on their original scale, or `None` for all ones
/// * `n_cells` - Number of cells
///
/// ### Returns
///
/// The log offsets and their summaries, or [`EdgeErrors`] if the offset length
/// disagrees with `n_cells`.
pub fn offset_summary(offset: Option<&[f64]>, n_cells: usize) -> Result<OffsetSummary, EdgeErrors> {
    if n_cells == 0 {
        return Err(EdgeErrors::MustBePositive("n_cells".to_string()));
    }
    if let Some(o) = offset
        && o.len() != n_cells
    {
        return Err(EdgeErrors::LengthMismatch {
            name: "offset",
            expected: n_cells,
            got: o.len(),
        });
    }

    let values: Vec<f64> = match offset {
        Some(o) => o.to_vec(),
        None => vec![1.0; n_cells],
    };
    let mean_offset = match offset {
        Some(_) => values.iter().sum::<f64>() / (n_cells as f64),
        None => 1.0,
    };

    let mut cv = 0.0;
    if mean_offset > 0.0 {
        let ss: f64 = values.iter().map(|v| (v - mean_offset).powi(2)).sum();
        cv = (ss / n_cells as f64).sqrt() / mean_offset;
    }

    let log_offset: Vec<f64> = values.iter().map(|v| v.ln()).collect();
    let mean_log_offset = match offset {
        Some(_) => log_offset.iter().sum::<f64>() / (n_cells as f64),
        None => 0.0,
    };

    Ok(OffsetSummary {
        log_offset,
        mean_log_offset,
        mean_offset,
        cv,
    })
}

/// Squared coefficient of variation of the fitted cell-level means.
///
/// Port of `get_cv`. Only cell-level columns contribute. This is the squared CV,
/// unlike [`offset_summary`]'s `cv` (the reference's own asymmetry).
///
/// ### Params
///
/// * `log_offset` - Log offset per cell, length `n_rows`
/// * `design` - Row-major `n_rows * n_cols` design matrix
/// * `n_rows` - Number of cells
/// * `n_cols` - Number of design columns
/// * `beta` - Coefficients, length `n_cols`
/// * `cell_columns` - Indices of the cell-level columns to include
///
/// ### Returns
///
/// `var(mu) / mean(mu)^2` over cells, or [`EdgeErrors`] if the shapes disagree
/// or a column index is out of range.
pub fn design_cv(
    log_offset: &[f64],
    design: &[f64],
    n_rows: usize,
    n_cols: usize,
    beta: &[f64],
    cell_columns: &[usize],
) -> Result<f64, EdgeErrors> {
    if n_rows == 0 || n_cols == 0 {
        return Err(EdgeErrors::MustBePositive("design shape".to_string()));
    }
    if design.len() != n_rows * n_cols {
        return Err(EdgeErrors::ShapeMismatch {
            expected: (n_rows, n_cols),
            got: (design.len(), 1),
        });
    }
    if log_offset.len() != n_rows {
        return Err(EdgeErrors::LengthMismatch {
            name: "log_offset",
            expected: n_rows,
            got: log_offset.len(),
        });
    }
    if beta.len() != n_cols {
        return Err(EdgeErrors::LengthMismatch {
            name: "beta",
            expected: n_cols,
            got: beta.len(),
        });
    }
    for &c in cell_columns {
        if c >= n_cols {
            return Err(EdgeErrors::CoefOutOfRange {
                index: c,
                n_coef: n_cols,
            });
        }
    }

    let mut mu = vec![0.0; n_rows];
    for i in 0..n_rows {
        let mut eta = log_offset[i];
        for &c in cell_columns {
            eta += design[i * n_cols + c] * beta[c];
        }
        mu[i] = eta.exp();
    }

    let mean = mu.iter().sum::<f64>() / (n_rows as f64);
    if mean <= 0.0 {
        return Ok(0.0);
    }
    let ss: f64 = mu.iter().map(|v| (v - mean).powi(2)).sum();

    Ok(ss / (n_rows as f64) / (mean * mean))
}

/// Flags the design columns that vary within a subject.
///
/// Port of `get_cell`. A column constant within every subject is subject-level
/// and absorbed by the random effect; the rest are cell-level.
///
/// ### Params
///
/// * `design` - Row-major `n_rows * n_cols` design matrix
/// * `n_rows` - Number of cells
/// * `n_cols` - Number of design columns
/// * `fid` - Subject boundaries, length `n_subjects + 1`
///
/// ### Returns
///
/// One flag per column, true where the column varies within at least one
/// subject, or [`EdgeErrors`] if the shapes disagree.
pub fn cell_level_columns(
    design: &[f64],
    n_rows: usize,
    n_cols: usize,
    fid: &[usize],
) -> Result<Vec<bool>, EdgeErrors> {
    if n_rows == 0 || n_cols == 0 {
        return Err(EdgeErrors::MustBePositive("design shape".to_string()));
    }
    if design.len() != n_rows * n_cols {
        return Err(EdgeErrors::ShapeMismatch {
            expected: (n_rows, n_cols),
            got: (design.len(), 1),
        });
    }
    if fid.len() < 2 {
        return Err(EdgeErrors::TooFewSubjects {
            required: 1,
            got: fid.len().saturating_sub(1),
        });
    }
    if fid[0] != 0 || fid[fid.len() - 1] != n_rows {
        return Err(EdgeErrors::InvalidArgument(format!(
            "fid must run from 0 to {n_rows}; got {} to {}.",
            fid[0],
            fid[fid.len() - 1]
        )));
    }

    let mut flags = vec![false; n_cols];
    for (j, flag) in flags.iter_mut().enumerate() {
        'subjects: for s in 0..fid.len() - 1 {
            let first = design[fid[s] * n_cols + j];
            for i in fid[s]..fid[s + 1] {
                if design[i * n_cols + j] != first {
                    *flag = true;
                    break 'subjects;
                }
            }
        }
    }

    Ok(flags)
}

///////////
// Tests //
///////////

// Reference literals are R's `%.17e` output, pasted verbatim.
#[cfg(test)]
#[allow(clippy::excessive_precision)]
mod tests {
    use super::*;
    use approx::assert_relative_eq;

    /// Fixture A: six cells over three subjects of sizes two, three and one, so
    /// the last subject is a single cell. Two design columns, and two cells with
    /// a zero count.
    ///
    /// ### Returns
    ///
    /// Design, log offsets, non-zero counts, their cells, subject totals, `fid`.
    #[allow(clippy::type_complexity)]
    fn fixture_a() -> (
        Vec<f64>,
        Vec<f64>,
        Vec<f64>,
        Vec<usize>,
        Vec<f64>,
        Vec<usize>,
    ) {
        let x = [0.0, 1.0, 0.5, -0.5, 0.25, -1.0];
        let design: Vec<f64> = x.iter().flat_map(|&v| [1.0, v]).collect();
        let log_offset = vec![0.0, 0.25, -0.5, 0.125, 0.75, -0.25];
        // Full counts 3, 0, 1, 2, 0, 5.
        let counts = vec![3.0, 1.0, 2.0, 5.0];
        let cells = vec![0, 2, 3, 5];
        let subject_totals = vec![3.0, 3.0, 5.0];
        let fid = vec![0, 2, 5, 6];
        (design, log_offset, counts, cells, subject_totals, fid)
    }

    /// Fixture B: eight cells over three subjects of sizes three, one and four.
    /// Three design columns, and the middle subject has a total count of zero.
    ///
    /// ### Returns
    ///
    /// Design, log offsets, non-zero counts, their cells, subject totals, `fid`.
    #[allow(clippy::type_complexity)]
    fn fixture_b() -> (
        Vec<f64>,
        Vec<f64>,
        Vec<f64>,
        Vec<usize>,
        Vec<f64>,
        Vec<usize>,
    ) {
        let x1 = [0.0, 1.0, 0.5, -0.5, 0.25, -1.0, 0.75, -0.25];
        let x2 = [0.5, -0.25, 0.0, 0.75, -1.0, 0.25, 1.0, -0.5];
        let design: Vec<f64> = x1
            .iter()
            .zip(x2.iter())
            .flat_map(|(&a, &b)| [1.0, a, b])
            .collect();
        let log_offset = vec![0.5, -0.25, 0.125, 0.0, 0.375, -0.75, 0.25, 1.0];
        // Full counts 0, 2, 1, 0, 4, 0, 1, 7.
        let counts = vec![2.0, 1.0, 4.0, 1.0, 7.0];
        let cells = vec![1, 2, 4, 6, 7];
        let subject_totals = vec![3.0, 0.0, 12.0];
        let fid = vec![0, 3, 4, 8];
        (design, log_offset, counts, cells, subject_totals, fid)
    }

    /// Checks all four entry points against one reference point.
    ///
    /// ### Params
    ///
    /// * `fixture` - Which fixture builder to use
    /// * `params` - `[beta, sigma, phi]`
    /// * `value` - Reference negative log-likelihood
    /// * `gradient` - Reference gradient
    /// * `hessian` - Reference Hessian, row-major
    /// * `gradient_tol` - Relative tolerance on the gradient
    /// * `hessian_tol` - Relative tolerance on the Hessian
    #[allow(clippy::type_complexity)]
    fn check(
        fixture: fn() -> (
            Vec<f64>,
            Vec<f64>,
            Vec<f64>,
            Vec<usize>,
            Vec<f64>,
            Vec<usize>,
        ),
        params: &[f64],
        value: f64,
        gradient: &[f64],
        hessian: &[f64],
        gradient_tol: f64,
        hessian_tol: f64,
    ) {
        let (design, log_offset, counts, cells, subject_totals, fid) = fixture();
        let n_coef = params.len() - 2;
        let data = GeneData::new(
            &design,
            &log_offset,
            &counts,
            &cells,
            &subject_totals,
            &fid,
            n_coef,
        )
        .unwrap();

        assert_relative_eq!(
            ptmg_neg_log_likelihood(&data, params),
            value,
            max_relative = 1e-10
        );

        let (v, g) = ptmg_value_and_gradient(&data, params);
        assert_relative_eq!(v, value, max_relative = 1e-10);
        for (got, want) in g.iter().zip(gradient.iter()) {
            assert_relative_eq!(got, want, max_relative = gradient_tol, epsilon = 1e-12);
        }

        let only_gradient = ptmg_gradient(&data, params);
        assert_eq!(only_gradient, g);

        let (v, g, h) = ptmg_value_gradient_hessian(&data, params);
        assert_relative_eq!(v, value, max_relative = 1e-10);
        for (got, want) in g.iter().zip(gradient.iter()) {
            assert_relative_eq!(got, want, max_relative = gradient_tol, epsilon = 1e-12);
        }
        for (got, want) in h.iter().zip(hessian.iter()) {
            assert_relative_eq!(got, want, max_relative = hessian_tol, epsilon = 1e-12);
        }
    }

    // Reference values from nebula 1.5.8, via
    //   Rscript -e 'library(nebula); nebula:::ptmg_ll_der_hes_eigen(X, offset, Y,
    //     fid-1, cumsumy, posind-1, posindy, nb, nind, k, beta, sigma)'
    // with the gamma-function corrections of R/ptmg.R applied (the arithmetic of
    // nebula:::ptmg_ll_der plus the Hessian lines of nebula:::ptmg_ll_der_hes2,
    // without its log-scale chain rule).
    //
    // Deliberate departure: the corrections call base R's lgamma, digamma and
    // trigamma, where R/ptmg.R calls Rfast's. Rfast::Digamma is ~5e-11 off base
    // R and Rfast::Trigamma ~1.4e-9, and `alpha_pr` multiplies that by up to 1e8
    // at the lower bound on sigma. gamma.rs tracks base R. Against the Rfast
    // path these references move the gradient by at most 1e-8 relative.
    //
    // Tolerances: 1e-10 on value and gradient, 1e-8 on the Hessian; at the
    // `sigma = 1e-4` lower bound 1e-8 and 1e-6. There `alpha_pr` is ~1e8 and
    // `alpha_pr^2` ~1e16, and the sigma derivatives cancel down to ~1e2 and 1e3,
    // so the last digits are set by summation order and the last bit of
    // `digamma`. Achieved: 8e-13 (value), 1.4e-14 (gradient), 9e-13 (Hessian)
    // away from the bound; 4e-9 and 2e-7 at it.

    #[test]
    fn test_matches_nebula_at_an_interior_point() {
        check(
            fixture_a,
            &[0.5, -0.25, 0.5, 2.0],
            5.68919227265053351e+00,
            &[
                1.23230096727342087e-01,
                2.14374578345666933e+00,
                1.96862959139235372e+00,
                -1.46930062054576371e-01,
            ],
            &[
                2.46157581206725151e+00,
                -7.12808069390344246e-01,
                2.39173560205083291e-01,
                -2.69050975196676243e-02,
                -7.12808069390344246e-01,
                2.05333290435219373e+00,
                -1.34438513739700594e+00,
                3.31302601132097918e-01,
                2.39173560205083291e-01,
                -1.34438513739700594e+00,
                3.91220510021810242e-02,
                -1.42705917336173937e-01,
                -2.69050975196676243e-02,
                3.31302601132097918e-01,
                -1.42705917336173937e-01,
                1.68849516491021756e-01,
            ],
            1e-10,
            1e-8,
        );
    }

    /// Both variance parameters at their lower bound, where `alpha` is of order
    /// `1e4` and the cell-level shape is of order `1e-4`.
    #[test]
    fn test_matches_nebula_at_the_lower_bounds() {
        check(
            fixture_a,
            &[1.25, 0.75, 1e-4, 1e-4],
            3.29896692291768900e+01,
            &[
                2.96406695818021149e-02,
                1.31693015791212176e-02,
                1.30305458759699832e+02,
                -3.99414489355896367e+04,
            ],
            &[
                6.86679893234001110e-02,
                2.79446317794878468e-02,
                2.95430767492699658e+02,
                1.82090109226938218e-01,
                2.79446317794878468e-02,
                2.14302885561563627e-02,
                1.27035501836412706e+02,
                4.34573972169806044e+00,
                2.95430767492699658e+02,
                1.27035501836412706e+02,
                -4.52776105356216431e+03,
                -4.46684522479116310e+01,
                1.82090109226938218e-01,
                4.34573972169806044e+00,
                -4.46684522479116310e+01,
                3.99940003972201765e+08,
            ],
            1e-8,
            1e-6,
        );
    }

    /// Near the upper bounds, `sigma = 8` and `phi = 512`.
    #[test]
    fn test_matches_nebula_near_the_upper_bounds() {
        check(
            fixture_a,
            &[-0.5, 0.125, 8.0, 512.0],
            2.65308778580696725e+01,
            &[
                9.12793250876120510e-04,
                2.53045547739533250e+00,
                2.99186987816233607e+00,
                -3.60045252223439860e-07,
            ],
            &[
                9.39311823971149956e-05,
                -6.57328696371346430e-05,
                -8.66127875657092474e-04,
                3.13400144628472963e-11,
                -6.57328696371346430e-05,
                1.13575813292771000e+00,
                9.98036672113045331e-06,
                1.32876795203876600e-05,
                -8.66127875657092474e-04,
                9.98036672113045331e-06,
                6.69606190840976723e-03,
                -1.05310011543468632e-09,
                3.13400144628472963e-11,
                1.32876795203876600e-05,
                -1.05310011543468632e-09,
                1.36270078923174387e-09,
            ],
            1e-10,
            1e-8,
        );
    }

    /// A zero coefficient vector, which is where the driver starts every gene.
    #[test]
    fn test_matches_nebula_at_zero_coefficients() {
        check(
            fixture_a,
            &[0.0, 0.0, 2.0, 0.25],
            1.12063514415980094e+01,
            &[
                -3.81456775084743072e-03,
                5.74173792151410289e-01,
                2.38393039589584355e+00,
                -7.26967190944776576e+00,
            ],
            &[
                4.23398202512259736e-01,
                -2.57145238578597779e-01,
                1.94645018901316025e-01,
                -1.15565351220650037e-02,
                -2.57145238578597779e-01,
                7.01861921458433624e-01,
                -3.59445802666915337e-01,
                1.18270172064454471e+00,
                1.94645018901316025e-01,
                -3.59445802666915337e-01,
                4.36597493935281022e-01,
                1.30165671415671189e-02,
                -1.15565351220650037e-02,
                1.18270172064454471e+00,
                1.30165671415671189e-02,
                4.50117573398768016e+01,
            ],
            1e-10,
            1e-8,
        );
    }

    /// Three design columns, and a subject whose counts are all zero.
    #[test]
    fn test_matches_nebula_with_three_columns() {
        check(
            fixture_b,
            &[0.25, -0.5, 0.75, 1.5, 3.0],
            7.46199076417621399e+00,
            &[
                5.01350696391867245e-01,
                -2.04257563318932167e+00,
                5.34153324797603979e+00,
                1.07011331098338225e+00,
                2.53578958023772838e-01,
            ],
            &[
                3.48740746684502056e-01,
                2.54876819316399161e-02,
                1.45533761085350766e-02,
                -4.41548536268330516e-01,
                3.91566084164468586e-03,
                2.54876819316399161e-02,
                2.09570270083503551e+00,
                2.64460193688231204e-01,
                1.02670366878276728e-01,
                -5.76636821621331688e-02,
                1.45533761085350766e-02,
                2.64460193688231204e-01,
                3.84459569945194923e+00,
                -2.78923375808989293e-01,
                6.39405403297957253e-01,
                -4.41548536268330516e-01,
                1.02670366878276728e-01,
                -2.78923375808989293e-01,
                3.66772893546357448e-01,
                -1.49218225719068356e-02,
                3.91566084164468586e-03,
                -5.76636821621331688e-02,
                6.39405403297957253e-01,
                -1.49218225719068356e-02,
                -7.22442940506229370e-02,
            ],
            1e-10,
            1e-8,
        );
    }

    /// Three columns with `sigma` at its lower bound and `phi` at its upper one.
    #[test]
    fn test_matches_nebula_with_three_columns_at_mixed_bounds() {
        check(
            fixture_b,
            &[-1.0, 0.5, 0.25, 1e-4, 1000.0],
            1.13584681060607693e+01,
            &[
                -1.07326878153205776e+01,
                -1.48756816066168707e+00,
                7.53628836395186852e+00,
                -4.62797267317655496e+01,
                1.95947859886434761e-05,
            ],
            &[
                4.26381285975427904e+00,
                1.01145653996075602e+00,
                5.35701406304531225e-01,
                2.61662160722283907e+01,
                -6.94406121181113076e-06,
                1.01145653996075602e+00,
                1.35186752241846908e+00,
                3.72773749626097717e-01,
                6.11253074542543828e+00,
                -1.69452414787726765e-07,
                5.35701406304531225e-01,
                3.72773749626097717e-01,
                1.95363290656862110e+00,
                9.21736461732613921e-01,
                4.48552067926583857e-06,
                2.61662160722283907e+01,
                6.11253074542543828e+00,
                9.21736461732613921e-01,
                4.35795623779296875e+02,
                -6.84104019316582210e-05,
                -6.94406121181113076e-06,
                -1.69452414787726765e-07,
                4.48552067926583857e-06,
                -6.84104019316582210e-05,
                -3.90898716648611179e-08,
            ],
            1e-8,
            1e-6,
        );
    }

    /// Three columns with `sigma` at its upper bound and `phi` at its lower one.
    #[test]
    fn test_matches_nebula_with_three_columns_at_the_other_bounds() {
        check(
            fixture_b,
            &[0.75, -0.25, 0.5, 10.0, 1e-4],
            5.76091855400661359e+01,
            &[
                1.35773392448967911e-04,
                -3.59794963305848370e-04,
                4.96443874810736929e-04,
                1.99820940875096320e+00,
                -4.99346186500725235e+04,
            ],
            &[
                4.32580556442951729e-07,
                4.91633557431160259e-08,
                2.83994566060920446e-08,
                -1.35563254538641632e-04,
                -9.85888077842389521e-08,
                4.91633557431160259e-08,
                3.29393259974663849e-04,
                -2.02106061975382121e-04,
                7.65773286491442306e-06,
                -3.52089986509287156e+00,
                2.83994566060920446e-08,
                -2.02106061975382121e-04,
                4.54276735397987554e-04,
                -4.52632990189595600e-05,
                4.51109987283559022e+00,
                -1.35563254538641632e-04,
                7.65773286491442306e-06,
                -4.52632990189595600e-05,
                1.57701505319085555e-03,
                -9.74640513908427847e-02,
                -9.85888077842389521e-08,
                -3.52089986509287156e+00,
                4.51109987283559022e+00,
                -9.74640513908427847e-02,
                4.99929026939662099e+08,
            ],
            1e-10,
            1e-8,
        );
    }

    /// The Hessian must be symmetric by construction, at every point.
    #[test]
    fn test_hessian_is_symmetric() {
        let (design, log_offset, counts, cells, subject_totals, fid) = fixture_b();
        let data = GeneData::new(
            &design,
            &log_offset,
            &counts,
            &cells,
            &subject_totals,
            &fid,
            3,
        )
        .unwrap();
        let params = [0.25, -0.5, 0.75, 1.5, 3.0];
        let (_, _, h) = ptmg_value_gradient_hessian(&data, &params);
        let n = 5;
        for i in 0..n {
            for j in 0..n {
                assert_relative_eq!(h[i * n + j], h[j * n + i], max_relative = 1e-14);
            }
        }
    }

    /// The analytic gradient must agree with a central difference of the value.
    #[test]
    fn test_gradient_agrees_with_finite_differences() {
        let (design, log_offset, counts, cells, subject_totals, fid) = fixture_a();
        let data = GeneData::new(
            &design,
            &log_offset,
            &counts,
            &cells,
            &subject_totals,
            &fid,
            2,
        )
        .unwrap();
        let params = [0.5, -0.25, 0.5, 2.0];
        let (_, gradient) = ptmg_value_and_gradient(&data, &params);

        for j in 0..params.len() {
            let step = 1e-6 * params[j].abs().max(1.0);
            let mut up = params;
            let mut down = params;
            up[j] += step;
            down[j] -= step;
            let fd = (ptmg_neg_log_likelihood(&data, &up) - ptmg_neg_log_likelihood(&data, &down))
                / (2.0 * step);
            assert_relative_eq!(gradient[j], fd, max_relative = 1e-6, epsilon = 1e-8);
        }
    }

    /// And the Hessian with a central difference of the gradient.
    #[test]
    fn test_hessian_agrees_with_finite_differences() {
        let (design, log_offset, counts, cells, subject_totals, fid) = fixture_a();
        let data = GeneData::new(
            &design,
            &log_offset,
            &counts,
            &cells,
            &subject_totals,
            &fid,
            2,
        )
        .unwrap();
        let params = [0.5, -0.25, 0.5, 2.0];
        let (_, _, hessian) = ptmg_value_gradient_hessian(&data, &params);
        let n = params.len();

        for j in 0..n {
            let step = 1e-6 * params[j].abs().max(1.0);
            let mut up = params;
            let mut down = params;
            up[j] += step;
            down[j] -= step;
            let g_up = ptmg_gradient(&data, &up);
            let g_down = ptmg_gradient(&data, &down);
            for i in 0..n {
                let fd = (g_up[i] - g_down[i]) / (2.0 * step);
                assert_relative_eq!(hessian[i * n + j], fd, max_relative = 1e-5, epsilon = 1e-7);
            }
        }
    }

    #[test]
    fn test_rejects_a_design_of_the_wrong_length() {
        let (_, log_offset, counts, cells, subject_totals, fid) = fixture_a();
        let design = vec![1.0; 5];
        let err = GeneData::new(
            &design,
            &log_offset,
            &counts,
            &cells,
            &subject_totals,
            &fid,
            2,
        )
        .unwrap_err();
        assert!(matches!(
            err,
            EdgeErrors::LengthMismatch { name: "design", .. }
        ));
    }

    #[test]
    fn test_rejects_a_fid_that_does_not_cover_the_cells() {
        let (design, log_offset, counts, cells, subject_totals, _) = fixture_a();
        let fid = vec![0, 2, 5, 7];
        let err = GeneData::new(
            &design,
            &log_offset,
            &counts,
            &cells,
            &subject_totals,
            &fid,
            2,
        )
        .unwrap_err();
        assert!(matches!(err, EdgeErrors::InvalidArgument(_)));
    }

    #[test]
    fn test_rejects_an_empty_subject() {
        let (design, log_offset, counts, cells, _, _) = fixture_a();
        let fid = vec![0, 2, 2, 6];
        let subject_totals = vec![3.0, 0.0, 8.0];
        let err = GeneData::new(
            &design,
            &log_offset,
            &counts,
            &cells,
            &subject_totals,
            &fid,
            2,
        )
        .unwrap_err();
        assert!(matches!(err, EdgeErrors::InvalidArgument(_)));
    }

    #[test]
    fn test_rejects_unsorted_cell_indices() {
        let (design, log_offset, counts, _, subject_totals, fid) = fixture_a();
        let cells = vec![0, 3, 2, 5];
        let err = GeneData::new(
            &design,
            &log_offset,
            &counts,
            &cells,
            &subject_totals,
            &fid,
            2,
        )
        .unwrap_err();
        assert!(matches!(err, EdgeErrors::InvalidArgument(_)));
    }

    // ------- //
    // helpers //
    // ------- //

    /// A three-gene, six-cell count matrix in gene-major CSR form.
    ///
    /// ### Returns
    ///
    /// The matrix, with dense rows `3 0 1 2 0 5`, `0 0 0 0 0 0` and
    /// `1 1 2 0 4 0`.
    fn count_fixture() -> CompressedSparse<f64> {
        let dense = [
            3.0, 0.0, 1.0, 2.0, 0.0, 5.0, //
            0.0, 0.0, 0.0, 0.0, 0.0, 0.0, //
            1.0, 1.0, 2.0, 0.0, 4.0, 0.0,
        ];
        CompressedSparse::from_dense(&dense, 3, 6, SparseFormat::Csr, |v: &f64| *v == 0.0).unwrap()
    }

    /// Reference from
    ///   Rscript -e 'library(nebula); library(Matrix);
    ///     m <- Matrix(c(3,0,1, 0,0,1, 1,0,2, 2,0,0, 0,0,4, 5,0,0), 3, 6, sparse=TRUE);
    ///     nebula:::call_cumsumy(m, c(0,2,5,6), 3, 3)'
    #[test]
    fn test_cumsum_y_matches_nebula() {
        let counts = count_fixture();
        let fid = vec![0, 2, 5, 6];
        let got = cumsum_y(&counts, &fid).unwrap();
        assert_eq!(got, vec![3.0, 3.0, 5.0, 0.0, 0.0, 0.0, 2.0, 6.0, 0.0]);
    }

    #[test]
    fn test_cumsum_y_rejects_cell_major_counts() {
        let counts = count_fixture().transpose();
        let fid = vec![0, 2, 5, 6];
        assert!(matches!(
            cumsum_y(&counts, &fid).unwrap_err(),
            EdgeErrors::MalformedSparse(_)
        ));
    }

    /// Reference from
    ///   Rscript -e 'library(nebula); library(Matrix);
    ///     m <- Matrix(c(3,0,1, 0,0,1, 1,0,2, 2,0,0, 0,0,4, 5,0,0), 3, 6, sparse=TRUE);
    ///     nebula:::call_posindy(t(m), 2, 6)'
    /// which is gene three, zero-based, over six cells.
    #[test]
    fn test_positive_indices_matches_nebula() {
        let counts = count_fixture();
        let got = positive_indices(&counts, 2).unwrap();
        assert_eq!(got.cells, vec![0, 1, 2, 4]);
        assert_eq!(got.counts, vec![1.0, 1.0, 2.0, 4.0]);
        assert_relative_eq!(got.mean_count, 8.0 / 6.0, max_relative = 1e-15);
        assert_eq!(got.n_one, 2);
        assert_eq!(got.n_two, 1);
    }

    #[test]
    fn test_positive_indices_on_an_all_zero_gene() {
        let counts = count_fixture();
        let got = positive_indices(&counts, 1).unwrap();
        assert!(got.cells.is_empty());
        assert_eq!(got.mean_count, 0.0);
        assert_eq!(got.n_one, 0);
        assert_eq!(got.n_two, 0);
    }

    #[test]
    fn test_positive_indices_rejects_a_gene_out_of_range() {
        let counts = count_fixture();
        assert!(matches!(
            positive_indices(&counts, 3).unwrap_err(),
            EdgeErrors::InvalidArgument(_)
        ));
    }

    /// Reference from
    ///   Rscript -e 'library(nebula);
    ///     X <- cbind(1, c(0,1,0.5,-0.5,0.25,-1), 0);
    ///     nebula:::center_m(X)'
    /// The intercept has no spread but a non-zero first entry, so it comes back
    /// as ones; the all-zero column is flagged with a standard deviation of -1.
    #[test]
    fn test_centre_design_matches_nebula() {
        let x = [0.0, 1.0, 0.5, -0.5, 0.25, -1.0];
        let design: Vec<f64> = x.iter().flat_map(|&v| [1.0, v, 0.0]).collect();
        let (centred, sds) = centre_design(&design, 6, 3).unwrap();

        assert_relative_eq!(sds[0], 0.0, epsilon = 1e-15);
        assert_relative_eq!(sds[1], 6.52186493437438730e-01, max_relative = 1e-14);
        assert_relative_eq!(sds[2], -1.0, max_relative = 1e-15);

        let mean = x.iter().sum::<f64>() / 6.0;
        for (i, &v) in x.iter().enumerate() {
            assert_relative_eq!(centred[i * 3], 1.0, max_relative = 1e-15);
            assert_relative_eq!(
                centred[i * 3 + 1],
                (v - mean) / sds[1],
                max_relative = 1e-13
            );
            assert_relative_eq!(centred[i * 3 + 2], 0.0, epsilon = 1e-15);
        }
    }

    #[test]
    fn test_centre_design_rejects_a_shape_mismatch() {
        assert!(matches!(
            centre_design(&[1.0, 2.0, 3.0], 2, 2).unwrap_err(),
            EdgeErrors::ShapeMismatch { .. }
        ));
    }

    /// Reference from
    ///   Rscript -e 'library(nebula);
    ///     nebula:::cv_offset(c(1,2,4,8), 1L, 4L)'
    #[test]
    fn test_offset_summary_matches_nebula() {
        let got = offset_summary(Some(&[1.0, 2.0, 4.0, 8.0]), 4).unwrap();
        assert_relative_eq!(got.mean_offset, 3.75, max_relative = 1e-15);
        assert_relative_eq!(got.cv, 7.14920352984240504e-01, max_relative = 1e-13);
        assert_relative_eq!(
            got.mean_log_offset,
            1.03972077083991787e+00,
            max_relative = 1e-13
        );
        for (got, want) in got
            .log_offset
            .iter()
            .zip([0.0, 2.0_f64.ln(), 4.0_f64.ln(), 8.0_f64.ln()].iter())
        {
            assert_relative_eq!(got, want, max_relative = 1e-15);
        }
    }

    /// With no offset every cell gets one, so the logs are zero and there is no
    /// spread. This is `input == 0` in the reference.
    #[test]
    fn test_offset_summary_without_an_offset() {
        let got = offset_summary(None, 4).unwrap();
        assert_eq!(got.log_offset, vec![0.0; 4]);
        assert_eq!(got.mean_log_offset, 0.0);
        assert_eq!(got.mean_offset, 1.0);
        assert_eq!(got.cv, 0.0);
    }

    #[test]
    fn test_offset_summary_rejects_a_length_mismatch() {
        assert!(matches!(
            offset_summary(Some(&[1.0, 2.0]), 3).unwrap_err(),
            EdgeErrors::LengthMismatch { .. }
        ));
    }

    /// Reference from
    ///   Rscript -e 'library(nebula);
    ///     X <- cbind(1, c(0,1,0.5,-0.5,0.25,-1));
    ///     nebula:::get_cv(rep(0,6), X, c(0.25,0.5), 2L, 1L, 6L)'
    /// which includes only column two, one-based.
    #[test]
    fn test_design_cv_matches_nebula() {
        let x = [0.0, 1.0, 0.5, -0.5, 0.25, -1.0];
        let design: Vec<f64> = x.iter().flat_map(|&v| [1.0, v]).collect();
        let got = design_cv(&[0.0; 6], &design, 6, 2, &[0.25, 0.5], &[1]).unwrap();
        assert_relative_eq!(got, 9.93386285703017347e-02, max_relative = 1e-12);
    }

    /// With no cell-level columns every fitted mean is the offset, so a constant
    /// offset has no spread at all.
    #[test]
    fn test_design_cv_with_no_cell_columns() {
        let design = vec![1.0; 6];
        let got = design_cv(&[0.0; 6], &design, 6, 1, &[0.5], &[]).unwrap();
        assert_relative_eq!(got, 0.0, epsilon = 1e-15);
    }

    #[test]
    fn test_design_cv_rejects_a_column_out_of_range() {
        let design = vec![1.0; 6];
        assert!(matches!(
            design_cv(&[0.0; 6], &design, 6, 1, &[0.5], &[1]).unwrap_err(),
            EdgeErrors::CoefOutOfRange { .. }
        ));
    }

    /// Reference from
    ///   Rscript -e 'library(nebula);
    ///     X <- cbind(1, c(0,1,0.5,-0.5,0.25,-1), c(1,1,2,2,2,3));
    ///     nebula:::get_cell(X, c(0,2,5,6), 3L, 3L)'
    /// The third column is constant within each subject, so it is subject-level.
    #[test]
    fn test_cell_level_columns_matches_nebula() {
        let x1 = [0.0, 1.0, 0.5, -0.5, 0.25, -1.0];
        let x2 = [1.0, 1.0, 2.0, 2.0, 2.0, 3.0];
        let design: Vec<f64> = x1
            .iter()
            .zip(x2.iter())
            .flat_map(|(&a, &b)| [1.0, a, b])
            .collect();
        let got = cell_level_columns(&design, 6, 3, &[0, 2, 5, 6]).unwrap();
        assert_eq!(got, vec![false, true, false]);
    }

    #[test]
    fn test_cell_level_columns_rejects_a_bad_fid() {
        let design = vec![1.0; 6];
        assert!(matches!(
            cell_level_columns(&design, 6, 1, &[0, 2, 5]).unwrap_err(),
            EdgeErrors::InvalidArgument(_)
        ));
    }

    #[test]
    fn test_tables_match_the_dense_value_and_gradient() {
        use rand::prelude::*;
        use rand::rngs::SmallRng;
        let mut rng = SmallRng::seed_from_u64(5);
        let fid = vec![0usize, 300, 700, 1000];
        let mut design = Vec::new();
        let mut log_offset = Vec::new();
        let mut counts = Vec::new();
        let mut cells = Vec::new();
        let mut totals = vec![0.0; 3];
        for c in 0..1000 {
            let s = fid.iter().rposition(|&f| f <= c).expect("in range");
            let kind = f64::from(rng.random_bool(0.3));
            design.extend([1.0, (s % 2) as f64 - 0.5, kind]);
            log_offset.push(rng.random_range(-2.0..2.5));
            if rng.random_bool(0.12) {
                let y = rng.random_range(1..8) as f64;
                counts.push(y);
                cells.push(c);
                totals[s] += y;
            }
        }
        let gene = GeneData::new(&design, &log_offset, &counts, &cells, &totals, &fid, 3)
            .expect("valid gene");
        let zeros = ZeroCells::build(&design, &log_offset, &fid, 3).expect("categorical");
        for params in [
            [-2.0, 0.3, 0.2, 0.5, 1.0],
            [0.5, -0.4, 1.0, 0.05, 0.02],
            [-4.0, 0.0, -1.0, 2.0, 30.0],
        ] {
            let (want_value, want_grad) = ptmg_value_and_gradient(&gene, &params);
            let mut scratch = PtmgScratch::tabled(&gene, &zeros);
            let (value, grad) = ptmg_value_and_gradient_with(&gene, &params, &mut scratch);
            assert_relative_eq!(value, want_value, max_relative = 1e-13);
            for (g, w) in grad.iter().zip(&want_grad) {
                assert_relative_eq!(*g, *w, max_relative = 1e-10, epsilon = 1e-9);
            }
        }
    }
}