limma/

optim.rs

//! Nelder–Mead simplex minimizer, a faithful port of R's `nmmin`
//! (`src/appl/optim.c`), which backs `optim(method = "Nelder-Mead")`.
//!
//! The algorithm is deterministic, so with the same starting point, objective
//! and control constants it reproduces R's iterate path — and therefore R's
//! `optim` result — bit for bit. This is what limma's `genas` relies on.
//!
//! Only the plain (unconstrained, unscaled) Nelder–Mead is ported; parameter
//! scaling (`parscale`) and the other `optim` methods are out of scope.

/// Result of [`nelder_mead`], mirroring the parts of R's `optim` return value
/// that callers need.
#[derive(Debug, Clone)]
pub struct NelderMead {
    /// Best parameter vector found (R's `$par`).
    pub par: Vec<f64>,
    /// Objective value at `par` (R's `$value`).
    pub value: f64,
    /// Number of objective evaluations (R's `$counts[1]`).
    pub fncount: usize,
    /// Convergence flag (R's `$convergence`): 0 = converged, 1 = `maxit`
    /// reached, 10 = simplex degenerated during shrink, 2 = objective was
    /// non-finite at the starting point.
    pub fail: i32,
}

// Sentinel substituted for non-finite objective values, matching R's `big`.
const BIG: f64 = 1.0e35;

/// Minimize `f` starting from `x0` using Nelder–Mead with R's default control
/// constants: `reltol = sqrt(f64::EPSILON)`, `abstol = -inf`, `maxit = 500`,
/// reflection `alpha = 1`, contraction `beta = 0.5`, expansion `gamma = 2`.
pub fn nelder_mead<F>(x0: &[f64], f: F) -> NelderMead
where
    F: FnMut(&[f64]) -> f64,
{
    nelder_mead_with(
        x0,
        f,
        f64::NEG_INFINITY,
        f64::EPSILON.sqrt(),
        500,
        1.0,
        0.5,
        2.0,
    )
}

/// Nelder–Mead with explicit control constants (see [`nelder_mead`] for the
/// defaults). `abstol`/`reltol` are R's `optim` `abstol`/`reltol`.
// The `for i in 0..n` / `for j in 0..n1` index loops below deliberately mirror
// the index arithmetic of R's `nmmin` C source (`p[i][col]`, `bvec[i]` updated
// in lockstep), so the iterate path stays bit-for-bit identical to R's.
#[allow(clippy::too_many_arguments, clippy::needless_range_loop)]
pub fn nelder_mead_with<F>(
    x0: &[f64],
    mut f: F,
    abstol: f64,
    reltol: f64,
    maxit: usize,
    alpha: f64,
    beta: f64,
    gamma: f64,
) -> NelderMead
where
    F: FnMut(&[f64]) -> f64,
{
    let n = x0.len();
    let mut bvec = x0.to_vec();

    if maxit == 0 {
        let value = f(&bvec);
        return NelderMead {
            par: bvec,
            value,
            fncount: 0,
            fail: 0,
        };
    }
    if n == 0 {
        let value = f(&bvec);
        return NelderMead {
            par: bvec,
            value,
            fncount: 1,
            fail: 0,
        };
    }

    // P has n+1 rows (n coordinate rows + one value row at index `n`) and n+2
    // columns (n+1 simplex vertices in columns 0..=n, plus a scratch/centroid
    // column at index n+1). Mirrors R's `matrix(n, n+1)`.
    let n1 = n + 1; // number of vertices
    let cidx = n + 1; // scratch column index (R's C-1)
    let vrow = n; // value row index (R's n1-1)
    let mut p = vec![vec![0.0_f64; n + 2]; n + 1];

    let f0 = f(&bvec);
    if !f0.is_finite() {
        return NelderMead {
            par: bvec,
            value: f0,
            fncount: 1,
            fail: 2,
        };
    }
    let mut funcount = 1usize;
    let convtol = reltol * (f0.abs() + reltol);

    p[vrow][0] = f0;
    for i in 0..n {
        p[i][0] = bvec[i];
    }

    let mut l = 1usize; // 1-indexed best vertex
    let mut size = 0.0;

    let mut step = 0.0;
    for i in 0..n {
        let s = 0.1 * bvec[i].abs();
        if s > step {
            step = s;
        }
    }
    if step == 0.0 {
        step = 0.1;
    }

    // Build the remaining n vertices by perturbing one coordinate each.
    for j in 2..=n1 {
        for i in 0..n {
            p[i][j - 1] = bvec[i];
        }
        let mut trystep = step;
        while p[j - 2][j - 1] == bvec[j - 2] {
            p[j - 2][j - 1] = bvec[j - 2] + trystep;
            trystep *= 10.0;
        }
        size += trystep;
    }
    let mut oldsize = size;
    let mut calcvert = true;
    let mut fail = 0i32;

    loop {
        if calcvert {
            for j in 0..n1 {
                if j + 1 != l {
                    for i in 0..n {
                        bvec[i] = p[i][j];
                    }
                    let mut fj = f(&bvec);
                    if !fj.is_finite() {
                        fj = BIG;
                    }
                    funcount += 1;
                    p[vrow][j] = fj;
                }
            }
            calcvert = false;
        }

        let mut vl = p[vrow][l - 1];
        let mut vh = vl;
        let mut h = l;
        for j in 1..=n1 {
            if j != l {
                let fj = p[vrow][j - 1];
                if fj < vl {
                    l = j;
                    vl = fj;
                }
                if fj > vh {
                    h = j;
                    vh = fj;
                }
            }
        }

        if vh <= vl + convtol || vl <= abstol {
            break;
        }

        // Centroid of all vertices except the worst (H), into the scratch col.
        for i in 0..n {
            let mut temp = -p[i][h - 1];
            for j in 0..n1 {
                temp += p[i][j];
            }
            p[i][cidx] = temp / n as f64;
        }
        // Reflect the worst vertex through the centroid.
        for i in 0..n {
            bvec[i] = (1.0 + alpha) * p[i][cidx] - alpha * p[i][h - 1];
        }
        let mut vr = f(&bvec);
        if !vr.is_finite() {
            vr = BIG;
        }
        funcount += 1;

        if vr < vl {
            // Reflection improved on the best: try to expand further.
            p[vrow][cidx] = vr;
            for i in 0..n {
                let fe = gamma * bvec[i] + (1.0 - gamma) * p[i][cidx];
                p[i][cidx] = bvec[i];
                bvec[i] = fe;
            }
            let mut fe = f(&bvec);
            if !fe.is_finite() {
                fe = BIG;
            }
            funcount += 1;
            if fe < vr {
                for i in 0..n {
                    p[i][h - 1] = bvec[i];
                }
                p[vrow][h - 1] = fe;
            } else {
                for i in 0..n {
                    p[i][h - 1] = p[i][cidx];
                }
                p[vrow][h - 1] = vr;
            }
        } else {
            // Reflection no better than the best.
            if vr < vh {
                for i in 0..n {
                    p[i][h - 1] = bvec[i];
                }
                p[vrow][h - 1] = vr;
            }
            // Contract.
            for i in 0..n {
                bvec[i] = (1.0 - beta) * p[i][h - 1] + beta * p[i][cidx];
            }
            let mut fc = f(&bvec);
            if !fc.is_finite() {
                fc = BIG;
            }
            funcount += 1;
            if fc < p[vrow][h - 1] {
                for i in 0..n {
                    p[i][h - 1] = bvec[i];
                }
                p[vrow][h - 1] = fc;
            } else if vr >= vh {
                // Shrink the whole simplex toward the best vertex.
                calcvert = true;
                size = 0.0;
                for j in 0..n1 {
                    if j + 1 != l {
                        for i in 0..n {
                            p[i][j] = beta * (p[i][j] - p[i][l - 1]) + p[i][l - 1];
                            size += (p[i][j] - p[i][l - 1]).abs();
                        }
                    }
                }
                if size < oldsize {
                    oldsize = size;
                } else {
                    fail = 10;
                    break;
                }
            }
        }

        if funcount > maxit {
            break;
        }
    }

    let value = p[vrow][l - 1];
    let par: Vec<f64> = (0..n).map(|i| p[i][l - 1]).collect();
    if funcount > maxit {
        fail = 1;
    }
    NelderMead {
        par,
        value,
        fncount: funcount,
        fail,
    }
}

#[cfg(test)]
#[allow(clippy::excessive_precision, clippy::approx_constant)]
mod tests {
    use super::*;

    fn rclose(a: f64, b: f64) -> bool {
        (a - b).abs() <= 1e-7 * (1.0 + b.abs())
    }

    // Rosenbrock: f(x,y) = 100*(y - x^2)^2 + (1 - x)^2, start (-1.2, 1).
    #[test]
    fn rosenbrock_matches_r_optim() {
        let res = nelder_mead(&[-1.2, 1.0], |x| {
            100.0 * (x[1] - x[0] * x[0]).powi(2) + (1.0 - x[0]).powi(2)
        });
        // optim(c(-1.2,1), fn, method="Nelder-Mead"):
        // par = (1.0002601387256695, 1.000505999303765), value = 8.8252410967e-08,
        // counts[1] = 195, convergence = 0.
        assert!(
            rclose(res.par[0], 1.0002601387256695),
            "par0 {}",
            res.par[0]
        );
        assert!(rclose(res.par[1], 1.000505999303765), "par1 {}", res.par[1]);
        assert!(
            rclose(res.value, 8.8252410967227472e-08),
            "val {}",
            res.value
        );
        assert_eq!(res.fncount, 195);
        assert_eq!(res.fail, 0);
    }

    // Rotated quadratic: f = (x1-1)^2 + 4*(x2-2)^2 + (x1-1)*(x2-2), start (0,0).
    #[test]
    fn rotated_quadratic_matches_r_optim() {
        let res = nelder_mead(&[0.0, 0.0], |x| {
            let a = x[0] - 1.0;
            let b = x[1] - 2.0;
            a * a + 4.0 * b * b + a * b
        });
        // optim(c(0,0), fn, method="Nelder-Mead"):
        // par = (1.000165999016468, 2.000030536283584), value = 3.6354524970e-08,
        // counts[1] = 65, convergence = 0.
        assert!(rclose(res.par[0], 1.000165999016468), "par0 {}", res.par[0]);
        assert!(rclose(res.par[1], 2.000030536283584), "par1 {}", res.par[1]);
        assert!(
            rclose(res.value, 3.6354524970336173e-08),
            "val {}",
            res.value
        );
        assert_eq!(res.fncount, 65);
        assert_eq!(res.fail, 0);
    }

    // Separable 3-D quadratic: f = sum (x - c)^2, c = (1,2,3), start (0,0,0).
    #[test]
    fn separable_3d_matches_r_optim() {
        let c = [1.0, 2.0, 3.0];
        let res = nelder_mead(&[0.0, 0.0, 0.0], |x| {
            (0..3).map(|i| (x[i] - c[i]).powi(2)).sum()
        });
        // optim(c(0,0,0), fn, method="Nelder-Mead"):
        // par = (1.0005383213703034, 1.9999848251163552, 2.9999188239843706),
        // value = 2.9660972033e-07, counts[1] = 112, convergence = 0.
        assert!(
            rclose(res.par[0], 1.0005383213703034),
            "par0 {}",
            res.par[0]
        );
        assert!(
            rclose(res.par[1], 1.9999848251163552),
            "par1 {}",
            res.par[1]
        );
        assert!(
            rclose(res.par[2], 2.9999188239843706),
            "par2 {}",
            res.par[2]
        );
        assert!(
            rclose(res.value, 2.9660972033243571e-07),
            "val {}",
            res.value
        );
        assert_eq!(res.fncount, 112);
        assert_eq!(res.fail, 0);
    }
}