use crate::core::math::{
AddDiagonalVectorInPlace, ComponentDivAssign, ComponentMaxAssign, ComponentMulAssign, Dot,
FloorZerosInPlace, GramMatrix, LinearSolveSpd, MatDiagonal, MatTransposeVec, NegInPlace,
NormInfinity, NormSquared, ScaleInPlace, ScaledAdd,
};
use crate::core::problem::{Jacobian, Residual};
use crate::core::solver::Solver;
use crate::core::state::BasicState;
use crate::core::termination::TerminationReason;
pub struct LevenbergMarquardt<V, M> {
tol_grad: f64,
tol_grad_rel: f64,
ftol: f64,
xtol: f64,
tau: f64,
max_inner_attempts: u32,
mu: Option<f64>,
nu: f64,
diag: Option<V>,
r_cache: Option<V>,
gram_cache: Option<M>,
jtr_cache: Option<V>,
}
impl<V, M> Default for LevenbergMarquardt<V, M> {
fn default() -> Self {
Self::new()
}
}
impl<V, M> LevenbergMarquardt<V, M> {
pub fn new() -> Self {
Self {
tol_grad: 1e-8,
tol_grad_rel: 0.0,
ftol: 0.0,
xtol: 0.0,
tau: 1e-3,
max_inner_attempts: 50,
mu: None,
nu: 2.0,
diag: None,
r_cache: None,
gram_cache: None,
jtr_cache: None,
}
}
pub fn tol_grad(mut self, tol: f64) -> Self {
assert!(tol >= 0.0, "tol_grad must be ≥ 0");
self.tol_grad = tol;
self
}
pub fn tol_grad_rel(mut self, tol: f64) -> Self {
assert!(tol >= 0.0, "tol_grad_rel must be ≥ 0");
self.tol_grad_rel = tol;
self
}
pub fn ftol(mut self, tol: f64) -> Self {
assert!(tol >= 0.0, "ftol must be ≥ 0");
self.ftol = tol;
self
}
pub fn xtol(mut self, tol: f64) -> Self {
assert!(tol >= 0.0, "xtol must be ≥ 0");
self.xtol = tol;
self
}
pub fn tau(mut self, tau: f64) -> Self {
assert!(tau > 0.0, "tau must be > 0");
self.tau = tau;
self
}
pub fn max_inner_attempts(mut self, n: u32) -> Self {
assert!(n > 0, "max_inner_attempts must be > 0");
self.max_inner_attempts = n;
self
}
}
impl<P, V, M> Solver<P, BasicState<V>> for LevenbergMarquardt<V, M>
where
P: Residual<Param = V, Output = V> + Jacobian<Param = V, Output = M>,
V: ScaledAdd<f64>
+ NormSquared
+ NormInfinity
+ NegInPlace
+ Dot
+ ScaleInPlace
+ ComponentMulAssign
+ ComponentDivAssign
+ ComponentMaxAssign
+ FloorZerosInPlace
+ Clone,
M: GramMatrix
+ MatTransposeVec<V>
+ LinearSolveSpd<V>
+ AddDiagonalVectorInPlace<V>
+ MatDiagonal<V>
+ Clone,
{
fn init(&mut self, problem: &P, mut state: BasicState<V>) -> BasicState<V> {
let r = problem.residual(&state.param);
let j = problem.jacobian(&state.param);
state.cost = Some(0.5 * r.norm_squared());
state.cost_evals += 1;
state.gradient_evals += 1;
let a = j.gram();
let mut d = a.diagonal();
d.floor_zeros_in_place(1.0);
self.diag = Some(d);
self.mu = Some(self.tau);
self.nu = 2.0;
self.jtr_cache = Some(j.mat_transpose_vec(&r));
self.gram_cache = Some(a);
self.r_cache = Some(r);
state
}
fn next_iter(
&mut self,
problem: &P,
mut state: BasicState<V>,
) -> (BasicState<V>, Option<TerminationReason>) {
let r = match self.r_cache.take() {
Some(r) => r,
None => {
state.cost_evals += 1;
problem.residual(&state.param)
}
};
let (a, g) = match (self.gram_cache.take(), self.jtr_cache.take()) {
(Some(a), Some(g)) => (a, g),
_ => {
state.gradient_evals += 1;
let j = problem.jacobian(&state.param);
(j.gram(), j.mat_transpose_vec(&r))
}
};
let diag_cur = a.diagonal();
let abs_converged = self.tol_grad > 0.0 && g.norm_infinity() <= self.tol_grad;
let rel_converged = self.tol_grad_rel > 0.0 && {
let mut cos_sq = g.clone();
cos_sq.component_mul_assign(&g);
let mut denom = diag_cur.clone();
denom.floor_zeros_in_place(1.0);
cos_sq.component_div_assign(&denom);
cos_sq.norm_infinity() <= self.tol_grad_rel * self.tol_grad_rel * r.norm_squared()
};
if abs_converged || rel_converged {
self.r_cache = Some(r);
self.gram_cache = Some(a);
self.jtr_cache = Some(g);
return (state, Some(TerminationReason::SolverConverged));
}
let mut neg_g = g.clone();
neg_g.neg_in_place();
let mut d = self
.diag
.take()
.expect("diag not set: Solver::init must run before next_iter");
d.component_max_assign(&diag_cur);
let mut mu = self
.mu
.expect("mu not set: Solver::init must run before next_iter");
let mut nu = self.nu;
let h;
let mut attempts: u32 = 0;
loop {
let mut a_damped = a.clone();
let mut damping = d.clone();
damping.scale_in_place(mu);
a_damped.add_diagonal_vector_in_place(&damping);
match a_damped.solve_spd(&neg_g) {
Ok(step) => {
h = step;
break;
}
Err(_) => {
attempts += 1;
if attempts >= self.max_inner_attempts || !mu.is_finite() {
self.mu = Some(mu);
self.nu = nu;
self.diag = Some(d);
self.r_cache = Some(r);
self.gram_cache = Some(a);
self.jtr_cache = Some(g);
return (state, Some(TerminationReason::SolverFailed));
}
mu *= nu;
nu *= 2.0;
}
}
}
let mut dh = d.clone();
dh.component_mul_assign(&h);
let l_diff = 0.5 * (mu * h.dot(&dh) - h.dot(&g));
let mut x_trial = state.param.clone();
x_trial.scaled_add(1.0, &h);
let r_trial = problem.residual(&x_trial);
state.cost_evals += 1;
let f_trial = 0.5 * r_trial.norm_squared();
let prev_cost = state
.cost
.expect("cost not set: Solver::init must run before next_iter");
let actual_diff = prev_cost - f_trial;
let rho = if l_diff > 0.0 {
actual_diff / l_diff
} else {
0.0
};
if rho > 0.0 {
state.param = x_trial;
state.cost = Some(f_trial);
let factor = 1.0 - (2.0 * rho - 1.0).powi(3);
mu *= factor.max(1.0 / 3.0);
nu = 2.0;
self.r_cache = Some(r_trial);
self.gram_cache = None;
self.jtr_cache = None;
} else {
mu *= nu;
nu *= 2.0;
self.r_cache = Some(r);
self.gram_cache = Some(a);
self.jtr_cache = Some(g);
}
self.mu = Some(mu);
self.nu = nu;
self.diag = Some(d);
let ftol_converged = self.ftol > 0.0
&& actual_diff.abs() <= self.ftol * prev_cost
&& l_diff <= self.ftol * prev_cost
&& rho <= 2.0;
let xtol_converged = self.xtol > 0.0
&& h.norm_squared() <= self.xtol * self.xtol * state.param.norm_squared();
if ftol_converged || xtol_converged {
return (state, Some(TerminationReason::SolverConverged));
}
(state, None)
}
}