regit-svi 2.0.0

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

//! Levenberg–Marquardt implementation.

use super::{OptimizerTermination, solve_spd};

/// Outcome of a crate-private Levenberg–Marquardt minimisation.
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct LevenbergMarquardtResult {
    pub(crate) params: Vec<f64>,
    pub(crate) cost: f64,
    pub(crate) iterations: usize,
    pub(crate) converged: bool,
    pub(crate) termination: OptimizerTermination,
}

/// Minimises a weighted nonlinear least-squares objective.
#[must_use]
#[allow(clippy::too_many_lines)]
pub(crate) fn levenberg_marquardt<F>(
    residual: F,
    start: &[f64],
    tolerance: f64,
    max_iterations: usize,
) -> LevenbergMarquardtResult
where
    F: Fn(&[f64]) -> Vec<(f64, f64, Vec<f64>)>,
{
    let dimensions = start.len();
    let mut parameters = start.to_vec();
    let mut damping = 1e-3;
    let damping_growth = 2.0;
    let cost_of = |observations: &[(f64, f64, Vec<f64>)]| -> f64 {
        observations
            .iter()
            .map(|(value, weight, _)| weight * value * value)
            .sum()
    };
    let mut observations = residual(&parameters);
    let mut cost = cost_of(&observations);
    let mut iterations = 0;
    let mut converged = false;
    let mut termination = OptimizerTermination::IterationLimit;
    if !cost.is_finite() {
        return LevenbergMarquardtResult {
            params: parameters,
            cost,
            iterations,
            converged,
            termination: OptimizerTermination::NonFiniteEvaluation,
        };
    }

    while iterations < max_iterations {
        iterations += 1;
        let mut normal_matrix = vec![0.0; dimensions * dimensions];
        let mut negative_gradient = vec![0.0; dimensions];
        for (value, weight, jacobian) in &observations {
            if jacobian.len() != dimensions
                || !value.is_finite()
                || !weight.is_finite()
                || *weight < 0.0
                || jacobian.iter().any(|entry| !entry.is_finite())
            {
                termination = OptimizerTermination::NonFiniteEvaluation;
                return LevenbergMarquardtResult {
                    params: parameters,
                    cost,
                    iterations,
                    converged,
                    termination,
                };
            }
            for row in 0..dimensions {
                negative_gradient[row] -= weight * jacobian[row] * value;
                for column in 0..dimensions {
                    normal_matrix[row * dimensions + column] +=
                        weight * jacobian[row] * jacobian[column];
                }
            }
        }

        let gradient_norm = negative_gradient
            .iter()
            .map(|gradient| gradient * gradient)
            .sum::<f64>()
            .sqrt();
        if gradient_norm <= tolerance {
            converged = true;
            termination = OptimizerTermination::GradientConverged;
            break;
        }
        if cost <= tolerance * tolerance {
            converged = true;
            termination = OptimizerTermination::ObjectiveConverged;
            break;
        }

        let mut accepted = false;
        let mut solved_system = false;
        for _ in 0..30 {
            let mut damped = normal_matrix.clone();
            for index in 0..dimensions {
                damped[index * dimensions + index] +=
                    damping * normal_matrix[index * dimensions + index].max(1e-12);
            }
            let Some(step) = solve_spd(&damped, &negative_gradient, dimensions) else {
                damping *= damping_growth;
                continue;
            };
            solved_system = true;
            let trial: Vec<f64> = parameters
                .iter()
                .zip(&step)
                .map(|(parameter, delta)| parameter + delta)
                .collect();
            let trial_observations = residual(&trial);
            let trial_cost = cost_of(&trial_observations);
            if !trial_cost.is_finite() {
                damping *= damping_growth;
                continue;
            }

            // The gain ratio compares the actual decrease with the local
            // damped-quadratic prediction; only positive improvement is used.
            let predicted_reduction = (0..dimensions)
                .map(|index| {
                    step[index]
                        * (damping
                            * normal_matrix[index * dimensions + index].max(1e-12)
                            * step[index]
                            + negative_gradient[index])
                })
                .sum::<f64>()
                / 2.0;
            let actual_reduction = cost - trial_cost;
            let gain = if predicted_reduction > 0.0 {
                actual_reduction / predicted_reduction
            } else {
                -1.0
            };
            if gain > 0.0 && trial_cost < cost {
                let step_norm = step.iter().map(|delta| delta * delta).sum::<f64>().sqrt();
                parameters = trial;
                observations = trial_observations;
                cost = trial_cost;
                let shrink = (1.0 - (2.0 * gain - 1.0).powi(3)).max(1.0 / 3.0);
                damping *= shrink;
                accepted = true;
                if step_norm <= tolerance {
                    converged = true;
                    termination = OptimizerTermination::StepConverged;
                }
                break;
            }
            damping *= damping_growth;
        }
        if !accepted {
            termination = if solved_system {
                OptimizerTermination::Stagnated
            } else {
                OptimizerTermination::SingularModel
            };
            break;
        }
        if converged {
            break;
        }
    }

    LevenbergMarquardtResult {
        params: parameters,
        cost,
        iterations,
        converged,
        termination,
    }
}

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

    #[test]
    fn fits_mean_and_reports_stagnation_at_tight_tolerance() {
        let data = [1.0_f64, 2.0, 3.0, 4.0];
        let result = levenberg_marquardt(
            |parameters| {
                data.iter()
                    .map(|&observation| (parameters[0] - observation, 1.0, vec![1.0]))
                    .collect()
            },
            &[0.0],
            1e-14,
            200,
        );
        assert!((result.params[0] - 2.5).abs() < 1e-8);
        assert!(!result.converged);
        assert_eq!(result.termination, OptimizerTermination::Stagnated);
    }

    #[test]
    fn fits_line() {
        let abscissas = [0.0_f64, 1.0, 2.0, 3.0, 4.0];
        let result = levenberg_marquardt(
            |parameters| {
                abscissas
                    .iter()
                    .map(|&x| {
                        let observation = 2.0 + 3.0 * x;
                        (
                            parameters[0] + parameters[1] * x - observation,
                            1.0,
                            vec![1.0, x],
                        )
                    })
                    .collect()
            },
            &[0.0, 0.0],
            1e-14,
            300,
        );
        assert!((result.params[0] - 2.0).abs() < 1e-6);
        assert!((result.params[1] - 3.0).abs() < 1e-6);
    }

    #[test]
    fn fits_nonlinear_exponential() {
        let abscissas = [0.0_f64, 0.5, 1.0, 1.5, 2.0];
        let result = levenberg_marquardt(
            |parameters| {
                abscissas
                    .iter()
                    .map(|&x| {
                        let model = (parameters[0] * x).exp();
                        let observation = (0.5 * x).exp();
                        (model - observation, 1.0, vec![x * model])
                    })
                    .collect()
            },
            &[0.1],
            1e-14,
            300,
        );
        assert!((result.params[0] - 0.5).abs() < 1e-5);
    }

    #[test]
    fn rejects_nonfinite_observations() {
        let result = levenberg_marquardt(|_| vec![(f64::NAN, 1.0, vec![1.0])], &[0.0], 1e-12, 5);
        assert_eq!(
            result.termination,
            OptimizerTermination::NonFiniteEvaluation
        );
    }
}