edge_rs/sc/

pml.rs

//! NEBULA's penalised maximum likelihood inner solver.
//!
//! Given fixed variance components, profiles out the fixed and subject random
//! effects jointly. The parameters are `(beta, log_w)`: `n_beta` fixed effects
//! and one random effect per subject. The `(n_beta + k)` Hessian is never
//! formed. The random-effect block is diagonal (independent subjects,
//! contiguous cells), so a Schur complement reduces each Newton step to an
//! `n_beta` system plus `k` divisions.
//!
//! * [`opt_pml`]: gamma prior on `w` (NBGMM). The penalty is
//!   `alpha log w - lambda w` with `alpha`, `lambda` derived from `sigma[0]`;
//!   `ord` extends the Laplace approximation past its leading term.
//! * [`opt_pml_nbm`]: normal prior on `log w` (NBLMM). The penalty is
//!   `-log_w'log_w / (2 alpha) - (k/2) log alpha`; no higher-order correction.
//!
//! ### Layout
//!
//! The design is row-major `n_cells * n_beta`. Counts are the positive entries
//! with their cell indices, and subjects are half-open `[start, end)` blocks.
//!
//! ### Parallelism
//!
//! Sequential by design: one call is one gene, and the caller parallelises over
//! genes. The outer loop also stops on an absolute improvement in the
//! log-likelihood, so a reassociating reduction can flip the last iteration and
//! move `beta` by more than the reassociation itself.
//!
//! ### Deviations from edgePython
//!
//! Ported from `nebula`'s `src/optimization.cpp`, not edgePython's
//! `_opt_pml_nb`, which adds a ridge, a `vw` floor, a linear-predictor clamp
//! and drops REML. See `UPSTREAM_DEVIATIONS.md` A21.
//!
//! ### References
//!
//! He et al., Communications Biology 4, 629, 2021

#![allow(clippy::needless_range_loop)]

use crate::errors::EdgeErrors;
use crate::sc::zeros::ZeroCells;

//////////////////////
// Tuning and codes //
//////////////////////

/// Newton iteration budget, matching nebula's hard-coded `maxstep`.
///
/// nebula flags a fit that reaches this as `-20`, so the value is fixed.
const DEFAULT_MAX_ITER: usize = 50;

/// Backtracking budget inside one Newton step, nebula's `maxstd`.
///
/// Exceeding it is what produces the `-10` and `-40` convergence codes.
const DEFAULT_MAX_BACKTRACK: usize = 10;

/// Default stopping tolerance on the improvement in the penalised
/// log-likelihood, nebula's `eps`.
const DEFAULT_EPS: f64 = 1e-6;

/// Largest gradient component nebula will still call a critical point.
///
/// Reached only when backtracking has been exhausted: a stalled step with every
/// gradient below this is `-10`, with any gradient above it `-40`.
const GRADIENT_TOLERANCE: f64 = 0.01;

/// Magnitude past which a Newton component is treated as unbounded rather than
/// damped.
///
/// nebula replaces such a step with a fixed displacement, seeded by the same
/// constant and halved once per backtrack, instead of halving it forty times.
const STEP_CUTOFF: f64 = 40.0;

/// Smallest eigenvalue of the information matrix nebula will invert.
///
/// Below it the fixed effects are not identifiable for this gene and the fit is
/// reported as `-25` with no standard errors.
const EIGENVALUE_CUTOFF: f64 = 1e-8;

/// Jacobi sweep budget for the smallest-eigenvalue check.
///
/// Cyclic Jacobi needs three or four sweeps on these small matrices; thirty is
/// slack.
const JACOBI_SWEEPS: usize = 30;

/// Converged, by a sufficiently small improvement in the objective.
pub const CONV_SUCCESS: i32 = 1;

/// Converged at a critical point: no step improved the objective, but every
/// gradient component is below nebula's `convd`, 0.01.
pub const CONV_CRITICAL_POINT: i32 = -10;

/// The Newton iteration budget was exhausted.
pub const CONV_MAX_ITER: i32 = -20;

/// The information matrix is numerically singular, so no standard errors.
pub const CONV_SINGULAR: i32 = -25;

/// The penalised log-likelihood was not a number at the reported optimum.
pub const CONV_NOT_FINITE: i32 = -30;

/// Backtracking stalled with a gradient still above nebula's `convd`, 0.01.
pub const CONV_STALLED: i32 = -40;

/// A variance component sat on its box constraint.
///
/// Not detectable from inside this module; pass `variance_at_bound` to
/// [`check_convergence`] to raise it.
pub const CONV_AT_BOUND: i32 = -60;

/////////////
// PmlData //
/////////////

/// One gene's data, in the layout NEBULA's C++ expects.
///
/// Cells must already be grouped by subject, which is what makes
/// `subject_start` a valid block partition. Only the positive counts are
/// stored: `counts[i]` is the count in cell `cell_index[i]`.
#[derive(Clone, Copy, Debug)]
pub struct PmlData<'a> {
    /// Design matrix, row-major `n_cells * n_beta`.
    pub design: &'a [f64],
    /// Log offset per cell, length `n_cells`.
    pub offset: &'a [f64],
    /// The positive counts, in increasing cell order.
    pub counts: &'a [f64],
    /// Cell index of each entry of `counts`, zero-based and increasing.
    pub cell_index: &'a [usize],
    /// Start of each subject's block plus a trailing `n_cells`, length `k + 1`.
    pub subject_start: &'a [usize],
    /// Total count per subject, length `k`. nebula's `cumsumy`.
    pub subject_total: &'a [f64],
}

impl PmlData<'_> {
    /// Number of cells.
    ///
    /// ### Returns
    ///
    /// The length of `offset`, which defines the cell axis.
    pub fn n_cells(&self) -> usize {
        self.offset.len()
    }

    /// Number of subjects.
    ///
    /// ### Returns
    ///
    /// One fewer than the length of `subject_start`, or zero if that is empty.
    pub fn n_subjects(&self) -> usize {
        self.subject_start.len().saturating_sub(1)
    }

    /// Number of fixed effects.
    ///
    /// ### Returns
    ///
    /// The design width implied by `design.len()` and `n_cells`, or zero when
    /// there are no cells.
    pub fn n_beta(&self) -> usize {
        let n = self.n_cells();
        self.design.len().checked_div(n).unwrap_or(0)
    }

    /// Checks every invariant the inner loops assume.
    ///
    /// ### Returns
    ///
    /// `Ok(())`, or the first violation found.
    fn validate(&self) -> Result<(), EdgeErrors> {
        let n_cells = self.n_cells();
        if n_cells == 0 {
            return Err(EdgeErrors::MustBePositive("n_cells".to_string()));
        }
        let k = self.n_subjects();
        if k == 0 {
            return Err(EdgeErrors::MustBePositive("n_subjects".to_string()));
        }
        let n_beta = self.n_beta();
        if n_beta == 0 || n_beta * n_cells != self.design.len() {
            return Err(EdgeErrors::LengthMismatch {
                name: "design",
                expected: n_cells * n_beta.max(1),
                got: self.design.len(),
            });
        }
        if self.counts.len() != self.cell_index.len() {
            return Err(EdgeErrors::LengthMismatch {
                name: "cell_index",
                expected: self.counts.len(),
                got: self.cell_index.len(),
            });
        }
        if self.subject_total.len() != k {
            return Err(EdgeErrors::LengthMismatch {
                name: "subject_total",
                expected: k,
                got: self.subject_total.len(),
            });
        }
        if self.subject_start[0] != 0 || self.subject_start[k] != n_cells {
            return Err(EdgeErrors::InvalidArgument(format!(
                "subject_start must run from 0 to {n_cells}; got {} to {}.",
                self.subject_start[0], self.subject_start[k]
            )));
        }
        for s in 0..k {
            if self.subject_start[s] >= self.subject_start[s + 1] {
                return Err(EdgeErrors::SubjectsNotContiguous { subject: s });
            }
        }
        let mut previous: Option<usize> = None;
        for &c in self.cell_index {
            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_index must be strictly increasing.".to_string(),
                ));
            }
            previous = Some(c);
        }
        Ok(())
    }
}

/////////////////
// PmlVariance //
/////////////////

/// The two variance components, held on nebula's own scale.
///
/// `subject` means different things in the two solvers: the C++ reads the same
/// `sigma[0]` slot differently in `opt_pml` and `opt_pml_nbm`.
#[derive(Clone, Copy, Debug)]
pub struct PmlVariance {
    /// nebula's `sigma[0]`.
    ///
    /// For [`opt_pml`] this is `log(1 + sigma^2)`, the log of one plus the
    /// squared coefficient of variation of the subject-level gamma random
    /// effect; the solver derives `alpha = 1 / (e^s - 1)` and
    /// `lambda = 1 / (sqrt(e^s) (e^s - 1))` from it. For [`opt_pml_nbm`] it is
    /// the variance of the normal random effect on the log scale, used as is.
    pub subject: f64,
    /// nebula's `sigma[1]`, its `gamma`: the cell-level negative binomial size,
    /// the reciprocal of the cell-level overdispersion.
    pub cell: f64,
}

///////////////
// PmlParams //
///////////////

/// Tuning knobs shared by both solvers.
#[derive(Clone, Copy, Debug)]
pub struct PmlParams {
    /// Add `log|det(information)|` to the log-determinant.
    ///
    /// nebula's `reml`: restricted instead of plain maximum likelihood for the
    /// outer overdispersion estimate. Off by default, as `nebula(..., reml = 0)`.
    pub reml: bool,
    /// Stop once a Newton step improves the penalised log-likelihood by less
    /// than this (absolute, as in nebula).
    pub eps: f64,
    /// Order of the Laplace expansion of the marginal likelihood.
    ///
    /// `1`: leading term only ([`PmlResult::second_order`] stays zero). `2` adds
    /// the third-derivative term, `3` the fourth. nebula uses `3` for genes with
    /// an expected count per subject below three. Ignored by [`opt_pml_nbm`].
    pub ord: u32,
    /// Newton iteration budget.
    pub max_iter: usize,
    /// Backtracking halvings allowed inside one Newton step.
    pub max_backtrack: usize,
}

impl Default for PmlParams {
    /// nebula's defaults as called from `nebula()`.
    fn default() -> Self {
        Self {
            reml: false,
            eps: DEFAULT_EPS,
            ord: 1,
            max_iter: DEFAULT_MAX_ITER,
            max_backtrack: DEFAULT_MAX_BACKTRACK,
        }
    }
}

///////////////
// PmlResult //
///////////////

/// The fitted penalised maximum likelihood solution for one gene.
#[derive(Clone, Debug)]
pub struct PmlResult {
    /// Fixed effect coefficients.
    pub beta: Vec<f64>,
    /// Subject random effects on the log scale.
    pub log_w: Vec<f64>,
    /// Observed information for the fixed effects, row-major `nb * nb`.
    ///
    /// nebula's `var`: the Schur complement of the joint observed information
    /// after eliminating the random-effect block. Its inverse is the covariance
    /// of `beta`.
    pub information: Vec<f64>,
    /// Value of the penalised log-likelihood at the optimum.
    pub log_likelihood: f64,
    /// Value at the previous Newton iterate, nebula's `loglikp`.
    ///
    /// nebula falls back to this when `log_likelihood` is not a number.
    pub log_likelihood_prev: f64,
    /// Log-determinant of the observed information.
    ///
    /// `sum log|vw|` over subjects, plus `log|det(information)|` when
    /// [`PmlParams::reml`] is set. The outer objective subtracts half of it.
    pub log_det: f64,
    /// Higher-order Laplace correction, nebula's `second`.
    ///
    /// Zero unless [`PmlParams::ord`] exceeds one. The outer objective adds
    /// `log(1 + second)`.
    pub second_order: f64,
    /// nebula's convergence code, one of [`CONV_SUCCESS`],
    /// [`CONV_CRITICAL_POINT`], [`CONV_MAX_ITER`], [`CONV_SINGULAR`],
    /// [`CONV_NOT_FINITE`] or [`CONV_STALLED`].
    ///
    /// [`CONV_AT_BOUND`] is not reachable from here; see [`check_convergence`].
    pub convergence: i32,
    /// Newton steps taken.
    pub iterations: usize,
    /// Backtracking halvings in the final Newton step, nebula's `damp`. Eleven:
    /// stalled at a critical point; twelve: stalled with a live gradient.
    pub backtracks: usize,
}

/// The penalised log-likelihood and its gradient at a given point.
///
/// Both are negated, as in nebula's `pml_ll_der_eigen`, for direct use by a
/// minimiser.
#[derive(Clone, Debug)]
pub struct PmlObjective {
    /// Negative penalised log-likelihood.
    pub objective: f64,
    /// Negative gradient: `n_beta` fixed-effect entries followed by
    /// `n_subjects` random-effect entries.
    pub gradient: Vec<f64>,
}

////////////////////
// Random effects //
////////////////////

/// The prior on the subject random effects, with its scalars already resolved.
#[derive(Clone, Copy, Debug)]
enum Penalty {
    /// Gamma prior on `w`, NEBULA's NBGMM.
    Gamma {
        /// Shape, `1 / (e^s - 1)`.
        alpha: f64,
        /// Rate, `1 / (sqrt(e^s) (e^s - 1))`.
        lambda: f64,
    },
    /// Normal prior on `log w`, NEBULA's NBLMM.
    LogNormal {
        /// Variance of the normal, used directly.
        alpha: f64,
    },
}

impl Penalty {
    /// Adds the prior's contribution to the log-likelihood.
    ///
    /// ### Params
    ///
    /// * `log_likelihood` - Data contribution accumulated so far
    /// * `log_w` - Random effects on the log scale
    /// * `w` - `exp(log_w)`, unused by the log-normal variant
    ///
    /// ### Returns
    ///
    /// The penalised log-likelihood.
    fn accumulate(&self, log_likelihood: f64, log_w: &[f64], w: &[f64]) -> f64 {
        match *self {
            Penalty::Gamma { alpha, lambda } => {
                let mut sum_log_w = 0.0;
                for &l in log_w {
                    sum_log_w += l;
                }
                let mut sum_w = 0.0;
                for &x in w {
                    sum_w += x;
                }
                log_likelihood + (alpha * sum_log_w - lambda * sum_w)
            }
            Penalty::LogNormal { alpha } => {
                let mut quad = 0.0;
                for &l in log_w {
                    quad += l * l;
                }
                log_likelihood - quad / alpha / 2.0 - (log_w.len() as f64) / 2.0 * alpha.ln()
            }
        }
    }

    /// The prior's contribution to the random-effect gradient.
    ///
    /// ### Params
    ///
    /// * `log_w_s` - Random effect for this subject
    /// * `w_s` - `exp(log_w_s)`
    ///
    /// ### Returns
    ///
    /// The term added to `cumsumy - sum(gstar_extb_phi)` for this subject.
    fn gradient(&self, log_w_s: f64, w_s: f64) -> f64 {
        match *self {
            Penalty::Gamma { alpha, lambda } => alpha - lambda * w_s,
            Penalty::LogNormal { alpha } => -log_w_s / alpha,
        }
    }

    /// The prior's contribution to the random-effect curvature.
    ///
    /// ### Params
    ///
    /// * `w_s` - `exp(log_w_s)`
    ///
    /// ### Returns
    ///
    /// The term added to the data curvature for this subject.
    fn curvature(&self, w_s: f64) -> f64 {
        match *self {
            Penalty::Gamma { lambda, .. } => lambda * w_s,
            Penalty::LogNormal { alpha } => 1.0 / alpha,
        }
    }
}

///////////////
// Front end //
///////////////

/// Penalised maximum likelihood fit with a gamma random effect, NEBULA's NBGMM.
///
/// Joint Newton iteration over `beta` and `log_w` (Schur complement,
/// per-coordinate backtracking, Laplace corrections per [`PmlParams::ord`]).
///
/// `w` has a `Gamma(alpha, lambda)` prior with `alpha = 1 / (e^s - 1)` and
/// `lambda = 1 / (sqrt(e^s) (e^s - 1))`, `s` being [`PmlVariance::subject`]: mean
/// `sqrt(e^s)`, squared coefficient of variation `e^s - 1`.
///
/// ### Params
///
/// * `data` - One gene's design, offsets, positive counts and subject blocks
/// * `beta_init` - Starting fixed effects, length `n_beta`. `log_w` always
///   starts at zero, as in the C++
/// * `variance` - The two variance components, held fixed throughout
/// * `params` - Tuning knobs, or `None` for nebula's defaults
///
/// ### Returns
///
/// The fitted coefficients, random effects, observed information and
/// convergence diagnostics.
///
/// ### References
///
/// He et al., Communications Biology 4, 629, 2021
pub fn opt_pml(
    data: &PmlData<'_>,
    beta_init: &[f64],
    variance: &PmlVariance,
    params: Option<PmlParams>,
) -> Result<PmlResult, EdgeErrors> {
    opt_pml_tabled(data, None, beta_init, variance, params)
}

/// [`opt_pml`], with the zero counts summed from the run's tables.
///
/// ### Params
///
/// * `data` - One gene's design, offsets, positive counts and subject blocks
/// * `zeros` - The tables built from this `data`'s design, offsets and
///   subjects, or `None` to sweep every cell
/// * `beta_init` - Starting fixed effects, length `n_beta`
/// * `variance` - The two variance components, held fixed throughout
/// * `params` - Tuning knobs, or `None` for nebula's defaults
///
/// ### Returns
///
/// As [`opt_pml`].
pub(crate) fn opt_pml_tabled(
    data: &PmlData<'_>,
    zeros: Option<&ZeroCells>,
    beta_init: &[f64],
    variance: &PmlVariance,
    params: Option<PmlParams>,
) -> Result<PmlResult, EdgeErrors> {
    let exps = variance.subject.exp();
    if !(exps.is_finite() && exps > 1.0) {
        return Err(EdgeErrors::InvalidArgument(format!(
            "PmlVariance::subject must be finite and strictly positive for opt_pml; got {}.",
            variance.subject
        )));
    }
    let penalty = Penalty::Gamma {
        alpha: 1.0 / (exps - 1.0),
        lambda: 1.0 / (exps.sqrt() * (exps - 1.0)),
    };
    optimise(
        data,
        zeros,
        beta_init,
        None,
        penalty,
        variance.cell,
        params.unwrap_or_default(),
    )
}

/// Penalised maximum likelihood fit with a log-normal random effect, NEBULA's
/// NBLMM.
///
/// The [`opt_pml`] loop with penalty `-log_w'log_w / (2 alpha) - (k/2) log alpha`,
/// `alpha` being [`PmlVariance::subject`] read as a variance.
/// [`PmlResult::second_order`] is always zero and [`PmlParams::ord`] is ignored.
///
/// ### Params
///
/// * `data` - One gene's design, offsets, positive counts and subject blocks
/// * `beta_init` - Starting fixed effects, length `n_beta`
/// * `variance` - The two variance components, held fixed throughout
/// * `params` - Tuning knobs, or `None` for nebula's defaults
///
/// ### Returns
///
/// The fitted coefficients, random effects, observed information and
/// convergence diagnostics.
pub fn opt_pml_nbm(
    data: &PmlData<'_>,
    beta_init: &[f64],
    variance: &PmlVariance,
    params: Option<PmlParams>,
) -> Result<PmlResult, EdgeErrors> {
    if !(variance.subject.is_finite() && variance.subject > 0.0) {
        return Err(EdgeErrors::InvalidArgument(format!(
            "PmlVariance::subject must be finite and strictly positive for opt_pml_nbm; got {}.",
            variance.subject
        )));
    }
    let penalty = Penalty::LogNormal {
        alpha: variance.subject,
    };
    optimise(
        data,
        None,
        beta_init,
        None,
        penalty,
        variance.cell,
        params.unwrap_or_default(),
    )
}

/// Penalised log-likelihood and gradient at a point, nebula's
/// `pml_ll_der_eigen`.
///
/// The entry the R side hands to `lbfgs` and `trust`. It reads
/// [`PmlVariance::subject`] as the shape of a mean-one `Gamma(alpha, alpha)`
/// prior (shape and rate both `alpha`), which is nebula's own inconsistency with
/// [`opt_pml`], reproduced.
///
/// ### Params
///
/// * `data` - One gene's design, offsets, positive counts and subject blocks
/// * `beta` - Fixed effects, length `n_beta`
/// * `log_w` - Random effects on the log scale, length `n_subjects`
/// * `variance` - The two variance components
///
/// ### Returns
///
/// The negated objective and the negated gradient, laid out as `n_beta` fixed
/// effect entries followed by `n_subjects` random effect entries.
pub fn pml_log_likelihood_gradient(
    data: &PmlData<'_>,
    beta: &[f64],
    log_w: &[f64],
    variance: &PmlVariance,
) -> Result<PmlObjective, EdgeErrors> {
    data.validate()?;
    let n_beta = data.n_beta();
    let k = data.n_subjects();
    if beta.len() != n_beta {
        return Err(EdgeErrors::LengthMismatch {
            name: "beta",
            expected: n_beta,
            got: beta.len(),
        });
    }
    if log_w.len() != k {
        return Err(EdgeErrors::LengthMismatch {
            name: "log_w",
            expected: k,
            got: log_w.len(),
        });
    }

    let alpha = variance.subject;
    let gamma = variance.cell;
    // Both the shape and the rate: `alpha*sum(log w) - alpha*sum(w)`.
    let penalty = Penalty::Gamma {
        alpha,
        lambda: alpha,
    };

    let mut work = Workspace::new(data, None);
    work.sweep_current(data, beta, log_w, gamma, penalty);

    let mut gradient = Vec::with_capacity(n_beta + k);
    gradient.extend(work.current.db.iter().map(|v| -v));
    gradient.extend(work.current.dw.iter().map(|v| -v));
    Ok(PmlObjective {
        objective: -work.current.log_likelihood,
        gradient,
    })
}

/// [`opt_pml`], started from a given point rather than from zero random effects.
///
/// The GPU path finds the optimum in `f32` and finishes it here in `f64`. From a
/// point at the optimum this takes one Newton step (about four passes over the
/// cells against fourteen for a cold fit) and returns what [`opt_pml`] returns,
/// including the log-determinant at the penultimate iterate.
///
/// ### Params
///
/// * `data` - One gene's design, offsets, positive counts and subject blocks
/// * `beta_init` - Starting fixed effects, length `n_beta`
/// * `log_w_init` - Starting random effects on the log scale, length
///   `n_subjects`
/// * `variance` - The two variance components, read as [`opt_pml`] reads them
/// * `params` - Tuning knobs, or `None` for nebula's defaults
///
/// ### Returns
///
/// As [`opt_pml`].
#[cfg(feature = "gpu")]
pub(crate) fn opt_pml_from(
    data: &PmlData<'_>,
    beta_init: &[f64],
    log_w_init: &[f64],
    variance: &PmlVariance,
    params: Option<PmlParams>,
) -> Result<PmlResult, EdgeErrors> {
    let exps = variance.subject.exp();
    if !(exps.is_finite() && exps > 1.0) {
        return Err(EdgeErrors::InvalidArgument(format!(
            "PmlVariance::subject must be finite and strictly positive for opt_pml; got {}.",
            variance.subject
        )));
    }
    let penalty = Penalty::Gamma {
        alpha: 1.0 / (exps - 1.0),
        lambda: 1.0 / (exps.sqrt() * (exps - 1.0)),
    };
    optimise(
        data,
        None,
        beta_init,
        Some(log_w_init),
        penalty,
        variance.cell,
        params.unwrap_or_default(),
    )
}

/// Largest move in a cell's linear predictor for which [`newton_finish`]
/// updates the stored `exp` by series instead of taking it again.
///
/// The second sweep is one converged Newton step from the first, so
/// `exp(eta + z) = exp(eta) exp(z)` with `z` of order `1e-4`. The fifth-order
/// truncation `z^6 / 720` is `1.4e-21` relative at this bound, under an ulp.
#[cfg(feature = "gpu")]
const EXP_SERIES_MAX: f64 = 1e-3;

/// Relative fall in the objective that [`newton_finish`] still reads as
/// rounding at a converged point rather than a step that went wrong.
///
/// From the device's `f32` optimum the `f64` step gains less than the rounding
/// of a sum over all cells, so `likdif` is a few ulp negative on about 15 per
/// cent of fits (S1: 10804 of 73650, between `1e-16` and `1e-15` relative).
/// [`opt_pml_from`] would backtrack to the same point. A real overshoot loses
/// orders of magnitude more.
#[cfg(feature = "gpu")]
const LIKDIF_ROUNDING: f64 = 1e-12;

/// `exp(z)` to fifth order, for `|z|` within [`EXP_SERIES_MAX`].
///
/// ### Params
///
/// * `z` - The exponent
///
/// ### Returns
///
/// `exp(z)`, to well under an ulp.
#[cfg(feature = "gpu")]
#[inline(always)]
fn exp_series(z: f64) -> f64 {
    1.0 + z * (1.0 + z * (0.5 + z * (1.0 / 6.0 + z * (1.0 / 24.0 + z * (1.0 / 120.0)))))
}

/// Terms multiplied together before [`LogSum`] takes one logarithm.
///
/// Each term is `extb + gamma`, far inside `1e-38..1e38`, so eight cannot leave
/// the `f64` range. An overflowing product gives an infinite sum, as the
/// per-term logarithm would.
#[cfg(feature = "gpu")]
const LOG_PRODUCT_GROUP: u32 = 8;

/// `sum ln(t)` taken as `ln` of running products, one logarithm per
/// [`LOG_PRODUCT_GROUP`] terms.
///
/// The logarithm is the dearest operation in the dense sweep of
/// [`newton_finish`]. Group rounding is `1e-15` absolute on the logarithm, the
/// same as the per-term form.
#[cfg(feature = "gpu")]
struct LogSum {
    /// Logarithms of the finished groups.
    sum: f64,
    /// Product of the group in flight.
    product: f64,
    /// Terms in the group in flight.
    count: u32,
}

#[cfg(feature = "gpu")]
impl LogSum {
    /// An empty sum.
    fn new() -> Self {
        Self {
            sum: 0.0,
            product: 1.0,
            count: 0,
        }
    }

    /// Adds `ln(t)`.
    ///
    /// ### Params
    ///
    /// * `t` - The term, strictly positive
    #[inline(always)]
    fn push(&mut self, t: f64) {
        self.product *= t;
        self.count += 1;
        if self.count == LOG_PRODUCT_GROUP {
            self.sum += self.product.ln();
            self.product = 1.0;
            self.count = 0;
        }
    }

    /// The sum, with the group in flight folded in.
    ///
    /// ### Returns
    ///
    /// `sum ln(t)` over everything pushed.
    fn finish(self) -> f64 {
        self.sum + self.product.ln()
    }
}

/// The device's curvature at, or one converged step behind, the point
/// [`newton_finish`] starts from.
#[cfg(feature = "gpu")]
#[derive(Clone, Copy, Debug)]
pub(crate) struct DeviceCurvature<'a> {
    /// Curvature of each random effect, `vw`, length `n_subjects`.
    pub(crate) subject: &'a [f64],
    /// Cross block, `vwb`, row-major `n_subjects * n_beta`.
    pub(crate) cross: &'a [f64],
    /// Schur complement, row-major `n_beta * n_beta`.
    pub(crate) schur: &'a [f64],
}

/// What [`newton_finish`] brings back: the three scalars the profile objective
/// reads off an order-one inner fit.
#[cfg(feature = "gpu")]
#[derive(Clone, Copy, Debug)]
pub(crate) struct NewtonFinish {
    /// Penalised log-likelihood after the step.
    pub(crate) log_likelihood: f64,
    /// Penalised log-likelihood at the starting point.
    pub(crate) log_likelihood_prev: f64,
    /// Log-determinant of the random-effect block at the final iterate.
    pub(crate) log_det: f64,
    /// Newton steps taken.
    pub(crate) iterations: usize,
}

/// Full `f64` Newton steps from a point already near the optimum, two fused
/// passes over the cells per step.
///
/// The common case of [`opt_pml_from`] on the GPU path: the device point is
/// converged to `f32`, one step settles it (two or three for a few per cent of
/// fits) and no step needs damping. [`opt_pml`] allocates
/// `(5 + n_beta) * n_cells` per call and walks it in a dozen passes; here each
/// step is two fused passes per subject with nothing stored per cell, split as
/// the device kernel splits it: a dense sweep over every cell at a count of
/// zero, then a sparse sweep adding what the positive counts change. Sums
/// associate differently from [`opt_pml`], so results agree to rounding, not
/// bitwise.
///
/// The log-determinant is taken at the final iterate, not the penultimate one as
/// in [`opt_pml`]: here the penultimate iterate is the device's `f32` point,
/// whose location error enters the log-determinant at first order, so the
/// second pass also accumulates the curvature weight.
///
/// ### Params
///
/// * `data` - One gene's design, offsets, positive counts and subject blocks
/// * `beta` - Starting fixed effects, length `n_beta`
/// * `log_w` - Starting random effects on the log scale, length `n_subjects`
/// * `variance` - The two variance components, read as [`opt_pml`] reads them
/// * `eps` - nebula's absolute stopping tolerance
/// * `max_iter` - Newton budget
/// * `curvature` - The device's curvature near the starting point. When given,
///   the first step takes its gradient in `f64` and its Hessian from here, so
///   its sweep skips the `n_beta^2` cross-products per cell. The Hessian only
///   steers the step: `f32` curvature and an `f64` gradient land as close as an
///   exact step from a point already converged to `f32`. Later steps, which a
///   few per cent of fits need, assemble their own
///
/// * `stored` - Scratch the first sweep leaves each cell's `exp` in for the
///   second; resized here, so one per thread serves every call
///
/// ### Returns
///
/// The finished scalars, or `None` when a full step worsens the objective by
/// more than [`LIKDIF_ROUNDING`] or leaves it not finite, or the budget runs
/// out. Those need the backtracking search and the convergence codes, so the
/// caller wants [`opt_pml_from`].
#[cfg(feature = "gpu")]
#[allow(clippy::too_many_arguments)]
pub(crate) fn newton_finish(
    data: &PmlData<'_>,
    beta: &[f64],
    log_w: &[f64],
    variance: &PmlVariance,
    eps: f64,
    max_iter: usize,
    curvature: Option<DeviceCurvature<'_>>,
    stored: &mut Vec<f64>,
) -> Option<NewtonFinish> {
    match data.n_beta() {
        1 => newton_finish_width::<1>(
            data, beta, log_w, variance, eps, max_iter, curvature, stored,
        ),
        2 => newton_finish_width::<2>(
            data, beta, log_w, variance, eps, max_iter, curvature, stored,
        ),
        3 => newton_finish_width::<3>(
            data, beta, log_w, variance, eps, max_iter, curvature, stored,
        ),
        4 => newton_finish_width::<4>(
            data, beta, log_w, variance, eps, max_iter, curvature, stored,
        ),
        5 => newton_finish_width::<5>(
            data, beta, log_w, variance, eps, max_iter, curvature, stored,
        ),
        6 => newton_finish_width::<6>(
            data, beta, log_w, variance, eps, max_iter, curvature, stored,
        ),
        7 => newton_finish_width::<7>(
            data, beta, log_w, variance, eps, max_iter, curvature, stored,
        ),
        8 => newton_finish_width::<8>(
            data, beta, log_w, variance, eps, max_iter, curvature, stored,
        ),
        // The device never fits a wider design, so nothing reaches this.
        _ => None,
    }
}

/// [`newton_finish`] for one design width.
///
/// With the width a constant, the loops over design columns unroll and a
/// subject's sums sit in `NB`-sized arrays held in registers, instead of heap
/// accumulators with a store-to-load round trip per cell.
///
/// The first sweep takes a subject's `exp` in its own loop, so the accumulating
/// loop has no call in it and the positive cells read the stored value. Together
/// with the registers: 15.6 to 12.1 ns per cell per step on a 20000-cell gene;
/// the split alone is slower.
#[cfg(feature = "gpu")]
#[allow(clippy::too_many_arguments)]
fn newton_finish_width<const NB: usize>(
    data: &PmlData<'_>,
    beta: &[f64],
    log_w: &[f64],
    variance: &PmlVariance,
    eps: f64,
    max_iter: usize,
    curvature: Option<DeviceCurvature<'_>>,
    stored: &mut Vec<f64>,
) -> Option<NewtonFinish> {
    let nb = NB;
    let k = data.n_subjects();
    let gamma = variance.cell;
    let exps = variance.subject.exp();
    if beta.len() != nb || log_w.len() != k || !(exps.is_finite() && exps > 1.0) {
        return None;
    }
    if curvature.is_some_and(|c| {
        c.subject.len() != k || c.cross.len() != k * nb || c.schur.len() != nb * nb
    }) {
        return None;
    }
    let alpha = 1.0 / (exps - 1.0);
    let lambda = 1.0 / (exps.sqrt() * (exps - 1.0));

    let run = positive_runs(data.cell_index, data.subject_start);

    let mut beta = beta.to_vec();
    let mut log_w = log_w.to_vec();
    let mut new_beta = vec![0.0; nb];
    let mut new_log_w = vec![0.0; k];
    let mut db = vec![0.0; nb];
    let mut dw = vec![0.0; k];
    let mut vw = vec![0.0; k];
    let mut vwb = vec![0.0; k * nb];
    let mut vb = vec![0.0; nb * nb];
    let mut factor = vec![0.0; nb * nb];
    let mut step_beta = vec![0.0; nb];
    let mut perm = vec![0usize; nb];
    let mut tmp = vec![0.0; nb];
    // Every entry is written by the first sweep before the second reads it.
    stored.resize(data.n_cells(), 0.0);

    for iterations in 1..=max_iter {
        let borrowed = curvature.filter(|_| iterations == 1);
        let assemble = borrowed.is_none();
        db.fill(0.0);
        vb.fill(0.0);
        let mut log_likelihood_prev = 0.0;
        for s in 0..k {
            let w_s = log_w[s].exp();
            let mut resid = 0.0;
            let mut weight = 0.0;
            let mut sum_log = LogSum::new();
            let mut linear = 0.0;
            let mut weighted_log = 0.0;
            let mut score = [0.0; NB];
            let mut first = [0.0; NB];
            let mut second = [[0.0; NB]; NB];

            let mut fold = |row: &[f64], d: f64, p: f64| {
                resid += d;
                for a in 0..nb {
                    score[a] += row[a] * d;
                }
                if assemble {
                    weight += p;
                    for a in 0..nb {
                        let pa = p * row[a];
                        first[a] += pa;
                        for b in a..nb {
                            second[a][b] += pa * row[b];
                        }
                    }
                }
            };

            for r in data.subject_start[s]..data.subject_start[s + 1] {
                let row = &data.design[r * nb..(r + 1) * nb];
                let mut eta = data.offset[r];
                for j in 0..nb {
                    eta += row[j] * beta[j];
                }
                stored[r] = (eta + log_w[s]).exp();
            }
            for r in data.subject_start[s]..data.subject_start[s + 1] {
                let row = &data.design[r * nb..(r + 1) * nb];
                let extb = stored[r];
                let t = extb + gamma;
                sum_log.push(t);
                let inv = 1.0 / t;
                let u = extb * inv;
                fold(row, -gamma * u, gamma * u * inv);
            }
            for i in run[s]..run[s + 1] {
                let c = data.cell_index[i];
                let y = data.counts[i];
                let row = &data.design[c * nb..(c + 1) * nb];
                let mut eta = data.offset[c];
                for j in 0..nb {
                    eta += row[j] * beta[j];
                }
                let extb = stored[c];
                let t = extb + gamma;
                linear += eta * y;
                weighted_log += y * t.ln();
                let inv = 1.0 / t;
                let u = extb * inv;
                fold(row, y * (1.0 - u), y * u * inv);
            }

            log_likelihood_prev +=
                linear + log_w[s] * data.subject_total[s] - gamma * sum_log.finish() - weighted_log
                    + (alpha * log_w[s] - lambda * w_s);
            dw[s] = resid + (alpha - lambda * w_s);
            for a in 0..nb {
                db[a] += score[a];
            }
            if assemble {
                vw[s] = gamma * weight + lambda * w_s;
                for a in 0..nb {
                    vwb[s * nb + a] = gamma * first[a];
                    for b in a..nb {
                        vb[a * nb + b] += gamma * second[a][b];
                    }
                }
            }
        }

        // Schur complement and the step, as in `optimise`.
        if let Some(device) = borrowed {
            vw.copy_from_slice(device.subject);
            vwb.copy_from_slice(device.cross);
            factor.copy_from_slice(device.schur);
        } else {
            for a in 0..nb {
                for b in a..nb {
                    let mut acc = 0.0;
                    for s in 0..k {
                        acc += vwb[s * nb + a] * vwb[s * nb + b] / vw[s];
                    }
                    let v = vb[a * nb + b] - acc;
                    factor[a * nb + b] = v;
                    factor[b * nb + a] = v;
                }
            }
        }
        for j in 0..nb {
            let mut acc = 0.0;
            for s in 0..k {
                acc += vwb[s * nb + j] * dw[s] / vw[s];
            }
            step_beta[j] = db[j] - acc;
        }
        ldlt_solve(&mut factor, &mut step_beta, nb, &mut perm, &mut tmp);

        for j in 0..nb {
            new_beta[j] = beta[j] + step_beta[j];
        }
        let mut log_likelihood = 0.0;
        let mut log_det = 0.0;
        for s in 0..k {
            let mut acc = 0.0;
            for j in 0..nb {
                acc += vwb[s * nb + j] * step_beta[j];
            }
            let new_log_w_s = log_w[s] + (dw[s] - acc) / vw[s];
            new_log_w[s] = new_log_w_s;
            let new_w_s = new_log_w_s.exp();
            let step_log_w_s = new_log_w_s - log_w[s];

            let mut sum_log = LogSum::new();
            let mut linear = 0.0;
            let mut weighted_log = 0.0;
            let mut weight = 0.0;
            for r in data.subject_start[s]..data.subject_start[s + 1] {
                let row = &data.design[r * nb..(r + 1) * nb];
                let mut z = step_log_w_s;
                for j in 0..nb {
                    z += row[j] * step_beta[j];
                }
                let extb = if z.abs() < EXP_SERIES_MAX {
                    stored[r] * exp_series(z)
                } else {
                    let mut eta = data.offset[r];
                    for j in 0..nb {
                        eta += row[j] * new_beta[j];
                    }
                    (eta + new_log_w_s).exp()
                };
                let t = extb + gamma;
                sum_log.push(t);
                let inv = 1.0 / t;
                weight += gamma * extb * inv * inv;
            }
            for i in run[s]..run[s + 1] {
                let c = data.cell_index[i];
                let y = data.counts[i];
                let row = &data.design[c * nb..(c + 1) * nb];
                let mut eta = data.offset[c];
                for j in 0..nb {
                    eta += row[j] * new_beta[j];
                }
                let extb = (eta + new_log_w_s).exp();
                let t = extb + gamma;
                linear += eta * y;
                weighted_log += y * t.ln();
                let inv = 1.0 / t;
                weight += y * extb * inv * inv;
            }
            log_likelihood += linear + new_log_w_s * data.subject_total[s]
                - gamma * sum_log.finish()
                - weighted_log
                + (alpha * new_log_w_s - lambda * new_w_s);
            log_det += (gamma * weight + lambda * new_w_s).abs().ln();
        }

        let likdif = log_likelihood - log_likelihood_prev;
        if likdif.is_nan() || likdif < -LIKDIF_ROUNDING * log_likelihood.abs() {
            return None;
        }
        if likdif <= eps {
            return Some(NewtonFinish {
                log_likelihood,
                log_likelihood_prev,
                log_det,
                iterations,
            });
        }
        beta.copy_from_slice(&new_beta);
        log_w.copy_from_slice(&new_log_w);
    }
    None
}

/// nebula's `check_conv`, applied to a finished fit.
///
/// [`opt_pml`] already runs this with `variance_at_bound` false and an initial
/// code of [`CONV_SUCCESS`]. The outer overdispersion loop re-runs it, as only
/// it knows whether a variance component hit its bound or an earlier optimiser
/// failed. The singular-information test overrides everything, including
/// [`CONV_AT_BOUND`] (nebula's ordering).
///
/// ### Params
///
/// * `result` - The fit to grade
/// * `params` - The knobs it was run with; only `max_iter` is consulted
/// * `variance_at_bound` - Whether a variance component sat on its bound
/// * `initial` - Code carried in from the outer optimiser, [`CONV_SUCCESS`]
///   when that succeeded and something negative when it did not
///
/// ### Returns
///
/// The convergence code.
pub fn check_convergence(
    result: &PmlResult,
    params: &PmlParams,
    variance_at_bound: bool,
    initial: i32,
) -> i32 {
    let mut code = initial;
    if code == CONV_SUCCESS {
        code = if variance_at_bound {
            CONV_AT_BOUND
        } else if result.log_likelihood.is_nan() {
            CONV_NOT_FINITE
        } else if result.iterations == params.max_iter {
            CONV_MAX_ITER
        } else if result.backtracks == params.max_backtrack + 1 {
            CONV_CRITICAL_POINT
        } else if result.backtracks == params.max_backtrack + 2 {
            CONV_STALLED
        } else {
            CONV_SUCCESS
        };
    }

    let n_beta = result.beta.len();
    if n_beta > 1 && min_eigenvalue(&result.information, n_beta) < EIGENVALUE_CUTOFF {
        code = CONV_SINGULAR;
    }
    code
}

///////////////
// Workspace //
///////////////

/// Where each subject's run of positive counts starts.
///
/// ### Params
///
/// * `cells` - Cell index of each positive count, increasing
/// * `starts` - Subject boundaries, length `k + 1`
///
/// ### Returns
///
/// For each boundary, the first positive count at or past it, length `k + 1`.
pub(crate) fn positive_runs(cells: &[usize], starts: &[usize]) -> Vec<usize> {
    starts
        .iter()
        .map(|&b| cells.partition_point(|&c| c < b))
        .collect()
}

/// Widest design whose per-subject sums sit in fixed-size arrays, so the loops
/// over design columns unroll and the sums stay in registers. Wider designs run
/// the same sweep with the width read at run time and the sums on the heap.
const MAX_UNROLLED_WIDTH: usize = 8;

/// Penalised log-likelihood, gradient and curvature at one point.
///
/// One [`sweep`] fills all of it, so a trial point that is accepted already
/// carries what the next Newton step needs.
struct Sweep {
    /// Penalised log-likelihood.
    log_likelihood: f64,
    /// `exp(log_w)`, one per subject.
    w: Vec<f64>,
    /// Fixed-effect gradient, length `n_beta`.
    db: Vec<f64>,
    /// Random-effect gradient, length `n_subjects`.
    dw: Vec<f64>,
    /// Random-effect curvature, length `n_subjects`.
    vw: Vec<f64>,
    /// Cross block of the information, row-major `n_subjects * n_beta`.
    vwb: Vec<f64>,
    /// Fixed-effect block of the information, row-major `n_beta * n_beta`.
    vb: Vec<f64>,
}

impl Sweep {
    /// Zeroed buffers for one gene's shape.
    ///
    /// ### Params
    ///
    /// * `nb` - Number of fixed effects
    /// * `k` - Number of subjects
    ///
    /// ### Returns
    ///
    /// The empty sweep.
    fn new(nb: usize, k: usize) -> Self {
        Self {
            log_likelihood: 0.0,
            w: vec![0.0; k],
            db: vec![0.0; nb],
            dw: vec![0.0; k],
            vw: vec![0.0; k],
            vwb: vec![0.0; k * nb],
            vb: vec![0.0; nb * nb],
        }
    }
}

/// Every buffer the Newton loop needs, allocated once per call.
struct Workspace<'z> {
    /// The run's zero-count tables, or `None` to sweep every cell.
    zeros: Option<&'z ZeroCells>,
    /// `exp(offset + design*beta + log_w)` per cell at the last swept point.
    /// With tables, only the positive and the loose cells are written.
    stored: Vec<f64>,
    /// Where each subject's run of positive counts starts, length `k + 1`.
    run: Vec<usize>,
    /// Per-subject sums for designs wider than [`MAX_UNROLLED_WIDTH`],
    /// length `2 * n_beta + n_beta^2`.
    wide: Vec<f64>,
    /// The sweep at the current iterate.
    current: Sweep,
    /// The sweep at the trial point.
    trial: Sweep,
    /// `vw` the last Newton step used, for the log-determinant.
    vw_step: Vec<f64>,
    /// Schur complement of the information, row-major `n_beta * n_beta`.
    vb2: Vec<f64>,
    /// Copy of `vb2`, destroyed by each factorisation.
    factor: Vec<f64>,
    /// Transposition record from the pivoted factorisation, length `n_beta`.
    perm: Vec<usize>,
    /// Scratch row for the factorisation, length `n_beta`.
    tmp: Vec<f64>,
    /// Newton step in `beta`, length `n_beta`.
    step_beta: Vec<f64>,
    /// Newton step in `log_w`, length `n_subjects`.
    step_log_w: Vec<f64>,
    /// `dw / vw`, length `n_subjects`.
    dwvw: Vec<f64>,
    /// Per-coordinate damping on `beta`, length `n_beta`.
    damp_beta: Vec<f64>,
    /// Per-coordinate damping on `log_w`, length `n_subjects`.
    damp_log_w: Vec<f64>,
    /// Trial fixed effects, length `n_beta`.
    new_beta: Vec<f64>,
    /// Trial random effects, length `n_subjects`.
    new_log_w: Vec<f64>,
}

impl<'z> Workspace<'z> {
    /// Allocates every buffer to the shape implied by `data`.
    ///
    /// ### Params
    ///
    /// * `data` - The gene whose shape sets the buffer sizes
    /// * `zeros` - The run's zero-count tables, or `None`
    ///
    /// ### Returns
    ///
    /// A zeroed workspace with the positive-count runs located.
    fn new(data: &PmlData<'_>, zeros: Option<&'z ZeroCells>) -> Self {
        let n_cells = data.n_cells();
        let k = data.n_subjects();
        let nb = data.n_beta();
        let run = positive_runs(data.cell_index, data.subject_start);

        Self {
            zeros,
            stored: vec![0.0; n_cells],
            run,
            wide: if nb > MAX_UNROLLED_WIDTH {
                vec![0.0; 2 * nb + nb * nb]
            } else {
                Vec::new()
            },
            current: Sweep::new(nb, k),
            trial: Sweep::new(nb, k),
            vw_step: vec![0.0; k],
            vb2: vec![0.0; nb * nb],
            factor: vec![0.0; nb * nb],
            perm: vec![0; nb],
            tmp: vec![0.0; nb],
            step_beta: vec![0.0; nb],
            step_log_w: vec![0.0; k],
            dwvw: vec![0.0; k],
            damp_beta: vec![0.0; nb],
            damp_log_w: vec![0.0; k],
            new_beta: vec![0.0; nb],
            new_log_w: vec![0.0; k],
        }
    }

    /// Sweeps the gene at `(beta, log_w)` into [`Self::current`].
    ///
    /// ### Params
    ///
    /// * `data` - The gene
    /// * `beta` - Fixed effects
    /// * `log_w` - Random effects on the log scale
    /// * `gamma` - Cell-level negative binomial size
    /// * `penalty` - Prior on the random effects
    fn sweep_current(
        &mut self,
        data: &PmlData<'_>,
        beta: &[f64],
        log_w: &[f64],
        gamma: f64,
        penalty: Penalty,
    ) {
        sweep(
            data,
            beta,
            log_w,
            gamma,
            penalty,
            self.zeros,
            &mut self.stored,
            &self.run,
            &mut self.wide,
            &mut self.current,
        );
    }

    /// Sweeps the gene at `(new_beta, new_log_w)` into [`Self::trial`].
    ///
    /// ### Params
    ///
    /// * `data` - The gene
    /// * `gamma` - Cell-level negative binomial size
    /// * `penalty` - Prior on the random effects
    ///
    /// ### Returns
    ///
    /// The penalised log-likelihood at the trial point.
    fn sweep_trial(&mut self, data: &PmlData<'_>, gamma: f64, penalty: Penalty) -> f64 {
        sweep(
            data,
            &self.new_beta,
            &self.new_log_w,
            gamma,
            penalty,
            self.zeros,
            &mut self.stored,
            &self.run,
            &mut self.wide,
            &mut self.trial,
        );
        self.trial.log_likelihood
    }
}

/// Penalised log-likelihood, gradient and curvature at one point, in one pass
/// per subject.
///
/// Each subject is a dense part over every cell as if its count were zero, then
/// a sparse sweep adding what the positive counts change. Every per-cell term is
/// linear in the count through `gstar = gamma + y`, so the split is exact.
/// Without tables the dense part is two loops, the `exp` of every cell into
/// `stored` and the accumulation; with them it is one lookup per group of
/// [`ZeroCells`] and the accumulation over the loose cells. Sums associate
/// differently from nebula's C++, so results agree to rounding, not bitwise.
///
/// ### Params
///
/// * `data` - The gene
/// * `beta` - Fixed effects
/// * `log_w` - Random effects on the log scale
/// * `gamma` - Cell-level negative binomial size
/// * `penalty` - Prior on the random effects
/// * `zeros` - The run's zero-count tables, or `None` to sweep every cell
/// * `stored` - Per-cell `exp`, overwritten
/// * `run` - Start of each subject's positive counts, from [`Workspace::new`]
/// * `wide` - Heap sums for designs wider than [`MAX_UNROLLED_WIDTH`]
/// * `out` - Filled with the sweep
#[allow(clippy::too_many_arguments)]
fn sweep(
    data: &PmlData<'_>,
    beta: &[f64],
    log_w: &[f64],
    gamma: f64,
    penalty: Penalty,
    zeros: Option<&ZeroCells>,
    stored: &mut [f64],
    run: &[usize],
    wide: &mut [f64],
    out: &mut Sweep,
) {
    macro_rules! width {
        ($nb:literal) => {
            sweep_width::<$nb>(
                data, beta, log_w, gamma, penalty, zeros, stored, run, wide, out,
            )
        };
    }
    match data.n_beta() {
        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),
    }
}

/// [`sweep`] for one design width, `NB == 0` meaning the width is read from
/// `data` at run time.
///
/// ### Params
///
/// As [`sweep`].
#[allow(clippy::too_many_arguments)]
fn sweep_width<const NB: usize>(
    data: &PmlData<'_>,
    beta: &[f64],
    log_w: &[f64],
    gamma: f64,
    penalty: Penalty,
    zeros: Option<&ZeroCells>,
    stored: &mut [f64],
    run: &[usize],
    wide: &mut [f64],
    out: &mut Sweep,
) {
    let nb = if NB == 0 { data.n_beta() } else { NB };
    let k = data.n_subjects();
    let mut local = [0.0; 2 * MAX_UNROLLED_WIDTH + MAX_UNROLLED_WIDTH * MAX_UNROLLED_WIDTH];
    let sums: &mut [f64] = if NB == 0 {
        wide
    } else {
        &mut local[..2 * nb + nb * nb]
    };

    out.db.fill(0.0);
    out.vb.fill(0.0);
    let ln_gamma = gamma.ln();
    let mut log_likelihood = 0.0;
    for s in 0..k {
        let cells = data.subject_start[s]..data.subject_start[s + 1];
        let log_w_s = log_w[s];
        let w_s = log_w_s.exp();
        sums.fill(0.0);
        let (score, rest) = sums.split_at_mut(nb);
        let (first, second) = rest.split_at_mut(nb);
        let mut resid = 0.0;
        let mut weight = 0.0;
        let mut sum_log = 0.0;
        let mut linear = 0.0;
        let mut weighted_log = 0.0;

        // `d` adds to the gradient, `p` to the curvature, the latter before the
        // factor of `gamma` the curvature blocks take at the end.
        let mut fold = |row: &[f64], d: f64, p: f64| {
            resid += d;
            weight += p;
            for a in 0..nb {
                score[a] += row[a] * d;
                let pa = p * row[a];
                first[a] += pa;
                for b in a..nb {
                    second[a * nb + b] += pa * row[b];
                }
            }
        };

        match zeros {
            None => {
                // The `exp` in its own loop keeps the accumulating loop free of
                // calls.
                for r in cells.clone() {
                    let row = &data.design[r * nb..(r + 1) * nb];
                    let mut eta = data.offset[r];
                    for j in 0..nb {
                        eta += row[j] * beta[j];
                    }
                    stored[r] = (eta + log_w_s).exp();
                }
                for r in cells {
                    let row = &data.design[r * nb..(r + 1) * nb];
                    let extb = stored[r];
                    let t = extb + gamma;
                    sum_log += t.ln();
                    let inv = 1.0 / t;
                    let u = extb * inv;
                    fold(row, -gamma * u, gamma * u * inv);
                }
            }
            Some(z) => {
                // Per group, `sum ln t = n ln gamma + L`, `sum u = A` and
                // `sum gamma u / t = B`.
                for g in &z.groups[z.subject_groups[s]..z.subject_groups[s + 1]] {
                    let row = &data.design[g.row * nb..(g.row + 1) * nb];
                    let mut log_tau = log_w_s - ln_gamma;
                    for j in 0..nb {
                        log_tau += row[j] * beta[j];
                    }
                    let [l, a, b] = g.eval::<0, 3>(log_tau);
                    sum_log += g.n * ln_gamma + l;
                    fold(row, -gamma * a, b);
                }
                for &r in &z.loose[z.subject_loose[s]..z.subject_loose[s + 1]] {
                    let row = &data.design[r * nb..(r + 1) * nb];
                    let mut eta = data.offset[r];
                    for j in 0..nb {
                        eta += row[j] * beta[j];
                    }
                    let extb = (eta + log_w_s).exp();
                    stored[r] = extb;
                    let t = extb + gamma;
                    sum_log += t.ln();
                    let inv = 1.0 / t;
                    let u = extb * inv;
                    fold(row, -gamma * u, gamma * u * inv);
                }
            }
        }
        for i in run[s]..run[s + 1] {
            let c = data.cell_index[i];
            let y = data.counts[i];
            let row = &data.design[c * nb..(c + 1) * nb];
            let mut eta = data.offset[c];
            for j in 0..nb {
                eta += row[j] * beta[j];
            }
            let extb = if zeros.is_some() {
                let e = (eta + log_w_s).exp();
                stored[c] = e;
                e
            } else {
                stored[c]
            };
            let t = extb + gamma;
            linear += eta * y;
            weighted_log += y * t.ln();
            let inv = 1.0 / t;
            let u = extb * inv;
            fold(row, y * (1.0 - u), y * u * inv);
        }

        log_likelihood += linear + log_w_s * data.subject_total[s] - gamma * sum_log - weighted_log;
        out.w[s] = w_s;
        out.dw[s] = resid + penalty.gradient(log_w_s, w_s);
        out.vw[s] = gamma * weight + penalty.curvature(w_s);
        for a in 0..nb {
            out.db[a] += score[a];
            out.vwb[s * nb + a] = gamma * first[a];
            for b in a..nb {
                out.vb[a * nb + b] += gamma * second[a * nb + b];
            }
        }
    }
    for a in 0..nb {
        for b in (a + 1)..nb {
            out.vb[b * nb + a] = out.vb[a * nb + b];
        }
    }
    out.log_likelihood = penalty.accumulate(log_likelihood, log_w, &out.w);
}

////////////
// Solver //
////////////

/// The shared Newton loop behind [`opt_pml`] and [`opt_pml_nbm`].
///
/// Each point is one [`sweep`]; an accepted trial point's sweep carries the
/// gradient and curvature of the next Newton step.
///
/// ### Params
///
/// * `data` - The gene
/// * `zeros` - The run's zero-count tables, or `None` to sweep every cell
/// * `beta_init` - Starting fixed effects
/// * `log_w_init` - Starting random effects, or `None` for zeros
/// * `penalty` - Prior on the random effects
/// * `gamma` - Cell-level negative binomial size
/// * `params` - Tuning knobs
///
/// ### Returns
///
/// The fitted solution and its diagnostics.
fn optimise(
    data: &PmlData<'_>,
    zeros: Option<&ZeroCells>,
    beta_init: &[f64],
    log_w_init: Option<&[f64]>,
    penalty: Penalty,
    gamma: f64,
    params: PmlParams,
) -> Result<PmlResult, EdgeErrors> {
    data.validate()?;
    let nb = data.n_beta();
    let k = data.n_subjects();
    if beta_init.len() != nb {
        return Err(EdgeErrors::LengthMismatch {
            name: "beta_init",
            expected: nb,
            got: beta_init.len(),
        });
    }
    if !(gamma.is_finite() && gamma > 0.0) {
        return Err(EdgeErrors::InvalidDispersion(gamma));
    }

    let mut beta = beta_init.to_vec();
    let mut log_w = match log_w_init {
        Some(init) if init.len() == k => init.to_vec(),
        Some(init) => {
            return Err(EdgeErrors::LengthMismatch {
                name: "log_w_init",
                expected: k,
                got: init.len(),
            });
        }
        None => vec![0.0; k],
    };
    let mut work = Workspace::new(data, zeros);
    work.sweep_current(data, &beta, &log_w, gamma, penalty);

    let mut log_likelihood = work.current.log_likelihood;
    let mut log_likelihood_prev = 0.0;
    let mut likdif = 0.0;
    let mut step = 0usize;
    let mut backtracks = 0usize;

    while step == 0 || (likdif > params.eps && step < params.max_iter) {
        step += 1;
        work.damp_beta.fill(1.0);
        work.damp_log_w.fill(1.0);

        let cur = &work.current;
        work.vw_step.copy_from_slice(&cur.vw);
        for s in 0..k {
            work.dwvw[s] = cur.dw[s] / cur.vw[s];
        }
        // Schur complement of the random-effect block, symmetric by
        // construction.
        for a in 0..nb {
            for b in a..nb {
                let mut acc = 0.0;
                for s in 0..k {
                    acc += cur.vwb[s * nb + a] * cur.vwb[s * nb + b] / cur.vw[s];
                }
                let v = cur.vb[a * nb + b] - acc;
                work.vb2[a * nb + b] = v;
                work.vb2[b * nb + a] = v;
            }
        }
        // `rhs = db - vwb' (dw/vw)`, then one `nb` solve gives the whole step.
        for j in 0..nb {
            let mut acc = 0.0;
            for s in 0..k {
                acc += cur.vwb[s * nb + j] * work.dwvw[s];
            }
            work.step_beta[j] = cur.db[j] - acc;
        }
        work.factor.copy_from_slice(&work.vb2);
        ldlt_solve(
            &mut work.factor,
            &mut work.step_beta,
            nb,
            &mut work.perm,
            &mut work.tmp,
        );
        for s in 0..k {
            let mut acc = 0.0;
            for j in 0..nb {
                acc += cur.vwb[s * nb + j] * work.step_beta[j];
            }
            work.step_log_w[s] = work.dwvw[s] - acc / cur.vw[s];
        }

        for j in 0..nb {
            work.new_beta[j] = beta[j] + work.step_beta[j];
        }
        for s in 0..k {
            work.new_log_w[s] = log_w[s] + work.step_log_w[s];
        }

        log_likelihood_prev = log_likelihood;
        log_likelihood = work.sweep_trial(data, gamma, penalty);
        likdif = log_likelihood - log_likelihood_prev;

        backtracks = 0;
        let mut min_step = STEP_CUTOFF;
        while likdif < 0.0 || log_likelihood.is_infinite() {
            backtracks += 1;
            min_step /= 2.0;

            if backtracks > params.max_backtrack {
                likdif = 0.0;
                log_likelihood = log_likelihood_prev;
                let grad_beta = work.current.db.iter().fold(0.0f64, |m, v| m.max(v.abs()));
                let grad_log_w = work.current.dw.iter().fold(0.0f64, |m, v| m.max(v.abs()));
                if grad_beta > GRADIENT_TOLERANCE || grad_log_w > GRADIENT_TOLERANCE {
                    backtracks += 1;
                }
                break;
            }

            // A step past the cutoff is replaced outright: halving would need
            // forty rounds and the budget is ten.
            for j in 0..nb {
                let d = work.step_beta[j];
                if d < STEP_CUTOFF && d > -STEP_CUTOFF {
                    work.damp_beta[j] /= 2.0;
                    work.new_beta[j] = beta[j] + d * work.damp_beta[j];
                } else if d > 0.0 {
                    work.new_beta[j] = beta[j] + min_step;
                } else {
                    work.new_beta[j] = beta[j] - min_step;
                }
            }
            for s in 0..k {
                let d = work.step_log_w[s];
                if d < STEP_CUTOFF && d > -STEP_CUTOFF {
                    work.damp_log_w[s] /= 2.0;
                    work.new_log_w[s] = log_w[s] + d * work.damp_log_w[s];
                } else if d > 0.0 {
                    work.new_log_w[s] = log_w[s] + min_step;
                } else {
                    work.new_log_w[s] = log_w[s] - min_step;
                }
            }

            log_likelihood = work.sweep_trial(data, gamma, penalty);
            likdif = log_likelihood - log_likelihood_prev;
        }

        beta.copy_from_slice(&work.new_beta);
        log_w.copy_from_slice(&work.new_log_w);
        std::mem::swap(&mut work.current, &mut work.trial);
    }

    // `vw` and `vb2` come from the previous iterate, as nebula reports; the
    // outer objective is calibrated against that.
    let mut log_det = 0.0;
    for s in 0..k {
        log_det += work.vw_step[s].abs().ln();
    }
    if params.reml {
        work.factor.copy_from_slice(&work.vb2);
        let det = determinant(&mut work.factor, nb, &mut work.perm);
        log_det += det.abs().ln();
    }

    let second_order = match penalty {
        Penalty::Gamma { lambda, .. } if params.ord > 1 => {
            laplace_correction(&work, data, &beta, &log_w, gamma, lambda, params.ord)
        }
        _ => 0.0,
    };

    let mut result = PmlResult {
        beta,
        log_w,
        information: work.vb2.clone(),
        log_likelihood,
        log_likelihood_prev,
        log_det,
        second_order,
        convergence: CONV_SUCCESS,
        iterations: step,
        backtracks,
    };
    result.convergence = check_convergence(&result, &params, false, CONV_SUCCESS);
    Ok(result)
}

/// Higher-order terms of the Laplace expansion of the marginal likelihood.
///
/// The leading Laplace term is `-0.5 log|V|`; this returns the next correction,
/// which the outer objective adds as `log(1 + second)`. `ord == 2` keeps only
/// the third-derivative term `5 m3^2 / (24 V^3)`; `ord == 3` adds the
/// fourth-derivative term `-m4 / (8 V^2)` and the mixed term
/// `7 m5 m3 / (48 V^4)`. Every derivative is diagonal in the subjects, and each
/// per-cell summand is `gstar` times a function of the cell's `exp`, so this is
/// one pass split dense and sparse as in [`sweep`].
///
/// Recomputes the curvature at the final iterate, not the one the
/// log-determinant used (as nebula does).
///
/// ### Params
///
/// * `work` - Workspace whose `stored` and `current.w` sit at the final iterate
/// * `data` - The gene
/// * `beta` - Fixed effects at the final iterate
/// * `log_w` - Random effects at the final iterate
/// * `gamma` - Cell-level negative binomial size
/// * `lambda` - Rate of the gamma prior
/// * `ord` - Expansion order, greater than one
///
/// ### Returns
///
/// The correction, nebula's `second`.
fn laplace_correction(
    work: &Workspace<'_>,
    data: &PmlData<'_>,
    beta: &[f64],
    log_w: &[f64],
    gamma: f64,
    lambda: f64,
    ord: u32,
) -> f64 {
    let k = data.n_subjects();
    let nb = data.n_beta();
    let ln_gamma = gamma.ln();
    let high = ord > 2;
    // The curvature and the three higher derivatives, per unit of `gstar`.
    let terms = |e: f64| {
        let t = e + gamma;
        let p2 = e / (t * t);
        let p3 = p2 / t;
        let p4 = p3 / t;
        let p5 = p4 / t;
        let e2 = e * e;
        [
            p2,
            p3 * (gamma - e),
            p4 * (gamma * gamma + e2 - (4.0 * gamma) * e),
            p5 * (gamma * gamma * gamma - (11.0 * gamma * gamma) * e + (11.0 * gamma) * e2
                - e2 * e),
        ]
    };

    let mut third_acc = 0.0;
    let mut fourth_acc = 0.0;
    let mut fifth_acc = 0.0;
    for s in 0..k {
        let mut sums = [0.0; 4];
        let mut dense = |e: f64| {
            let f = terms(e);
            for j in 0..4 {
                sums[j] += gamma * f[j];
            }
        };
        match work.zeros {
            None => {
                for &e in &work.stored[data.subject_start[s]..data.subject_start[s + 1]] {
                    dense(e);
                }
            }
            Some(z) => {
                for &r in &z.loose[z.subject_loose[s]..z.subject_loose[s + 1]] {
                    dense(work.stored[r]);
                }
                // `gamma` times each factor sums to the group's `B`..`E`.
                for g in &z.groups[z.subject_groups[s]..z.subject_groups[s + 1]] {
                    let row = &data.design[g.row * nb..(g.row + 1) * nb];
                    let mut log_tau = log_w[s] - ln_gamma;
                    for j in 0..nb {
                        log_tau += row[j] * beta[j];
                    }
                    let f = g.eval::<2, 4>(log_tau);
                    for j in 0..4 {
                        sums[j] += f[j];
                    }
                }
            }
        }
        for i in work.run[s]..work.run[s + 1] {
            let y = data.counts[i];
            let f = terms(work.stored[data.cell_index[i]]);
            for j in 0..4 {
                sums[j] += y * f[j];
            }
        }
        let prior = lambda * work.current.w[s];
        let vw = gamma * sums[0] + prior;
        let third = gamma * sums[1] + prior;
        third_acc += third * third / (vw * vw * vw);
        if high {
            let fourth = gamma * sums[2] + prior;
            let fifth = gamma * sums[3] + prior;
            let vs = vw * vw;
            fourth_acc += fourth / vs;
            fifth_acc += fifth * third / (vs * vs);
        }
    }

    let mut second = 5.0 * third_acc / 24.0;
    if high {
        second -= fourth_acc / 8.0;
        second += 7.0 * fifth_acc / 48.0;
    }
    second
}

////////////////////
// Linear algebra //
////////////////////

/// Solves a small symmetric system in place by pivoted `LDL'`.
///
/// Mirrors Eigen's `LDLT` as the C++ calls it: symmetric pivoting on the largest
/// remaining diagonal, and a solve that zeroes a vanishing pivot instead of
/// dividing. Backtracking visits indefinite points where plain Cholesky fails,
/// and a rank-deficient design must still give a finite step. See
/// [`crate::glm::levenberg`] for the unpivoted Cholesky.
///
/// ### Params
///
/// * `a` - Row-major `n * n` symmetric matrix, overwritten by the factor
/// * `b` - Right-hand side, overwritten by the solution
/// * `n` - System size
/// * `perm` - Scratch for the transposition record, length `n`
/// * `tmp` - Scratch row, length `n`
fn ldlt_solve(a: &mut [f64], b: &mut [f64], n: usize, perm: &mut [usize], tmp: &mut [f64]) {
    for step in 0..n {
        let mut pivot = step;
        let mut best = a[step * n + step].abs();
        for i in (step + 1)..n {
            let v = a[i * n + i].abs();
            if v > best {
                best = v;
                pivot = i;
            }
        }
        perm[step] = pivot;
        if pivot != step {
            for j in 0..step {
                a.swap(step * n + j, pivot * n + j);
            }
            for i in (pivot + 1)..n {
                a.swap(i * n + step, i * n + pivot);
            }
            a.swap(step * n + step, pivot * n + pivot);
            for i in (step + 1)..pivot {
                a.swap(i * n + step, pivot * n + i);
            }
        }

        if step > 0 {
            for j in 0..step {
                tmp[j] = a[j * n + j] * a[step * n + j];
            }
            let mut acc = 0.0;
            for j in 0..step {
                acc += a[step * n + j] * tmp[j];
            }
            a[step * n + step] -= acc;
            for i in (step + 1)..n {
                let mut acc = 0.0;
                for j in 0..step {
                    acc += a[i * n + j] * tmp[j];
                }
                a[i * n + step] -= acc;
            }
        }

        let d = a[step * n + step];
        if d != 0.0 {
            for i in (step + 1)..n {
                a[i * n + step] /= d;
            }
        }
    }

    for step in 0..n {
        b.swap(step, perm[step]);
    }
    for i in 0..n {
        let mut acc = b[i];
        for j in 0..i {
            acc -= a[i * n + j] * b[j];
        }
        b[i] = acc;
    }
    for i in 0..n {
        let d = a[i * n + i];
        // Eigen's tolerance: a pivot below the smallest normal double is an
        // exact zero and the component is dropped.
        b[i] = if d.abs() > f64::MIN_POSITIVE {
            b[i] / d
        } else {
            0.0
        };
    }
    for i in (0..n).rev() {
        let mut acc = b[i];
        for j in (i + 1)..n {
            acc -= a[j * n + i] * b[j];
        }
        b[i] = acc;
    }
    for step in (0..n).rev() {
        b.swap(step, perm[step]);
    }
}

/// Determinant of a small square matrix by partial-pivot `LU`.
///
/// A product, not a sum of logs, as in the C++, so under- and overflow match.
///
/// ### Params
///
/// * `a` - Row-major `n * n` matrix, overwritten by the factor
/// * `n` - Matrix order
/// * `perm` - Scratch, length `n`
///
/// ### Returns
///
/// The signed determinant.
fn determinant(a: &mut [f64], n: usize, perm: &mut [usize]) -> f64 {
    let mut det = 1.0;
    perm.fill(0);
    for step in 0..n {
        let mut pivot = step;
        let mut best = a[step * n + step].abs();
        for i in (step + 1)..n {
            let v = a[i * n + step].abs();
            if v > best {
                best = v;
                pivot = i;
            }
        }
        if pivot != step {
            for j in 0..n {
                a.swap(step * n + j, pivot * n + j);
            }
            det = -det;
        }
        let d = a[step * n + step];
        det *= d;
        if d == 0.0 {
            return 0.0;
        }
        for i in (step + 1)..n {
            let factor = a[i * n + step] / d;
            a[i * n + step] = factor;
            for j in (step + 1)..n {
                a[i * n + j] -= factor * a[step * n + j];
            }
        }
    }
    det
}

/// Smallest eigenvalue of a small symmetric matrix, by cyclic Jacobi.
///
/// Used only for the identifiability test against [`EIGENVALUE_CUTOFF`]. Jacobi
/// beats a faer decomposition on matrices a few columns wide.
///
/// ### Params
///
/// * `matrix` - Row-major `n * n` symmetric matrix, not modified
/// * `n` - Matrix order
///
/// ### Returns
///
/// The smallest eigenvalue, or the single entry when `n` is one.
fn min_eigenvalue(matrix: &[f64], n: usize) -> f64 {
    if n == 0 {
        return f64::NAN;
    }
    if n == 1 {
        return matrix[0];
    }
    let mut a = matrix.to_vec();
    let mut frobenius = 0.0;
    for &v in &a {
        frobenius += v * v;
    }
    let tolerance = (frobenius * f64::EPSILON * f64::EPSILON).max(f64::MIN_POSITIVE);

    for _ in 0..JACOBI_SWEEPS {
        let mut off = 0.0;
        for p in 0..n {
            for q in (p + 1)..n {
                off += a[p * n + q] * a[p * n + q];
            }
        }
        if off <= tolerance {
            break;
        }
        for p in 0..n {
            for q in (p + 1)..n {
                let apq = a[p * n + q];
                if apq == 0.0 {
                    continue;
                }
                let theta = (a[q * n + q] - a[p * n + p]) / (2.0 * apq);
                let t = theta.signum() / (theta.abs() + (theta * theta + 1.0).sqrt());
                let c = 1.0 / (t * t + 1.0).sqrt();
                let s = t * c;
                for r in 0..n {
                    let arp = a[r * n + p];
                    let arq = a[r * n + q];
                    a[r * n + p] = c * arp - s * arq;
                    a[r * n + q] = s * arp + c * arq;
                }
                for r in 0..n {
                    let apr = a[p * n + r];
                    let aqr = a[q * n + r];
                    a[p * n + r] = c * apr - s * aqr;
                    a[q * n + r] = s * apr + c * aqr;
                }
            }
        }
    }
    (0..n).map(|i| a[i * n + i]).fold(f64::INFINITY, f64::min)
}

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

#[cfg(test)]
// Reference values are pasted verbatim from R's 17-digit print.
#[allow(clippy::excessive_precision)]
// Fixture builders return the five parallel arrays `PmlData` borrows.
#[allow(clippy::type_complexity)]
mod tests {
    use super::*;
    use approx::assert_relative_eq;

    /// Design of fixture A, row-major `12 * 2`: intercept plus a cell-level
    /// binary covariate.
    const DESIGN: [f64; 24] = [
        1.0, 0.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0,
        1.0, 0.0, 1.0, 1.0, 1.0, 1.0,
    ];

    /// Offsets of fixture A, all dyadic so they round-trip exactly through R.
    const OFFSET: [f64; 12] = [
        0.25, -0.5, 0.75, 0.0, 0.5, -0.25, 0.125, 0.375, -0.125, 0.25, 0.0, -0.375,
    ];

    /// Subject block boundaries: three subjects of four, five and three cells.
    const SUBJECT_START: [usize; 4] = [0, 4, 9, 12];

    /// Positive counts of fixture A.
    const COUNTS: [f64; 8] = [3.0, 5.0, 2.0, 4.0, 1.0, 7.0, 2.0, 6.0];

    /// Cells holding those counts.
    const CELL_INDEX: [usize; 8] = [0, 2, 3, 5, 6, 7, 9, 10];

    /// Per-subject totals of fixture A.
    const SUBJECT_TOTAL: [f64; 3] = [10.0, 12.0, 8.0];

    /// Positive counts of fixture B, where subject two is entirely zero.
    const COUNTS_B: [f64; 5] = [3.0, 5.0, 2.0, 2.0, 6.0];

    /// Cells holding those counts.
    const CELL_INDEX_B: [usize; 5] = [0, 2, 3, 9, 10];

    /// Per-subject totals of fixture B.
    const SUBJECT_TOTAL_B: [f64; 3] = [10.0, 0.0, 8.0];

    /// Relative tolerance against the R package.
    ///
    /// Three orders tighter than 1e-8, reached on every fixture below. The
    /// binding constraint is nebula's own absolute `eps`: at a near-tie Newton
    /// step, a one-ulp shift in summation adds or drops a whole damped step and
    /// moves the answer by 1e-7. Fixtures avoid such ties.
    const TOL: f64 = 1e-11;

    /// Fixture A: a well-behaved gene.
    fn fixture_a() -> PmlData<'static> {
        PmlData {
            design: &DESIGN,
            offset: &OFFSET,
            counts: &COUNTS,
            cell_index: &CELL_INDEX,
            subject_start: &SUBJECT_START,
            subject_total: &SUBJECT_TOTAL,
        }
    }

    /// Fixture B: the same gene with every count in subject two set to zero.
    fn fixture_b() -> PmlData<'static> {
        PmlData {
            design: &DESIGN,
            offset: &OFFSET,
            counts: &COUNTS_B,
            cell_index: &CELL_INDEX_B,
            subject_start: &SUBJECT_START,
            subject_total: &SUBJECT_TOTAL_B,
        }
    }

    /// Compares two slices elementwise at [`TOL`].
    ///
    /// ### Params
    ///
    /// * `got` - Values produced here
    /// * `want` - Values produced by nebula 1.5.8
    fn assert_slice(got: &[f64], want: &[f64]) {
        assert_eq!(got.len(), want.len());
        for (g, w) in got.iter().zip(want) {
            assert_relative_eq!(*g, *w, max_relative = TOL, epsilon = 1e-12);
        }
    }

    // Reference values from nebula 1.5.8:
    //   Rscript -e 'library(nebula); nebula:::opt_pml(X, offset, Y, fid, cumsumy,
    //     posind, posindy, nb, nind, k, beta, sigma, reml, eps, ord)'
    // with X, offset, Y, fid, cumsumy, posindy as the constants above.

    #[test]
    fn test_opt_pml_well_behaved_matches_nebula() {
        let data = fixture_a();
        let variance = PmlVariance {
            subject: 0.25,
            cell: 2.0,
        };
        let fit = opt_pml(&data, &[0.5, -0.25], &variance, None).expect("fit");

        assert_slice(&fit.beta, &[1.1055943907637618, -0.98348970166026206]);
        assert_slice(
            &fit.log_w,
            &[
                -0.025429869309964156,
                0.022419629678247836,
                0.33781545877640862,
            ],
        );
        assert_slice(
            &fit.information,
            &[
                5.4708370512237092,
                2.6193993778249318,
                2.6193993778249318,
                4.0223986651682804,
            ],
        );
        assert_relative_eq!(fit.log_likelihood, -60.378642466146786, max_relative = TOL);
        assert_relative_eq!(fit.log_det, 6.0421923783970968, max_relative = TOL);
        assert_eq!(fit.second_order, 0.0);
        assert_eq!(fit.iterations, 4);
        assert_eq!(fit.backtracks, 0);
        assert_eq!(fit.convergence, CONV_SUCCESS);
    }

    #[test]
    fn test_opt_pml_third_order_laplace_correction() {
        let data = fixture_a();
        let variance = PmlVariance {
            subject: 0.25,
            cell: 2.0,
        };
        let params = PmlParams {
            ord: 3,
            ..Default::default()
        };
        let fit = opt_pml(&data, &[0.5, -0.25], &variance, Some(params)).expect("fit");

        // Same optimum, only the correction differs from ord = 1.
        assert_slice(&fit.beta, &[1.1055943907637618, -0.98348970166026206]);
        assert_relative_eq!(fit.second_order, 0.0028301625845031939, max_relative = TOL);
    }

    #[test]
    fn test_opt_pml_reml_log_determinant() {
        let data = fixture_a();
        let variance = PmlVariance {
            subject: 0.25,
            cell: 2.0,
        };
        let params = PmlParams {
            reml: true,
            ..Default::default()
        };
        let fit = opt_pml(&data, &[0.5, -0.25], &variance, Some(params)).expect("fit");

        assert_relative_eq!(fit.log_det, 8.759838692617901, max_relative = TOL);
        assert_slice(&fit.beta, &[1.1055943907637618, -0.98348970166026206]);
    }

    #[test]
    fn test_opt_pml_subject_with_all_zero_counts() {
        let data = fixture_b();
        let variance = PmlVariance {
            subject: 0.25,
            cell: 2.0,
        };
        let fit = opt_pml(&data, &[0.5, -0.25], &variance, None).expect("fit");

        assert_slice(&fit.beta, &[0.15781792039265752, -0.098178037869394941]);
        assert_slice(
            &fit.log_w,
            &[
                0.31781101972092152,
                -0.83477809302857631,
                0.46458776408134972,
            ],
        );
        assert_slice(
            &fit.information,
            &[
                4.9151627959404793,
                2.59325292186585,
                2.59325292186585,
                3.5781777249175555,
            ],
        );
        assert_relative_eq!(fit.log_likelihood, -50.769290856943599, max_relative = TOL);
        assert_relative_eq!(fit.log_det, 5.3802869340559738, max_relative = TOL);
        assert_eq!(fit.convergence, CONV_SUCCESS);
    }

    #[test]
    fn test_opt_pml_zero_subject_third_order() {
        let data = fixture_b();
        let variance = PmlVariance {
            subject: 0.25,
            cell: 2.0,
        };
        let params = PmlParams {
            ord: 3,
            ..Default::default()
        };
        let fit = opt_pml(&data, &[0.5, -0.25], &variance, Some(params)).expect("fit");
        assert_relative_eq!(fit.second_order, 0.029026901671004347, max_relative = TOL);
    }

    #[test]
    fn test_opt_pml_start_near_bound_backtracks() {
        // beta[0] = 90 is just inside nebula's +-100 box: the first Newton step
        // overshoots and per-coordinate damping runs.
        let data = fixture_a();
        let variance = PmlVariance {
            subject: 0.25,
            cell: 2.0,
        };
        let fit = opt_pml(&data, &[90.0, 0.0], &variance, None).expect("fit");

        assert_slice(&fit.beta, &[1.1055943759446161, -0.98348968952057669]);
        assert_slice(
            &fit.log_w,
            &[
                -0.025429869876909816,
                0.022419639320108459,
                0.3378154594265026,
            ],
        );
        assert_slice(
            &fit.information,
            &[
                5.4711020927179739,
                2.6193969725187061,
                2.6193969725187061,
                4.0224446400576319,
            ],
        );
        assert_eq!(fit.iterations, 9);
        assert_eq!(fit.convergence, CONV_SUCCESS);
    }

    #[test]
    fn test_opt_pml_far_outside_bounds_uses_fixed_displacement() {
        // Both coordinates start well outside the box: the Newton step exceeds
        // the cutoff and the fixed-displacement branch runs.
        let data = fixture_a();
        let variance = PmlVariance {
            subject: 0.25,
            cell: 2.0,
        };
        let fit = opt_pml(&data, &[-90.0, 60.0], &variance, None).expect("fit");

        assert_slice(&fit.beta, &[1.1055943725487707, -0.98348968650982715]);
        assert_relative_eq!(fit.log_det, 6.0423215051590002, max_relative = TOL);
        assert_eq!(fit.iterations, 12);
    }

    #[test]
    fn test_opt_pml_critical_point_reports_minus_ten() {
        // Backtracking is exhausted with every gradient under 0.01: a critical
        // point, not a failure. The stall is decisive (orders of magnitude from
        // improving); near-tie stalls are not testable to this tolerance.
        let data = fixture_a();
        let variance = PmlVariance {
            subject: 8.0,
            cell: 0.001,
        };
        let fit = opt_pml(&data, &[15.0, 5.0], &variance, None).expect("fit");

        assert_slice(&fit.beta, &[4.9609375, -4.9609375]);
        assert_slice(&fit.log_w, &[0.0390625, 0.0390625, 4.5234136840078936]);
        assert_slice(
            &fit.information,
            &[
                7.0733722800021132e-05,
                7.0105348391562988e-05,
                7.0105348391562988e-05,
                0.00016310178559274507,
            ],
        );
        assert_relative_eq!(fit.log_det, -20.515549220945573, max_relative = TOL);
        assert_eq!(fit.iterations, 3);
        assert_eq!(fit.backtracks, 11);
        assert_eq!(fit.convergence, CONV_CRITICAL_POINT);
    }

    #[test]
    fn test_opt_pml_stalled_reports_minus_forty() {
        let data = fixture_a();
        let variance = PmlVariance {
            subject: 8.0,
            cell: 0.001,
        };
        let fit = opt_pml(&data, &[20.0, 0.0], &variance, None).expect("fit");

        assert_slice(&fit.beta, &[14.9609375, -4.9609375]);
        assert_slice(&fit.log_w, &[-4.9609375, -4.9609375, 4.9962571608576489]);
        assert_eq!(fit.iterations, 2);
        assert_eq!(fit.backtracks, 12);
        assert_eq!(fit.convergence, CONV_STALLED);
    }

    #[test]
    fn test_opt_pml_iteration_budget_reports_minus_twenty() {
        // A negative eps is unsatisfiable: the loop exhausts its budget.
        let data = fixture_a();
        let variance = PmlVariance {
            subject: 0.25,
            cell: 2.0,
        };
        let params = PmlParams {
            eps: -1e-300,
            ..Default::default()
        };
        let fit = opt_pml(&data, &[0.5, -0.25], &variance, Some(params)).expect("fit");

        assert_eq!(fit.iterations, 50);
        assert_eq!(fit.convergence, CONV_MAX_ITER);
        assert_slice(&fit.beta, &[1.1055943907637629, -0.98348970166026151]);
        assert_slice(
            &fit.information,
            &[
                5.4708368635191373,
                2.619399280826618,
                2.619399280826618,
                4.0223986123865263,
            ],
        );
    }

    #[test]
    fn test_opt_pml_collinear_design_reports_minus_twentyfive() {
        // The covariate is duplicated, so the information matrix is singular.
        let design: Vec<f64> = DESIGN
            .chunks_exact(2)
            .flat_map(|r| [r[0], r[1], r[1]])
            .collect();
        let data = PmlData {
            design: &design,
            offset: &OFFSET,
            counts: &COUNTS,
            cell_index: &CELL_INDEX,
            subject_start: &SUBJECT_START,
            subject_total: &SUBJECT_TOTAL,
        };
        let variance = PmlVariance {
            subject: 0.25,
            cell: 2.0,
        };
        let fit = opt_pml(&data, &[0.5, -0.25, 0.1], &variance, None).expect("fit");

        assert_eq!(fit.convergence, CONV_SINGULAR);
        assert_slice(
            &fit.beta,
            &[1.1055943330380458, -1.0834897078722583, 0.10000000000000001],
        );
        assert_relative_eq!(fit.log_det, 6.042689573294485, max_relative = TOL);
    }

    #[test]
    fn test_opt_pml_single_coefficient() {
        let design = [1.0; 12];
        let data = PmlData {
            design: &design,
            offset: &OFFSET,
            counts: &COUNTS,
            cell_index: &CELL_INDEX,
            subject_start: &SUBJECT_START,
            subject_total: &SUBJECT_TOTAL,
        };
        let variance = PmlVariance {
            subject: 0.25,
            cell: 2.0,
        };
        let params = PmlParams {
            reml: true,
            ord: 3,
            ..Default::default()
        };
        let fit = opt_pml(&data, &[0.0], &variance, Some(params)).expect("fit");

        assert_slice(&fit.beta, &[0.64526354365336647]);
        assert_slice(
            &fit.log_w,
            &[
                0.040002950937536336,
                0.10805651074009694,
                0.21875064209166895,
            ],
        );
        assert_slice(&fit.information, &[5.7533150000267472]);
        assert_relative_eq!(fit.log_det, 7.9235375511695514, max_relative = TOL);
        assert_relative_eq!(fit.second_order, 0.0044134970888637676, max_relative = TOL);
    }

    #[test]
    fn test_opt_pml_nbm_matches_nebula() {
        let data = fixture_a();
        let variance = PmlVariance {
            subject: 0.3,
            cell: 2.0,
        };
        let params = PmlParams {
            reml: true,
            ..Default::default()
        };
        let fit = opt_pml_nbm(&data, &[0.5, -0.25], &variance, Some(params)).expect("fit");

        assert_slice(&fit.beta, &[1.2213403122060453, -0.98985988766009869]);
        assert_slice(
            &fit.log_w,
            &[
                -0.14301073703498665,
                -0.095894342874540239,
                0.23890507990952689,
            ],
        );
        assert_slice(
            &fit.information,
            &[
                5.4025646031315233,
                2.5076353884976541,
                2.5076353884976541,
                3.9141513741254204,
            ],
        );
        assert_relative_eq!(fit.log_likelihood, -49.318101708775217, max_relative = TOL);
        assert_relative_eq!(fit.log_det, 8.6579114426780528, max_relative = TOL);
        assert_eq!(fit.second_order, 0.0);
        assert_eq!(fit.iterations, 3);
        assert_eq!(fit.convergence, CONV_SUCCESS);
    }

    #[test]
    fn test_pml_log_likelihood_gradient_matches_nebula() {
        // Rscript -e 'library(nebula); nebula:::pml_ll_der_eigen(X, offset, Y,
        //   fid, cumsumy, posind, posindy, 2, 12, 3, c(0.5,-0.25),
        //   c(0.1,-0.2,0.05), c(0.25,2))'
        let data = fixture_a();
        let variance = PmlVariance {
            subject: 0.25,
            cell: 2.0,
        };
        let out = pml_log_likelihood_gradient(&data, &[0.5, -0.25], &[0.1, -0.2, 0.05], &variance)
            .expect("gradient");

        assert_relative_eq!(out.objective, 53.926203132942852, max_relative = TOL);
        assert_slice(
            &out.gradient,
            &[
                -5.2597994101747254,
                -0.19846989257393055,
                -0.40624529144311583,
                -2.8367693486273806,
                -2.0229915782218155,
            ],
        );
    }

    /// Builds the wider fixture: 64 cells, 4 subjects, 3 coefficients.
    ///
    /// Generated by an integer recurrence so R reproduces it with the same three
    /// expressions. Every value is dyadic.
    ///
    /// ### Returns
    ///
    /// Design, offsets, counts, cell indices and subject totals.
    fn fixture_large() -> (Vec<f64>, Vec<f64>, Vec<f64>, Vec<usize>, Vec<f64>) {
        let n_cells = 64usize;
        let mut design = Vec::with_capacity(n_cells * 3);
        let mut offset = Vec::with_capacity(n_cells);
        let mut counts = Vec::new();
        let mut cell_index = Vec::new();
        let mut subject_total = vec![0.0; 4];
        let starts = [0usize, 12, 32, 48, 64];
        for i in 0..n_cells {
            design.push(1.0);
            design.push(if i % 3 == 0 { 1.0 } else { 0.0 });
            design.push((i % 7) as f64 / 8.0);
            offset.push(((i % 9) as f64 - 4.0) / 8.0);
            let y = ((i * 7 + 3) % 11) as f64;
            if y > 0.0 {
                counts.push(y);
                cell_index.push(i);
            }
            let s = starts.iter().rposition(|&b| b <= i).unwrap_or(0).min(3);
            subject_total[s] += y;
        }
        (design, offset, counts, cell_index, subject_total)
    }

    #[test]
    fn test_opt_pml_wider_design_and_more_cells() {
        let (design, offset, counts, cell_index, subject_total) = fixture_large();
        assert_eq!(subject_total, vec![58.0, 100.0, 83.0, 82.0]);
        let starts = [0usize, 12, 32, 48, 64];
        let data = PmlData {
            design: &design,
            offset: &offset,
            counts: &counts,
            cell_index: &cell_index,
            subject_start: &starts,
            subject_total: &subject_total,
        };
        let variance = PmlVariance {
            subject: 0.35,
            cell: 1.5,
        };
        let params = PmlParams {
            reml: true,
            ord: 3,
            ..Default::default()
        };
        let fit = opt_pml(&data, &[0.5, -0.25, 0.75], &variance, Some(params)).expect("fit");

        assert_slice(
            &fit.beta,
            &[
                1.4272053191863423,
                0.20743588590645803,
                -0.017982095723719935,
            ],
        );
        assert_slice(
            &fit.log_w,
            &[
                0.21944093016930771,
                0.12415265795703513,
                0.20913336580790309,
                0.14393280971073011,
            ],
        );
        assert_slice(
            &fit.information,
            &[
                8.398261243665857,
                2.7062153540842324,
                3.0637064599168475,
                2.7062153540842324,
                16.226178650135836,
                1.2821611476475319,
                3.0637064599168475,
                1.2821611476475319,
                5.5864643706193053,
            ],
        );
        assert_relative_eq!(fit.log_likelihood, -276.5244116459354, max_relative = TOL);
        assert_relative_eq!(fit.log_det, 18.407453390744379, max_relative = TOL);
        assert_relative_eq!(fit.second_order, 0.0029993806316514779, max_relative = TOL);
        assert_eq!(fit.iterations, 4);
        assert_eq!(fit.convergence, CONV_SUCCESS);

        // Only the REML term of the log-determinant should move.
        let plain = opt_pml(&data, &[0.5, -0.25, 0.75], &variance, None).expect("fit");
        assert_relative_eq!(plain.log_det, 12.052188343414912, max_relative = TOL);
        assert_eq!(plain.second_order, 0.0);
        assert_slice(&plain.beta, &fit.beta);
    }

    #[test]
    fn test_check_convergence_bound_and_nan_branches() {
        let base = PmlResult {
            beta: vec![0.0, 0.0],
            log_w: vec![0.0],
            information: vec![4.0, 0.0, 0.0, 4.0],
            log_likelihood: -1.0,
            log_likelihood_prev: -1.0,
            log_det: 0.0,
            second_order: 0.0,
            convergence: CONV_SUCCESS,
            iterations: 3,
            backtracks: 0,
        };
        let params = PmlParams::default();

        assert_eq!(
            check_convergence(&base, &params, false, CONV_SUCCESS),
            CONV_SUCCESS
        );
        assert_eq!(
            check_convergence(&base, &params, true, CONV_SUCCESS),
            CONV_AT_BOUND
        );

        let nan = PmlResult {
            log_likelihood: f64::NAN,
            ..base.clone()
        };
        assert_eq!(
            check_convergence(&nan, &params, false, CONV_SUCCESS),
            CONV_NOT_FINITE
        );

        // An incoming failure from the outer optimiser is carried through.
        assert_eq!(check_convergence(&base, &params, false, -50), -50);

        // The singular test overrides even that.
        let singular = PmlResult {
            information: vec![1.0, 1.0, 1.0, 1.0],
            ..base.clone()
        };
        assert_eq!(
            check_convergence(&singular, &params, true, CONV_SUCCESS),
            CONV_SINGULAR
        );
    }

    #[test]
    fn test_min_eigenvalue_known_matrix() {
        // Eigenvalues 1 and 3.
        let a = [2.0, 1.0, 1.0, 2.0];
        assert_relative_eq!(min_eigenvalue(&a, 2), 1.0, max_relative = 1e-13);
        assert_relative_eq!(min_eigenvalue(&[7.0], 1), 7.0, max_relative = 1e-15);
    }

    #[test]
    fn test_ldlt_solve_indefinite_system() {
        // Eigenvalues 3 and -1, so a plain Cholesky would fail here.
        let mut a = [1.0, 2.0, 2.0, 1.0];
        let mut b = [3.0, 3.0];
        let mut perm = [0usize; 2];
        let mut tmp = [0.0; 2];
        ldlt_solve(&mut a, &mut b, 2, &mut perm, &mut tmp);
        assert_relative_eq!(b[0], 1.0, max_relative = 1e-13);
        assert_relative_eq!(b[1], 1.0, max_relative = 1e-13);
    }

    #[test]
    fn test_determinant_matches_closed_form() {
        let mut a = [4.0, 1.0, 2.0, 1.0, 5.0, 3.0, 2.0, 3.0, 6.0];
        let mut perm = [0usize; 3];
        // 4(30-9) - 1(6-6) + 2(3-10) = 84 - 0 - 14 = 70.
        assert_relative_eq!(
            determinant(&mut a, 3, &mut perm),
            70.0,
            max_relative = 1e-13
        );
    }

    #[test]
    fn test_rejects_malformed_input() {
        let data = fixture_a();
        let variance = PmlVariance {
            subject: 0.25,
            cell: 2.0,
        };
        assert!(opt_pml(&data, &[0.5], &variance, None).is_err());

        let bad_gamma = PmlVariance {
            subject: 0.25,
            cell: 0.0,
        };
        assert!(opt_pml(&data, &[0.5, -0.25], &bad_gamma, None).is_err());

        let bad_subject = PmlVariance {
            subject: 0.0,
            cell: 2.0,
        };
        assert!(opt_pml(&data, &[0.5, -0.25], &bad_subject, None).is_err());

        let short = PmlData {
            subject_total: &SUBJECT_TOTAL[..2],
            ..data
        };
        assert!(opt_pml(&short, &[0.5, -0.25], &variance, None).is_err());
    }

    /// Three subjects of 400 cells, a subject-level group and a cell-level
    /// indicator, offsets over two orders of magnitude, about a tenth of the
    /// cells expressed.
    fn categorical_gene() -> (
        Vec<f64>,
        Vec<f64>,
        Vec<f64>,
        Vec<usize>,
        Vec<f64>,
        Vec<usize>,
    ) {
        use rand::prelude::*;
        use rand::rngs::SmallRng;
        let mut rng = SmallRng::seed_from_u64(11);
        let starts = vec![0usize, 400, 800, 1200];
        let mut design = Vec::new();
        let mut offset = Vec::new();
        let mut counts = Vec::new();
        let mut cell_index = Vec::new();
        let mut subject_total = vec![0.0; 3];
        for c in 0..1200 {
            let s = c / 400;
            let kind = f64::from(rng.random_bool(0.4));
            design.extend([1.0, (s % 2) as f64, kind]);
            let o: f64 = rng.random_range(-2.3..2.3);
            offset.push(o);
            if rng.random_bool(0.1 + 0.05 * kind) {
                let y = rng.random_range(1..6) as f64;
                counts.push(y);
                cell_index.push(c);
                subject_total[s] += y;
            }
        }
        (design, offset, counts, cell_index, subject_total, starts)
    }

    #[test]
    fn test_opt_pml_tables_match_the_dense_sweep() {
        let (design, offset, counts, cell_index, subject_total, starts) = categorical_gene();
        let data = PmlData {
            design: &design,
            offset: &offset,
            counts: &counts,
            cell_index: &cell_index,
            subject_start: &starts,
            subject_total: &subject_total,
        };
        let zeros = ZeroCells::build(&design, &offset, &starts, 3).expect("categorical");
        for (subject, cell) in [(0.3, 0.05), (0.8, 2.0), (0.05, 40.0)] {
            let variance = PmlVariance { subject, cell };
            let params = PmlParams {
                reml: true,
                ord: 3,
                ..Default::default()
            };
            let start = [-1.5, 0.2, 0.1];
            let dense = opt_pml(&data, &start, &variance, Some(params)).expect("dense");
            let tabled = opt_pml_tabled(&data, Some(&zeros), &start, &variance, Some(params))
                .expect("tabled");
            assert_eq!(dense.iterations, tabled.iterations);
            for (a, b) in dense.beta.iter().zip(&tabled.beta) {
                assert_relative_eq!(*a, *b, max_relative = 1e-11);
            }
            for (a, b) in dense.information.iter().zip(&tabled.information) {
                assert_relative_eq!(*a, *b, max_relative = 1e-11);
            }
            assert_relative_eq!(
                dense.log_likelihood,
                tabled.log_likelihood,
                max_relative = 1e-13
            );
            assert_relative_eq!(dense.log_det, tabled.log_det, max_relative = 1e-12);
            assert_relative_eq!(
                dense.second_order,
                tabled.second_order,
                max_relative = 1e-10
            );
        }
    }
}