use crate::error::ErrorLogging;
use crate::{core::problem, error, linalg, optimizer};
use apex_manifolds as manifold;
use std::{collections, time};
use tracing::debug;
use crate::linalg::{
DenseCholeskySolver, DenseMode, DenseQRSolver, JacobianMode, LinearSolver, LinearSolverType,
SparseCholeskySolver, SparseMode, SparseQRSolver,
};
use crate::optimizer::{AssemblyBackend, IterationStats};
#[derive(Clone)]
pub struct GaussNewtonConfig {
pub linear_solver_type: linalg::LinearSolverType,
pub max_iterations: usize,
pub cost_tolerance: f64,
pub parameter_tolerance: f64,
pub gradient_tolerance: f64,
pub timeout: Option<time::Duration>,
pub use_jacobi_scaling: bool,
pub min_diagonal: f64,
pub min_cost_threshold: Option<f64>,
pub max_condition_number: Option<f64>,
pub compute_covariances: bool,
#[cfg(feature = "visualization")]
pub enable_visualization: bool,
}
impl Default for GaussNewtonConfig {
fn default() -> Self {
Self {
linear_solver_type: linalg::LinearSolverType::default(),
max_iterations: 50,
cost_tolerance: 1e-6,
parameter_tolerance: 1e-8,
gradient_tolerance: 1e-10,
timeout: None,
use_jacobi_scaling: false,
min_diagonal: 1e-10,
min_cost_threshold: None,
max_condition_number: None,
compute_covariances: false,
#[cfg(feature = "visualization")]
enable_visualization: false,
}
}
}
impl GaussNewtonConfig {
pub fn new() -> Self {
Self::default()
}
pub fn with_linear_solver_type(mut self, linear_solver_type: linalg::LinearSolverType) -> Self {
self.linear_solver_type = linear_solver_type;
self
}
pub fn with_max_iterations(mut self, max_iterations: usize) -> Self {
self.max_iterations = max_iterations;
self
}
pub fn with_cost_tolerance(mut self, cost_tolerance: f64) -> Self {
self.cost_tolerance = cost_tolerance;
self
}
pub fn with_parameter_tolerance(mut self, parameter_tolerance: f64) -> Self {
self.parameter_tolerance = parameter_tolerance;
self
}
pub fn with_gradient_tolerance(mut self, gradient_tolerance: f64) -> Self {
self.gradient_tolerance = gradient_tolerance;
self
}
pub fn with_timeout(mut self, timeout: time::Duration) -> Self {
self.timeout = Some(timeout);
self
}
pub fn with_jacobi_scaling(mut self, use_jacobi_scaling: bool) -> Self {
self.use_jacobi_scaling = use_jacobi_scaling;
self
}
pub fn with_min_diagonal(mut self, min_diagonal: f64) -> Self {
self.min_diagonal = min_diagonal;
self
}
pub fn with_min_cost_threshold(mut self, min_cost: f64) -> Self {
self.min_cost_threshold = Some(min_cost);
self
}
pub fn with_max_condition_number(mut self, max_cond: f64) -> Self {
self.max_condition_number = Some(max_cond);
self
}
pub fn with_compute_covariances(mut self, compute_covariances: bool) -> Self {
self.compute_covariances = compute_covariances;
self
}
#[cfg(feature = "visualization")]
pub fn with_visualization(mut self, enable: bool) -> Self {
self.enable_visualization = enable;
self
}
pub fn print_configuration(&self) {
debug!(
"\nConfiguration:\n Solver: Gauss-Newton\n Linear solver: {:?}\n Convergence Criteria:\n Max iterations: {}\n Cost tolerance: {:.2e}\n Parameter tolerance: {:.2e}\n Gradient tolerance: {:.2e}\n Timeout: {:?}\n Numerical Settings:\n Jacobi scaling: {}\n Compute covariances: {}",
self.linear_solver_type,
self.max_iterations,
self.cost_tolerance,
self.parameter_tolerance,
self.gradient_tolerance,
self.timeout,
if self.use_jacobi_scaling {
"enabled"
} else {
"disabled"
},
if self.compute_covariances {
"enabled"
} else {
"disabled"
}
);
}
}
struct StepResult {
step: faer::Mat<f64>,
gradient_norm: f64,
}
struct CostEvaluation {
new_cost: f64,
cost_reduction: f64,
}
pub struct GaussNewton {
config: GaussNewtonConfig,
jacobi_scaling: Option<Vec<f64>>,
observers: optimizer::OptObserverVec,
}
impl Default for GaussNewton {
fn default() -> Self {
Self::new()
}
}
impl GaussNewton {
pub fn new() -> Self {
Self::with_config(GaussNewtonConfig::default())
}
pub fn with_config(config: GaussNewtonConfig) -> Self {
Self {
config,
jacobi_scaling: None,
observers: optimizer::OptObserverVec::new(),
}
}
pub fn add_observer(&mut self, observer: impl optimizer::OptObserver + 'static) {
self.observers.add(observer);
}
fn compute_step_generic<M: AssemblyBackend>(
&self,
residuals: &faer::Mat<f64>,
scaled_jacobian: &M::Jacobian,
linear_solver: &mut dyn LinearSolver<M>,
) -> Result<StepResult, optimizer::OptimizerError> {
let residuals_owned = residuals.as_ref().to_owned();
let scaled_step = linear_solver
.solve_normal_equation(&residuals_owned, scaled_jacobian)
.map_err(|e| {
optimizer::OptimizerError::LinearSolveFailed(e.to_string()).log_with_source(e)
})?;
let gradient = linear_solver.get_gradient().ok_or_else(|| {
optimizer::OptimizerError::NumericalInstability("Gradient not available".into()).log()
})?;
let gradient_norm = gradient.norm_l2();
let step = if self.config.use_jacobi_scaling {
let scaling = self
.jacobi_scaling
.as_ref()
.ok_or_else(|| optimizer::OptimizerError::JacobiScalingNotInitialized.log())?;
M::apply_inverse_scaling(&scaled_step, scaling)
} else {
scaled_step
};
Ok(StepResult {
step,
gradient_norm,
})
}
fn apply_step_and_evaluate_cost(
&self,
step_result: &StepResult,
state: &mut optimizer::InitializedState,
problem: &problem::Problem,
) -> error::ApexSolverResult<CostEvaluation> {
let _step_norm = optimizer::apply_parameter_step(
&mut state.variables,
step_result.step.as_ref(),
&state.sorted_vars,
);
let new_residual = problem.compute_residual_sparse(&state.variables)?;
let new_cost = optimizer::compute_cost(&new_residual);
let cost_reduction = state.current_cost - new_cost;
state.current_cost = new_cost;
Ok(CostEvaluation {
new_cost,
cost_reduction,
})
}
fn optimize_with_mode<M: AssemblyBackend>(
&mut self,
problem: &problem::Problem,
initial_params: &collections::HashMap<
String,
(manifold::ManifoldType, nalgebra::DVector<f64>),
>,
linear_solver: &mut dyn LinearSolver<M>,
) -> Result<
optimizer::SolverResult<collections::HashMap<String, problem::VariableEnum>>,
error::ApexSolverError,
> {
let start_time = time::Instant::now();
let mut iteration = 0;
let mut cost_evaluations = 1; let mut jacobian_evaluations = 0;
let mut state = optimizer::initialize_optimization_state(problem, initial_params)?;
let mut max_gradient_norm: f64 = 0.0;
let mut max_parameter_update_norm: f64 = 0.0;
let mut total_cost_reduction = 0.0;
let mut final_gradient_norm;
let mut final_parameter_update_norm;
let mut iteration_stats = Vec::with_capacity(self.config.max_iterations);
let mut previous_cost = state.current_cost;
if tracing::enabled!(tracing::Level::DEBUG) {
self.config.print_configuration();
IterationStats::print_header();
}
loop {
let iter_start = time::Instant::now();
let (residuals, jacobian) = M::assemble(
problem,
&state.variables,
&state.variable_index_map,
state.symbolic_structure.as_ref(),
state.total_dof,
)?;
jacobian_evaluations += 1;
let scaled_jacobian = if self.config.use_jacobi_scaling {
optimizer::process_jacobian_generic::<M>(
&jacobian,
&mut self.jacobi_scaling,
iteration,
)?
} else {
jacobian
};
let step_result =
self.compute_step_generic::<M>(&residuals, &scaled_jacobian, linear_solver)?;
max_gradient_norm = max_gradient_norm.max(step_result.gradient_norm);
final_gradient_norm = step_result.gradient_norm;
let step_norm = step_result.step.norm_l2();
max_parameter_update_norm = max_parameter_update_norm.max(step_norm);
final_parameter_update_norm = step_norm;
let cost_before_step = state.current_cost;
let cost_eval = self.apply_step_and_evaluate_cost(&step_result, &mut state, problem)?;
cost_evaluations += 1;
total_cost_reduction += cost_eval.cost_reduction;
if tracing::enabled!(tracing::Level::DEBUG) {
let iter_elapsed_ms = iter_start.elapsed().as_secs_f64() * 1000.0;
let total_elapsed_ms = start_time.elapsed().as_secs_f64() * 1000.0;
let stats = IterationStats {
iteration,
cost: state.current_cost,
cost_change: previous_cost - state.current_cost,
gradient_norm: step_result.gradient_norm,
step_norm,
tr_ratio: 0.0, tr_radius: 0.0, ls_iter: 0, iter_time_ms: iter_elapsed_ms,
total_time_ms: total_elapsed_ms,
accepted: true, };
iteration_stats.push(stats.clone());
stats.print_line();
}
previous_cost = state.current_cost;
optimizer::notify_observers_generic::<M>(
&mut self.observers,
&state.variables,
iteration,
state.current_cost,
step_result.gradient_norm,
None, step_norm,
None, linear_solver,
);
let parameter_norm = optimizer::compute_parameter_norm(&state.variables);
let elapsed = start_time.elapsed();
if let Some(status) = optimizer::check_convergence(&optimizer::ConvergenceParams {
iteration,
current_cost: cost_before_step,
new_cost: cost_eval.new_cost,
parameter_norm,
parameter_update_norm: step_norm,
gradient_norm: step_result.gradient_norm,
elapsed,
step_accepted: true, max_iterations: self.config.max_iterations,
gradient_tolerance: self.config.gradient_tolerance,
parameter_tolerance: self.config.parameter_tolerance,
cost_tolerance: self.config.cost_tolerance,
min_cost_threshold: self.config.min_cost_threshold,
timeout: self.config.timeout,
trust_region_radius: None,
min_trust_region_radius: None,
}) {
if tracing::enabled!(tracing::Level::DEBUG) {
let summary = optimizer::create_optimizer_summary(
"Gauss-Newton",
state.initial_cost,
state.current_cost,
iteration + 1,
None,
None,
max_gradient_norm,
final_gradient_norm,
max_parameter_update_norm,
final_parameter_update_norm,
total_cost_reduction,
elapsed,
iteration_stats.clone(),
status.clone(),
None,
None,
None,
);
debug!("{}", summary);
}
let covariances = if self.config.compute_covariances {
problem.compute_and_set_covariances_generic::<M>(
linear_solver,
&mut state.variables,
&state.variable_index_map,
)
} else {
None
};
return Ok(optimizer::build_solver_result(
status,
iteration + 1,
state,
elapsed,
final_gradient_norm,
final_parameter_update_norm,
cost_evaluations,
jacobian_evaluations,
covariances,
));
}
iteration += 1;
}
}
pub fn optimize(
&mut self,
problem: &problem::Problem,
initial_params: &collections::HashMap<
String,
(manifold::ManifoldType, nalgebra::DVector<f64>),
>,
) -> Result<
optimizer::SolverResult<collections::HashMap<String, problem::VariableEnum>>,
error::ApexSolverError,
> {
match problem.jacobian_mode {
JacobianMode::Dense => match self.config.linear_solver_type {
LinearSolverType::DenseQR => {
let mut solver = DenseQRSolver::new();
self.optimize_with_mode::<DenseMode>(problem, initial_params, &mut solver)
}
_ => {
let mut solver = DenseCholeskySolver::new();
self.optimize_with_mode::<DenseMode>(problem, initial_params, &mut solver)
}
},
JacobianMode::Sparse => match self.config.linear_solver_type {
linalg::LinearSolverType::SparseQR => {
let mut solver = SparseQRSolver::new();
self.optimize_with_mode::<SparseMode>(problem, initial_params, &mut solver)
}
_ => {
let mut solver = SparseCholeskySolver::new();
self.optimize_with_mode::<SparseMode>(problem, initial_params, &mut solver)
}
},
}
}
}
impl optimizer::Optimizer for GaussNewton {
fn optimize(
&mut self,
problem: &problem::Problem,
initial_params: &std::collections::HashMap<
String,
(manifold::ManifoldType, nalgebra::DVector<f64>),
>,
) -> Result<
optimizer::SolverResult<std::collections::HashMap<String, problem::VariableEnum>>,
crate::error::ApexSolverError,
> {
self.optimize(problem, initial_params)
}
}
#[cfg(test)]
mod tests {
use crate::{core::problem, factors, linalg::JacobianMode, optimizer};
use apex_manifolds as manifold;
use nalgebra::dvector;
use std::collections;
type TestResult = Result<(), Box<dyn std::error::Error>>;
#[derive(Debug, Clone)]
struct RosenbrockFactor1;
impl factors::Factor for RosenbrockFactor1 {
fn linearize(
&self,
params: &[nalgebra::DVector<f64>],
compute_jacobian: bool,
) -> (nalgebra::DVector<f64>, Option<nalgebra::DMatrix<f64>>) {
let x1 = params[0][0];
let x2 = params[1][0];
let residual = nalgebra::dvector![10.0 * (x2 - x1 * x1)];
let jacobian = if compute_jacobian {
let mut jac = nalgebra::DMatrix::zeros(1, 2);
jac[(0, 0)] = -20.0 * x1;
jac[(0, 1)] = 10.0;
Some(jac)
} else {
None
};
(residual, jacobian)
}
fn get_dimension(&self) -> usize {
1
}
}
#[derive(Debug, Clone)]
struct RosenbrockFactor2;
impl factors::Factor for RosenbrockFactor2 {
fn linearize(
&self,
params: &[nalgebra::DVector<f64>],
compute_jacobian: bool,
) -> (nalgebra::DVector<f64>, Option<nalgebra::DMatrix<f64>>) {
let x1 = params[0][0];
let residual = nalgebra::dvector![1.0 - x1];
let jacobian = if compute_jacobian {
Some(nalgebra::DMatrix::from_element(1, 1, -1.0))
} else {
None
};
(residual, jacobian)
}
fn get_dimension(&self) -> usize {
1
}
}
#[test]
fn test_rosenbrock_optimization() -> TestResult {
let mut problem = problem::Problem::new(JacobianMode::Sparse);
let mut initial_values = collections::HashMap::new();
initial_values.insert(
"x1".to_string(),
(manifold::ManifoldType::RN, dvector![-1.2]),
);
initial_values.insert(
"x2".to_string(),
(manifold::ManifoldType::RN, dvector![1.0]),
);
problem.add_residual_block(&["x1", "x2"], Box::new(RosenbrockFactor1), None);
problem.add_residual_block(&["x1"], Box::new(RosenbrockFactor2), None);
let config = optimizer::gauss_newton::GaussNewtonConfig::new()
.with_max_iterations(100)
.with_cost_tolerance(1e-8)
.with_parameter_tolerance(1e-8)
.with_gradient_tolerance(1e-10);
let mut solver = optimizer::GaussNewton::with_config(config);
let result = solver.optimize(&problem, &initial_values)?;
let x1_final = result
.parameters
.get("x1")
.ok_or("x1 not found")?
.to_vector()[0];
let x2_final = result
.parameters
.get("x2")
.ok_or("x2 not found")?
.to_vector()[0];
assert!(
matches!(
result.status,
optimizer::OptimizationStatus::Converged
| optimizer::OptimizationStatus::CostToleranceReached
| optimizer::OptimizationStatus::ParameterToleranceReached
| optimizer::OptimizationStatus::GradientToleranceReached
),
"Optimization should converge"
);
assert!(
(x1_final - 1.0).abs() < 1e-4,
"x1 should converge to 1.0, got {}",
x1_final
);
assert!(
(x2_final - 1.0).abs() < 1e-4,
"x2 should converge to 1.0, got {}",
x2_final
);
assert!(
result.final_cost < 1e-6,
"Final cost should be near zero, got {}",
result.final_cost
);
Ok(())
}
struct LinearFactor {
target: f64,
}
impl factors::Factor for LinearFactor {
fn linearize(
&self,
params: &[nalgebra::DVector<f64>],
compute_jacobian: bool,
) -> (nalgebra::DVector<f64>, Option<nalgebra::DMatrix<f64>>) {
let residual = nalgebra::dvector![params[0][0] - self.target];
let jacobian = if compute_jacobian {
Some(nalgebra::DMatrix::from_element(1, 1, 1.0))
} else {
None
};
(residual, jacobian)
}
fn get_dimension(&self) -> usize {
1
}
}
fn rosenbrock_problem() -> (
problem::Problem,
collections::HashMap<String, (manifold::ManifoldType, nalgebra::DVector<f64>)>,
) {
let mut prob = problem::Problem::new(JacobianMode::Sparse);
let mut init = collections::HashMap::new();
init.insert(
"x1".to_string(),
(manifold::ManifoldType::RN, nalgebra::dvector![-1.2]),
);
init.insert(
"x2".to_string(),
(manifold::ManifoldType::RN, nalgebra::dvector![1.0]),
);
prob.add_residual_block(&["x1", "x2"], Box::new(RosenbrockFactor1), None);
prob.add_residual_block(&["x1"], Box::new(RosenbrockFactor2), None);
(prob, init)
}
fn linear_problem(
start: f64,
) -> (
problem::Problem,
collections::HashMap<String, (manifold::ManifoldType, nalgebra::DVector<f64>)>,
) {
let mut prob = problem::Problem::new(JacobianMode::Sparse);
let mut init = collections::HashMap::new();
init.insert(
"x".to_string(),
(manifold::ManifoldType::RN, nalgebra::dvector![start]),
);
prob.add_residual_block(&["x"], Box::new(LinearFactor { target: 0.0 }), None);
(prob, init)
}
#[test]
fn test_gn_config_default() {
let cfg = optimizer::gauss_newton::GaussNewtonConfig::default();
assert_eq!(cfg.max_iterations, 50);
assert!((cfg.cost_tolerance - 1e-6).abs() < 1e-15);
assert!(!cfg.use_jacobi_scaling);
assert!(!cfg.compute_covariances);
}
#[test]
fn test_gn_config_builders() {
use crate::linalg::LinearSolverType;
let cfg = optimizer::gauss_newton::GaussNewtonConfig::new()
.with_max_iterations(15)
.with_cost_tolerance(1e-5)
.with_parameter_tolerance(1e-6)
.with_gradient_tolerance(1e-7)
.with_jacobi_scaling(true)
.with_min_diagonal(1e-8)
.with_min_cost_threshold(1e-10)
.with_compute_covariances(true)
.with_linear_solver_type(LinearSolverType::SparseQR);
assert_eq!(cfg.max_iterations, 15);
assert!((cfg.cost_tolerance - 1e-5).abs() < 1e-20);
assert!(cfg.use_jacobi_scaling);
assert!(cfg.min_cost_threshold.is_some());
assert!(cfg.compute_covariances);
assert!(matches!(cfg.linear_solver_type, LinearSolverType::SparseQR));
}
#[test]
fn test_gn_print_configuration_no_panic() {
optimizer::gauss_newton::GaussNewtonConfig::default().print_configuration();
}
#[test]
fn test_gn_default_equals_new() {
let _a = optimizer::GaussNewton::new();
let _b = optimizer::GaussNewton::default();
}
#[test]
fn test_gn_with_config_method() {
let cfg = optimizer::gauss_newton::GaussNewtonConfig::new().with_max_iterations(3);
let _solver = optimizer::GaussNewton::with_config(cfg);
}
#[test]
fn test_gn_max_iterations_termination() -> TestResult {
let (problem, initial_values) = rosenbrock_problem();
let cfg = optimizer::gauss_newton::GaussNewtonConfig::new().with_max_iterations(2);
let mut solver = optimizer::GaussNewton::with_config(cfg);
let result = solver.optimize(&problem, &initial_values)?;
assert_eq!(
result.status,
optimizer::OptimizationStatus::MaxIterationsReached
);
Ok(())
}
#[test]
fn test_gn_gradient_tolerance_convergence() -> TestResult {
let (problem, initial_values) = linear_problem(1.0);
let cfg = optimizer::gauss_newton::GaussNewtonConfig::new()
.with_gradient_tolerance(1e3)
.with_cost_tolerance(1e-20)
.with_parameter_tolerance(1e-20);
let mut solver = optimizer::GaussNewton::with_config(cfg);
let result = solver.optimize(&problem, &initial_values)?;
assert_eq!(
result.status,
optimizer::OptimizationStatus::GradientToleranceReached
);
Ok(())
}
#[test]
fn test_gn_cost_tolerance_convergence() -> TestResult {
let (problem, initial_values) = rosenbrock_problem();
let cfg = optimizer::gauss_newton::GaussNewtonConfig::new()
.with_cost_tolerance(1e2) .with_gradient_tolerance(1e-20)
.with_parameter_tolerance(1e-20);
let mut solver = optimizer::GaussNewton::with_config(cfg);
let result = solver.optimize(&problem, &initial_values)?;
assert!(matches!(
result.status,
optimizer::OptimizationStatus::CostToleranceReached
| optimizer::OptimizationStatus::GradientToleranceReached
| optimizer::OptimizationStatus::ParameterToleranceReached
| optimizer::OptimizationStatus::Converged
));
Ok(())
}
#[test]
fn test_gn_qr_solver() -> TestResult {
use crate::linalg::LinearSolverType;
let (problem, initial_values) = rosenbrock_problem();
let cfg = optimizer::gauss_newton::GaussNewtonConfig::new()
.with_linear_solver_type(LinearSolverType::SparseQR)
.with_max_iterations(100);
let mut solver = optimizer::GaussNewton::with_config(cfg);
let result = solver.optimize(&problem, &initial_values)?;
assert!(result.final_cost < 1e-6);
Ok(())
}
#[test]
fn test_gn_jacobi_scaling_enabled() -> TestResult {
let (problem, initial_values) = rosenbrock_problem();
let cfg = optimizer::gauss_newton::GaussNewtonConfig::new()
.with_jacobi_scaling(true)
.with_max_iterations(100);
let mut solver = optimizer::GaussNewton::with_config(cfg);
let result = solver.optimize(&problem, &initial_values)?;
assert!(result.final_cost < 1e-6);
Ok(())
}
#[test]
fn test_gn_min_cost_threshold() -> TestResult {
let (problem, initial_values) = rosenbrock_problem();
let cfg = optimizer::gauss_newton::GaussNewtonConfig::new()
.with_min_cost_threshold(1e10)
.with_cost_tolerance(1e-20)
.with_gradient_tolerance(1e-20)
.with_parameter_tolerance(1e-20);
let mut solver = optimizer::GaussNewton::with_config(cfg);
let result = solver.optimize(&problem, &initial_values)?;
assert_eq!(
result.status,
optimizer::OptimizationStatus::MinCostThresholdReached
);
Ok(())
}
#[test]
fn test_gn_result_fields() -> TestResult {
let (problem, initial_values) = rosenbrock_problem();
let mut solver = optimizer::GaussNewton::new();
let result = solver.optimize(&problem, &initial_values)?;
assert!(result.initial_cost > result.final_cost);
assert!(result.iterations > 0);
Ok(())
}
#[test]
fn test_gn_convergence_info_populated() -> TestResult {
let (problem, initial_values) = rosenbrock_problem();
let mut solver = optimizer::GaussNewton::new();
let result = solver.optimize(&problem, &initial_values)?;
assert!(result.convergence_info.is_some());
Ok(())
}
#[test]
fn test_gn_timeout_config() {
let cfg = optimizer::gauss_newton::GaussNewtonConfig::new()
.with_timeout(std::time::Duration::from_secs(60));
assert!(cfg.timeout.is_some());
}
}