use crate::dynmatrix::{DynMatrix, DynVector};
use crate::traits::FloatScalar;
use super::line_search::backtracking_armijo_dyn;
use super::{BfgsSettings, GaussNewtonSettings, LmSettings, OptimError};
#[derive(Debug, Clone)]
pub struct MinimizeResultDyn<T: FloatScalar> {
pub x: DynVector<T>,
pub fx: T,
pub grad_norm: T,
pub iterations: usize,
pub f_evals: usize,
pub grad_evals: usize,
}
#[derive(Debug, Clone)]
pub struct LeastSquaresResultDyn<T: FloatScalar> {
pub x: DynVector<T>,
pub cost: T,
pub grad_norm: T,
pub iterations: usize,
pub r_evals: usize,
pub j_evals: usize,
}
#[inline]
fn vec_add<T: FloatScalar>(a: &DynVector<T>, b: &DynVector<T>) -> DynVector<T> {
debug_assert_eq!(a.len(), b.len());
let n = a.len();
let mut out = DynVector::zeros(n);
for i in 0..n {
out[i] = a[i] + b[i];
}
out
}
#[inline]
fn vec_sub<T: FloatScalar>(a: &DynVector<T>, b: &DynVector<T>) -> DynVector<T> {
debug_assert_eq!(a.len(), b.len());
let n = a.len();
let mut out = DynVector::zeros(n);
for i in 0..n {
out[i] = a[i] - b[i];
}
out
}
#[inline]
fn vec_neg<T: FloatScalar>(a: &DynVector<T>) -> DynVector<T> {
let n = a.len();
let mut out = DynVector::zeros(n);
for i in 0..n {
out[i] = T::zero() - a[i];
}
out
}
#[inline]
fn vec_scale<T: FloatScalar>(a: &DynVector<T>, s: T) -> DynVector<T> {
let n = a.len();
let mut out = DynVector::zeros(n);
for i in 0..n {
out[i] = a[i] * s;
}
out
}
#[inline]
fn mat_vec<T: FloatScalar>(m: &DynMatrix<T>, v: &DynVector<T>) -> DynVector<T> {
let nr = m.nrows();
let nc = m.ncols();
debug_assert_eq!(nc, v.len());
let mut out = DynVector::zeros(nr);
for j in 0..nc {
let vj = v[j];
for i in 0..nr {
out[i] = out[i] + m[(i, j)] * vj;
}
}
out
}
#[inline]
fn matt_vec<T: FloatScalar>(m: &DynMatrix<T>, v: &DynVector<T>) -> DynVector<T> {
let nr = m.nrows();
let nc = m.ncols();
debug_assert_eq!(nr, v.len());
let mut out = DynVector::zeros(nc);
for j in 0..nc {
let mut s = T::zero();
for i in 0..nr {
s = s + m[(i, j)] * v[i];
}
out[j] = s;
}
out
}
pub fn finite_difference_jacobian_dyn<T: FloatScalar>(
mut f: impl FnMut(&DynVector<T>) -> DynVector<T>,
x: &DynVector<T>,
) -> DynMatrix<T> {
let sqrt_eps = T::epsilon().sqrt();
let f0 = f(x);
let n = x.len();
let m = f0.len();
let mut jac = DynMatrix::zeros(m, n);
let mut x_pert = x.clone();
for j in 0..n {
let h = sqrt_eps * x[j].abs().max(T::one());
x_pert[j] = x[j] + h;
let f_pert = f(&x_pert);
x_pert[j] = x[j];
for i in 0..m {
jac[(i, j)] = (f_pert[i] - f0[i]) / h;
}
}
jac
}
pub fn finite_difference_gradient_dyn<T: FloatScalar>(
mut f: impl FnMut(&DynVector<T>) -> T,
x: &DynVector<T>,
) -> DynVector<T> {
let sqrt_eps = T::epsilon().sqrt();
let f0 = f(x);
let n = x.len();
let mut grad = DynVector::zeros(n);
let mut x_pert = x.clone();
for j in 0..n {
let h = sqrt_eps * x[j].abs().max(T::one());
x_pert[j] = x[j] + h;
grad[j] = (f(&x_pert) - f0) / h;
x_pert[j] = x[j];
}
grad
}
#[cfg(feature = "rayon")]
const FD_PAR_MIN_COLS: usize = 8;
#[cfg(feature = "rayon")]
pub fn finite_difference_jacobian_dyn_par<T: FloatScalar + Send + Sync>(
f: impl Fn(&DynVector<T>) -> DynVector<T> + Sync + Send,
x: &DynVector<T>,
) -> DynMatrix<T> {
let sqrt_eps = T::epsilon().sqrt();
let f0 = f(x);
let m = f0.len();
let f0 = f0.as_slice();
let mut jac = DynMatrix::zeros(m, x.len());
crate::par::for_each_chunk_mut(jac.as_mut_slice(), m, FD_PAR_MIN_COLS, |j, col| {
let h = sqrt_eps * x[j].abs().max(T::one());
let mut x_pert = x.clone();
x_pert[j] = x_pert[j] + h;
let f_pert = f(&x_pert);
for i in 0..m {
col[i] = (f_pert[i] - f0[i]) / h;
}
});
jac
}
#[cfg(feature = "rayon")]
pub fn finite_difference_gradient_dyn_par<T: FloatScalar + Send + Sync>(
f: impl Fn(&DynVector<T>) -> T + Sync + Send,
x: &DynVector<T>,
) -> DynVector<T> {
let sqrt_eps = T::epsilon().sqrt();
let f0 = f(x);
let mut grad = DynVector::zeros(x.len());
crate::par::for_each_chunk_mut(grad.as_mut_slice(), 1, FD_PAR_MIN_COLS, |j, cell| {
let h = sqrt_eps * x[j].abs().max(T::one());
let mut x_pert = x.clone();
x_pert[j] = x_pert[j] + h;
cell[0] = (f(&x_pert) - f0) / h;
});
grad
}
pub fn minimize_bfgs_dyn<T: FloatScalar>(
mut f: impl FnMut(&DynVector<T>) -> T,
mut grad: impl FnMut(&DynVector<T>) -> DynVector<T>,
x0: &DynVector<T>,
settings: &BfgsSettings<T>,
) -> Result<MinimizeResultDyn<T>, OptimError> {
let n = x0.len();
let mut x = x0.clone();
let mut fx = f(&x);
let mut g = grad(&x);
let mut f_evals = 1usize;
let mut grad_evals = 1usize;
assert_eq!(g.len(), n, "gradient length must match x length");
let mut h = DynMatrix::<T>::zeros(n, n);
for i in 0..n {
h[(i, i)] = T::one();
}
for iter in 0..settings.max_iter {
let g_norm = g.norm();
if g_norm < settings.grad_tol {
return Ok(MinimizeResultDyn {
x,
fx,
grad_norm: g_norm,
iterations: iter,
f_evals,
grad_evals,
});
}
let hg = mat_vec(&h, &g);
let p_init = vec_neg(&hg);
let grad_dot_p = g.dot(&p_init);
let (p, h_reset) = if grad_dot_p >= T::zero() {
(vec_neg(&g), true)
} else {
(p_init, false)
};
let grad_dot_p = if h_reset { -g.dot(&g) } else { grad_dot_p };
let (alpha, f_new, ls_evals) = backtracking_armijo_dyn(
fx,
grad_dot_p,
&x,
&p,
&mut f,
settings.armijo_c1,
settings.armijo_rho,
settings.max_ls_iter,
)?;
f_evals += ls_evals;
let s = vec_scale(&p, alpha);
let x_new = vec_add(&x, &s);
let g_new = grad(&x_new);
grad_evals += 1;
let y = vec_sub(&g_new, &g);
let ys = y.dot(&s);
if s.norm() < settings.x_tol * (T::one() + x.norm()) {
return Ok(MinimizeResultDyn {
x: x_new,
fx: f_new,
grad_norm: g_new.norm(),
iterations: iter + 1,
f_evals,
grad_evals,
});
}
if (fx - f_new).abs() < settings.f_tol * (T::one() + fx.abs()) {
return Ok(MinimizeResultDyn {
x: x_new,
fx: f_new,
grad_norm: g_new.norm(),
iterations: iter + 1,
f_evals,
grad_evals,
});
}
if h_reset {
for i in 0..n {
for j in 0..n {
h[(i, j)] = if i == j { T::one() } else { T::zero() };
}
}
}
if ys > T::epsilon() {
let rho = T::one() / ys;
let hy = mat_vec(&h, &y);
let yhy = y.dot(&hy);
let factor = (T::one() + rho * yhy) * rho;
for i in 0..n {
for j in 0..n {
h[(i, j)] =
h[(i, j)] + factor * s[i] * s[j] - rho * (hy[i] * s[j] + s[i] * hy[j]);
}
}
}
x = x_new;
fx = f_new;
g = g_new;
}
Err(OptimError::MaxIterations)
}
pub fn least_squares_gn_dyn<T: FloatScalar>(
mut residual: impl FnMut(&DynVector<T>) -> DynVector<T>,
mut jacobian: impl FnMut(&DynVector<T>) -> DynMatrix<T>,
x0: &DynVector<T>,
settings: &GaussNewtonSettings<T>,
) -> Result<LeastSquaresResultDyn<T>, OptimError> {
let n = x0.len();
let mut x = x0.clone();
let mut r = residual(&x);
let mut r_evals = 1usize;
let mut j_evals = 0usize;
let half = T::one() / (T::one() + T::one());
let mut cost = r.dot(&r) * half;
for iter in 0..settings.max_iter {
let j = jacobian(&x);
j_evals += 1;
assert_eq!(
j.nrows(),
r.len(),
"jacobian rows must match residual length"
);
assert_eq!(j.ncols(), n, "jacobian cols must match x length");
let g = matt_vec(&j, &r);
let g_norm = g.norm();
if g_norm < settings.grad_tol {
return Ok(LeastSquaresResultDyn {
x,
cost,
grad_norm: g_norm,
iterations: iter,
r_evals,
j_evals,
});
}
let neg_r = vec_neg(&r);
let qr = j.qr().map_err(|_| OptimError::Singular)?;
let delta = qr.solve(&neg_r);
let x_new = vec_add(&x, &delta);
r = residual(&x_new);
r_evals += 1;
let cost_new = r.dot(&r) * half;
if delta.norm() < settings.x_tol * (T::one() + x.norm()) {
return Ok(LeastSquaresResultDyn {
x: x_new,
cost: cost_new,
grad_norm: g_norm,
iterations: iter + 1,
r_evals,
j_evals,
});
}
if (cost - cost_new).abs() < settings.f_tol * (T::one() + cost.abs()) {
return Ok(LeastSquaresResultDyn {
x: x_new,
cost: cost_new,
grad_norm: g_norm,
iterations: iter + 1,
r_evals,
j_evals,
});
}
x = x_new;
cost = cost_new;
}
Err(OptimError::MaxIterations)
}
pub fn least_squares_lm_dyn<T: FloatScalar>(
mut residual: impl FnMut(&DynVector<T>) -> DynVector<T>,
mut jacobian: impl FnMut(&DynVector<T>) -> DynMatrix<T>,
x0: &DynVector<T>,
settings: &LmSettings<T>,
) -> Result<LeastSquaresResultDyn<T>, OptimError> {
let n = x0.len();
let mut x = x0.clone();
let mut r = residual(&x);
let mut r_evals = 1usize;
let mut j_evals = 0usize;
let half = T::one() / (T::one() + T::one());
let mut cost = r.dot(&r) * half;
let mut mu = settings.mu_init;
for iter in 0..settings.max_iter {
let j = jacobian(&x);
j_evals += 1;
assert_eq!(
j.nrows(),
r.len(),
"jacobian rows must match residual length"
);
assert_eq!(j.ncols(), n, "jacobian cols must match x length");
let jt = j.transpose();
let jtj = &jt * &j;
let g = matt_vec(&j, &r);
let g_norm = g.norm();
if g_norm < settings.grad_tol {
return Ok(LeastSquaresResultDyn {
x,
cost,
grad_norm: g_norm,
iterations: iter,
r_evals,
j_evals,
});
}
let neg_g = vec_neg(&g);
let delta = solve_damped_dyn(&jtj, mu, &neg_g)?;
let x_new = vec_add(&x, &delta);
let r_new = residual(&x_new);
r_evals += 1;
let cost_new = r_new.dot(&r_new) * half;
let mu_delta_minus_g = vec_sub(&vec_scale(&delta, mu), &g);
let predicted = delta.dot(&mu_delta_minus_g);
let eps = T::epsilon();
if predicted > T::zero() && predicted.abs() >= eps * cost.abs() {
let actual = cost - cost_new;
let gain_ratio = actual / predicted;
if gain_ratio > T::zero() {
x = x_new;
r = r_new;
cost = cost_new;
mu = (mu * settings.mu_decrease).max(settings.mu_min);
if delta.norm() < settings.x_tol * (T::one() + x.norm()) {
return Ok(LeastSquaresResultDyn {
x,
cost,
grad_norm: g_norm,
iterations: iter + 1,
r_evals,
j_evals,
});
}
if actual.abs() < settings.f_tol * (T::one() + cost.abs()) {
return Ok(LeastSquaresResultDyn {
x,
cost,
grad_norm: g_norm,
iterations: iter + 1,
r_evals,
j_evals,
});
}
} else {
mu = (mu * settings.mu_increase).min(settings.mu_max);
}
} else {
mu = (mu * settings.mu_increase).min(settings.mu_max);
}
}
Err(OptimError::MaxIterations)
}
fn solve_damped_dyn<T: FloatScalar>(
a: &DynMatrix<T>,
mu: T,
b: &DynVector<T>,
) -> Result<DynVector<T>, OptimError> {
let n = a.nrows();
let mut damped = a.clone();
for i in 0..n {
damped[(i, i)] = damped[(i, i)] + mu;
}
if let Ok(chol) = damped.cholesky() {
return Ok(chol.solve(b));
}
let lu = damped.lu().map_err(|_| OptimError::Singular)?;
Ok(lu.solve(b))
}