regit-svi 2.0.0

// Copyright 2026 Regit.io — Nicolas Koenig
// SPDX-License-Identifier: Apache-2.0

//! Nelder–Mead implementation.

use super::{OptimizerTermination, index_to_f64};

const REFLECT: f64 = 1.0;
const EXPAND: f64 = 2.0;
const CONTRACT: f64 = 0.5;
const SHRINK: f64 = 0.5;

/// Outcome of a [`nelder_mead`] minimisation.
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct NelderMeadResult {
    pub(crate) x: Vec<f64>,
    pub(crate) fx: f64,
    pub(crate) iterations: usize,
    pub(crate) converged: bool,
    pub(crate) termination: OptimizerTermination,
}

/// Minimises a scalar objective by the Nelder–Mead downhill simplex method.
#[must_use]
#[allow(clippy::too_many_lines)]
pub(crate) fn nelder_mead<F>(
    objective: F,
    start: &[f64],
    tolerance: f64,
    max_iterations: usize,
) -> NelderMeadResult
where
    F: Fn(&[f64]) -> f64,
{
    let dimensions = start.len();
    if dimensions == 0 {
        return NelderMeadResult {
            x: Vec::new(),
            fx: objective(&[]),
            iterations: 0,
            converged: true,
            termination: OptimizerTermination::ObjectiveConverged,
        };
    }

    // A nonzero absolute step keeps a zero coordinate from collapsing the
    // initial simplex; nonzero coordinates use the standard relative step.
    let mut simplex: Vec<Vec<f64>> = Vec::with_capacity(dimensions + 1);
    simplex.push(start.to_vec());
    for axis in 0..dimensions {
        let mut vertex = start.to_vec();
        let step = if vertex[axis].abs() > 1e-8 {
            0.05 * vertex[axis].abs()
        } else {
            0.000_25
        };
        vertex[axis] += step;
        simplex.push(vertex);
    }

    let mut values: Vec<f64> = simplex
        .iter()
        .map(|vertex| finite_objective(&objective, vertex))
        .collect();
    let mut order: Vec<usize> = (0..=dimensions).collect();
    let mut iterations = 0;
    let mut converged = false;

    while iterations < max_iterations {
        iterations += 1;
        order.sort_by(|&left, &right| values[left].total_cmp(&values[right]));
        let best = order[0];
        let worst = order[dimensions];
        let second_worst = order[dimensions - 1];

        // Objective spread alone can stop with geometrically distant points
        // of equal height, so termination also requires a small diameter.
        let spread = (values[worst] - values[best]).abs();
        let mut diameter = 0.0_f64;
        for vertex in &simplex {
            let distance_squared: f64 = vertex
                .iter()
                .zip(&simplex[best])
                .map(|(coordinate, best_coordinate)| {
                    let difference = coordinate - best_coordinate;
                    difference * difference
                })
                .sum();
            diameter = diameter.max(distance_squared.sqrt());
        }
        if spread <= tolerance && diameter <= tolerance.sqrt().max(tolerance) {
            converged = true;
            break;
        }

        let mut centroid = vec![0.0; dimensions];
        for &vertex in &order[..dimensions] {
            for (coordinate, value) in centroid.iter_mut().zip(&simplex[vertex]) {
                *coordinate += value;
            }
        }
        let inverse_dimensions = 1.0 / index_to_f64(dimensions);
        for coordinate in &mut centroid {
            *coordinate *= inverse_dimensions;
        }

        let reflected = affine_step(&centroid, REFLECT, &centroid, &simplex[worst]);
        let reflected_value = finite_objective(&objective, &reflected);
        if reflected_value < values[best] {
            let expanded = affine_step(&centroid, EXPAND, &reflected, &centroid);
            let expanded_value = finite_objective(&objective, &expanded);
            if expanded_value < reflected_value {
                simplex[worst] = expanded;
                values[worst] = expanded_value;
            } else {
                simplex[worst] = reflected;
                values[worst] = reflected_value;
            }
        } else if reflected_value < values[second_worst] {
            simplex[worst] = reflected;
            values[worst] = reflected_value;
        } else {
            let (contracted, contracted_value) = if reflected_value < values[worst] {
                let point = affine_step(&centroid, CONTRACT, &reflected, &centroid);
                let value = finite_objective(&objective, &point);
                (point, value)
            } else {
                let point = affine_step(&centroid, CONTRACT, &simplex[worst], &centroid);
                let value = finite_objective(&objective, &point);
                (point, value)
            };
            if contracted_value < reflected_value.min(values[worst]) {
                simplex[worst] = contracted;
                values[worst] = contracted_value;
            } else {
                let best_point = simplex[best].clone();
                for &vertex in &order[1..] {
                    for (coordinate, best_coordinate) in simplex[vertex].iter_mut().zip(&best_point)
                    {
                        *coordinate = best_coordinate + SHRINK * (*coordinate - best_coordinate);
                    }
                    values[vertex] = finite_objective(&objective, &simplex[vertex]);
                }
            }
        }
    }

    order.sort_by(|&left, &right| values[left].total_cmp(&values[right]));
    let best = order[0];
    let finite_best = values[best].is_finite();
    NelderMeadResult {
        x: simplex[best].clone(),
        fx: values[best],
        iterations,
        converged: converged && finite_best,
        termination: if !finite_best {
            OptimizerTermination::NonFiniteEvaluation
        } else if converged {
            OptimizerTermination::ObjectiveConverged
        } else {
            OptimizerTermination::IterationLimit
        },
    }
}

fn finite_objective(function: &impl Fn(&[f64]) -> f64, point: &[f64]) -> f64 {
    let value = function(point);
    if value.is_finite() {
        value
    } else {
        f64::INFINITY
    }
}

fn affine_step(centroid: &[f64], coefficient: f64, point: &[f64], reference: &[f64]) -> Vec<f64> {
    centroid
        .iter()
        .zip(point)
        .zip(reference)
        .map(|((&center, &value), &base)| center + coefficient * (value - base))
        .collect()
}

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

    fn rosenbrock(point: &[f64]) -> f64 {
        let first = 1.0 - point[0];
        let second = point[1] - point[0] * point[0];
        first * first + 100.0 * second * second
    }

    #[test]
    fn minimises_sphere() {
        let result = nelder_mead(
            |point| point[0] * point[0] + point[1] * point[1],
            &[3.0, -2.0],
            1e-14,
            1_000,
        );
        assert!(result.fx < 1e-10);
        assert!(result.x[0].abs() < 1e-5);
        assert!(result.x[1].abs() < 1e-5);
        assert!(result.converged);
    }

    #[test]
    fn minimises_rosenbrock() {
        let result = nelder_mead(rosenbrock, &[-1.2, 1.0], 1e-14, 5_000);
        assert!((result.x[0] - 1.0).abs() < 1e-3);
        assert!((result.x[1] - 1.0).abs() < 1e-3);
        assert!(result.fx < 1e-6);
    }

    #[test]
    fn handles_one_dimension() {
        let result = nelder_mead(
            |point| (point[0] - 4.0) * (point[0] - 4.0),
            &[0.0],
            1e-14,
            1_000,
        );
        assert!((result.x[0] - 4.0).abs() < 1e-5);
    }

    #[test]
    fn handles_empty_start() {
        let result = nelder_mead(|_| 7.0, &[], 1e-12, 10);
        assert!((result.fx - 7.0).abs() < 1e-15);
        assert!(result.converged);
    }
}