use crate::{core::problem, error, linalg, optimizer};
use apex_manifolds as manifold;
use faer::sparse;
use std::{collections, time};
use tracing::debug;
use crate::optimizer::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<sparse::SparseColMat<usize, 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_gauss_newton_step(
&self,
residuals: &faer::Mat<f64>,
scaled_jacobian: &sparse::SparseColMat<usize, f64>,
linear_solver: &mut Box<dyn linalg::SparseLinearSolver>,
) -> Option<StepResult> {
let residuals_owned = residuals.as_ref().to_owned();
let scaled_step = linear_solver
.solve_normal_equation(&residuals_owned, scaled_jacobian)
.ok()?;
let gradient = linear_solver.get_gradient()?;
let gradient_norm = gradient.norm_l2();
let step = if self.config.use_jacobi_scaling {
let scaling = self.jacobi_scaling.as_ref()?;
&scaled_step * scaling
} else {
scaled_step
};
Some(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,
})
}
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,
> {
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 linear_solver = optimizer::create_linear_solver(&self.config.linear_solver_type);
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) = problem.compute_residual_and_jacobian_sparse(
&state.variables,
&state.variable_index_map,
&state.symbolic_structure,
)?;
jacobian_evaluations += 1;
let scaled_jacobian = if self.config.use_jacobi_scaling {
optimizer::process_jacobian(&jacobian, &mut self.jacobi_scaling, iteration)?
} else {
jacobian
};
let step_result = match self.compute_gauss_newton_step(
&residuals,
&scaled_jacobian,
&mut linear_solver,
) {
Some(result) => result,
None => {
return Err(optimizer::OptimizerError::LinearSolveFailed(
"Linear solver failed to solve Gauss-Newton system".to_string(),
)
.into());
}
};
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(
&mut self.observers,
&state.variables,
iteration,
state.current_cost,
step_result.gradient_norm,
None, step_norm,
None, linear_solver.as_ref(),
);
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(
&mut 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;
}
}
}
impl optimizer::Solver for GaussNewton {
type Config = GaussNewtonConfig;
type Error = error::ApexSolverError;
fn new() -> Self {
Self::default()
}
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>>,
Self::Error,
> {
self.optimize(problem, initial_params)
}
}
#[cfg(test)]
mod tests {
use crate::{core::problem, factors, optimizer};
use apex_manifolds as manifold;
use nalgebra::dvector;
use std::collections;
#[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() -> Result<(), Box<dyn std::error::Error>> {
let mut problem = problem::Problem::new();
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(())
}
}