use crate::core::problem::{Problem, SymbolicStructure, VariableEnum};
use crate::error;
use crate::linalg::{self, SparseCholeskySolver, SparseLinearSolver, SparseQRSolver};
use apex_manifolds::ManifoldType;
use faer::sparse::{SparseColMat, Triplet};
use faer::{Mat, MatRef};
use nalgebra::DVector;
use std::collections::HashMap;
use std::time::{self, Duration};
use std::{
fmt,
fmt::{Display, Formatter},
};
use thiserror::Error;
use tracing::{debug, error};
pub mod dog_leg;
pub mod gauss_newton;
pub mod levenberg_marquardt;
pub use dog_leg::DogLeg;
pub use gauss_newton::GaussNewton;
pub use levenberg_marquardt::LevenbergMarquardt;
pub use crate::observers::{OptObserver, OptObserverVec};
#[derive(Default, Clone, Copy, PartialEq, Eq)]
pub enum OptimizerType {
#[default]
LevenbergMarquardt,
GaussNewton,
DogLeg,
}
impl Display for OptimizerType {
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
match self {
OptimizerType::LevenbergMarquardt => write!(f, "Levenberg-Marquardt"),
OptimizerType::GaussNewton => write!(f, "Gauss-Newton"),
OptimizerType::DogLeg => write!(f, "Dog Leg"),
}
}
}
#[derive(Debug, Clone, Error)]
pub enum OptimizerError {
#[error("Linear system solve failed: {0}")]
LinearSolveFailed(String),
#[error("Maximum iterations ({max_iters}) reached without convergence")]
MaxIterationsReached { max_iters: usize },
#[error("Trust region radius became too small: {radius:.6e} < {min_radius:.6e}")]
TrustRegionFailure { radius: f64, min_radius: f64 },
#[error("Damping parameter became too large: {damping:.6e} > {max_damping:.6e}")]
DampingFailure { damping: f64, max_damping: f64 },
#[error("Cost increased unexpectedly: {old_cost:.6e} -> {new_cost:.6e}")]
CostIncrease { old_cost: f64, new_cost: f64 },
#[error("Jacobian computation failed: {0}")]
JacobianFailed(String),
#[error("Invalid optimization parameters: {0}")]
InvalidParameters(String),
#[error("Numerical instability detected: {0}")]
NumericalInstability(String),
#[error("Linear algebra error: {0}")]
LinAlg(#[from] linalg::LinAlgError),
#[error("Problem has no variables to optimize")]
EmptyProblem,
#[error("Problem has no residual blocks")]
NoResidualBlocks,
#[error("Failed to create Jacobi scaling matrix: {0}")]
JacobiScalingCreation(String),
#[error("Jacobi scaling not initialized")]
JacobiScalingNotInitialized,
#[error("Unknown linear solver type: {0}")]
UnknownLinearSolver(String),
}
impl OptimizerError {
#[must_use]
pub fn log(self) -> Self {
error!("{}", self);
self
}
#[must_use]
pub fn log_with_source<E: std::fmt::Debug>(self, source_error: E) -> Self {
error!("{} | Source: {:?}", self, source_error);
self
}
}
pub type OptimizerResult<T> = Result<T, OptimizerError>;
#[derive(Debug, Clone)]
pub struct ConvergenceInfo {
pub final_gradient_norm: f64,
pub final_parameter_update_norm: f64,
pub cost_evaluations: usize,
pub jacobian_evaluations: usize,
}
impl Display for ConvergenceInfo {
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
write!(
f,
"Final gradient norm: {:.2e}, Final parameter update norm: {:.2e}, Cost evaluations: {}, Jacobian evaluations: {}",
self.final_gradient_norm,
self.final_parameter_update_norm,
self.cost_evaluations,
self.jacobian_evaluations
)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum OptimizationStatus {
Converged,
MaxIterationsReached,
CostToleranceReached,
ParameterToleranceReached,
GradientToleranceReached,
NumericalFailure,
UserTerminated,
Timeout,
TrustRegionRadiusTooSmall,
MinCostThresholdReached,
IllConditionedJacobian,
InvalidNumericalValues,
Failed(String),
}
impl Display for OptimizationStatus {
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
match self {
OptimizationStatus::Converged => write!(f, "Converged"),
OptimizationStatus::MaxIterationsReached => write!(f, "Maximum iterations reached"),
OptimizationStatus::CostToleranceReached => write!(f, "Cost tolerance reached"),
OptimizationStatus::ParameterToleranceReached => {
write!(f, "Parameter tolerance reached")
}
OptimizationStatus::GradientToleranceReached => write!(f, "Gradient tolerance reached"),
OptimizationStatus::NumericalFailure => write!(f, "Numerical failure"),
OptimizationStatus::UserTerminated => write!(f, "User terminated"),
OptimizationStatus::Timeout => write!(f, "Timeout"),
OptimizationStatus::TrustRegionRadiusTooSmall => {
write!(f, "Trust region radius too small")
}
OptimizationStatus::MinCostThresholdReached => {
write!(f, "Minimum cost threshold reached")
}
OptimizationStatus::IllConditionedJacobian => {
write!(f, "Ill-conditioned Jacobian matrix")
}
OptimizationStatus::InvalidNumericalValues => {
write!(f, "Invalid numerical values (NaN/Inf) detected")
}
OptimizationStatus::Failed(msg) => write!(f, "Failed: {msg}"),
}
}
}
#[derive(Clone)]
pub struct SolverResult<T> {
pub parameters: T,
pub status: OptimizationStatus,
pub initial_cost: f64,
pub final_cost: f64,
pub iterations: usize,
pub elapsed_time: time::Duration,
pub convergence_info: Option<ConvergenceInfo>,
pub covariances: Option<HashMap<String, Mat<f64>>>,
}
pub trait Solver {
type Config;
type Error;
fn new() -> Self;
fn optimize(
&mut self,
problem: &Problem,
initial_params: &HashMap<String, (ManifoldType, DVector<f64>)>,
) -> Result<SolverResult<HashMap<String, VariableEnum>>, Self::Error>;
}
pub fn apply_parameter_step(
variables: &mut HashMap<String, VariableEnum>,
step: MatRef<f64>,
variable_order: &[String],
) -> f64 {
let mut step_offset = 0;
for var_name in variable_order {
if let Some(var) = variables.get_mut(var_name) {
let var_size = var.get_size();
let var_step = step.subrows(step_offset, var_size);
var.apply_tangent_step(var_step);
step_offset += var_size;
}
}
step.norm_l2()
}
pub fn apply_negative_parameter_step(
variables: &mut HashMap<String, VariableEnum>,
step: MatRef<f64>,
variable_order: &[String],
) {
let mut negative_step = Mat::zeros(step.nrows(), 1);
for i in 0..step.nrows() {
negative_step[(i, 0)] = -step[(i, 0)];
}
apply_parameter_step(variables, negative_step.as_ref(), variable_order);
}
pub fn compute_cost(residual: &Mat<f64>) -> f64 {
let cost = residual.norm_l2();
0.5 * cost * cost
}
#[derive(Debug, Clone)]
pub struct IterationStats {
pub iteration: usize,
pub cost: f64,
pub cost_change: f64,
pub gradient_norm: f64,
pub step_norm: f64,
pub tr_ratio: f64,
pub tr_radius: f64,
pub ls_iter: usize,
pub iter_time_ms: f64,
pub total_time_ms: f64,
pub accepted: bool,
}
impl IterationStats {
pub fn print_header() {
debug!(
"{:>4} {:>13} {:>13} {:>13} {:>13} {:>11} {:>11} {:>7} {:>11} {:>13} {:>6}",
"iter",
"cost",
"cost_change",
"|gradient|",
"|step|",
"tr_ratio",
"tr_radius",
"ls_iter",
"iter_time",
"total_time",
"status"
);
}
pub fn print_line(&self) {
let status = if self.iteration == 0 {
"-"
} else if self.accepted {
"✓"
} else {
"✗"
};
debug!(
"{:>4} {:>13.6e} {:>13.2e} {:>13.2e} {:>13.2e} {:>11.2e} {:>11.2e} {:>7} {:>9.2}ms {:>11.2}ms {:>6}",
self.iteration,
self.cost,
self.cost_change,
self.gradient_norm,
self.step_norm,
self.tr_ratio,
self.tr_radius,
self.ls_iter,
self.iter_time_ms,
self.total_time_ms,
status
);
}
}
pub struct InitializedState {
pub variables: HashMap<String, VariableEnum>,
pub variable_index_map: HashMap<String, usize>,
pub sorted_vars: Vec<String>,
pub symbolic_structure: SymbolicStructure,
pub current_cost: f64,
pub initial_cost: f64,
}
pub fn compute_parameter_norm(variables: &HashMap<String, VariableEnum>) -> f64 {
variables
.values()
.map(|v| {
let vec = v.to_vector();
vec.norm_squared()
})
.sum::<f64>()
.sqrt()
}
pub fn create_jacobi_scaling(
jacobian: &SparseColMat<usize, f64>,
) -> Result<SparseColMat<usize, f64>, OptimizerError> {
let cols = jacobian.ncols();
let jacobi_scaling_vec: Vec<Triplet<usize, usize, f64>> = (0..cols)
.map(|c| {
let col_norm_squared: f64 = jacobian
.triplet_iter()
.filter(|t| t.col == c)
.map(|t| t.val * t.val)
.sum();
let col_norm = col_norm_squared.sqrt();
let scaling = 1.0 / (1.0 + col_norm);
Triplet::new(c, c, scaling)
})
.collect();
SparseColMat::try_new_from_triplets(cols, cols, &jacobi_scaling_vec)
.map_err(|e| OptimizerError::JacobiScalingCreation(e.to_string()).log_with_source(e))
}
pub fn process_jacobian(
jacobian: &SparseColMat<usize, f64>,
jacobi_scaling: &mut Option<SparseColMat<usize, f64>>,
iteration: usize,
) -> Result<SparseColMat<usize, f64>, OptimizerError> {
if iteration == 0 {
let scaling = create_jacobi_scaling(jacobian)?;
*jacobi_scaling = Some(scaling);
}
let scaling = jacobi_scaling
.as_ref()
.ok_or_else(|| OptimizerError::JacobiScalingNotInitialized.log())?;
Ok(jacobian * scaling)
}
pub fn initialize_optimization_state(
problem: &Problem,
initial_params: &HashMap<String, (ManifoldType, DVector<f64>)>,
) -> Result<InitializedState, error::ApexSolverError> {
let variables = problem.initialize_variables(initial_params);
let mut variable_index_map = HashMap::new();
let mut col_offset = 0;
let mut sorted_vars: Vec<String> = variables.keys().cloned().collect();
sorted_vars.sort();
for var_name in &sorted_vars {
variable_index_map.insert(var_name.clone(), col_offset);
col_offset += variables[var_name].get_size();
}
let symbolic_structure =
problem.build_symbolic_structure(&variables, &variable_index_map, col_offset)?;
let residual = problem.compute_residual_sparse(&variables)?;
let current_cost = compute_cost(&residual);
let initial_cost = current_cost;
Ok(InitializedState {
variables,
variable_index_map,
sorted_vars,
symbolic_structure,
current_cost,
initial_cost,
})
}
pub struct ConvergenceParams {
pub iteration: usize,
pub current_cost: f64,
pub new_cost: f64,
pub parameter_norm: f64,
pub parameter_update_norm: f64,
pub gradient_norm: f64,
pub elapsed: Duration,
pub step_accepted: bool,
pub max_iterations: usize,
pub gradient_tolerance: f64,
pub parameter_tolerance: f64,
pub cost_tolerance: f64,
pub min_cost_threshold: Option<f64>,
pub timeout: Option<Duration>,
pub trust_region_radius: Option<f64>,
pub min_trust_region_radius: Option<f64>,
}
pub fn check_convergence(params: &ConvergenceParams) -> Option<OptimizationStatus> {
if !params.new_cost.is_finite()
|| !params.parameter_update_norm.is_finite()
|| !params.gradient_norm.is_finite()
{
return Some(OptimizationStatus::InvalidNumericalValues);
}
if let Some(timeout) = params.timeout {
if params.elapsed >= timeout {
return Some(OptimizationStatus::Timeout);
}
}
if params.iteration >= params.max_iterations {
return Some(OptimizationStatus::MaxIterationsReached);
}
if !params.step_accepted {
return None;
}
if params.gradient_norm < params.gradient_tolerance {
return Some(OptimizationStatus::GradientToleranceReached);
}
if params.iteration > 0 {
let relative_step_tolerance =
params.parameter_tolerance * (params.parameter_norm + params.parameter_tolerance);
if params.parameter_update_norm <= relative_step_tolerance {
return Some(OptimizationStatus::ParameterToleranceReached);
}
let cost_change = (params.current_cost - params.new_cost).abs();
let relative_cost_change = cost_change / params.current_cost.max(1e-10);
if relative_cost_change < params.cost_tolerance {
return Some(OptimizationStatus::CostToleranceReached);
}
}
if let Some(min_cost) = params.min_cost_threshold {
if params.new_cost < min_cost {
return Some(OptimizationStatus::MinCostThresholdReached);
}
}
if let (Some(radius), Some(min_radius)) =
(params.trust_region_radius, params.min_trust_region_radius)
{
if radius < min_radius {
return Some(OptimizationStatus::TrustRegionRadiusTooSmall);
}
}
None
}
pub fn compute_step_quality(current_cost: f64, new_cost: f64, predicted_reduction: f64) -> f64 {
let actual_reduction = current_cost - new_cost;
if predicted_reduction.abs() < 1e-15 {
if actual_reduction > 0.0 { 1.0 } else { 0.0 }
} else {
actual_reduction / predicted_reduction
}
}
pub fn create_linear_solver(solver_type: &linalg::LinearSolverType) -> Box<dyn SparseLinearSolver> {
match solver_type {
linalg::LinearSolverType::SparseCholesky => Box::new(SparseCholeskySolver::new()),
linalg::LinearSolverType::SparseQR => Box::new(SparseQRSolver::new()),
linalg::LinearSolverType::SparseSchurComplement => {
Box::new(SparseCholeskySolver::new())
}
}
}
#[allow(clippy::too_many_arguments)]
pub fn notify_observers(
observers: &mut OptObserverVec,
variables: &HashMap<String, VariableEnum>,
iteration: usize,
cost: f64,
gradient_norm: f64,
damping: Option<f64>,
step_norm: f64,
step_quality: Option<f64>,
linear_solver: &dyn SparseLinearSolver,
) {
observers.set_iteration_metrics(cost, gradient_norm, damping, step_norm, step_quality);
if !observers.is_empty() {
if let (Some(hessian), Some(gradient)) =
(linear_solver.get_hessian(), linear_solver.get_gradient())
{
observers.set_matrix_data(Some(hessian.clone()), Some(gradient.clone()));
}
}
observers.notify(variables, iteration);
}
#[allow(clippy::too_many_arguments)]
pub fn build_solver_result(
status: OptimizationStatus,
iterations: usize,
state: InitializedState,
elapsed: Duration,
final_gradient_norm: f64,
final_parameter_update_norm: f64,
cost_evaluations: usize,
jacobian_evaluations: usize,
covariances: Option<HashMap<String, Mat<f64>>>,
) -> SolverResult<HashMap<String, VariableEnum>> {
SolverResult {
status,
iterations,
initial_cost: state.initial_cost,
final_cost: state.current_cost,
parameters: state.variables.into_iter().collect(),
elapsed_time: elapsed,
convergence_info: Some(ConvergenceInfo {
final_gradient_norm,
final_parameter_update_norm,
cost_evaluations,
jacobian_evaluations,
}),
covariances,
}
}
#[derive(Debug, Clone)]
pub struct OptimizerSummary {
pub optimizer_name: &'static str,
pub initial_cost: f64,
pub final_cost: f64,
pub iterations: usize,
pub successful_steps: Option<usize>,
pub unsuccessful_steps: Option<usize>,
pub average_cost_reduction: f64,
pub max_gradient_norm: f64,
pub final_gradient_norm: f64,
pub max_parameter_update_norm: f64,
pub final_parameter_update_norm: f64,
pub total_time: Duration,
pub average_time_per_iteration: Duration,
pub iteration_history: Vec<IterationStats>,
pub convergence_status: OptimizationStatus,
pub final_damping: Option<f64>,
pub final_trust_region_radius: Option<f64>,
pub rho: Option<f64>,
}
impl Display for OptimizerSummary {
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
let converged = matches!(
self.convergence_status,
OptimizationStatus::Converged
| OptimizationStatus::CostToleranceReached
| OptimizationStatus::GradientToleranceReached
| OptimizationStatus::ParameterToleranceReached
);
writeln!(f, "{} Final Result", self.optimizer_name)?;
if converged {
writeln!(f, "CONVERGED ({:?})", self.convergence_status)?;
} else {
writeln!(f, "DIVERGED ({:?})", self.convergence_status)?;
}
writeln!(f)?;
writeln!(f, "Cost:")?;
writeln!(f, " Initial: {:.6e}", self.initial_cost)?;
writeln!(f, " Final: {:.6e}", self.final_cost)?;
writeln!(
f,
" Reduction: {:.6e} ({:.2}%)",
self.initial_cost - self.final_cost,
100.0 * (self.initial_cost - self.final_cost) / self.initial_cost.max(1e-12)
)?;
writeln!(f)?;
writeln!(f, "Iterations:")?;
writeln!(f, " Total: {}", self.iterations)?;
if let (Some(successful), Some(unsuccessful)) =
(self.successful_steps, self.unsuccessful_steps)
{
writeln!(
f,
" Successful steps: {} ({:.1}%)",
successful,
100.0 * successful as f64 / self.iterations.max(1) as f64
)?;
writeln!(
f,
" Unsuccessful steps: {} ({:.1}%)",
unsuccessful,
100.0 * unsuccessful as f64 / self.iterations.max(1) as f64
)?;
}
if let Some(radius) = self.final_trust_region_radius {
writeln!(f)?;
writeln!(f, "Trust Region:")?;
writeln!(f, " Final radius: {:.6e}", radius)?;
}
writeln!(f)?;
writeln!(f, "Gradient:")?;
writeln!(f, " Max norm: {:.2e}", self.max_gradient_norm)?;
writeln!(f, " Final norm: {:.2e}", self.final_gradient_norm)?;
writeln!(f)?;
writeln!(f, "Parameter Update:")?;
writeln!(f, " Max norm: {:.2e}", self.max_parameter_update_norm)?;
writeln!(f, " Final norm: {:.2e}", self.final_parameter_update_norm)?;
writeln!(f)?;
writeln!(f, "Performance:")?;
writeln!(
f,
" Total time: {:.2}ms",
self.total_time.as_secs_f64() * 1000.0
)?;
writeln!(
f,
" Average per iteration: {:.2}ms",
self.average_time_per_iteration.as_secs_f64() * 1000.0
)?;
Ok(())
}
}
#[allow(clippy::too_many_arguments)]
pub fn create_optimizer_summary(
optimizer_name: &'static str,
initial_cost: f64,
final_cost: f64,
iterations: usize,
successful_steps: Option<usize>,
unsuccessful_steps: Option<usize>,
max_gradient_norm: f64,
final_gradient_norm: f64,
max_parameter_update_norm: f64,
final_parameter_update_norm: f64,
total_cost_reduction: f64,
total_time: Duration,
iteration_history: Vec<IterationStats>,
convergence_status: OptimizationStatus,
final_damping: Option<f64>,
final_trust_region_radius: Option<f64>,
rho: Option<f64>,
) -> OptimizerSummary {
OptimizerSummary {
optimizer_name,
initial_cost,
final_cost,
iterations,
successful_steps,
unsuccessful_steps,
average_cost_reduction: if iterations > 0 {
total_cost_reduction / iterations as f64
} else {
0.0
},
max_gradient_norm,
final_gradient_norm,
max_parameter_update_norm,
final_parameter_update_norm,
total_time,
average_time_per_iteration: if iterations > 0 {
total_time / iterations as u32
} else {
Duration::from_secs(0)
},
iteration_history,
convergence_status,
final_damping,
final_trust_region_radius,
rho,
}
}