use crate::numerical::Nonlinear_systems::engine::nonlinear_log_info;
use nalgebra::{DMatrix, DVector};
use strum_macros::Display;
#[derive(Debug, Clone, PartialEq, Display, Copy)]
pub enum ScalingMethod {
Levenberg,
Marquardt,
More,
}
pub struct TrustRegionScaling;
impl TrustRegionScaling {
pub fn init_scaling(jacobian: &DMatrix<f64>, method: &ScalingMethod) -> DVector<f64> {
match method {
ScalingMethod::Levenberg => Self::init_levenberg_scaling(jacobian),
ScalingMethod::Marquardt => Self::init_marquardt_scaling(jacobian),
ScalingMethod::More => Self::init_more_scaling(jacobian),
}
}
pub fn update_scaling(
jacobian: &DMatrix<f64>,
method: &ScalingMethod,
current_diag: &mut DVector<f64>,
) {
match method {
ScalingMethod::Levenberg => {
}
ScalingMethod::Marquardt => {
Self::update_marquardt_scaling(jacobian, current_diag);
}
ScalingMethod::More => {
Self::update_more_scaling(jacobian, current_diag);
}
}
}
fn init_levenberg_scaling(jacobian: &DMatrix<f64>) -> DVector<f64> {
let n_params = jacobian.ncols();
DVector::from_element(n_params, 1.0)
}
fn init_marquardt_scaling(jacobian: &DMatrix<f64>) -> DVector<f64> {
let n_params = jacobian.ncols();
let mut diag = DVector::zeros(n_params);
Self::update_marquardt_scaling(jacobian, &mut diag);
nonlinear_log_info!("\n Diagonal scaling vector created!: {}", diag);
diag
}
fn update_marquardt_scaling(jacobian: &DMatrix<f64>, diag: &mut DVector<f64>) {
let n_params = jacobian.ncols();
for j in 0..n_params {
let column = jacobian.column(j);
let norm = column.norm();
let scaling_value = if norm < 1e-12 { 1.0 } else { norm };
diag[j] = scaling_value;
}
}
fn init_more_scaling(jacobian: &DMatrix<f64>) -> DVector<f64> {
Self::init_marquardt_scaling(jacobian)
}
fn update_more_scaling(jacobian: &DMatrix<f64>, diag: &mut DVector<f64>) {
let n_params = jacobian.ncols();
for j in 0..n_params {
let column = jacobian.column(j);
let norm = column.norm();
let scaling_value = if norm < 1e-12 { 1.0 } else { norm };
diag[j] = scaling_value.max(diag[j]);
}
}
pub fn create_scaled_regularization(
mu: f64,
diag: &DVector<f64>,
method: &ScalingMethod,
) -> DMatrix<f64> {
let n = diag.len();
let mut regularization = DMatrix::zeros(n, n);
match method {
ScalingMethod::Levenberg => {
for i in 0..n {
regularization[(i, i)] = mu;
}
}
ScalingMethod::Marquardt | ScalingMethod::More => {
for i in 0..n {
regularization[(i, i)] = mu * diag[i] * diag[i];
}
}
}
regularization
}
pub(crate) fn add_scaled_regularization_in_place(
matrix: &mut DMatrix<f64>,
mu: f64,
diag: &DVector<f64>,
method: &ScalingMethod,
) {
debug_assert_eq!(matrix.nrows(), matrix.ncols());
debug_assert_eq!(matrix.nrows(), diag.len());
for index in 0..diag.len() {
let increment = match method {
ScalingMethod::Levenberg => mu,
ScalingMethod::Marquardt | ScalingMethod::More => mu * diag[index] * diag[index],
};
matrix[(index, index)] += increment;
}
}
pub(crate) fn add_squared_diagonal_regularization_in_place(
matrix: &mut DMatrix<f64>,
mu: f64,
diag: &DVector<f64>,
) {
debug_assert_eq!(matrix.nrows(), matrix.ncols());
debug_assert_eq!(matrix.nrows(), diag.len());
for index in 0..diag.len() {
matrix[(index, index)] += mu * diag[index] * diag[index];
}
}
pub(crate) fn add_jtj_diagonal_regularization_in_place(matrix: &mut DMatrix<f64>, lambda: f64) {
debug_assert_eq!(matrix.nrows(), matrix.ncols());
for index in 0..matrix.nrows() {
let diagonal = matrix[(index, index)];
matrix[(index, index)] += lambda * diagonal;
}
}
pub(crate) fn add_identity_regularization_in_place(matrix: &mut DMatrix<f64>, lambda: f64) {
debug_assert_eq!(matrix.nrows(), matrix.ncols());
for index in 0..matrix.nrows() {
matrix[(index, index)] += lambda;
}
}
}
#[derive(Debug, Clone, PartialEq, Display, Copy)]
pub enum ReductionRatio {
Nielsen, More, }
pub struct ReductionRatioSolver {}
impl ReductionRatioSolver {
pub fn solve_reduction_ratio(
Jy: &DMatrix<f64>,
Fy: &DVector<f64>,
Fy_new: &DVector<f64>,
scaling: &DVector<f64>,
p: &DVector<f64>,
mu: f64,
reduction_ratio_method: ReductionRatio,
) -> f64 {
match reduction_ratio_method {
ReductionRatio::Nielsen => {
let rr = Self::nielsen_reduction_ratio(Jy, Fy, Fy_new, p);
rr
}
ReductionRatio::More => {
let rr = Self::gsl_reduction_ratio(Jy, Fy, Fy_new, p, scaling, mu);
rr
}
}
}
fn nielsen_reduction_ratio(
Jy: &DMatrix<f64>,
Fy: &DVector<f64>,
Fy_new: &DVector<f64>,
p: &DVector<f64>,
) -> f64 {
let actual_reduction = Fy.norm_squared() - Fy_new.norm_squared();
let predicted_residual = Fy + Jy * p;
let predicted_reduction = Fy.norm_squared() - predicted_residual.norm_squared();
let rho = if predicted_reduction.abs() > 1e-12 {
actual_reduction / predicted_reduction
} else {
0.0
};
nonlinear_log_info!(
"rho = {}, actual_reduction = {}, predicted_reduction = {}",
rho,
actual_reduction,
predicted_reduction
);
rho
}
fn gsl_reduction_ratio(
Jy: &DMatrix<f64>,
Fy: &DVector<f64>,
Fy_new: &DVector<f64>,
p: &DVector<f64>,
scaling: &DVector<f64>,
mu: f64,
) -> f64 {
let norm_f = Fy.norm_squared();
let norm_f_trial = Fy_new.norm_squared();
if norm_f_trial >= norm_f {
nonlinear_log_info!("Cost increased, rejecting step");
return -1.0; }
let actual_reduction = 1.0 - norm_f_trial / norm_f;
let predicted_reduction = Self::predicted_reduction(Fy, Jy, p, scaling, mu);
if predicted_reduction > 0.0 {
let rho = actual_reduction / predicted_reduction;
nonlinear_log_info!(
"rho = {}, actual_reduction = {}, predicted_reduction = {}",
rho,
actual_reduction,
predicted_reduction
);
return rho;
} else {
nonlinear_log_info!(
"rho = {}, actual_reduction = {}, predicted_reduction = {}",
-1.0,
actual_reduction,
predicted_reduction
);
return -1.0; }
}
fn predicted_reduction(
Fy: &DVector<f64>,
Jy: &DMatrix<f64>,
p: &DVector<f64>,
scaling: &DVector<f64>,
mu: f64,
) -> f64 {
let norm_f = Fy.norm();
let norm_Jp = (Jy * p).norm();
let norm_Dp = (scaling.component_mul(p)).norm();
let u = norm_Jp / norm_f;
let v = norm_Dp / norm_f;
let predicted_reduction = u.powf(2.0) + 2.0 * mu * v.powf(2.0);
predicted_reduction
}
}
#[derive(Debug, Clone, PartialEq, Display)]
pub enum ConvergenceCriteria {
GSL,
SimpleScaled,
}
pub struct ConvergenceCriteriaSolver {}
impl ConvergenceCriteriaSolver {
pub fn test_convergence(
x: &DVector<f64>, dx: &DVector<f64>, g: &DVector<f64>, f: &DVector<f64>, diag: &DVector<f64>, xtol: f64, gtol: Option<f64>, ftol: Option<f64>, convergence_criteria: ConvergenceCriteria,
) -> (bool, i32) {
match convergence_criteria {
ConvergenceCriteria::GSL => test_convergence_gsl(x, dx, g, f, xtol, gtol, ftol),
ConvergenceCriteria::SimpleScaled => test_convergence_simple_scaled(dx, diag, xtol),
}
}
}
fn test_convergence_simple_scaled(
dx: &DVector<f64>,
diag: &DVector<f64>,
xtol: f64,
) -> (bool, i32) {
let scaled_norm = scaled_norm_common(diag, dx);
let converged = scaled_norm < xtol;
let info = if converged { 1 } else { 0 };
(converged, info)
}
pub fn scaled_norm_common(diag: &DVector<f64>, vector: &DVector<f64>) -> f64 {
let mut sum_squares = 0.0;
for i in 0..vector.len() {
let scaled_component = vector[i] * diag[i];
sum_squares += scaled_component * scaled_component;
}
let scaled_norm = sum_squares.sqrt();
nonlinear_log_info!("scaled norm ||step*D|| {}", scaled_norm);
scaled_norm
}
pub fn test_convergence_gsl(
x: &DVector<f64>, dx: &DVector<f64>, g: &DVector<f64>, f: &DVector<f64>, xtol: f64, gtol: Option<f64>, ftol: Option<f64>, ) -> (bool, i32) {
let gtol = gtol.unwrap_or(1e-6);
let ftol = ftol.unwrap_or(1e-6);
nonlinear_log_info!("Checking convergence conditions:");
nonlinear_log_info!(" xtol = {}, gtol = {}, ftol = {}", xtol, gtol, ftol);
nonlinear_log_info!(
" ||dx|| = {}, ||g|| = {}, ||f|| = {}",
dx.norm(),
g.norm(),
f.norm()
);
if test_parameter_change(dx, x, xtol) {
nonlinear_log_info!("CONVERGED: Parameter change test passed");
nonlinear_log_info!(" ||dx|| components are all < xtol = {}", xtol);
return (true, 1);
}
let gnorm = scaled_infnorm(x, g);
let fnorm = f.norm();
let phi = 0.5 * fnorm * fnorm;
let gradient_threshold = gtol * phi.max(1.0);
nonlinear_log_info!("Gradient test:");
nonlinear_log_info!(" gnorm = {}", gnorm);
nonlinear_log_info!(" phi = 0.5 * ||f||² = {}", phi);
nonlinear_log_info!(" threshold = gtol * max(phi, 1) = {}", gradient_threshold);
if gnorm <= gradient_threshold {
nonlinear_log_info!("CONVERGED: Gradient test passed");
nonlinear_log_info!(
" Scaled gradient norm {} <= threshold {}",
gnorm,
gradient_threshold
);
return (true, 2);
}
let _ftol_unused = ftol;
nonlinear_log_info!("No convergence detected:");
nonlinear_log_info!(" Parameter change test: FAILED");
nonlinear_log_info!(
" Gradient test: gnorm = {} > threshold = {}",
gnorm,
gradient_threshold
);
nonlinear_log_info!(" Function change test: DISABLED (following GSL)");
(false, 0)
}
fn test_parameter_change(dx: &DVector<f64>, x: &DVector<f64>, xtol: f64) -> bool {
let n = dx.len();
let xtol_squared = xtol * xtol;
nonlinear_log_info!("\n Parameter change test (component-wise):");
for i in 0..n {
let dx_i = dx[i];
let x_i = x[i];
if dx_i.abs() > xtol {
nonlinear_log_info!(
" Component {}: |dx| = {} > xtol = {} (FAIL absolute)",
i,
dx_i.abs(),
xtol
);
return false;
}
if dx_i.abs() > xtol_squared * x_i.abs() {
nonlinear_log_info!(
" Component {}: |dx| = {} > xtol² * |x| = {} (FAIL relative)",
i,
dx_i.abs(),
xtol_squared * x_i.abs()
);
return false;
}
nonlinear_log_info!(" Component {}: |dx| = {} (PASS both tests)", i, dx_i.abs());
}
nonlinear_log_info!("All components passed parameter change test");
true
}
fn scaled_infnorm(x: &DVector<f64>, g: &DVector<f64>) -> f64 {
let n = x.len();
let mut max_scaled_grad: f64 = 0.0;
nonlinear_log_info!("\n Computing scaled gradient norm:");
for i in 0..n {
let x_i = x[i];
let g_i = g[i];
let scale_factor = x_i.abs().max(1.0);
let scaled_grad_i: f64 = g_i.abs() * scale_factor;
max_scaled_grad = max_scaled_grad.max(scaled_grad_i);
nonlinear_log_info!(
" Component {}: |g| = {}, scale = {}, scaled = {}",
i,
g_i.abs(),
scale_factor,
scaled_grad_i
);
}
nonlinear_log_info!("Maximum scaled gradient component: {}", max_scaled_grad);
max_scaled_grad
}
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum ConvergenceInfo {
NoConvergence = 0,
ParameterChange = 1,
GradientSmall = 2,
FunctionChange = 3, }
impl From<i32> for ConvergenceInfo {
fn from(code: i32) -> Self {
match code {
1 => ConvergenceInfo::ParameterChange,
2 => ConvergenceInfo::GradientSmall,
3 => ConvergenceInfo::FunctionChange,
_ => ConvergenceInfo::NoConvergence,
}
}
}
impl std::fmt::Display for ConvergenceInfo {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
ConvergenceInfo::NoConvergence => write!(f, "No convergence"),
ConvergenceInfo::ParameterChange => write!(f, "Converged: parameter change small"),
ConvergenceInfo::GradientSmall => write!(f, "Converged: gradient small"),
ConvergenceInfo::FunctionChange => write!(f, "Converged: function change small"),
}
}
}
#[allow(dead_code)]
fn compute_scaled_gradient_norm(
jacobian: &DMatrix<f64>,
residuals: &DVector<f64>,
residuals_norm: f64,
) -> f64 {
let mut max_scaled_grad: f64 = 0.0;
for j in 0..jacobian.ncols() {
let col = jacobian.column(j);
let col_norm = col.norm();
if col_norm > 0.0 && residuals_norm > 0.0 {
let grad_component = col.dot(residuals).abs();
let scaled_grad: f64 = grad_component / (col_norm * residuals_norm);
max_scaled_grad = max_scaled_grad.max(scaled_grad);
}
}
max_scaled_grad
}
pub fn compute_scaled_gradient_norm2(
jacobian: &DMatrix<f64>, residuals: &DVector<f64>, ) -> Option<f64> {
let residual_norm = residuals.norm();
if residual_norm == 0.0 {
return Some(0.0);
}
let mut max_scaled_gradient: f64 = 0.0;
let n = jacobian.ncols();
for j in 0..n {
let jacobian_column = jacobian.column(j);
let mut column_norm = jacobian_column.norm();
if column_norm.abs() < 1e-40 {
column_norm = 1.0;
}
let gradient_component = jacobian_column.dot(residuals);
let scaled_gradient = gradient_component.abs() / (column_norm * residual_norm);
if scaled_gradient.is_nan() {
return None;
}
max_scaled_gradient = max_scaled_gradient.max(scaled_gradient);
}
Some(max_scaled_gradient)
}
#[cfg(test)]
mod tests {
use super::{ScalingMethod, TrustRegionScaling};
use nalgebra::{DMatrix, DVector};
#[test]
fn in_place_regularization_matches_materialized_matrix() {
let base =
DMatrix::from_row_slice(3, 3, &[4.0, 0.5, -0.25, 0.5, 9.0, 1.25, -0.25, 1.25, 16.0]);
let scaling = DVector::from_vec(vec![2.0, 3.0, 4.0]);
let mu = 0.125;
for method in [ScalingMethod::Levenberg, ScalingMethod::Marquardt] {
let mut actual = base.clone();
TrustRegionScaling::add_scaled_regularization_in_place(
&mut actual,
mu,
&scaling,
&method,
);
let regularization =
TrustRegionScaling::create_scaled_regularization(mu, &scaling, &method);
let expected = &base + regularization;
assert!((actual - expected).norm() < 1e-14);
}
let mut actual = base.clone();
TrustRegionScaling::add_jtj_diagonal_regularization_in_place(&mut actual, mu);
let diagonal = DMatrix::from_diagonal(&base.diagonal());
let expected = &base + mu * diagonal;
assert!((actual - expected).norm() < 1e-14);
}
#[test]
fn identity_regularization_updates_only_the_diagonal() {
let base = DMatrix::from_row_slice(2, 2, &[2.0, 3.0, 5.0, 7.0]);
let mut actual = base.clone();
TrustRegionScaling::add_identity_regularization_in_place(&mut actual, 0.75);
assert_eq!(actual[(0, 0)], 2.75);
assert_eq!(actual[(0, 1)], 3.0);
assert_eq!(actual[(1, 0)], 5.0);
assert_eq!(actual[(1, 1)], 7.75);
}
}