edge_rs/numeric/

optimise.rs

//! Scalar minimisation, root finding and derivative-free simplex search.
//!
//! Covers the `scipy.optimize` calls in edgePython except the bounded
//! quasi-Newton, which lives in [`crate::numeric::lbfgsb`]:
//!
//! * `brentq` is [`brentq`], used by `disp_pearson`, `disp_deviance` and
//!   limma's `fitFDist`
//! * `minimize(method = 'Nelder-Mead')` is [`nelder_mead`], used by the spline
//!   and power trend fitters
//!
//! [`brent_fmin`] is R's `Brent_fmin` (not a `scipy.optimize` port), used by
//! `disp_cox_reid` and limma's `squeezeVar`, where stopping at R's point
//! matters.
//!
//! [`nelder_mead`] reproduces scipy's initial simplex exactly: the trend fits
//! minimise a nearly flat objective and are sensitive to it.

use std::cmp::Ordering;

use crate::prelude::*;

////////////
// Consts //
////////////

/// Default `tol` of R's `optimize`, `.Machine$double.eps^0.25`, which is exactly
/// `2^-13`.
pub const OPTIMIZE_TOL: f64 = 1.220_703_125e-4;

/// `sqrt(DBL_EPSILON)`, exactly `2^-26`, the relative half of `Brent_fmin`'s
/// convergence test.
const BRENT_SQRT_EPS: f64 = 1.490_116_119_384_765_6e-8;

/// Golden section step `(3 - sqrt(5)) / 2`, as `Brent_fmin` spells it.
const BRENT_GOLDEN: f64 = 0.381_966_011_250_105_15;

/// Relative step used to build scipy's initial Nelder-Mead simplex.
///
/// scipy perturbs each coordinate by 5% of its value.
const NELDER_MEAD_STEP: f64 = 0.05;

/// Absolute perturbation used where a starting coordinate is exactly zero.
///
/// A relative step cannot move a zero.
const NELDER_MEAD_ZERO_STEP: f64 = 0.000_25;

/// Simplex extent, in units of `f64::EPSILON` times the best vertex's largest
/// coordinate, below which Nelder-Mead stops whatever the objective says.
///
/// scipy needs both tolerances to pass, which a discontinuous objective can
/// prevent: NEBULA's profile likelihood jumps by a few `1e-7` wherever an inner
/// Newton loop changes its step count, and a simplex shrunk onto such a jump
/// cycles until the iteration cap. Measured on one gene in 500: vertices two ulp
/// apart, a `4.5e-7` jump against an `fatol` of `1e-7`, 1650 wasted evaluations.
/// Stopping returns the same minimiser to the last few bits.
const NELDER_MEAD_COLLAPSED_ULPS: f64 = 16.0;

//////////////////
// Root finding //
//////////////////

/// Finds a root of a continuous function on a bracketing interval.
///
/// Port of `scipy.optimize.brentq`: bisection, accelerated by inverse quadratic
/// interpolation and the secant rule.
///
/// ### Params
///
/// * `f` - Function whose root is wanted
/// * `lower` - Lower end of the bracket
/// * `upper` - Upper end of the bracket
/// * `xtol` - Absolute tolerance on the root
/// * `rtol` - Relative tolerance on the root
/// * `max_iter` - Iteration budget
///
/// ### Returns
///
/// The root, or [`EdgeErrors::RootNotBracketed`] if the endpoints share a sign,
/// or [`EdgeErrors::NoConvergence`] if the budget runs out.
///
/// ### References
///
/// Brent, Algorithms for Minimization without Derivatives, 1973, chapter 4
pub fn brentq<F>(
    mut f: F,
    lower: f64,
    upper: f64,
    xtol: f64,
    rtol: f64,
    max_iter: usize,
) -> Result<f64, EdgeErrors>
where
    F: FnMut(f64) -> f64,
{
    if lower.partial_cmp(&upper) != Some(Ordering::Less) {
        return Err(EdgeErrors::InvalidArgument(format!(
            "brentq needs lower < upper, got [{lower}, {upper}]"
        )));
    }
    if xtol <= 0.0 {
        return Err(EdgeErrors::MustBePositive("xtol".to_string()));
    }

    let (mut xpre, mut xcur) = (lower, upper);
    let (mut fpre, mut fcur) = (f(xpre), f(xcur));

    if fpre == 0.0 {
        return Ok(xpre);
    }
    if fcur == 0.0 {
        return Ok(xcur);
    }
    if fpre.signum() == fcur.signum() {
        return Err(EdgeErrors::RootNotBracketed {
            lower,
            upper,
            f_lower: fpre,
            f_upper: fcur,
        });
    }

    let (mut xblk, mut fblk) = (0.0, 0.0);
    let (mut spre, mut scur) = (0.0, 0.0);

    for _ in 0..max_iter {
        if fpre != 0.0 && fcur != 0.0 && fpre.signum() != fcur.signum() {
            xblk = xpre;
            fblk = fpre;
            spre = xcur - xpre;
            scur = spre;
        }
        if fblk.abs() < fcur.abs() {
            xpre = xcur;
            xcur = xblk;
            xblk = xpre;
            fpre = fcur;
            fcur = fblk;
            fblk = fpre;
        }

        let delta = (xtol + rtol * xcur.abs()) / 2.0;
        let sbis = (xblk - xcur) / 2.0;
        if fcur == 0.0 || sbis.abs() < delta {
            return Ok(xcur);
        }

        if spre.abs() > delta && fcur.abs() < fpre.abs() {
            let stry = if xpre == xblk {
                // Secant.
                -fcur * (xcur - xpre) / (fcur - fpre)
            } else {
                // Inverse quadratic interpolation.
                let dpre = (fpre - fcur) / (xpre - xcur);
                let dblk = (fblk - fcur) / (xblk - xcur);
                -fcur * (fblk * dblk - fpre * dpre) / (dblk * dpre * (fblk - fpre))
            };
            if 2.0 * stry.abs() < (spre.abs()).min(3.0 * sbis.abs() - delta) {
                spre = scur;
                scur = stry;
            } else {
                spre = sbis;
                scur = sbis;
            }
        } else {
            spre = sbis;
            scur = sbis;
        }

        xpre = xcur;
        fpre = fcur;
        if scur.abs() > delta {
            xcur += scur;
        } else {
            xcur += if sbis > 0.0 { delta } else { -delta };
        }
        fcur = f(xcur);
    }

    Err(EdgeErrors::NoConvergence {
        routine: "brentq",
        iterations: max_iter,
        last_delta: (xblk - xcur).abs(),
    })
}

/// Minimises a scalar function on a closed interval, as R's `optimize` does.
///
/// Line-by-line port of R's `Brent_fmin`, including its convergence test
/// (`sqrt(eps) * |x| + tol / 3`), so it stops where limma stops.
///
/// ### Params
///
/// * `ax` - Lower end of the interval
/// * `bx` - Upper end of the interval
/// * `f` - Objective
/// * `tol` - Convergence tolerance, R's `optimize(tol = )`
///
/// ### Returns
///
/// The minimiser.
///
/// ### References
///
/// Brent, Algorithms for Minimization without Derivatives, 1973, chapter 5
pub fn brent_fmin<F>(ax: f64, bx: f64, mut f: F, tol: f64) -> f64
where
    F: FnMut(f64) -> f64,
{
    let (mut a, mut b) = (ax, bx);
    let mut v = a + BRENT_GOLDEN * (b - a);
    let (mut w, mut x) = (v, v);
    let mut d = 0.0_f64;
    let mut e = 0.0_f64;
    let mut fx = f(x);
    let (mut fv, mut fw) = (fx, fx);
    let tol3 = tol / 3.0;

    loop {
        let xm = 0.5 * (a + b);
        let tol1 = BRENT_SQRT_EPS * x.abs() + tol3;
        let t2 = 2.0 * tol1;
        if (x - xm).abs() <= t2 - 0.5 * (b - a) {
            return x;
        }

        let (mut p, mut q, mut r) = (0.0_f64, 0.0_f64, 0.0_f64);
        if e.abs() > tol1 {
            r = (x - w) * (fx - fv);
            q = (x - v) * (fx - fw);
            p = (x - v) * q - (x - w) * r;
            q = 2.0 * (q - r);
            if q > 0.0 {
                p = -p;
            } else {
                q = -q;
            }
            r = e;
            e = d;
        }

        if p.abs() >= (0.5 * q * r).abs() || p <= q * (a - x) || p >= q * (b - x) {
            e = if x < xm { b - x } else { a - x };
            d = BRENT_GOLDEN * e;
        } else {
            d = p / q;
            let u = x + d;
            // f must not be evaluated too close to either end of the interval.
            if u - a < t2 || b - u < t2 {
                d = if x >= xm { -tol1 } else { tol1 };
            }
        }

        // Nor too close to the incumbent.
        let u = if d.abs() >= tol1 {
            x + d
        } else if d > 0.0 {
            x + tol1
        } else {
            x - tol1
        };
        let fu = f(u);

        if fu <= fx {
            if u < x {
                b = x;
            } else {
                a = x;
            }
            v = w;
            fv = fw;
            w = x;
            fw = fx;
            x = u;
            fx = fu;
        } else {
            if u < x {
                a = u;
            } else {
                b = u;
            }
            if fu <= fw || w == x {
                v = w;
                fv = fw;
                w = u;
                fw = fu;
            } else if fu <= fv || v == x || v == w {
                v = u;
                fv = fu;
            }
        }
    }
}

/////////////////
// Nelder-Mead //
/////////////////

/// Tuning knobs for [`nelder_mead`].
#[derive(Clone, Copy, Debug)]
pub struct NelderMeadParams {
    /// Absolute tolerance on the spread of the simplex vertices.
    pub xatol: f64,
    /// Absolute tolerance on the spread of the objective across the simplex.
    pub fatol: f64,
    /// Iteration budget.
    pub max_iter: usize,
}

impl Default for NelderMeadParams {
    /// `xatol`/`fatol` match scipy's defaults. `max_iter` is a fixed fallback,
    /// not scipy's rule (scipy defaults to `N * 200` for `N` parameters).
    fn default() -> Self {
        Self {
            xatol: 1e-4,
            fatol: 1e-4,
            max_iter: 1000,
        }
    }
}

/// Outcome of a simplex search.
#[derive(Clone, Debug)]
pub struct NelderMeadResult {
    /// Best vertex found.
    pub x: Vec<f64>,
    /// Objective at that vertex.
    pub f: f64,
    /// Iterations performed.
    pub iterations: usize,
    /// Whether both tolerances were met before the budget ran out.
    pub converged: bool,
}

/// Derivative-free minimisation by the Nelder-Mead simplex method.
///
/// The initial simplex is scipy's: the start plus one vertex per coordinate,
/// perturbed by 5% or a small absolute step at zero. Reflection, expansion,
/// contraction and shrink use the coefficients 1, 2, 0.5 and 0.5.
///
/// ### Params
///
/// * `f` - Objective
/// * `x0` - Starting point
/// * `params` - Tuning knobs, or [`NelderMeadParams::default`]
///
/// ### Returns
///
/// The best vertex and its value, or [`EdgeErrors::MustBePositive`] if `x0` is
/// empty.
pub fn nelder_mead<F>(
    mut f: F,
    x0: &[f64],
    params: Option<NelderMeadParams>,
) -> Result<NelderMeadResult, EdgeErrors>
where
    F: FnMut(&[f64]) -> f64,
{
    let mut stepper = NelderMeadStepper::new(x0, params)?;
    while let Some(x) = stepper.ask() {
        let value = f(x);
        stepper.tell(value);
    }
    Ok(stepper.result())
}

/// Where a [`NelderMeadStepper`] is waiting for a value.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum NelderMeadPhase {
    /// Evaluating vertex `i` of the initial simplex.
    Init(usize),
    /// Evaluating the reflected point.
    Reflect,
    /// Evaluating the expanded point.
    Expand,
    /// Evaluating the contracted point.
    Contract,
    /// Evaluating the shrunk vertex at position `i` of the sorted order.
    Shrink(usize),
    /// Finished.
    Done,
}

/// [`nelder_mead`] as a reverse-communication state machine.
///
/// The caller asks for the next point, evaluates it, and reports the value, so
/// many independent searches can advance in lockstep with batched evaluations.
/// Points and arithmetic are identical to the closure-driven form:
/// [`nelder_mead`] is this stepper driven by a closure.
#[derive(Clone, Debug)]
pub struct NelderMeadStepper {
    /// Tuning knobs.
    params: NelderMeadParams,
    /// Dimension.
    n: usize,
    /// Vertices, `n + 1` of them.
    simplex: Vec<Vec<f64>>,
    /// Objective at each vertex.
    values: Vec<f64>,
    /// Vertex indices, sorted by value at the top of every iteration. Kept
    /// across iterations so the stable sort keeps the previous order on ties.
    order: Vec<usize>,
    /// Index of the next iteration to run.
    next_iter: usize,
    /// Iterations performed so far.
    iterations: usize,
    /// Whether both tolerances were met.
    converged: bool,
    /// What the pending value is for.
    phase: NelderMeadPhase,
    /// The point awaiting a value.
    pending: Vec<f64>,
    /// Centroid of all but the worst vertex, for the current iteration.
    centroid: Vec<f64>,
    /// The reflected point of the current iteration.
    reflected: Vec<f64>,
    /// Its value.
    f_reflected: f64,
    /// Whether the current contraction is the inside one.
    inside: bool,
    /// Best vertex of the current iteration.
    best: usize,
    /// Worst vertex of the current iteration.
    worst: usize,
    /// Second worst vertex of the current iteration.
    second_worst: usize,
}

impl NelderMeadStepper {
    /// Builds the initial simplex, scipy's, and asks for its first vertex.
    ///
    /// ### Params
    ///
    /// * `x0` - Starting point
    /// * `params` - Tuning knobs, or [`NelderMeadParams::default`]
    ///
    /// ### Returns
    ///
    /// The stepper, or [`EdgeErrors::MustBePositive`] if `x0` is empty.
    pub fn new(x0: &[f64], params: Option<NelderMeadParams>) -> Result<Self, EdgeErrors> {
        let params = params.unwrap_or_default();
        let n = x0.len();
        if n == 0 {
            return Err(EdgeErrors::MustBePositive("x0.len()".to_string()));
        }
        let mut simplex: Vec<Vec<f64>> = Vec::with_capacity(n + 1);
        simplex.push(x0.to_vec());
        for i in 0..n {
            let mut vertex = x0.to_vec();
            if vertex[i] != 0.0 {
                vertex[i] *= 1.0 + NELDER_MEAD_STEP;
            } else {
                vertex[i] = NELDER_MEAD_ZERO_STEP;
            }
            simplex.push(vertex);
        }
        let pending = simplex[0].clone();
        Ok(Self {
            params,
            n,
            simplex,
            values: vec![0.0; n + 1],
            order: (0..=n).collect(),
            next_iter: 0,
            iterations: 0,
            converged: false,
            phase: NelderMeadPhase::Init(0),
            pending,
            centroid: vec![0.0; n],
            reflected: vec![0.0; n],
            f_reflected: 0.0,
            inside: false,
            best: 0,
            worst: 0,
            second_worst: 0,
        })
    }

    /// The point the search wants evaluated next.
    ///
    /// ### Returns
    ///
    /// The point, or `None` once the search has finished.
    pub fn ask(&self) -> Option<&[f64]> {
        if self.phase == NelderMeadPhase::Done {
            None
        } else {
            Some(&self.pending)
        }
    }

    /// Hands the stepper the value at the point [`Self::ask`] returned.
    ///
    /// ### Params
    ///
    /// * `value` - The objective at the pending point
    pub fn tell(&mut self, value: f64) {
        let n = self.n;
        match self.phase {
            NelderMeadPhase::Init(i) => {
                self.values[i] = value;
                if i < n {
                    self.phase = NelderMeadPhase::Init(i + 1);
                    self.pending.clone_from(&self.simplex[i + 1]);
                } else {
                    self.begin_iteration();
                }
            }
            NelderMeadPhase::Reflect => {
                self.f_reflected = value;
                let worst = self.worst;
                if value < self.values[self.best] {
                    self.pending = (0..n)
                        .map(|j| {
                            self.centroid[j] + 2.0 * (self.centroid[j] - self.simplex[worst][j])
                        })
                        .collect();
                    self.phase = NelderMeadPhase::Expand;
                } else if value < self.values[self.second_worst] {
                    self.simplex[worst].clone_from(&self.reflected);
                    self.values[worst] = value;
                    self.begin_iteration();
                } else {
                    self.inside = value >= self.values[worst];
                    self.pending = if self.inside {
                        (0..n)
                            .map(|j| {
                                self.centroid[j] - 0.5 * (self.centroid[j] - self.simplex[worst][j])
                            })
                            .collect()
                    } else {
                        (0..n)
                            .map(|j| {
                                self.centroid[j] + 0.5 * (self.centroid[j] - self.simplex[worst][j])
                            })
                            .collect()
                    };
                    self.phase = NelderMeadPhase::Contract;
                }
            }
            NelderMeadPhase::Expand => {
                let worst = self.worst;
                if value < self.f_reflected {
                    self.simplex[worst] = std::mem::take(&mut self.pending);
                    self.values[worst] = value;
                } else {
                    self.simplex[worst].clone_from(&self.reflected);
                    self.values[worst] = self.f_reflected;
                }
                self.begin_iteration();
            }
            NelderMeadPhase::Contract => {
                let worst = self.worst;
                let accept = if self.inside {
                    value < self.values[worst]
                } else {
                    value <= self.f_reflected
                };
                if accept {
                    self.simplex[worst] = std::mem::take(&mut self.pending);
                    self.values[worst] = value;
                    self.begin_iteration();
                } else {
                    // Shrink every vertex towards the best one, one at a time in
                    // sorted order; each only touches itself.
                    self.shrink_vertex(1);
                }
            }
            NelderMeadPhase::Shrink(pos) => {
                let i = self.order[pos];
                self.values[i] = value;
                if pos < n {
                    self.shrink_vertex(pos + 1);
                } else {
                    self.begin_iteration();
                }
            }
            NelderMeadPhase::Done => {}
        }
    }

    /// Moves the vertex at sorted position `pos` halfway to the best one and
    /// asks for its value.
    ///
    /// ### Params
    ///
    /// * `pos` - Position in the sorted order, from `1` to `n`
    fn shrink_vertex(&mut self, pos: usize) {
        let i = self.order[pos];
        let best = self.best;
        for j in 0..self.n {
            let anchor = self.simplex[best][j];
            self.simplex[i][j] = anchor + 0.5 * (self.simplex[i][j] - anchor);
        }
        self.pending.clone_from(&self.simplex[i]);
        self.phase = NelderMeadPhase::Shrink(pos);
    }

    /// Runs the top of an iteration: sort, test convergence, reflect.
    fn begin_iteration(&mut self) {
        let n = self.n;
        if self.next_iter >= self.params.max_iter {
            self.finish();
            return;
        }
        self.iterations = self.next_iter + 1;
        self.next_iter += 1;

        let values = &self.values;
        self.order.sort_by(|&a, &b| values[a].total_cmp(&values[b]));
        self.best = self.order[0];
        self.worst = self.order[n];
        self.second_worst = self.order[n - 1];
        let best = self.best;

        let spread_x = self.order[1..]
            .iter()
            .map(|&i| {
                self.simplex[i]
                    .iter()
                    .zip(self.simplex[best].iter())
                    .map(|(a, b)| (a - b).abs())
                    .fold(0.0_f64, f64::max)
            })
            .fold(0.0_f64, f64::max);
        let spread_f = self.order[1..]
            .iter()
            .map(|&i| (self.values[i] - self.values[best]).abs())
            .fold(0.0_f64, f64::max);
        let scale = self.simplex[best]
            .iter()
            .fold(f64::MIN_POSITIVE, |m, v| m.max(v.abs()));
        let collapsed = spread_x <= NELDER_MEAD_COLLAPSED_ULPS * f64::EPSILON * scale;
        if collapsed || (spread_x <= self.params.xatol && spread_f <= self.params.fatol) {
            self.converged = true;
            self.finish();
            return;
        }

        self.centroid.iter_mut().for_each(|c| *c = 0.0);
        for &i in &self.order[..n] {
            for (c, v) in self.centroid.iter_mut().zip(self.simplex[i].iter()) {
                *c += *v;
            }
        }
        for c in self.centroid.iter_mut() {
            *c /= n as f64;
        }
        let worst = self.worst;
        self.reflected = (0..n)
            .map(|j| self.centroid[j] + (self.centroid[j] - self.simplex[worst][j]))
            .collect();
        self.pending.clone_from(&self.reflected);
        self.phase = NelderMeadPhase::Reflect;
    }

    /// Sorts once more and stops.
    fn finish(&mut self) {
        let values = &self.values;
        self.order.sort_by(|&a, &b| values[a].total_cmp(&values[b]));
        self.phase = NelderMeadPhase::Done;
    }

    /// The best vertex found so far.
    ///
    /// ### Returns
    ///
    /// The result; meaningful once [`Self::ask`] returns `None`.
    pub fn result(&self) -> NelderMeadResult {
        let best = self.order[0];
        NelderMeadResult {
            x: self.simplex[best].clone(),
            f: self.values[best],
            iterations: self.iterations,
            converged: self.converged,
        }
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use approx::assert_relative_eq;

    #[test]
    fn test_brentq_finds_a_polynomial_root() {
        // x^3 - 2x - 5 has its real root near 2.0945514815423265.
        let root = brentq(
            |x: f64| x.powi(3) - 2.0 * x - 5.0,
            1.0,
            3.0,
            1e-14,
            1e-15,
            200,
        )
        .unwrap();
        assert_relative_eq!(root, 2.094_551_481_542_326_5, max_relative = 1e-12);
    }

    #[test]
    fn test_brentq_handles_a_root_at_an_endpoint() {
        let root = brentq(|x: f64| x - 1.0, 1.0, 3.0, 1e-14, 1e-15, 200).unwrap();
        assert_relative_eq!(root, 1.0, max_relative = 1e-14);
    }

    #[test]
    fn test_brentq_rejects_an_unbracketed_interval() {
        let err = brentq(|x: f64| x * x + 1.0, -1.0, 1.0, 1e-12, 1e-14, 100).unwrap_err();
        assert!(matches!(err, EdgeErrors::RootNotBracketed { .. }));
    }

    /// A steep function where bisection alone would be slow, so the
    /// interpolation branches actually run.
    #[test]
    fn test_brentq_converges_on_a_steep_function() {
        let root = brentq(|x: f64| x.exp() - 1e6, 0.0, 30.0, 1e-12, 1e-14, 200).unwrap();
        assert_relative_eq!(root, 1e6_f64.ln(), max_relative = 1e-10);
    }

    #[test]
    fn test_nelder_mead_minimises_rosenbrock() {
        let res = nelder_mead(
            |x: &[f64]| 100.0 * (x[1] - x[0] * x[0]).powi(2) + (1.0 - x[0]).powi(2),
            &[-1.2, 1.0],
            Some(NelderMeadParams {
                xatol: 1e-10,
                fatol: 1e-12,
                max_iter: 5000,
            }),
        )
        .unwrap();
        assert!(res.converged);
        assert_relative_eq!(res.x[0], 1.0, max_relative = 1e-4);
        assert_relative_eq!(res.x[1], 1.0, max_relative = 1e-4);
    }

    #[test]
    fn test_nelder_mead_handles_a_zero_starting_coordinate() {
        // The zero coordinate exercises the absolute-step branch of the simplex.
        let res = nelder_mead(
            |x: &[f64]| (x[0] - 3.0).powi(2) + (x[1] + 2.0).powi(2),
            &[0.0, 0.0],
            Some(NelderMeadParams {
                xatol: 1e-10,
                fatol: 1e-12,
                max_iter: 5000,
            }),
        )
        .unwrap();
        assert_relative_eq!(res.x[0], 3.0, max_relative = 1e-5);
        assert_relative_eq!(res.x[1], -2.0, max_relative = 1e-5);
    }

    #[test]
    fn test_nelder_mead_rejects_empty_start() {
        let err = nelder_mead(|_: &[f64]| 0.0, &[], None).unwrap_err();
        assert!(matches!(err, EdgeErrors::MustBePositive(_)));
    }

    /// Closure-driven Nelder-Mead, the specification [`NelderMeadStepper`] is
    /// checked against bit for bit.
    fn nelder_mead_reference<F>(
        mut f: F,
        x0: &[f64],
        params: Option<NelderMeadParams>,
    ) -> Result<NelderMeadResult, EdgeErrors>
    where
        F: FnMut(&[f64]) -> f64,
    {
        let params = params.unwrap_or_default();
        let n = x0.len();
        if n == 0 {
            return Err(EdgeErrors::MustBePositive("x0.len()".to_string()));
        }

        // scipy's initial simplex.
        let mut simplex: Vec<Vec<f64>> = Vec::with_capacity(n + 1);
        simplex.push(x0.to_vec());
        for i in 0..n {
            let mut vertex = x0.to_vec();
            if vertex[i] != 0.0 {
                vertex[i] *= 1.0 + NELDER_MEAD_STEP;
            } else {
                vertex[i] = NELDER_MEAD_ZERO_STEP;
            }
            simplex.push(vertex);
        }

        let mut values: Vec<f64> = simplex.iter().map(|v| f(v)).collect();
        let mut order: Vec<usize> = (0..=n).collect();
        let mut iterations = 0;
        let mut converged = false;

        for iter in 0..params.max_iter {
            iterations = iter + 1;

            order.sort_by(|&a, &b| values[a].total_cmp(&values[b]));
            let best = order[0];
            let worst = order[n];
            let second_worst = order[n - 1];

            // Convergence: both the simplex and the objective values have collapsed.
            let spread_x = order[1..]
                .iter()
                .map(|&i| {
                    simplex[i]
                        .iter()
                        .zip(simplex[best].iter())
                        .map(|(a, b)| (a - b).abs())
                        .fold(0.0_f64, f64::max)
                })
                .fold(0.0_f64, f64::max);
            let spread_f = order[1..]
                .iter()
                .map(|&i| (values[i] - values[best]).abs())
                .fold(0.0_f64, f64::max);
            if spread_x <= params.xatol && spread_f <= params.fatol {
                converged = true;
                break;
            }

            // Centroid of everything but the worst vertex.
            let mut centroid = vec![0.0; n];
            for &i in &order[..n] {
                for (c, v) in centroid.iter_mut().zip(simplex[i].iter()) {
                    *c += *v;
                }
            }
            for c in centroid.iter_mut() {
                *c /= n as f64;
            }

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

            if f_reflected < values[best] {
                let expanded: Vec<f64> = (0..n)
                    .map(|j| centroid[j] + 2.0 * (centroid[j] - simplex[worst][j]))
                    .collect();
                let f_expanded = f(&expanded);
                if f_expanded < f_reflected {
                    simplex[worst] = expanded;
                    values[worst] = f_expanded;
                } else {
                    simplex[worst] = reflected;
                    values[worst] = f_reflected;
                }
            } else if f_reflected < values[second_worst] {
                simplex[worst] = reflected;
                values[worst] = f_reflected;
            } else {
                let inside = f_reflected >= values[worst];
                let contracted: Vec<f64> = if inside {
                    (0..n)
                        .map(|j| centroid[j] - 0.5 * (centroid[j] - simplex[worst][j]))
                        .collect()
                } else {
                    (0..n)
                        .map(|j| centroid[j] + 0.5 * (centroid[j] - simplex[worst][j]))
                        .collect()
                };
                let f_contracted = f(&contracted);
                let accept = if inside {
                    f_contracted < values[worst]
                } else {
                    f_contracted <= f_reflected
                };

                if accept {
                    simplex[worst] = contracted;
                    values[worst] = f_contracted;
                } else {
                    // Shrink every vertex towards the best one.
                    let anchor = simplex[best].clone();
                    for &i in &order[1..] {
                        for j in 0..n {
                            simplex[i][j] = anchor[j] + 0.5 * (simplex[i][j] - anchor[j]);
                        }
                        values[i] = f(&simplex[i]);
                    }
                }
            }
        }

        order.sort_by(|&a, &b| values[a].total_cmp(&values[b]));
        let best = order[0];
        Ok(NelderMeadResult {
            x: simplex[best].clone(),
            f: values[best],
            iterations,
            converged,
        })
    }

    /// Drives a [`NelderMeadStepper`] with a closure, recording every point.
    fn drive_stepper(
        f: &dyn Fn(&[f64]) -> f64,
        x0: &[f64],
        params: NelderMeadParams,
    ) -> (NelderMeadResult, Vec<Vec<f64>>) {
        let mut stepper = NelderMeadStepper::new(x0, Some(params)).unwrap();
        let mut asked = Vec::new();
        while let Some(x) = stepper.ask() {
            asked.push(x.to_vec());
            let v = f(x);
            stepper.tell(v);
        }
        (stepper.result(), asked)
    }

    #[test]
    fn test_nelder_mead_stepper_is_bit_identical_to_the_closure_form() {
        let rosenbrock = |x: &[f64]| (1.0 - x[0]).powi(2) + 100.0 * (x[1] - x[0] * x[0]).powi(2);
        let quadratic_1d = |x: &[f64]| (x[0] - 0.3).powi(2) + 1.0;
        // Plateaus force ties, so the stable sort's ordering has to match.
        let staircase = |x: &[f64]| (x[0] * 4.0).floor().abs() + (x[1] * 4.0).floor().abs();
        // Non-convex with a narrow valley: exercises shrink steps.
        let bumpy =
            |x: &[f64]| (x[0] * 3.0).sin() * (x[1] * 5.0).cos() + 0.1 * (x[0] * x[0] + x[1] * x[1]);
        let three_d = |x: &[f64]| {
            x.iter()
                .enumerate()
                .map(|(i, v)| (i as f64 + 1.0) * v * v)
                .sum()
        };

        type Case<'a> = (&'a dyn Fn(&[f64]) -> f64, Vec<f64>);
        let cases: Vec<Case<'_>> = vec![
            (&rosenbrock, vec![-1.2, 1.0]),
            (&quadratic_1d, vec![2.0]),
            (&quadratic_1d, vec![0.0]),
            (&staircase, vec![1.3, -0.7]),
            (&bumpy, vec![0.4, 0.9]),
            (&three_d, vec![1.0, -2.0, 0.5]),
        ];

        for max_iter in [3, 17, 1000] {
            let params = NelderMeadParams {
                xatol: 1e-9,
                fatol: 1e-9,
                max_iter,
            };
            for (f, x0) in &cases {
                let mut expected_points = Vec::new();
                let expected = nelder_mead_reference(
                    |x: &[f64]| {
                        expected_points.push(x.to_vec());
                        f(x)
                    },
                    x0,
                    Some(params),
                )
                .unwrap();
                let (got, got_points) = drive_stepper(*f, x0, params);

                assert_eq!(got_points, expected_points, "evaluation sequence differs");
                assert_eq!(got.x, expected.x);
                assert_eq!(got.f.to_bits(), expected.f.to_bits());
                assert_eq!(got.iterations, expected.iterations);
                assert_eq!(got.converged, expected.converged);
            }
        }
    }
}