use ndarray::{Array1, Array2};
use thiserror::Error;
#[derive(Debug, Error)]
pub enum LMError {
#[error("bounds/x0 dimension mismatch: x0 has {x0_len}, bounds has {bounds_len}")]
DimensionMismatch {
x0_len: usize,
bounds_len: usize,
},
#[error("invalid bounds at index {index}: lower ({lower}) > upper ({upper})")]
InvalidBounds {
index: usize,
lower: f64,
upper: f64,
},
#[error("residual dimension changed: was {expected}, now {got}")]
ResidualDimensionChanged {
expected: usize,
got: usize,
},
#[error("singular Jacobian — cannot compute step")]
SingularJacobian,
}
pub type LMResult<T> = std::result::Result<T, LMError>;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LMCallbackAction {
Continue,
Stop,
}
pub struct LMIntermediate {
pub x: Array1<f64>,
pub fun: f64,
pub lambda: f64,
pub iter: usize,
}
pub type LMCallback = Box<dyn FnMut(&LMIntermediate) -> LMCallbackAction>;
pub struct LMConfig {
pub maxiter: usize,
pub tol: f64,
pub atol: f64,
pub lambda_init: f64,
pub jacobian_epsilon: f64,
pub x0: Array1<f64>,
pub weights: Option<Array1<f64>>,
pub disp: bool,
pub callback: Option<LMCallback>,
}
pub struct LMConfigBuilder {
maxiter: usize,
tol: f64,
atol: f64,
lambda_init: f64,
jacobian_epsilon: f64,
x0: Option<Array1<f64>>,
weights: Option<Array1<f64>>,
disp: bool,
callback: Option<LMCallback>,
}
impl LMConfigBuilder {
pub fn new() -> Self {
Self {
maxiter: 100,
tol: 1e-10,
atol: 1e-14,
lambda_init: 1.0,
jacobian_epsilon: 1e-8,
x0: None,
weights: None,
disp: false,
callback: None,
}
}
pub fn x0(mut self, x0: Array1<f64>) -> Self {
self.x0 = Some(x0);
self
}
pub fn maxiter(mut self, n: usize) -> Self {
self.maxiter = n;
self
}
pub fn tol(mut self, t: f64) -> Self {
self.tol = t;
self
}
pub fn atol(mut self, t: f64) -> Self {
self.atol = t;
self
}
pub fn lambda_init(mut self, l: f64) -> Self {
self.lambda_init = l;
self
}
pub fn jacobian_epsilon(mut self, eps: f64) -> Self {
self.jacobian_epsilon = eps;
self
}
pub fn weights(mut self, w: Array1<f64>) -> Self {
self.weights = Some(w);
self
}
pub fn disp(mut self, d: bool) -> Self {
self.disp = d;
self
}
pub fn callback(mut self, cb: Box<dyn FnMut(&LMIntermediate) -> LMCallbackAction>) -> Self {
self.callback = Some(cb);
self
}
pub fn build(self) -> LMConfig {
LMConfig {
maxiter: self.maxiter,
tol: self.tol,
atol: self.atol,
lambda_init: self.lambda_init,
jacobian_epsilon: self.jacobian_epsilon,
x0: self.x0.expect("LMConfigBuilder: x0 is required"),
weights: self.weights,
disp: self.disp,
callback: self.callback,
}
}
}
impl Default for LMConfigBuilder {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug)]
pub struct LMReport {
pub x: Array1<f64>,
pub fun: f64,
pub residuals: Array1<f64>,
pub success: bool,
pub message: String,
pub nit: usize,
pub nfev: usize,
}
pub fn levenberg_marquardt<F>(
residual_fn: &F,
bounds: &[(f64, f64)],
config: LMConfig,
) -> LMResult<LMReport>
where
F: Fn(&Array1<f64>) -> Array1<f64>,
{
let n_params = config.x0.len();
if bounds.len() != n_params {
return Err(LMError::DimensionMismatch {
x0_len: n_params,
bounds_len: bounds.len(),
});
}
for (i, &(lo, hi)) in bounds.iter().enumerate() {
if lo > hi {
return Err(LMError::InvalidBounds {
index: i,
lower: lo,
upper: hi,
});
}
}
let mut x = project(&config.x0, bounds);
let mut r = residual_fn(&x);
let n_residuals = r.len();
let mut nfev: usize = 1;
let w = config.weights.unwrap_or_else(|| Array1::ones(n_residuals));
let mut f_val = weighted_sos(&r, &w);
let mut lambda = config.lambda_init;
let eps = config.jacobian_epsilon;
let mut success = false;
let mut message = format!("max iterations ({}) reached", config.maxiter);
let mut callback = config.callback;
let mut nit: usize = 0;
for iter in 0..config.maxiter {
nit = iter + 1;
if let Some(ref mut cb) = callback {
let intermediate = LMIntermediate {
x: x.clone(),
fun: f_val,
lambda,
iter,
};
if cb(&intermediate) == LMCallbackAction::Stop {
message = "stopped by callback".to_string();
break;
}
}
if f_val <= config.atol {
success = true;
message = format!(
"converged: objective {:.2e} <= atol {:.2e}",
f_val, config.atol
);
break;
}
let jac = compute_jacobian(residual_fn, &x, &r, n_residuals, eps, &mut nfev)?;
let jtwj = jtw_j(&jac, &w);
let jtwr = jtw_r(&jac, &w, &r);
let diag: Array1<f64> = (0..n_params).map(|j| jtwj[[j, j]].max(1e-20)).collect();
let mut h = jtwj.clone();
for j in 0..n_params {
h[[j, j]] += lambda * diag[j];
}
let neg_jtwr = jtwr.mapv(|v| -v);
let delta = match solve_linear_system(&h, &neg_jtwr) {
Some(d) => d,
None => {
lambda *= 4.0;
continue;
}
};
let x_trial = project(&(&x + &delta), bounds);
let r_trial = residual_fn(&x_trial);
nfev += 1;
if r_trial.len() != n_residuals {
return Err(LMError::ResidualDimensionChanged {
expected: n_residuals,
got: r_trial.len(),
});
}
let f_trial = weighted_sos(&r_trial, &w);
let pred: f64 = {
let mut p = 0.0;
for j in 0..n_params {
let mut jtwj_delta_j = 0.0;
for k in 0..n_params {
jtwj_delta_j += jtwj[[j, k]] * delta[k];
}
p += delta[j] * jtwj_delta_j + 2.0 * lambda * diag[j] * delta[j] * delta[j];
}
p
};
let actual = f_val - f_trial;
let rho = if pred.abs() > 1e-30 {
actual / pred
} else {
0.0
};
if rho > 0.0 && f_trial < f_val {
let f_old = f_val;
x = x_trial;
r = r_trial;
f_val = f_trial;
lambda = (lambda / 2.0).max(1e-15);
if (f_old - f_val).abs() < config.tol * f_old + config.atol {
success = true;
message = format!(
"converged: |df|={:.2e} < tol*f+atol={:.2e}",
(f_old - f_val).abs(),
config.tol * f_old + config.atol
);
break;
}
} else {
lambda = (lambda * 2.0).min(1e15);
}
if config.disp && iter % 10 == 0 {
eprintln!(
"LM iter {}: f={:.6e}, lambda={:.2e}, rho={:.3}",
iter, f_val, lambda, rho
);
}
}
Ok(LMReport {
x,
fun: f_val,
residuals: r,
success,
message,
nit,
nfev,
})
}
fn project(x: &Array1<f64>, bounds: &[(f64, f64)]) -> Array1<f64> {
Array1::from(
x.iter()
.zip(bounds.iter())
.map(|(&xi, &(lo, hi))| xi.clamp(lo, hi))
.collect::<Vec<_>>(),
)
}
fn weighted_sos(r: &Array1<f64>, w: &Array1<f64>) -> f64 {
r.iter().zip(w.iter()).map(|(&ri, &wi)| wi * ri * ri).sum()
}
fn compute_jacobian<F>(
residual_fn: &F,
x: &Array1<f64>,
_r0: &Array1<f64>,
n_residuals: usize,
eps: f64,
nfev: &mut usize,
) -> LMResult<Array2<f64>>
where
F: Fn(&Array1<f64>) -> Array1<f64>,
{
let n_params = x.len();
let mut jac = Array2::zeros((n_residuals, n_params));
for j in 0..n_params {
let mut x_plus = x.clone();
let mut x_minus = x.clone();
let h = jacobian_step(eps, x[j]);
x_plus[j] += h;
x_minus[j] -= h;
let r_plus = residual_fn(&x_plus);
let r_minus = residual_fn(&x_minus);
*nfev += 2;
if r_plus.len() != n_residuals {
return Err(LMError::ResidualDimensionChanged {
expected: n_residuals,
got: r_plus.len(),
});
}
if r_minus.len() != n_residuals {
return Err(LMError::ResidualDimensionChanged {
expected: n_residuals,
got: r_minus.len(),
});
}
let inv_2h = 1.0 / (2.0 * h);
for i in 0..n_residuals {
jac[[i, j]] = (r_plus[i] - r_minus[i]) * inv_2h;
}
}
Ok(jac)
}
fn jacobian_step(eps: f64, x: f64) -> f64 {
(eps * (1.0 + x.abs())).max(eps)
}
fn jtw_j(jac: &Array2<f64>, w: &Array1<f64>) -> Array2<f64> {
let n_params = jac.ncols();
let n_res = jac.nrows();
let mut result = Array2::zeros((n_params, n_params));
for j in 0..n_params {
for k in j..n_params {
let mut val = 0.0;
for i in 0..n_res {
val += w[i] * jac[[i, j]] * jac[[i, k]];
}
result[[j, k]] = val;
result[[k, j]] = val;
}
}
result
}
fn jtw_r(jac: &Array2<f64>, w: &Array1<f64>, r: &Array1<f64>) -> Array1<f64> {
let n_params = jac.ncols();
let n_res = jac.nrows();
let mut result = Array1::zeros(n_params);
for j in 0..n_params {
let mut val = 0.0;
for i in 0..n_res {
val += w[i] * jac[[i, j]] * r[i];
}
result[j] = val;
}
result
}
fn solve_linear_system(a: &Array2<f64>, b: &Array1<f64>) -> Option<Array1<f64>> {
let n = b.len();
debug_assert_eq!(a.nrows(), n);
debug_assert_eq!(a.ncols(), n);
let mut aug = Array2::zeros((n, n + 1));
for i in 0..n {
for j in 0..n {
aug[[i, j]] = a[[i, j]];
}
aug[[i, n]] = b[i];
}
for col in 0..n {
let mut max_val = aug[[col, col]].abs();
let mut max_row = col;
for row in (col + 1)..n {
let val = aug[[row, col]].abs();
if val > max_val {
max_val = val;
max_row = row;
}
}
if max_val < 1e-12 {
return None; }
if max_row != col {
for j in 0..=n {
let tmp = aug[[col, j]];
aug[[col, j]] = aug[[max_row, j]];
aug[[max_row, j]] = tmp;
}
}
let pivot = aug[[col, col]];
for row in (col + 1)..n {
let factor = aug[[row, col]] / pivot;
for j in col..=n {
aug[[row, j]] -= factor * aug[[col, j]];
}
}
}
let mut x = Array1::zeros(n);
for col in (0..n).rev() {
let mut sum = aug[[col, n]];
for j in (col + 1)..n {
sum -= aug[[col, j]] * x[j];
}
x[col] = sum / aug[[col, col]];
}
if x.iter().any(|v| !v.is_finite()) {
return None;
}
Some(x)
}
#[cfg(test)]
mod tests {
use super::*;
use ndarray::array;
#[test]
fn test_sphere() {
let residual = |x: &Array1<f64>| x.clone();
let bounds = vec![(-10.0, 10.0); 3];
let config = LMConfigBuilder::new()
.x0(array![3.0, -2.0, 1.0])
.maxiter(50)
.build();
let report = levenberg_marquardt(&residual, &bounds, config).unwrap();
assert!(report.success, "should converge: {}", report.message);
assert!(
report.fun < 1e-12,
"objective should be ~0, got {}",
report.fun
);
for &xi in report.x.iter() {
assert!(xi.abs() < 1e-6, "x should be ~0, got {}", xi);
}
}
#[test]
fn test_rosenbrock_residual() {
let residual = |x: &Array1<f64>| array![10.0 * (x[1] - x[0] * x[0]), 1.0 - x[0]];
let bounds = vec![(-5.0, 5.0); 2];
let config = LMConfigBuilder::new()
.x0(array![-1.0, 1.0])
.maxiter(200)
.tol(1e-12)
.build();
let report = levenberg_marquardt(&residual, &bounds, config).unwrap();
assert!(report.success, "should converge: {}", report.message);
assert!(
(report.x[0] - 1.0).abs() < 1e-4,
"x0 should be ~1, got {}",
report.x[0]
);
assert!(
(report.x[1] - 1.0).abs() < 1e-4,
"x1 should be ~1, got {}",
report.x[1]
);
}
#[test]
fn test_bounded_solution() {
let residual = |x: &Array1<f64>| array![x[0] - 5.0];
let bounds = vec![(-10.0, 3.0)];
let config = LMConfigBuilder::new().x0(array![0.0]).maxiter(50).build();
let report = levenberg_marquardt(&residual, &bounds, config).unwrap();
assert!(
(report.x[0] - 3.0).abs() < 1e-6,
"x should be at bound 3.0, got {}",
report.x[0]
);
}
#[test]
fn test_nan_residual_handled() {
let residual = |x: &Array1<f64>| {
if x[0].abs() > 100.0 {
array![f64::NAN]
} else {
array![x[0] - 1.0]
}
};
let bounds = vec![(-200.0, 200.0)];
let config = LMConfigBuilder::new().x0(array![0.0]).maxiter(50).build();
let result = levenberg_marquardt(&residual, &bounds, config);
assert!(result.is_ok());
}
#[test]
fn test_zero_residual() {
let residual = |x: &Array1<f64>| array![x[0], x[1]];
let bounds = vec![(-10.0, 10.0); 2];
let config = LMConfigBuilder::new()
.x0(array![0.0, 0.0])
.maxiter(10)
.build();
let report = levenberg_marquardt(&residual, &bounds, config).unwrap();
assert!(report.success, "already at optimum: {}", report.message);
assert!(report.fun < 1e-14);
}
#[test]
fn test_callback_stop() {
let residual = |x: &Array1<f64>| x.clone();
let bounds = vec![(-10.0, 10.0); 2];
let config = LMConfigBuilder::new()
.x0(array![5.0, 5.0])
.maxiter(1000)
.callback(Box::new(|inter| {
if inter.iter >= 3 {
LMCallbackAction::Stop
} else {
LMCallbackAction::Continue
}
}))
.build();
let report = levenberg_marquardt(&residual, &bounds, config).unwrap();
assert_eq!(report.message, "stopped by callback");
}
#[test]
fn test_weighted_residuals() {
let residual = |x: &Array1<f64>| array![x[0] - 1.0, x[0] - 3.0];
let bounds = vec![(-10.0, 10.0)];
let config_unw = LMConfigBuilder::new().x0(array![0.0]).maxiter(50).build();
let report_unw = levenberg_marquardt(&residual, &bounds, config_unw).unwrap();
let config_w = LMConfigBuilder::new()
.x0(array![0.0])
.maxiter(50)
.weights(array![10.0, 1.0])
.build();
let report_w = levenberg_marquardt(&residual, &bounds, config_w).unwrap();
assert!(
(report_unw.x[0] - 2.0).abs() < 0.01,
"unweighted x should be ~2, got {}",
report_unw.x[0]
);
assert!(
report_w.x[0] < report_unw.x[0],
"weighted x ({}) should be less than unweighted ({})",
report_w.x[0],
report_unw.x[0]
);
assert!(
(report_w.x[0] - 1.0).abs() < 0.5,
"weighted x should be near 1.0, got {}",
report_w.x[0]
);
}
#[test]
fn test_dimension_mismatch() {
let residual = |x: &Array1<f64>| x.clone();
let bounds = vec![(-1.0, 1.0); 3]; let config = LMConfigBuilder::new()
.x0(array![0.0, 0.0]) .build();
let err = levenberg_marquardt(&residual, &bounds, config).unwrap_err();
assert!(matches!(err, LMError::DimensionMismatch { .. }));
}
#[test]
fn test_nit_tracks_iterations_not_nfev() {
let residual = |x: &Array1<f64>| x.clone();
let bounds = vec![(-10.0, 10.0); 2];
let config = LMConfigBuilder::new()
.x0(array![5.0, 5.0])
.maxiter(5)
.tol(1e-20) .atol(0.0)
.build();
let report = levenberg_marquardt(&residual, &bounds, config).unwrap();
assert_eq!(
report.nit, 5,
"nit should be maxiter (5), got {}",
report.nit
);
assert!(
report.nfev > report.nit,
"nfev ({}) should be much larger than nit ({})",
report.nfev,
report.nit
);
}
#[test]
fn test_invalid_bounds() {
let residual = |x: &Array1<f64>| x.clone();
let bounds = vec![(5.0, 1.0)]; let config = LMConfigBuilder::new().x0(array![0.0]).build();
let err = levenberg_marquardt(&residual, &bounds, config).unwrap_err();
assert!(matches!(err, LMError::InvalidBounds { .. }));
}
#[test]
fn test_jacobian_step_size_uses_relative_scale() {
let eps = 1e-8;
assert!((jacobian_step(eps, 0.0) - eps).abs() < 1e-15);
assert!((jacobian_step(eps, 1.0) - 2.0e-8).abs() < 1e-15);
let h_large = jacobian_step(eps, 1.0e12);
assert!(h_large >= eps);
assert!((h_large - eps * (1.0 + 1.0e12)).abs() < 1.0);
}
#[test]
fn test_large_x0_lm_converges() {
let residual = |x: &Array1<f64>| array![x[0] - 1.0e6];
let bounds = vec![(0.0, 2.0e6)];
let config = LMConfigBuilder::new()
.x0(array![1.0e8])
.maxiter(100)
.build();
let report = levenberg_marquardt(&residual, &bounds, config).unwrap();
assert!(report.success, "should converge: {}", report.message);
assert!(
(report.x[0] - 1.0e6).abs() < 1.0,
"x should be ~1e6, got {}",
report.x[0]
);
}
#[test]
fn test_linear_solver_rejects_near_singular() {
let a = Array2::from_shape_vec((2, 2), vec![1.0, 1.0, 1.0, 1.0 + 1e-15]).unwrap();
let b = Array1::from(vec![1.0, 1.0]);
assert!(
solve_linear_system(&a, &b).is_none(),
"near-singular matrix should be rejected"
);
}
}