regression-diagnostics 0.2.0

//! Small pure-Rust optimizers used by the maximum-likelihood models that have no
//! closed form (parametric survival, general mixed models, GLMM).
//!
//! Two tools, both dependency-free:
//!
//! * [`nelder_mead`] — a derivative-free simplex minimizer, for objectives whose
//!   gradient is awkward (the profiled REML criterion, the Laplace marginal).
//! * [`numerical_hessian`] — a central-difference Hessian of a supplied gradient,
//!   for turning a converged MLE into a covariance matrix.
//!
//! Neither appears in the public API.

/// Minimize `f` by the Nelder–Mead downhill simplex method.
///
/// Starts from `x0`, builds the initial simplex by stepping `step` along each
/// axis, and iterates the standard reflect/expand/contract/shrink moves until
/// the simplex spread (in both parameters and function value) falls below `tol`
/// or `max_iter` iterations elapse. Returns the best vertex found.
pub(crate) fn nelder_mead(
    f: impl Fn(&[f64]) -> f64,
    x0: &[f64],
    step: f64,
    tol: f64,
    max_iter: usize,
) -> Vec<f64> {
    let n = x0.len();
    if n == 0 {
        return Vec::new();
    }
    // Simplex of n+1 vertices.
    let mut simplex: Vec<Vec<f64>> = Vec::with_capacity(n + 1);
    simplex.push(x0.to_vec());
    for i in 0..n {
        let mut v = x0.to_vec();
        let s = if x0[i].abs() > 1e-8 {
            step * x0[i].abs()
        } else {
            step
        };
        v[i] += s;
        simplex.push(v);
    }
    let mut fvals: Vec<f64> = simplex.iter().map(|v| f(v)).collect();

    // Coefficients: reflection, expansion, contraction, shrink.
    let (alpha, gamma, rho, sigma) = (1.0, 2.0, 0.5, 0.5);

    for _ in 0..max_iter {
        // Order vertices by function value (best first).
        let mut order: Vec<usize> = (0..=n).collect();
        order.sort_by(|&a, &b| fvals[a].partial_cmp(&fvals[b]).unwrap());
        let best = order[0];
        let worst = order[n];
        let second_worst = order[n - 1];

        // Convergence: spread in f and in the simplex.
        let fspread = (fvals[worst] - fvals[best]).abs();
        let mut xspread = 0.0_f64;
        for (a, b) in simplex[worst].iter().zip(simplex[best].iter()) {
            xspread = xspread.max((a - b).abs());
        }
        if fspread < tol && xspread < tol {
            break;
        }

        // Centroid of all but the worst vertex.
        let mut centroid = vec![0.0; n];
        for (i, v) in simplex.iter().enumerate() {
            if i == worst {
                continue;
            }
            for j in 0..n {
                centroid[j] += v[j] / n as f64;
            }
        }

        // Reflection.
        let reflected: Vec<f64> = (0..n)
            .map(|j| centroid[j] + alpha * (centroid[j] - simplex[worst][j]))
            .collect();
        let fr = f(&reflected);

        if fr < fvals[best] {
            // Expansion.
            let expanded: Vec<f64> = (0..n)
                .map(|j| centroid[j] + gamma * (reflected[j] - centroid[j]))
                .collect();
            let fe = f(&expanded);
            if fe < fr {
                simplex[worst] = expanded;
                fvals[worst] = fe;
            } else {
                simplex[worst] = reflected;
                fvals[worst] = fr;
            }
        } else if fr < fvals[second_worst] {
            simplex[worst] = reflected;
            fvals[worst] = fr;
        } else {
            // Contraction.
            let contracted: Vec<f64> = (0..n)
                .map(|j| centroid[j] + rho * (simplex[worst][j] - centroid[j]))
                .collect();
            let fc = f(&contracted);
            if fc < fvals[worst] {
                simplex[worst] = contracted;
                fvals[worst] = fc;
            } else {
                // Shrink toward the best vertex.
                let best_v = simplex[best].clone();
                for (i, v) in simplex.iter_mut().enumerate() {
                    if i == best {
                        continue;
                    }
                    for j in 0..n {
                        v[j] = best_v[j] + sigma * (v[j] - best_v[j]);
                    }
                    fvals[i] = f(v);
                }
            }
        }
    }

    let best = (0..=n)
        .min_by(|&a, &b| fvals[a].partial_cmp(&fvals[b]).unwrap())
        .unwrap();
    simplex[best].clone()
}

/// Central-difference Hessian of a function whose analytic `grad` is supplied,
/// evaluated at `x`. Symmetrized to counter finite-difference asymmetry.
pub(crate) fn numerical_hessian(
    grad: impl Fn(&[f64]) -> Vec<f64>,
    x: &[f64],
) -> Vec<Vec<f64>> {
    let m = x.len();
    let mut h = vec![vec![0.0; m]; m];
    for j in 0..m {
        let step = 1e-6 * x[j].abs().max(1.0);
        let mut xp = x.to_vec();
        let mut xm = x.to_vec();
        xp[j] += step;
        xm[j] -= step;
        let gp = grad(&xp);
        let gm = grad(&xm);
        for (i, hij) in h.iter_mut().enumerate() {
            hij[j] = (gp[i] - gm[i]) / (2.0 * step);
        }
    }
    // Symmetrize; the transposed writes make an iterator form awkward.
    #[allow(clippy::needless_range_loop)]
    for i in 0..m {
        for j in (i + 1)..m {
            let avg = 0.5 * (h[i][j] + h[j][i]);
            h[i][j] = avg;
            h[j][i] = avg;
        }
    }
    h
}