1use crate::core::problem::{Problem, SymbolicStructure, VariableEnum};
10use crate::error;
11use crate::linalg::{self, SparseCholeskySolver, SparseLinearSolver, SparseQRSolver};
12use apex_manifolds::ManifoldType;
13use faer::sparse::{SparseColMat, Triplet};
14use faer::{Mat, MatRef};
15use nalgebra::DVector;
16use std::collections::HashMap;
17use std::time::{self, Duration};
18use std::{
19 fmt,
20 fmt::{Display, Formatter},
21};
22use thiserror::Error;
23use tracing::{debug, error};
24
25pub mod dog_leg;
26pub mod gauss_newton;
27pub mod levenberg_marquardt;
28
29pub use dog_leg::DogLeg;
30pub use gauss_newton::GaussNewton;
31pub use levenberg_marquardt::LevenbergMarquardt;
32
33pub use crate::observers::{OptObserver, OptObserverVec};
35
36#[derive(Default, Clone, Copy, PartialEq, Eq)]
38pub enum OptimizerType {
39 #[default]
41 LevenbergMarquardt,
42 GaussNewton,
44 DogLeg,
46}
47
48impl Display for OptimizerType {
49 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
50 match self {
51 OptimizerType::LevenbergMarquardt => write!(f, "Levenberg-Marquardt"),
52 OptimizerType::GaussNewton => write!(f, "Gauss-Newton"),
53 OptimizerType::DogLeg => write!(f, "Dog Leg"),
54 }
55 }
56}
57
58#[derive(Debug, Clone, Error)]
60pub enum OptimizerError {
61 #[error("Linear system solve failed: {0}")]
63 LinearSolveFailed(String),
64
65 #[error("Maximum iterations ({max_iters}) reached without convergence")]
67 MaxIterationsReached { max_iters: usize },
68
69 #[error("Trust region radius became too small: {radius:.6e} < {min_radius:.6e}")]
71 TrustRegionFailure { radius: f64, min_radius: f64 },
72
73 #[error("Damping parameter became too large: {damping:.6e} > {max_damping:.6e}")]
75 DampingFailure { damping: f64, max_damping: f64 },
76
77 #[error("Cost increased unexpectedly: {old_cost:.6e} -> {new_cost:.6e}")]
79 CostIncrease { old_cost: f64, new_cost: f64 },
80
81 #[error("Jacobian computation failed: {0}")]
83 JacobianFailed(String),
84
85 #[error("Invalid optimization parameters: {0}")]
87 InvalidParameters(String),
88
89 #[error("Numerical instability detected: {0}")]
91 NumericalInstability(String),
92
93 #[error("Linear algebra error: {0}")]
95 LinAlg(#[from] linalg::LinAlgError),
96
97 #[error("Problem has no variables to optimize")]
99 EmptyProblem,
100
101 #[error("Problem has no residual blocks")]
103 NoResidualBlocks,
104
105 #[error("Failed to create Jacobi scaling matrix: {0}")]
107 JacobiScalingCreation(String),
108
109 #[error("Jacobi scaling not initialized")]
111 JacobiScalingNotInitialized,
112
113 #[error("Unknown linear solver type: {0}")]
115 UnknownLinearSolver(String),
116}
117
118impl OptimizerError {
119 #[must_use]
135 pub fn log(self) -> Self {
136 error!("{}", self);
137 self
138 }
139
140 #[must_use]
163 pub fn log_with_source<E: std::fmt::Debug>(self, source_error: E) -> Self {
164 error!("{} | Source: {:?}", self, source_error);
165 self
166 }
167}
168
169pub type OptimizerResult<T> = Result<T, OptimizerError>;
171
172#[derive(Debug, Clone)]
189pub struct ConvergenceInfo {
190 pub final_gradient_norm: f64,
192 pub final_parameter_update_norm: f64,
194 pub cost_evaluations: usize,
196 pub jacobian_evaluations: usize,
198}
199
200impl Display for ConvergenceInfo {
201 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
202 write!(
203 f,
204 "Final gradient norm: {:.2e}, Final parameter update norm: {:.2e}, Cost evaluations: {}, Jacobian evaluations: {}",
205 self.final_gradient_norm,
206 self.final_parameter_update_norm,
207 self.cost_evaluations,
208 self.jacobian_evaluations
209 )
210 }
211}
212
213#[derive(Debug, Clone, PartialEq, Eq)]
215pub enum OptimizationStatus {
216 Converged,
218 MaxIterationsReached,
220 CostToleranceReached,
222 ParameterToleranceReached,
224 GradientToleranceReached,
226 NumericalFailure,
228 UserTerminated,
230 Timeout,
232 TrustRegionRadiusTooSmall,
234 MinCostThresholdReached,
236 IllConditionedJacobian,
238 InvalidNumericalValues,
240 Failed(String),
242}
243
244impl Display for OptimizationStatus {
245 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
246 match self {
247 OptimizationStatus::Converged => write!(f, "Converged"),
248 OptimizationStatus::MaxIterationsReached => write!(f, "Maximum iterations reached"),
249 OptimizationStatus::CostToleranceReached => write!(f, "Cost tolerance reached"),
250 OptimizationStatus::ParameterToleranceReached => {
251 write!(f, "Parameter tolerance reached")
252 }
253 OptimizationStatus::GradientToleranceReached => write!(f, "Gradient tolerance reached"),
254 OptimizationStatus::NumericalFailure => write!(f, "Numerical failure"),
255 OptimizationStatus::UserTerminated => write!(f, "User terminated"),
256 OptimizationStatus::Timeout => write!(f, "Timeout"),
257 OptimizationStatus::TrustRegionRadiusTooSmall => {
258 write!(f, "Trust region radius too small")
259 }
260 OptimizationStatus::MinCostThresholdReached => {
261 write!(f, "Minimum cost threshold reached")
262 }
263 OptimizationStatus::IllConditionedJacobian => {
264 write!(f, "Ill-conditioned Jacobian matrix")
265 }
266 OptimizationStatus::InvalidNumericalValues => {
267 write!(f, "Invalid numerical values (NaN/Inf) detected")
268 }
269 OptimizationStatus::Failed(msg) => write!(f, "Failed: {msg}"),
270 }
271 }
272}
273
274#[derive(Clone)]
276pub struct SolverResult<T> {
277 pub parameters: T,
279 pub status: OptimizationStatus,
281 pub initial_cost: f64,
283 pub final_cost: f64,
285 pub iterations: usize,
287 pub elapsed_time: time::Duration,
289 pub convergence_info: Option<ConvergenceInfo>,
291 pub covariances: Option<HashMap<String, Mat<f64>>>,
299}
300
301pub trait Solver {
303 type Config;
305 type Error;
307
308 fn new() -> Self;
310
311 fn optimize(
313 &mut self,
314 problem: &Problem,
315 initial_params: &HashMap<String, (ManifoldType, DVector<f64>)>,
316 ) -> Result<SolverResult<HashMap<String, VariableEnum>>, Self::Error>;
317}
318
319pub fn apply_parameter_step(
338 variables: &mut HashMap<String, VariableEnum>,
339 step: MatRef<f64>,
340 variable_order: &[String],
341) -> f64 {
342 let mut step_offset = 0;
343
344 for var_name in variable_order {
345 if let Some(var) = variables.get_mut(var_name) {
346 let var_size = var.get_size();
347 let var_step = step.subrows(step_offset, var_size);
348
349 var.apply_tangent_step(var_step);
352
353 step_offset += var_size;
354 }
355 }
356
357 step.norm_l2()
359}
360
361pub fn apply_negative_parameter_step(
372 variables: &mut HashMap<String, VariableEnum>,
373 step: MatRef<f64>,
374 variable_order: &[String],
375) {
376 let mut negative_step = Mat::zeros(step.nrows(), 1);
378 for i in 0..step.nrows() {
379 negative_step[(i, 0)] = -step[(i, 0)];
380 }
381
382 apply_parameter_step(variables, negative_step.as_ref(), variable_order);
384}
385
386pub fn compute_cost(residual: &Mat<f64>) -> f64 {
387 let cost = residual.norm_l2();
388 0.5 * cost * cost
389}
390
391#[derive(Debug, Clone)]
403pub struct IterationStats {
404 pub iteration: usize,
406 pub cost: f64,
408 pub cost_change: f64,
410 pub gradient_norm: f64,
412 pub step_norm: f64,
414 pub tr_ratio: f64,
416 pub tr_radius: f64,
418 pub ls_iter: usize,
420 pub iter_time_ms: f64,
422 pub total_time_ms: f64,
424 pub accepted: bool,
426}
427
428impl IterationStats {
429 pub fn print_header() {
431 debug!(
432 "{:>4} {:>13} {:>13} {:>13} {:>13} {:>11} {:>11} {:>7} {:>11} {:>13} {:>6}",
433 "iter",
434 "cost",
435 "cost_change",
436 "|gradient|",
437 "|step|",
438 "tr_ratio",
439 "tr_radius",
440 "ls_iter",
441 "iter_time",
442 "total_time",
443 "status"
444 );
445 }
446
447 pub fn print_line(&self) {
449 let status = if self.iteration == 0 {
450 "-"
451 } else if self.accepted {
452 "✓"
453 } else {
454 "✗"
455 };
456
457 debug!(
458 "{:>4} {:>13.6e} {:>13.2e} {:>13.2e} {:>13.2e} {:>11.2e} {:>11.2e} {:>7} {:>9.2}ms {:>11.2}ms {:>6}",
459 self.iteration,
460 self.cost,
461 self.cost_change,
462 self.gradient_norm,
463 self.step_norm,
464 self.tr_ratio,
465 self.tr_radius,
466 self.ls_iter,
467 self.iter_time_ms,
468 self.total_time_ms,
469 status
470 );
471 }
472}
473
474pub struct InitializedState {
476 pub variables: HashMap<String, VariableEnum>,
477 pub variable_index_map: HashMap<String, usize>,
478 pub sorted_vars: Vec<String>,
479 pub symbolic_structure: SymbolicStructure,
480 pub current_cost: f64,
481 pub initial_cost: f64,
482}
483
484pub fn compute_parameter_norm(variables: &HashMap<String, VariableEnum>) -> f64 {
486 variables
487 .values()
488 .map(|v| {
489 let vec = v.to_vector();
490 vec.norm_squared()
491 })
492 .sum::<f64>()
493 .sqrt()
494}
495
496pub fn create_jacobi_scaling(
501 jacobian: &SparseColMat<usize, f64>,
502) -> Result<SparseColMat<usize, f64>, OptimizerError> {
503 let cols = jacobian.ncols();
504 let jacobi_scaling_vec: Vec<Triplet<usize, usize, f64>> = (0..cols)
505 .map(|c| {
506 let col_norm_squared: f64 = jacobian
507 .triplet_iter()
508 .filter(|t| t.col == c)
509 .map(|t| t.val * t.val)
510 .sum();
511 let col_norm = col_norm_squared.sqrt();
512 let scaling = 1.0 / (1.0 + col_norm);
513 Triplet::new(c, c, scaling)
514 })
515 .collect();
516
517 SparseColMat::try_new_from_triplets(cols, cols, &jacobi_scaling_vec)
518 .map_err(|e| OptimizerError::JacobiScalingCreation(e.to_string()).log_with_source(e))
519}
520
521pub fn process_jacobian(
526 jacobian: &SparseColMat<usize, f64>,
527 jacobi_scaling: &mut Option<SparseColMat<usize, f64>>,
528 iteration: usize,
529) -> Result<SparseColMat<usize, f64>, OptimizerError> {
530 if iteration == 0 {
531 let scaling = create_jacobi_scaling(jacobian)?;
532 *jacobi_scaling = Some(scaling);
533 }
534 let scaling = jacobi_scaling
535 .as_ref()
536 .ok_or_else(|| OptimizerError::JacobiScalingNotInitialized.log())?;
537 Ok(jacobian * scaling)
538}
539
540pub fn initialize_optimization_state(
548 problem: &Problem,
549 initial_params: &HashMap<String, (ManifoldType, DVector<f64>)>,
550) -> Result<InitializedState, error::ApexSolverError> {
551 let variables = problem.initialize_variables(initial_params);
552
553 let mut variable_index_map = HashMap::new();
554 let mut col_offset = 0;
555 let mut sorted_vars: Vec<String> = variables.keys().cloned().collect();
556 sorted_vars.sort();
557
558 for var_name in &sorted_vars {
559 variable_index_map.insert(var_name.clone(), col_offset);
560 col_offset += variables[var_name].get_size();
561 }
562
563 let symbolic_structure =
564 problem.build_symbolic_structure(&variables, &variable_index_map, col_offset)?;
565
566 let residual = problem.compute_residual_sparse(&variables)?;
567 let current_cost = compute_cost(&residual);
568 let initial_cost = current_cost;
569
570 Ok(InitializedState {
571 variables,
572 variable_index_map,
573 sorted_vars,
574 symbolic_structure,
575 current_cost,
576 initial_cost,
577 })
578}
579
580pub struct ConvergenceParams {
582 pub iteration: usize,
583 pub current_cost: f64,
584 pub new_cost: f64,
585 pub parameter_norm: f64,
586 pub parameter_update_norm: f64,
587 pub gradient_norm: f64,
588 pub elapsed: Duration,
589 pub step_accepted: bool,
590 pub max_iterations: usize,
592 pub gradient_tolerance: f64,
593 pub parameter_tolerance: f64,
594 pub cost_tolerance: f64,
595 pub min_cost_threshold: Option<f64>,
596 pub timeout: Option<Duration>,
597 pub trust_region_radius: Option<f64>,
599 pub min_trust_region_radius: Option<f64>,
601}
602
603pub fn check_convergence(params: &ConvergenceParams) -> Option<OptimizationStatus> {
607 if !params.new_cost.is_finite()
611 || !params.parameter_update_norm.is_finite()
612 || !params.gradient_norm.is_finite()
613 {
614 return Some(OptimizationStatus::InvalidNumericalValues);
615 }
616
617 if let Some(timeout) = params.timeout {
619 if params.elapsed >= timeout {
620 return Some(OptimizationStatus::Timeout);
621 }
622 }
623
624 if params.iteration >= params.max_iterations {
626 return Some(OptimizationStatus::MaxIterationsReached);
627 }
628
629 if !params.step_accepted {
631 return None;
632 }
633
634 if params.gradient_norm < params.gradient_tolerance {
636 return Some(OptimizationStatus::GradientToleranceReached);
637 }
638
639 if params.iteration > 0 {
641 let relative_step_tolerance =
643 params.parameter_tolerance * (params.parameter_norm + params.parameter_tolerance);
644 if params.parameter_update_norm <= relative_step_tolerance {
645 return Some(OptimizationStatus::ParameterToleranceReached);
646 }
647
648 let cost_change = (params.current_cost - params.new_cost).abs();
650 let relative_cost_change = cost_change / params.current_cost.max(1e-10);
651 if relative_cost_change < params.cost_tolerance {
652 return Some(OptimizationStatus::CostToleranceReached);
653 }
654 }
655
656 if let Some(min_cost) = params.min_cost_threshold {
658 if params.new_cost < min_cost {
659 return Some(OptimizationStatus::MinCostThresholdReached);
660 }
661 }
662
663 if let (Some(radius), Some(min_radius)) =
665 (params.trust_region_radius, params.min_trust_region_radius)
666 {
667 if radius < min_radius {
668 return Some(OptimizationStatus::TrustRegionRadiusTooSmall);
669 }
670 }
671
672 None
673}
674
675pub fn compute_step_quality(current_cost: f64, new_cost: f64, predicted_reduction: f64) -> f64 {
684 let actual_reduction = current_cost - new_cost;
685 if predicted_reduction.abs() < 1e-15 {
686 if actual_reduction > 0.0 { 1.0 } else { 0.0 }
687 } else {
688 actual_reduction / predicted_reduction
689 }
690}
691
692pub fn create_linear_solver(solver_type: &linalg::LinearSolverType) -> Box<dyn SparseLinearSolver> {
697 match solver_type {
698 linalg::LinearSolverType::SparseCholesky => Box::new(SparseCholeskySolver::new()),
699 linalg::LinearSolverType::SparseQR => Box::new(SparseQRSolver::new()),
700 linalg::LinearSolverType::SparseSchurComplement => {
701 Box::new(SparseCholeskySolver::new())
703 }
704 }
705}
706
707#[allow(clippy::too_many_arguments)]
711pub fn notify_observers(
712 observers: &mut OptObserverVec,
713 variables: &HashMap<String, VariableEnum>,
714 iteration: usize,
715 cost: f64,
716 gradient_norm: f64,
717 damping: Option<f64>,
718 step_norm: f64,
719 step_quality: Option<f64>,
720 linear_solver: &dyn SparseLinearSolver,
721) {
722 observers.set_iteration_metrics(cost, gradient_norm, damping, step_norm, step_quality);
723
724 if !observers.is_empty() {
725 if let (Some(hessian), Some(gradient)) =
726 (linear_solver.get_hessian(), linear_solver.get_gradient())
727 {
728 observers.set_matrix_data(Some(hessian.clone()), Some(gradient.clone()));
729 }
730 }
731
732 observers.notify(variables, iteration);
733}
734
735#[allow(clippy::too_many_arguments)]
739pub fn build_solver_result(
740 status: OptimizationStatus,
741 iterations: usize,
742 state: InitializedState,
743 elapsed: Duration,
744 final_gradient_norm: f64,
745 final_parameter_update_norm: f64,
746 cost_evaluations: usize,
747 jacobian_evaluations: usize,
748 covariances: Option<HashMap<String, Mat<f64>>>,
749) -> SolverResult<HashMap<String, VariableEnum>> {
750 SolverResult {
751 status,
752 iterations,
753 initial_cost: state.initial_cost,
754 final_cost: state.current_cost,
755 parameters: state.variables.into_iter().collect(),
756 elapsed_time: elapsed,
757 convergence_info: Some(ConvergenceInfo {
758 final_gradient_norm,
759 final_parameter_update_norm,
760 cost_evaluations,
761 jacobian_evaluations,
762 }),
763 covariances,
764 }
765}
766
767#[derive(Debug, Clone)]
773pub struct OptimizerSummary {
774 pub optimizer_name: &'static str,
776 pub initial_cost: f64,
778 pub final_cost: f64,
780 pub iterations: usize,
782 pub successful_steps: Option<usize>,
784 pub unsuccessful_steps: Option<usize>,
786 pub average_cost_reduction: f64,
788 pub max_gradient_norm: f64,
790 pub final_gradient_norm: f64,
792 pub max_parameter_update_norm: f64,
794 pub final_parameter_update_norm: f64,
796 pub total_time: Duration,
798 pub average_time_per_iteration: Duration,
800 pub iteration_history: Vec<IterationStats>,
802 pub convergence_status: OptimizationStatus,
804 pub final_damping: Option<f64>,
806 pub final_trust_region_radius: Option<f64>,
808 pub rho: Option<f64>,
810}
811
812impl Display for OptimizerSummary {
813 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
814 let converged = matches!(
815 self.convergence_status,
816 OptimizationStatus::Converged
817 | OptimizationStatus::CostToleranceReached
818 | OptimizationStatus::GradientToleranceReached
819 | OptimizationStatus::ParameterToleranceReached
820 );
821
822 writeln!(f, "{} Final Result", self.optimizer_name)?;
823
824 if converged {
825 writeln!(f, "CONVERGED ({:?})", self.convergence_status)?;
826 } else {
827 writeln!(f, "DIVERGED ({:?})", self.convergence_status)?;
828 }
829
830 writeln!(f)?;
831 writeln!(f, "Cost:")?;
832 writeln!(f, " Initial: {:.6e}", self.initial_cost)?;
833 writeln!(f, " Final: {:.6e}", self.final_cost)?;
834 writeln!(
835 f,
836 " Reduction: {:.6e} ({:.2}%)",
837 self.initial_cost - self.final_cost,
838 100.0 * (self.initial_cost - self.final_cost) / self.initial_cost.max(1e-12)
839 )?;
840 writeln!(f)?;
841 writeln!(f, "Iterations:")?;
842 writeln!(f, " Total: {}", self.iterations)?;
843 if let (Some(successful), Some(unsuccessful)) =
844 (self.successful_steps, self.unsuccessful_steps)
845 {
846 writeln!(
847 f,
848 " Successful steps: {} ({:.1}%)",
849 successful,
850 100.0 * successful as f64 / self.iterations.max(1) as f64
851 )?;
852 writeln!(
853 f,
854 " Unsuccessful steps: {} ({:.1}%)",
855 unsuccessful,
856 100.0 * unsuccessful as f64 / self.iterations.max(1) as f64
857 )?;
858 }
859 if let Some(radius) = self.final_trust_region_radius {
860 writeln!(f)?;
861 writeln!(f, "Trust Region:")?;
862 writeln!(f, " Final radius: {:.6e}", radius)?;
863 }
864 writeln!(f)?;
865 writeln!(f, "Gradient:")?;
866 writeln!(f, " Max norm: {:.2e}", self.max_gradient_norm)?;
867 writeln!(f, " Final norm: {:.2e}", self.final_gradient_norm)?;
868 writeln!(f)?;
869 writeln!(f, "Parameter Update:")?;
870 writeln!(f, " Max norm: {:.2e}", self.max_parameter_update_norm)?;
871 writeln!(f, " Final norm: {:.2e}", self.final_parameter_update_norm)?;
872 writeln!(f)?;
873 writeln!(f, "Performance:")?;
874 writeln!(
875 f,
876 " Total time: {:.2}ms",
877 self.total_time.as_secs_f64() * 1000.0
878 )?;
879 writeln!(
880 f,
881 " Average per iteration: {:.2}ms",
882 self.average_time_per_iteration.as_secs_f64() * 1000.0
883 )?;
884
885 Ok(())
886 }
887}
888
889#[allow(clippy::too_many_arguments)]
891pub fn create_optimizer_summary(
892 optimizer_name: &'static str,
893 initial_cost: f64,
894 final_cost: f64,
895 iterations: usize,
896 successful_steps: Option<usize>,
897 unsuccessful_steps: Option<usize>,
898 max_gradient_norm: f64,
899 final_gradient_norm: f64,
900 max_parameter_update_norm: f64,
901 final_parameter_update_norm: f64,
902 total_cost_reduction: f64,
903 total_time: Duration,
904 iteration_history: Vec<IterationStats>,
905 convergence_status: OptimizationStatus,
906 final_damping: Option<f64>,
907 final_trust_region_radius: Option<f64>,
908 rho: Option<f64>,
909) -> OptimizerSummary {
910 OptimizerSummary {
911 optimizer_name,
912 initial_cost,
913 final_cost,
914 iterations,
915 successful_steps,
916 unsuccessful_steps,
917 average_cost_reduction: if iterations > 0 {
918 total_cost_reduction / iterations as f64
919 } else {
920 0.0
921 },
922 max_gradient_norm,
923 final_gradient_norm,
924 max_parameter_update_norm,
925 final_parameter_update_norm,
926 total_time,
927 average_time_per_iteration: if iterations > 0 {
928 total_time / iterations as u32
929 } else {
930 Duration::from_secs(0)
931 },
932 iteration_history,
933 convergence_status,
934 final_damping,
935 final_trust_region_radius,
936 rho,
937 }
938}