From 31e0b93801d217cebcf6d1d01a213fc45b5c0bb0 Mon Sep 17 00:00:00 2001 From: Hauke Strasdat Date: Mon, 6 Jan 2025 23:20:34 -0800 Subject: [PATCH] linear solvers --- crates/sophus_opt/Cargo.toml | 1 + crates/sophus_opt/src/block.rs | 210 +-------------- crates/sophus_opt/src/block/block_gradient.rs | 92 +++++++ crates/sophus_opt/src/block/block_hessian.rs | 117 +++++++++ crates/sophus_opt/src/cost.rs | 186 ------------- .../src/example_problems/cam_calib.rs | 34 +-- .../cost_fn/isometry2_prior.rs | 18 +- .../cost_fn/isometry3_prior.rs | 18 +- .../example_problems/cost_fn/pose_graph.rs | 18 +- .../example_problems/cost_fn/reprojection.rs | 18 +- .../src/example_problems/pose_circle.rs | 15 +- .../src/example_problems/simple_prior.rs | 30 ++- crates/sophus_opt/src/lib.rs | 17 +- crates/sophus_opt/src/nlls.rs | 58 ++++- crates/sophus_opt/src/quadratic_cost.rs | 12 + .../compare_idx.rs} | 5 +- .../src/{ => quadratic_cost}/cost_fn.rs | 202 +++++--------- .../evaluated_term.rs} | 62 ++--- .../src/quadratic_cost/residual_fn.rs | 28 ++ crates/sophus_opt/src/solvers.rs | 72 +---- .../sophus_opt/src/solvers/normal_equation.rs | 246 ++++++++++++++++++ .../src/{ldlt.rs => solvers/sparse_ldlt.rs} | 153 +++++------ 22 files changed, 799 insertions(+), 813 deletions(-) create mode 100644 crates/sophus_opt/src/block/block_gradient.rs create mode 100644 crates/sophus_opt/src/block/block_hessian.rs delete mode 100644 crates/sophus_opt/src/cost.rs create mode 100644 crates/sophus_opt/src/quadratic_cost.rs rename crates/sophus_opt/src/{cost_args.rs => quadratic_cost/compare_idx.rs} (98%) rename crates/sophus_opt/src/{ => quadratic_cost}/cost_fn.rs (62%) rename crates/sophus_opt/src/{term.rs => quadratic_cost/evaluated_term.rs} (88%) create mode 100644 crates/sophus_opt/src/quadratic_cost/residual_fn.rs create mode 100644 crates/sophus_opt/src/solvers/normal_equation.rs rename crates/sophus_opt/src/{ldlt.rs => solvers/sparse_ldlt.rs} (69%) diff --git a/crates/sophus_opt/Cargo.toml b/crates/sophus_opt/Cargo.toml index 17b53a8..d4f5a1e 100644 --- a/crates/sophus_opt/Cargo.toml +++ b/crates/sophus_opt/Cargo.toml @@ -25,6 +25,7 @@ nalgebra.workspace = true ndarray.workspace = true rand.workspace = true rayon = "1.10" +snafu.workspace = true [features] simd = [ diff --git a/crates/sophus_opt/src/block.rs b/crates/sophus_opt/src/block.rs index f8c70ca..f425131 100644 --- a/crates/sophus_opt/src/block.rs +++ b/crates/sophus_opt/src/block.rs @@ -1,6 +1,7 @@ -use nalgebra::Const; -use sophus_autodiff::linalg::MatF64; -use sophus_autodiff::linalg::VecF64; +/// Block Hessian matrix +pub mod block_hessian; +/// Block gradient vector +pub mod block_gradient; /// Range of a block #[derive(Clone, Debug, Copy, Default)] @@ -10,206 +11,3 @@ pub struct BlockRange { /// Dimension of the block pub dim: usize, } - -/// Block vector -#[derive(Debug, Clone)] -pub struct BlockVector { - /// vector storage - pub vec: nalgebra::SVector, - /// ranges, one for each block - pub ranges: [BlockRange; NUM_ARGS], -} - -impl BlockVector { - /// create a new block vector - pub fn new(dims: &[usize; NUM_ARGS]) -> Self { - debug_assert!(!dims.is_empty()); - - let num_blocks = dims.len(); - - let mut ranges = [BlockRange::default(); NUM_ARGS]; - - let mut num_rows: usize = 0; - - for i in 0..num_blocks { - let dim = dims[i]; - ranges[i] = BlockRange { - index: num_rows as i64, - dim, - }; - num_rows += dim; - } - Self { - vec: nalgebra::SVector::zeros(), - ranges, - } - } - - /// Number of blocks - pub fn num_blocks(&self) -> usize { - self.ranges.len() - } - - /// set the ith block - pub fn set_block(&mut self, ith: usize, v: VecF64) { - debug_assert!(ith < self.num_blocks()); - debug_assert_eq!(R, self.ranges[ith].dim); - self.mut_block::(ith).copy_from(&v); - } - - /// get the ith block - pub fn block( - &self, - ith: usize, - ) -> nalgebra::Matrix< - f64, - nalgebra::Dyn, - nalgebra::Const<1>, - nalgebra::ViewStorage< - '_, - f64, - nalgebra::Dyn, - nalgebra::Const<1>, - nalgebra::Const<1>, - Const, - >, - > { - let idx = self.ranges[ith].index as usize; - self.vec.rows(idx, self.ranges[ith].dim) - } - - /// mutable reference to the ith block - pub fn mut_block( - &mut self, - ith: usize, - ) -> nalgebra::Matrix< - f64, - nalgebra::Const, - nalgebra::Const<1>, - nalgebra::ViewStorageMut< - '_, - f64, - nalgebra::Const, - nalgebra::Const<1>, - nalgebra::Const<1>, - Const, - >, - > { - let idx = self.ranges[ith].index as usize; - self.vec.fixed_rows_mut::(idx) - } -} - -/// Block matrix -#[derive(Clone, Debug)] -pub struct BlockMatrix { - /// matrix storage - pub mat: nalgebra::SMatrix, - /// ranges, one for each block - pub ranges: [BlockRange; NUM_ARGS], -} - -impl BlockMatrix { - /// create a new block matrix - pub fn new(dims: &[usize]) -> Self { - debug_assert!(!dims.is_empty()); - - let num_blocks = dims.len(); - - let mut ranges = [BlockRange::default(); NUM_ARGS]; - let mut num_rows: usize = 0; - - for i in 0..num_blocks { - let dim = dims[i]; - ranges[i] = BlockRange { - index: num_rows as i64, - dim, - }; - num_rows += dim; - } - Self { - mat: nalgebra::SMatrix::zeros(), - ranges, - } - } - - /// Number of blocks - pub fn num_blocks(&self) -> usize { - self.ranges.len() - } - - /// set block (i, j) - pub fn set_block( - &mut self, - ith: usize, - jth: usize, - m: MatF64, - ) { - debug_assert!(ith < self.num_blocks()); - debug_assert!(jth < self.num_blocks()); - debug_assert_eq!(R, self.ranges[ith].dim); - debug_assert_eq!(C, self.ranges[jth].dim); - - if ith == jth { - debug_assert_eq!(R, C); - - let mut mut_block = self.mut_block::(ith, jth); - mut_block.copy_from(&m); - } else { - debug_assert!(ith < jth); - { - self.mut_block::(ith, jth).copy_from(&m); - } - { - self.mut_block::(jth, ith).copy_from(&m.transpose()); - } - } - } - - /// get block (i, j) - pub fn block( - &self, - ith: usize, - jth: usize, - ) -> nalgebra::Matrix< - f64, - nalgebra::Dyn, - nalgebra::Dyn, - nalgebra::ViewStorage< - '_, - f64, - nalgebra::Dyn, - nalgebra::Dyn, - nalgebra::Const<1>, - nalgebra::Const, - >, - > { - let idx_i = self.ranges[ith].index as usize; - let idx_j = self.ranges[jth].index as usize; - self.mat - .view((idx_i, idx_j), (self.ranges[ith].dim, self.ranges[jth].dim)) - } - - /// mutable reference to block (i, j) - pub fn mut_block( - &mut self, - ith: usize, - jth: usize, - ) -> nalgebra::Matrix< - f64, - nalgebra::Const, - nalgebra::Const, - nalgebra::ViewStorageMut< - '_, - f64, - nalgebra::Const, - nalgebra::Const, - nalgebra::Const<1>, - nalgebra::Const, - >, - > { - let idx_i = self.ranges[ith].index as usize; - let idx_j = self.ranges[jth].index as usize; - self.mat.fixed_view_mut::(idx_i, idx_j) - } -} diff --git a/crates/sophus_opt/src/block/block_gradient.rs b/crates/sophus_opt/src/block/block_gradient.rs new file mode 100644 index 0000000..63688e3 --- /dev/null +++ b/crates/sophus_opt/src/block/block_gradient.rs @@ -0,0 +1,92 @@ +use super::BlockRange; +use nalgebra::Const; +use sophus_autodiff::linalg::VecF64; + +/// Block gradient vector +#[derive(Debug, Clone)] +pub struct BlockGradient { + /// vector storage + pub vec: nalgebra::SVector, + /// ranges, one for each block + pub ranges: [BlockRange; NUM_ARGS], +} + +impl BlockGradient { + /// create a new block vector + pub fn new(dims: &[usize; NUM_ARGS]) -> Self { + debug_assert!(!dims.is_empty()); + + let num_blocks = dims.len(); + + let mut ranges = [BlockRange::default(); NUM_ARGS]; + + let mut num_rows: usize = 0; + + for i in 0..num_blocks { + let dim = dims[i]; + ranges[i] = BlockRange { + index: num_rows as i64, + dim, + }; + num_rows += dim; + } + Self { + vec: nalgebra::SVector::zeros(), + ranges, + } + } + + /// Number of blocks + pub fn num_blocks(&self) -> usize { + self.ranges.len() + } + + /// set the ith block + pub fn set_block(&mut self, ith: usize, v: VecF64) { + debug_assert!(ith < self.num_blocks()); + debug_assert_eq!(R, self.ranges[ith].dim); + self.mut_block::(ith).copy_from(&v); + } + + /// get the ith block + pub fn block( + &self, + ith: usize, + ) -> nalgebra::Matrix< + f64, + nalgebra::Dyn, + nalgebra::Const<1>, + nalgebra::ViewStorage< + '_, + f64, + nalgebra::Dyn, + nalgebra::Const<1>, + nalgebra::Const<1>, + Const, + >, + > { + let idx = self.ranges[ith].index as usize; + self.vec.rows(idx, self.ranges[ith].dim) + } + + /// mutable reference to the ith block + pub fn mut_block( + &mut self, + ith: usize, + ) -> nalgebra::Matrix< + f64, + nalgebra::Const, + nalgebra::Const<1>, + nalgebra::ViewStorageMut< + '_, + f64, + nalgebra::Const, + nalgebra::Const<1>, + nalgebra::Const<1>, + Const, + >, + > { + let idx = self.ranges[ith].index as usize; + self.vec.fixed_rows_mut::(idx) + } +} diff --git a/crates/sophus_opt/src/block/block_hessian.rs b/crates/sophus_opt/src/block/block_hessian.rs new file mode 100644 index 0000000..d183746 --- /dev/null +++ b/crates/sophus_opt/src/block/block_hessian.rs @@ -0,0 +1,117 @@ +use sophus_autodiff::linalg::MatF64; + +use super::BlockRange; + +/// Block Hessian matrix +#[derive(Clone, Debug)] +pub struct BlockHessian { + /// matrix storage + pub mat: nalgebra::SMatrix, + /// ranges, one for each block + pub ranges: [BlockRange; NUM_ARGS], +} + +impl BlockHessian { + /// create a new block matrix + pub fn new(dims: &[usize]) -> Self { + debug_assert!(!dims.is_empty()); + + let num_blocks = dims.len(); + + let mut ranges = [BlockRange::default(); NUM_ARGS]; + let mut num_rows: usize = 0; + + for i in 0..num_blocks { + let dim = dims[i]; + ranges[i] = BlockRange { + index: num_rows as i64, + dim, + }; + num_rows += dim; + } + Self { + mat: nalgebra::SMatrix::zeros(), + ranges, + } + } + + /// Number of blocks + pub fn num_blocks(&self) -> usize { + self.ranges.len() + } + + /// set block (i, j) + pub fn set_block( + &mut self, + ith: usize, + jth: usize, + m: MatF64, + ) { + debug_assert!(ith < self.num_blocks()); + debug_assert!(jth < self.num_blocks()); + debug_assert_eq!(R, self.ranges[ith].dim); + debug_assert_eq!(C, self.ranges[jth].dim); + + if ith == jth { + debug_assert_eq!(R, C); + + let mut mut_block = self.mut_block::(ith, jth); + mut_block.copy_from(&m); + } else { + debug_assert!(ith < jth); + { + self.mut_block::(ith, jth).copy_from(&m); + } + { + self.mut_block::(jth, ith).copy_from(&m.transpose()); + } + } + } + + /// get block (i, j) + pub fn block( + &self, + ith: usize, + jth: usize, + ) -> nalgebra::Matrix< + f64, + nalgebra::Dyn, + nalgebra::Dyn, + nalgebra::ViewStorage< + '_, + f64, + nalgebra::Dyn, + nalgebra::Dyn, + nalgebra::Const<1>, + nalgebra::Const, + >, + > { + let idx_i = self.ranges[ith].index as usize; + let idx_j = self.ranges[jth].index as usize; + self.mat + .view((idx_i, idx_j), (self.ranges[ith].dim, self.ranges[jth].dim)) + } + + /// mutable reference to block (i, j) + pub fn mut_block( + &mut self, + ith: usize, + jth: usize, + ) -> nalgebra::Matrix< + f64, + nalgebra::Const, + nalgebra::Const, + nalgebra::ViewStorageMut< + '_, + f64, + nalgebra::Const, + nalgebra::Const, + nalgebra::Const<1>, + nalgebra::Const, + >, + > { + let idx_i = self.ranges[ith].index as usize; + let idx_j = self.ranges[jth].index as usize; + self.mat.fixed_view_mut::(idx_i, idx_j) + } +} diff --git a/crates/sophus_opt/src/cost.rs b/crates/sophus_opt/src/cost.rs deleted file mode 100644 index 2bf6408..0000000 --- a/crates/sophus_opt/src/cost.rs +++ /dev/null @@ -1,186 +0,0 @@ -use crate::term::Term; -use crate::variables::VarKind; -use crate::variables::VarPool; -use core::fmt::Debug; -use core::ops::AddAssign; -use dyn_clone::DynClone; - -extern crate alloc; - -/// Evaluated cost -pub trait IsCost: Debug + DynClone { - /// squared error - fn calc_square_error(&self) -> f64; - - /// populate the normal equation - fn populate_normal_equation( - &self, - variables: &VarPool, - nu: f64, - upper_hessian_triplet: &mut alloc::vec::Vec<(usize, usize, f64)>, - neg_grad: &mut nalgebra::DVector, - ); -} - -/// Generic evaluated cost -#[derive(Debug, Clone)] -pub struct Cost { - /// one name (of the corresponding variable family) for each argument (of the cost function - pub family_names: [alloc::string::String; NUM_ARGS], - /// evaluated terms of the cost function - pub terms: alloc::vec::Vec>, - /// degrees of freedom for each argument - pub dof_tuple: [i64; NUM_ARGS], -} - -impl Cost { - /// create a new evaluated cost - pub fn new( - family_names: [alloc::string::String; NUM_ARGS], - dof_tuple: [i64; NUM_ARGS], - ) -> Self { - Cost { - family_names, - terms: alloc::vec::Vec::new(), - dof_tuple, - } - } - - /// add a term to the evaluated cost - pub fn num_of_kind(&self, kind: VarKind, pool: &VarPool) -> usize { - let mut c = 0; - - for name in self.family_names.iter() { - c += if pool.families.get(name).unwrap().get_var_kind() == kind { - 1 - } else { - 0 - }; - } - c - } -} - -impl IsCost for Cost { - fn calc_square_error(&self) -> f64 { - let mut error = 0.0; - for term in self.terms.iter() { - error += term.cost; - } - error - } - - fn populate_normal_equation( - &self, - variables: &VarPool, - nu: f64, - upper_hessian_triplet: &mut alloc::vec::Vec<(usize, usize, f64)>, - neg_grad: &mut nalgebra::DVector, - ) { - let num_args = self.family_names.len(); - - let mut start_indices_per_arg = alloc::vec::Vec::new(); - - let mut dof_per_arg = alloc::vec::Vec::new(); - for name in self.family_names.iter() { - let family = variables.families.get(name).unwrap(); - start_indices_per_arg.push(family.get_start_indices().clone()); - dof_per_arg.push(family.free_or_marg_dof()); - } - - for evaluated_term in self.terms.iter() { - let idx = evaluated_term.idx; - assert_eq!(idx.len(), num_args); - - for arg_id_alpha in 0..num_args { - let dof_alpha = dof_per_arg[arg_id_alpha]; - if dof_alpha == 0 { - continue; - } - - let var_idx_alpha = idx[arg_id_alpha]; - let start_idx_alpha = start_indices_per_arg[arg_id_alpha][var_idx_alpha]; - - if start_idx_alpha == -1 { - continue; - } - - let grad_block = evaluated_term.gradient.block(arg_id_alpha); - let start_idx_alpha = start_idx_alpha as usize; - assert_eq!(dof_alpha, grad_block.nrows()); - - neg_grad - .rows_mut(start_idx_alpha, dof_alpha) - .add_assign(-grad_block); - - let hessian_block = evaluated_term.hessian.block(arg_id_alpha, arg_id_alpha); - assert_eq!(dof_alpha, hessian_block.nrows()); - assert_eq!(dof_alpha, hessian_block.ncols()); - - // block diagonal - for r in 0..dof_alpha { - for c in 0..dof_alpha { - let mut d = 0.0; - if r == c { - d = nu; - } - - if r <= c { - // upper triangular - upper_hessian_triplet.push(( - start_idx_alpha + r, - start_idx_alpha + c, - hessian_block[(r, c)] + d, - )); - } - } - } - - // off diagonal hessian - for arg_id_beta in 0..num_args { - // skip diagonal blocks - if arg_id_alpha == arg_id_beta { - continue; - } - let dof_beta = dof_per_arg[arg_id_beta]; - if dof_beta == 0 { - continue; - } - - let var_idx_beta = idx[arg_id_beta]; - let start_idx_beta = start_indices_per_arg[arg_id_beta][var_idx_beta]; - if start_idx_beta == -1 { - continue; - } - let start_idx_beta = start_idx_beta as usize; - if start_idx_beta < start_idx_alpha { - // upper triangular, skip lower triangular - continue; - } - - let hessian_block_alpha_beta = - evaluated_term.hessian.block(arg_id_alpha, arg_id_beta); - let hessian_block_beta_alpha = - evaluated_term.hessian.block(arg_id_beta, arg_id_alpha); - - assert_eq!(dof_alpha, hessian_block_alpha_beta.nrows()); - assert_eq!(dof_beta, hessian_block_alpha_beta.ncols()); - assert_eq!(dof_beta, hessian_block_beta_alpha.nrows()); - assert_eq!(dof_alpha, hessian_block_beta_alpha.ncols()); - - // alpha-beta off-diagonal - for r in 0..dof_alpha { - for c in 0..dof_beta { - // upper triangular - upper_hessian_triplet.push(( - start_idx_alpha + r, - start_idx_beta + c, - hessian_block_alpha_beta[(r, c)], - )); - } - } - } - } - } - } -} diff --git a/crates/sophus_opt/src/example_problems/cam_calib.rs b/crates/sophus_opt/src/example_problems/cam_calib.rs index f737c7b..bfb61be 100644 --- a/crates/sophus_opt/src/example_problems/cam_calib.rs +++ b/crates/sophus_opt/src/example_problems/cam_calib.rs @@ -1,12 +1,11 @@ -use crate::cost_fn::CostFn; -use crate::cost_fn::CostSignature; use crate::example_problems::cost_fn::isometry3_prior::Isometry3PriorCostFn; -use crate::example_problems::cost_fn::isometry3_prior::Isometry3PriorTermSignature; -use crate::example_problems::cost_fn::reprojection::ReprojTermSignature; +use crate::example_problems::cost_fn::isometry3_prior::Isometry3PriorTerm; use crate::example_problems::cost_fn::reprojection::ReprojectionCostFn; use crate::nlls::optimize; use crate::nlls::OptParams; use crate::prelude::*; +use crate::quadratic_cost::cost_fn::CostFn; +use crate::quadratic_cost::cost_term::CostTerms; use crate::robust_kernel::HuberKernel; use crate::variables::VarFamily; use crate::variables::VarKind; @@ -18,6 +17,7 @@ use sophus_lie::Isometry3; use sophus_lie::Isometry3F64; use sophus_lie::Rotation3; use sophus_sensor::camera_enum::perspective_camera::PinholeCameraF64; +use super::cost_fn::reprojection::ReprojTerm; extern crate alloc; @@ -31,7 +31,7 @@ pub struct CamCalibProblem { /// points in world pub points_in_world: alloc::vec::Vec>, /// observations - pub observations: alloc::vec::Vec, + pub observations: alloc::vec::Vec, /// true intrinsics pub true_intrinsics: PinholeCameraF64, @@ -76,7 +76,7 @@ impl CamCalibProblem { let v = rng.gen::() * (image_size.height as f64 - 1.0); let true_uv_in_img0 = VecF64::<2>::new(u, v); let img_noise = VecF64::<2>::new(rng.gen::() - 0.5, rng.gen::() - 0.5); - observations.push(ReprojTermSignature { + observations.push(ReprojTerm { uv_in_image: true_uv_in_img0 + img_noise, entity_indices: [0, 0, i], }); @@ -95,12 +95,12 @@ impl CamCalibProblem { if spurious_matches && i == 0 { let u = rng.gen::() * (image_size.width as f64 - 1.0); let v = rng.gen::() * (image_size.height as f64 - 1.0); - observations.push(ReprojTermSignature { + observations.push(ReprojTerm { uv_in_image: VecF64::<2>::new(u, v), entity_indices: [0, 1, i], }); } else { - observations.push(ReprojTermSignature { + observations.push(ReprojTerm { uv_in_image: true_uv_in_img1 + img_noise, entity_indices: [0, 1, i], }); @@ -111,7 +111,7 @@ impl CamCalibProblem { let true_point_in_cam2 = true_cam2_from_cam0.transform(true_point_in_cam0); let true_uv_in_img2 = true_intrinsics.cam_proj(true_point_in_cam2); let img_noise = VecF64::<2>::new(rng.gen::() - 0.5, rng.gen::() - 0.5); - observations.push(ReprojTermSignature { + observations.push(ReprojTerm { uv_in_image: true_uv_in_img2 + img_noise, entity_indices: [0, 2, i], }); @@ -135,7 +135,7 @@ impl CamCalibProblem { /// optimize with two poses fixed pub fn optimize_with_two_poses_fixed(&self, intrinsics_var_kind: VarKind) { - let reproj_obs = CostSignature::<3, VecF64<2>, ReprojTermSignature>::new( + let reproj_obs = CostTerms::<3, VecF64<2>, ReprojTerm>::new( ["cams".into(), "poses".into(), "points".into()], self.observations.clone(), ); @@ -177,8 +177,9 @@ impl CamCalibProblem { num_iter: 25, // should converge in single iteration initial_lm_nu: 1.0, // if lm prior param is tiny parallelize: true, + linear_solver: crate::nlls::LinearSolverType::NalgebraDenseFullPiVLu, }, - ); + ).unwrap(); let refined_world_from_robot = up_var_pool.get_members::("poses".into()); @@ -192,17 +193,17 @@ impl CamCalibProblem { /// optimize with priors pub fn optimize_with_priors(&self) { let priors = - CostSignature::<1, (Isometry3F64, MatF64<6, 6>), Isometry3PriorTermSignature>::new( + CostTerms::<1, (Isometry3F64, MatF64<6, 6>), Isometry3PriorTerm>::new( ["poses".into()], alloc::vec![ - Isometry3PriorTermSignature { + Isometry3PriorTerm { entity_indices: [0], isometry_prior: ( self.true_world_from_cameras[0], MatF64::<6, 6>::new_scaling(10000.0), ), }, - Isometry3PriorTermSignature { + Isometry3PriorTerm { entity_indices: [1], isometry_prior: ( self.true_world_from_cameras[1], @@ -212,7 +213,7 @@ impl CamCalibProblem { ], ); - let reproj_obs = CostSignature::<3, VecF64<2>, ReprojTermSignature>::new( + let reproj_obs = CostTerms::<3, VecF64<2>, ReprojTerm>::new( ["cams".into(), "poses".into(), "points".into()], self.observations.clone(), ); @@ -242,8 +243,9 @@ impl CamCalibProblem { num_iter: 5, initial_lm_nu: 1.0, parallelize: true, + linear_solver: crate::nlls::LinearSolverType::NalgebraDenseFullPiVLu, }, - ); + ).unwrap(); let refined_world_from_robot = up_var_pool.get_members::("poses".into()); diff --git a/crates/sophus_opt/src/example_problems/cost_fn/isometry2_prior.rs b/crates/sophus_opt/src/example_problems/cost_fn/isometry2_prior.rs index c31ff97..167d7a4 100644 --- a/crates/sophus_opt/src/example_problems/cost_fn/isometry2_prior.rs +++ b/crates/sophus_opt/src/example_problems/cost_fn/isometry2_prior.rs @@ -1,9 +1,9 @@ -use crate::cost_fn::IsResidualFn; -use crate::cost_fn::IsTermSignature; use crate::prelude::*; +use crate::quadratic_cost::evaluated_term::EvaluatedCostTerm; +use crate::quadratic_cost::evaluated_term::MakeEvaluatedCostTerm; +use crate::quadratic_cost::residual_fn::IsResidualFn; +use crate::quadratic_cost::cost_term::IsTerm; use crate::robust_kernel; -use crate::term::MakeTerm; -use crate::term::Term; use crate::variables::VarKind; use sophus_autodiff::dual::DualScalar; use sophus_autodiff::dual::DualVector; @@ -16,16 +16,16 @@ use sophus_lie::Isometry2F64; #[derive(Copy, Clone)] pub struct Isometry2PriorCostFn {} -/// Isometry2 prior term signature +/// Isometry2 prior term #[derive(Clone)] -pub struct Isometry2PriorTermSignature { +pub struct Isometry2PriorTerm { /// prior mean pub isometry_prior_mean: Isometry2F64, /// entity index pub entity_indices: [usize; 1], } -impl IsTermSignature<1> for Isometry2PriorTermSignature { +impl IsTerm<1> for Isometry2PriorTerm { type Constants = Isometry2F64; fn c_ref(&self) -> &Self::Constants { @@ -55,7 +55,7 @@ impl IsResidualFn<3, 1, (), Isometry2F64, Isometry2F64> for Isometry2PriorCostFn var_kinds: [VarKind; 1], robust_kernel: Option, isometry_prior_mean: &Isometry2F64, - ) -> Term<3, 1> { + ) -> EvaluatedCostTerm<3, 1> { let isometry: Isometry2F64 = args; let residual = res_fn(isometry, *isometry_prior_mean); @@ -74,6 +74,6 @@ impl IsResidualFn<3, 1, (), Isometry2F64, Isometry2F64> for Isometry2PriorCostFn dx_res_fn, zeros, ) },) - .make_term(idx, var_kinds, residual, robust_kernel, None) + .make(idx, var_kinds, residual, robust_kernel, None) } } diff --git a/crates/sophus_opt/src/example_problems/cost_fn/isometry3_prior.rs b/crates/sophus_opt/src/example_problems/cost_fn/isometry3_prior.rs index 7b92945..e3c353e 100644 --- a/crates/sophus_opt/src/example_problems/cost_fn/isometry3_prior.rs +++ b/crates/sophus_opt/src/example_problems/cost_fn/isometry3_prior.rs @@ -1,9 +1,9 @@ -use crate::cost_fn::IsResidualFn; -use crate::cost_fn::IsTermSignature; use crate::prelude::*; +use crate::quadratic_cost::evaluated_term::EvaluatedCostTerm; +use crate::quadratic_cost::residual_fn::IsResidualFn; +use crate::quadratic_cost::evaluated_term::MakeEvaluatedCostTerm; +use crate::quadratic_cost::cost_term::IsTerm; use crate::robust_kernel; -use crate::term::MakeTerm; -use crate::term::Term; use crate::variables::VarKind; use sophus_autodiff::dual::DualScalar; use sophus_autodiff::dual::DualVector; @@ -17,16 +17,16 @@ use sophus_lie::Isometry3F64; #[derive(Copy, Clone)] pub struct Isometry3PriorCostFn {} -/// Isometry3 prior term signature +/// Isometry3 prior term #[derive(Clone)] -pub struct Isometry3PriorTermSignature { +pub struct Isometry3PriorTerm { /// prior mean pub isometry_prior: (Isometry3F64, MatF64<6, 6>), /// entity index pub entity_indices: [usize; 1], } -impl IsTermSignature<1> for Isometry3PriorTermSignature { +impl IsTerm<1> for Isometry3PriorTerm { type Constants = (Isometry3F64, MatF64<6, 6>); fn c_ref(&self) -> &Self::Constants { @@ -56,7 +56,7 @@ impl IsResidualFn<6, 1, (), Isometry3F64, (Isometry3F64, MatF64<6, 6>)> for Isom var_kinds: [VarKind; 1], robust_kernel: Option, isometry_prior: &(Isometry3F64, MatF64<6, 6>), - ) -> Term<6, 1> { + ) -> EvaluatedCostTerm<6, 1> { let isometry: Isometry3F64 = args; let residual = res_fn(isometry, isometry_prior.0); @@ -75,7 +75,7 @@ impl IsResidualFn<6, 1, (), Isometry3F64, (Isometry3F64, MatF64<6, 6>)> for Isom dx_res_fn, zeros, ) },) - .make_term( + .make( idx, var_kinds, residual, diff --git a/crates/sophus_opt/src/example_problems/cost_fn/pose_graph.rs b/crates/sophus_opt/src/example_problems/cost_fn/pose_graph.rs index fd31379..b1cc2b8 100644 --- a/crates/sophus_opt/src/example_problems/cost_fn/pose_graph.rs +++ b/crates/sophus_opt/src/example_problems/cost_fn/pose_graph.rs @@ -1,9 +1,9 @@ -use crate::cost_fn::IsResidualFn; -use crate::cost_fn::IsTermSignature; use crate::prelude::*; +use crate::quadratic_cost::evaluated_term::EvaluatedCostTerm; +use crate::quadratic_cost::residual_fn::IsResidualFn; +use crate::quadratic_cost::evaluated_term::MakeEvaluatedCostTerm; +use crate::quadratic_cost::cost_term::IsTerm; use crate::robust_kernel; -use crate::term::MakeTerm; -use crate::term::Term; use crate::variables::VarKind; use sophus_lie::Isometry2; use sophus_lie::Isometry2F64; @@ -30,7 +30,7 @@ impl IsResidualFn<12, 2, (), (Isometry2F64, Isometry2F64), Isometry2F64> for Pos var_kinds: [VarKind; 2], robust_kernel: Option, obs: &Isometry2F64, - ) -> Term<12, 2> { + ) -> EvaluatedCostTerm<12, 2> { let world_from_pose_a = world_from_pose_x.0; let world_from_pose_b = world_from_pose_x.1; @@ -50,20 +50,20 @@ impl IsResidualFn<12, 2, (), (Isometry2F64, Isometry2F64), Isometry2F64> for Pos ) }, ) - .make_term(idx, var_kinds, residual, robust_kernel, None) + .make(idx, var_kinds, residual, robust_kernel, None) } } -/// Pose graph term signature +/// Pose graph term #[derive(Debug, Clone)] -pub struct PoseGraphCostTermSignature { +pub struct PoseGraphCostTerm { /// 2d relative pose constraint pub pose_a_from_pose_b: Isometry2F64, /// ids of the two poses pub entity_indices: [usize; 2], } -impl IsTermSignature<2> for PoseGraphCostTermSignature { +impl IsTerm<2> for PoseGraphCostTerm { type Constants = Isometry2F64; fn c_ref(&self) -> &Self::Constants { diff --git a/crates/sophus_opt/src/example_problems/cost_fn/reprojection.rs b/crates/sophus_opt/src/example_problems/cost_fn/reprojection.rs index f258b54..5b70317 100644 --- a/crates/sophus_opt/src/example_problems/cost_fn/reprojection.rs +++ b/crates/sophus_opt/src/example_problems/cost_fn/reprojection.rs @@ -1,9 +1,9 @@ -use crate::cost_fn::IsResidualFn; -use crate::cost_fn::IsTermSignature; use crate::prelude::*; +use crate::quadratic_cost::evaluated_term::EvaluatedCostTerm; +use crate::quadratic_cost::evaluated_term::MakeEvaluatedCostTerm; +use crate::quadratic_cost::residual_fn::IsResidualFn; +use crate::quadratic_cost::cost_term::IsTerm; use crate::robust_kernel; -use crate::term::MakeTerm; -use crate::term::Term; use crate::variables::VarKind; use sophus_autodiff::dual::DualScalar; use sophus_autodiff::dual::DualVector; @@ -28,16 +28,16 @@ fn res_fn, const DM: usize, const DN: usize>( uv_in_image - intrinscs.cam_proj(point_in_cam) } -/// Reprojection term signature +/// Reprojection term #[derive(Clone)] -pub struct ReprojTermSignature { +pub struct ReprojTerm { /// Pixel measurement pub uv_in_image: VecF64<2>, /// camera/intrinsics index, pose index, point index pub entity_indices: [usize; 3], } -impl IsTermSignature<3> for ReprojTermSignature { +impl IsTerm<3> for ReprojTerm { type Constants = VecF64<2>; fn c_ref(&self) -> &Self::Constants { @@ -66,7 +66,7 @@ impl IsResidualFn<13, 3, (), (PinholeCameraF64, Isometry3F64, VecF64<3>), VecF64 var_kinds: [VarKind; 3], robust_kernel: Option, uv_in_image: &VecF64<2>, - ) -> Term<13, 3> { + ) -> EvaluatedCostTerm<13, 3> { // calculate residual let residual = res_fn( intrinsics, @@ -132,6 +132,6 @@ impl IsResidualFn<13, 3, (), (PinholeCameraF64, Isometry3F64, VecF64<3>), VecF64 ) }, ) - .make_term(idx, var_kinds, residual, robust_kernel, None) + .make(idx, var_kinds, residual, robust_kernel, None) } } diff --git a/crates/sophus_opt/src/example_problems/pose_circle.rs b/crates/sophus_opt/src/example_problems/pose_circle.rs index ac9a33d..1ecc945 100644 --- a/crates/sophus_opt/src/example_problems/pose_circle.rs +++ b/crates/sophus_opt/src/example_problems/pose_circle.rs @@ -1,10 +1,10 @@ use super::cost_fn::pose_graph::PoseGraphCostFn; -use crate::cost_fn::CostFn; -use crate::cost_fn::CostSignature; -use crate::example_problems::cost_fn::pose_graph::PoseGraphCostTermSignature; +use crate::example_problems::cost_fn::pose_graph::PoseGraphCostTerm; use crate::nlls::optimize; use crate::nlls::OptParams; use crate::prelude::*; +use crate::quadratic_cost::cost_fn::CostFn; +use crate::quadratic_cost::cost_term::CostTerms; use crate::variables::VarFamily; use crate::variables::VarKind; use crate::variables::VarPool; @@ -23,7 +23,7 @@ pub struct PoseCircleProblem { /// estimated poses pub est_world_from_robot: alloc::vec::Vec, /// pose-pose constraints - pub obs_pose_a_from_pose_b_poses: CostSignature<2, Isometry2F64, PoseGraphCostTermSignature>, + pub obs_pose_a_from_pose_b_poses: CostTerms<2, Isometry2F64, PoseGraphCostTerm>, } impl Default for PoseCircleProblem { @@ -38,7 +38,7 @@ impl PoseCircleProblem { let mut true_world_from_robot_poses = alloc::vec![]; let mut est_world_from_robot_poses = alloc::vec![]; let mut obs_pose_a_from_pose_b_poses = - CostSignature::<2, Isometry2F64, PoseGraphCostTermSignature>::new( + CostTerms::<2, Isometry2F64, PoseGraphCostTerm>::new( ["poses".into(), "poses".into()], alloc::vec![], ); @@ -66,7 +66,7 @@ impl PoseCircleProblem { obs_pose_a_from_pose_b_poses .terms - .push(PoseGraphCostTermSignature { + .push(PoseGraphCostTerm { pose_a_from_pose_b, entity_indices: [a_idx, b_idx], }); @@ -140,8 +140,9 @@ impl PoseCircleProblem { num_iter: 5, initial_lm_nu: 1.0, parallelize: true, + ..Default::default() }, - ) + ).unwrap() } } diff --git a/crates/sophus_opt/src/example_problems/simple_prior.rs b/crates/sophus_opt/src/example_problems/simple_prior.rs index 4dce4f7..3de81e2 100644 --- a/crates/sophus_opt/src/example_problems/simple_prior.rs +++ b/crates/sophus_opt/src/example_problems/simple_prior.rs @@ -1,12 +1,12 @@ -use crate::cost_fn::CostFn; -use crate::cost_fn::CostSignature; use crate::example_problems::cost_fn::isometry2_prior::Isometry2PriorCostFn; -use crate::example_problems::cost_fn::isometry2_prior::Isometry2PriorTermSignature; +use crate::example_problems::cost_fn::isometry2_prior::Isometry2PriorTerm; use crate::example_problems::cost_fn::isometry3_prior::Isometry3PriorCostFn; -use crate::example_problems::cost_fn::isometry3_prior::Isometry3PriorTermSignature; +use crate::example_problems::cost_fn::isometry3_prior::Isometry3PriorTerm; use crate::nlls::optimize; use crate::nlls::OptParams; use crate::prelude::*; +use crate::quadratic_cost::cost_fn::CostFn; +use crate::quadratic_cost::cost_term::CostTerms; use crate::variables::VarFamily; use crate::variables::VarKind; use crate::variables::VarPoolBuilder; @@ -44,16 +44,16 @@ impl SimpleIso2PriorProblem { /// Test the simple 2D isometry prior problem pub fn test(&self) { use sophus_autodiff::linalg::EPS_F64; - let cost_signature = alloc::vec![Isometry2PriorTermSignature { + let cost_term = alloc::vec![Isometry2PriorTerm { isometry_prior_mean: self.true_world_from_robot, entity_indices: [0], }]; - let obs_pose_a_from_pose_b_poses = CostSignature::< + let obs_pose_a_from_pose_b_poses = CostTerms::< 1, Isometry2F64, - Isometry2PriorTermSignature, - >::new(["poses".into()], cost_signature); + Isometry2PriorTerm, + >::new(["poses".into()], cost_term); let family: VarFamily = VarFamily::new(VarKind::Free, alloc::vec![self.est_world_from_robot]); @@ -77,8 +77,9 @@ impl SimpleIso2PriorProblem { num_iter: 1, // should converge in single iteration initial_lm_nu: EPS_F64, // if lm prior param is tiny parallelize: true, + ..Default::default() }, - ); + ).unwrap(); let refined_world_from_robot = up_families.get_members::("poses".into()); approx::assert_abs_diff_eq!( @@ -117,16 +118,16 @@ impl SimpleIso3PriorProblem { pub fn test(&self) { use sophus_autodiff::linalg::EPS_F64; - let cost_signature = alloc::vec![Isometry3PriorTermSignature { + let cost_term = alloc::vec![Isometry3PriorTerm { isometry_prior: (self.true_world_from_robot, MatF64::<6, 6>::identity()), entity_indices: [0], }]; - let obs_pose_a_from_pose_b_poses = CostSignature::< + let obs_pose_a_from_pose_b_poses = CostTerms::< 1, (Isometry3F64, MatF64<6, 6>), - Isometry3PriorTermSignature, - >::new(["poses".into()], cost_signature); + Isometry3PriorTerm, + >::new(["poses".into()], cost_term); let family: VarFamily = VarFamily::new(VarKind::Free, alloc::vec![self.est_world_from_robot]); @@ -150,8 +151,9 @@ impl SimpleIso3PriorProblem { num_iter: 1, // should converge in single iteration initial_lm_nu: EPS_F64, // if lm prior param is tiny parallelize: true, + linear_solver: crate::nlls::LinearSolverType::NalgebraDenseFullPiVLu, }, - ); + ).unwrap(); let refined_world_from_robot = up_families.get_members::("poses".into()); approx::assert_abs_diff_eq!( diff --git a/crates/sophus_opt/src/lib.rs b/crates/sophus_opt/src/lib.rs index 68fdad4..191feb9 100644 --- a/crates/sophus_opt/src/lib.rs +++ b/crates/sophus_opt/src/lib.rs @@ -1,30 +1,21 @@ #![cfg_attr(feature = "simd", feature(portable_simd))] #![deny(missing_docs)] #![allow(clippy::needless_range_loop)] -#![no_std] //! # Non-linear least squares optimization crate - part of the sophus-rs project /// Block vector and matrix operations pub mod block; -/// Evaluated costs -pub mod cost; -/// Cost function arguments -pub mod cost_args; -/// Cost functions -pub mod cost_fn; /// Example problems pub mod example_problems; -/// LDLt Cholesky factorization -pub mod ldlt; +/// Linear solvers +pub mod solvers; /// Entry point for the non-linear least squares optimization pub mod nlls; +/// Cost functions, terms, residuals etc. +pub mod quadratic_cost; /// Robust kernel functions pub mod robust_kernel; -/// Linear solvers -pub mod solvers; -/// Evaluated cost terms -pub mod term; /// Decision variables pub mod variables; diff --git a/crates/sophus_opt/src/nlls.rs b/crates/sophus_opt/src/nlls.rs index dfb6873..fabb497 100644 --- a/crates/sophus_opt/src/nlls.rs +++ b/crates/sophus_opt/src/nlls.rs @@ -1,6 +1,9 @@ -use crate::cost::IsCost; -use crate::cost_fn::IsCostFn; -use crate::solvers::solve; +use crate::quadratic_cost::evaluated_cost::IsEvaluatedCost; +use crate::quadratic_cost::cost_fn::IsCostFn; +use crate::solvers::normal_equation::solve; +use crate::solvers::normal_equation::SolveError; +use crate::solvers::normal_equation::TripletType; +use crate::solvers::sparse_ldlt::SparseLdltParams; use crate::variables::VarPool; use core::fmt::Debug; use log::debug; @@ -8,6 +11,40 @@ use log::info; extern crate alloc; +/// Linear solver type +#[derive(Copy, Clone, Debug)] +pub enum LinearSolverType { + /// Sparse LDLT solver (using faer crate) + FaerSparseLdlt(SparseLdltParams), + /// Sparse partial pivoting LU solver (using faer crate) + FaerSparsePartialPivLu, + /// Sparse QR solver (using faer crate) + FaerSparseQR, + /// Dense cholesky (using nalgebra crate) + NalgebraDenseCholesky, + /// Dense full-piovt. LU solver (using nalgebra crate) + NalgebraDenseFullPiVLu, +} + +impl Default for LinearSolverType { + fn default() -> Self { + LinearSolverType::FaerSparseLdlt(Default::default()) + } +} + +impl LinearSolverType { + /// Get the triplet type + pub fn triplet_type(&self) -> TripletType { + match self { + LinearSolverType::FaerSparseLdlt(_) => TripletType::Upper, + LinearSolverType::FaerSparsePartialPivLu => TripletType::Full, + LinearSolverType::FaerSparseQR => TripletType::Full, + LinearSolverType::NalgebraDenseCholesky => TripletType::Full, + LinearSolverType::NalgebraDenseFullPiVLu => TripletType::Full, + } + } +} + /// Optimization parameters #[derive(Copy, Clone, Debug)] pub struct OptParams { @@ -17,6 +54,8 @@ pub struct OptParams { pub initial_lm_nu: f64, /// whether to use parallelization pub parallelize: bool, + /// linear solver type + pub linear_solver: LinearSolverType, } impl Default for OptParams { @@ -25,6 +64,7 @@ impl Default for OptParams { num_iter: 20, initial_lm_nu: 10.0, parallelize: true, + linear_solver: Default::default(), } } } @@ -34,8 +74,8 @@ pub fn optimize( mut variables: VarPool, mut cost_fns: alloc::vec::Vec>, params: OptParams, -) -> VarPool { - let mut init_costs: alloc::vec::Vec> = alloc::vec::Vec::new(); +) -> Result { + let mut init_costs: alloc::vec::Vec> = alloc::vec::Vec::new(); for cost_fn in cost_fns.iter_mut() { // sort to achieve more efficient evaluation and reduction @@ -52,15 +92,15 @@ pub fn optimize( let initial_mse = mse; for i in 0..params.num_iter { - let mut evaluated_costs: alloc::vec::Vec> = + let mut evaluated_costs: alloc::vec::Vec> = alloc::vec::Vec::new(); for cost_fn in cost_fns.iter_mut() { evaluated_costs.push(cost_fn.eval(&variables, true, params.parallelize)); } - let updated_families = solve(&variables, evaluated_costs, nu); + let updated_families = solve(&variables, evaluated_costs, nu, params.linear_solver)?; - let mut new_costs: alloc::vec::Vec> = alloc::vec::Vec::new(); + let mut new_costs: alloc::vec::Vec> = alloc::vec::Vec::new(); for cost_fn in cost_fns.iter_mut() { new_costs.push(cost_fn.eval(&updated_families, true, params.parallelize)); } @@ -84,5 +124,5 @@ pub fn optimize( } info!("e^2: {:?} -> {:?}", initial_mse, mse); - variables + Ok(variables) } diff --git a/crates/sophus_opt/src/quadratic_cost.rs b/crates/sophus_opt/src/quadratic_cost.rs new file mode 100644 index 0000000..ba74fb0 --- /dev/null +++ b/crates/sophus_opt/src/quadratic_cost.rs @@ -0,0 +1,12 @@ +/// Cost function arguments +pub mod compare_idx; +/// evaluated cost +pub mod evaluated_cost; +/// Cost functions +pub mod cost_fn; +/// Cost terms +pub mod cost_term; +/// Residual fn - trait of the main customization point +pub mod residual_fn; +/// EEvaluated cost term +pub mod evaluated_term; diff --git a/crates/sophus_opt/src/cost_args.rs b/crates/sophus_opt/src/quadratic_cost/compare_idx.rs similarity index 98% rename from crates/sophus_opt/src/cost_args.rs rename to crates/sophus_opt/src/quadratic_cost/compare_idx.rs index 0cc4d67..0210fd0 100644 --- a/crates/sophus_opt/src/cost_args.rs +++ b/crates/sophus_opt/src/quadratic_cost/compare_idx.rs @@ -129,7 +129,6 @@ where } mod test { - use crate::cost_args::CompareIdx; #[allow(dead_code)] fn le_than( @@ -137,7 +136,7 @@ mod test { lhs: [usize; N], rhs: [usize; N], ) -> core::cmp::Ordering { - CompareIdx::new(c).le_than(lhs, rhs) + super::CompareIdx::new(c).le_than(lhs, rhs) } #[test] @@ -401,7 +400,7 @@ mod test { [3, 2, 1], ]; - let less = CompareIdx::new(c); + let less = crate::quadratic_cost::compare_idx::CompareIdx::new(c); l.sort_by(|a, b| less.le_than(*a, *b)); } } diff --git a/crates/sophus_opt/src/cost_fn.rs b/crates/sophus_opt/src/quadratic_cost/cost_fn.rs similarity index 62% rename from crates/sophus_opt/src/cost_fn.rs rename to crates/sophus_opt/src/quadratic_cost/cost_fn.rs index b62d0a3..6b33f98 100644 --- a/crates/sophus_opt/src/cost_fn.rs +++ b/crates/sophus_opt/src/quadratic_cost/cost_fn.rs @@ -1,8 +1,12 @@ -use crate::cost::Cost; -use crate::cost::IsCost; -use crate::cost_args::c_from_var_kind; +use super::cost_term::CostTerms; +use super::cost_term::IsTerm; +use super::evaluated_cost::IsEvaluatedCost; +use super::residual_fn::IsResidualFn; +use crate::quadratic_cost::compare_idx::c_from_var_kind; +use crate::quadratic_cost::compare_idx::CompareIdx; +use crate::quadratic_cost::evaluated_cost::EvaluatedCost; +use crate::quadratic_cost::evaluated_term::EvaluatedCostTerm; use crate::robust_kernel::RobustKernel; -use crate::term::Term; use crate::variables::IsVarTuple; use crate::variables::VarKind; use crate::variables::VarPool; @@ -11,84 +15,17 @@ use core::ops::Range; extern crate alloc; -/// Signature of a term of a cost function -pub trait IsTermSignature: Send + Sync + 'static { - /// associated constants such as measurements, etc. - type Constants; - - /// one DOF for each argument - const DOF_TUPLE: [i64; N]; - - /// reference to the constants - fn c_ref(&self) -> &Self::Constants; - - /// one index (into the variable family) for each argument - fn idx_ref(&self) -> &[usize; N]; -} - -/// Signature of a cost function -#[derive(Debug, Clone)] -pub struct CostSignature< - const NUM_ARGS: usize, - Constants, - TermSignature: IsTermSignature, -> { - /// one variable family name for each argument - pub family_names: [alloc::string::String; NUM_ARGS], - /// terms of the unevaluated cost function - pub terms: alloc::vec::Vec, - pub(crate) reduction_ranges: Option>>, -} - -impl< - const NUM_ARGS: usize, - Constants, - TermSignature: IsTermSignature, - > CostSignature -{ - /// Create a new cost signature - pub fn new( - family_names: [alloc::string::String; NUM_ARGS], - terms: alloc::vec::Vec, - ) -> Self { - CostSignature { - family_names, - terms, - reduction_ranges: None, - } - } -} - -/// Residual function -pub trait IsResidualFn< - const NUM: usize, - const NUM_ARGS: usize, - GlobalConstants: 'static + Send + Sync, - Args: IsVarTuple, - Constants, ->: Copy + Send + Sync + 'static -{ - /// evaluate the residual function which shall be defined by the user - fn eval( - &self, - global_constants: &GlobalConstants, - idx: [usize; NUM_ARGS], - args: Args, - derivatives: [VarKind; NUM_ARGS], - robust_kernel: Option, - constants: &Constants, - ) -> Term; -} - -/// Quadratic cost function of the non-linear least squares problem +/// Quadratic cost function of the non-linear least squares problem. +/// +/// This is producing an evaluated cost: Box which is a sum of squared residuals. pub trait IsCostFn { - /// evaluate the cost function + /// Evaluate the cost function. fn eval( &self, var_pool: &VarPool, calc_derivatives: bool, parallelize: bool, - ) -> alloc::boxed::Box; + ) -> alloc::boxed::Box; /// sort the terms of the cost function (to ensure more efficient evaluation and reduction over /// conditioned variables) @@ -98,21 +35,23 @@ pub trait IsCostFn { fn robust_kernel(&self) -> Option; } -/// Generic cost function of the non-linear least squares problem +/// Generic cost function of the non-linear least squares problem. +/// +/// This struct is passed as a Box to the optimizer. #[derive(Debug, Clone)] pub struct CostFn< const NUM: usize, const NUM_ARGS: usize, GlobalConstants: 'static + Send + Sync, Constants, - TermSignature: IsTermSignature, + Term: IsTerm, ResidualFn, VarTuple: IsVarTuple + 'static, > where ResidualFn: IsResidualFn, { global_constants: GlobalConstants, - signature: CostSignature, + terms: CostTerms, residual_fn: ResidualFn, robust_kernel: Option, phantom: PhantomData, @@ -123,38 +62,38 @@ impl< const NUM_ARGS: usize, GlobalConstants: 'static + Send + Sync, Constants: 'static, - TermSignature: IsTermSignature + 'static, + Term: IsTerm + 'static, ResidualFn, VarTuple: IsVarTuple + 'static, - > CostFn + > CostFn where ResidualFn: IsResidualFn + 'static, { - /// create a new cost function from a signature and a residual function + /// create a new cost function from the cost terms and a residual function pub fn new_box( global_constants: GlobalConstants, - signature: CostSignature, + terms: CostTerms, residual_fn: ResidualFn, ) -> alloc::boxed::Box { alloc::boxed::Box::new(Self { global_constants, - signature, + terms, residual_fn, robust_kernel: None, phantom: PhantomData, }) } - /// create a new robust cost function from a signature, a residual function and a robust kernel + /// create a new robust cost function from the cost terms, a residual function and a robust kernel pub fn new_robust( global_constants: GlobalConstants, - signature: CostSignature, + terms: CostTerms, residual_fn: ResidualFn, robust_kernel: RobustKernel, ) -> alloc::boxed::Box { alloc::boxed::Box::new(Self { global_constants, - signature, + terms, residual_fn, robust_kernel: Some(robust_kernel), phantom: PhantomData, @@ -167,11 +106,10 @@ impl< const NUM_ARGS: usize, GlobalConstants: 'static + Send + Sync, Constants, - TermSignature: IsTermSignature, + Term: IsTerm, ResidualFn, VarTuple: IsVarTuple + 'static, - > IsCostFn - for CostFn + > IsCostFn for CostFn where ResidualFn: IsResidualFn, { @@ -180,34 +118,32 @@ where var_pool: &VarPool, calc_derivatives: bool, parallelize: bool, - ) -> alloc::boxed::Box { + ) -> alloc::boxed::Box { let mut var_kind_array = - VarTuple::var_kind_array(var_pool, self.signature.family_names.clone()); + VarTuple::var_kind_array(var_pool, self.terms.family_names.clone()); if !calc_derivatives { var_kind_array = var_kind_array.map(|_x| VarKind::Conditioned) } - let mut evaluated_terms = Cost::new( - self.signature.family_names.clone(), - TermSignature::DOF_TUPLE, - ); + let mut evaluated_terms = + EvaluatedCost::new(self.terms.family_names.clone(), Term::DOF_TUPLE); let var_family_tuple = - VarTuple::ref_var_family_tuple(var_pool, self.signature.family_names.clone()); + VarTuple::ref_var_family_tuple(var_pool, self.terms.family_names.clone()); - let eval_res = |term_signature: &TermSignature| { + let eval_res = |term: &Term| { self.residual_fn.eval( &self.global_constants, - *term_signature.idx_ref(), - VarTuple::extract(&var_family_tuple, *term_signature.idx_ref()), + *term.idx_ref(), + VarTuple::extract(&var_family_tuple, *term.idx_ref()), var_kind_array, self.robust_kernel, - term_signature.c_ref(), + term.c_ref(), ) }; - let reduction_ranges = self.signature.reduction_ranges.as_ref().unwrap(); + let reduction_ranges = self.terms.reduction_ranges.as_ref().unwrap(); #[derive(Debug)] enum ParallelizationStrategy { @@ -219,8 +155,7 @@ where const INNER_LOOP_THRESHOLD: f64 = 100.0; const REDUCTION_RATIO_THRESHOLD: f64 = 1.0; - let average_inner_loop_size = - self.signature.terms.len() as f64 / reduction_ranges.len() as f64; + let average_inner_loop_size = self.terms.terms.len() as f64 / reduction_ranges.len() as f64; let reduction_ratio = average_inner_loop_size / reduction_ranges.len() as f64; let parallelization_strategy = match parallelize { @@ -248,7 +183,7 @@ where // evaluated_terms.terms = reduction_ranges // .iter() // sequential outer loop // .map(|range| { - // let evaluated_term_sum = self.signature.terms[range.start..range.end] + // let evaluated_term_sum = self.terms.terms[range.start..range.end] // .iter() // sequential inner loop // .fold(None, |acc: Option>, term| { // let evaluated_term = eval_res(term); @@ -267,9 +202,9 @@ where evaluated_terms.terms.reserve(reduction_ranges.len()); for range in reduction_ranges.iter() { - let mut evaluated_term_sum: Option> = None; + let mut evaluated_term_sum: Option> = None; - for term in self.signature.terms[range.start..range.end].iter() { + for term in self.terms.terms[range.start..range.end].iter() { match evaluated_term_sum { Some(mut sum) => { sum.reduce(eval_res(term)); @@ -288,18 +223,21 @@ where evaluated_terms.terms = reduction_ranges .par_iter() // parallelize over the outer terms .map(|range| { - let evaluated_term_sum = self.signature.terms[range.start..range.end] + let evaluated_term_sum = self.terms.terms[range.start..range.end] .iter() // sequential inner loop - .fold(None, |acc: Option>, term| { - let evaluated_term = eval_res(term); - match acc { - Some(mut sum) => { - sum.reduce(evaluated_term); - Some(sum) + .fold( + None, + |acc: Option>, term| { + let evaluated_term = eval_res(term); + match acc { + Some(mut sum) => { + sum.reduce(evaluated_term); + Some(sum) + } + None => Some(evaluated_term), } - None => Some(evaluated_term), - } - }); + }, + ); evaluated_term_sum.unwrap() }) @@ -316,11 +254,11 @@ where // // todo: Consider adding an if statement here and only parallelize the // inner loop if the range length is greater than some threshold. - let evaluated_term_sum = self.signature.terms[range.start..range.end] + let evaluated_term_sum = self.terms.terms[range.start..range.end] .par_iter() // parallelize over the inner terms .fold( || None, - |acc: Option>, term| { + |acc: Option>, term| { let evaluated_term = eval_res(term); match acc { Some(mut sum) => { @@ -353,38 +291,36 @@ where } fn sort(&mut self, variables: &VarPool) { - let var_kind_array = - &VarTuple::var_kind_array(variables, self.signature.family_names.clone()); - use crate::cost_args::CompareIdx; + let var_kind_array = &VarTuple::var_kind_array(variables, self.terms.family_names.clone()); let c_array = c_from_var_kind(var_kind_array); let less = CompareIdx::new(&c_array); - assert!(!self.signature.terms.is_empty()); + assert!(!self.terms.terms.is_empty()); - self.signature + self.terms .terms .sort_by(|a, b| less.le_than(*a.idx_ref(), *b.idx_ref())); - for t in 0..self.signature.terms.len() - 1 { + for t in 0..self.terms.terms.len() - 1 { assert!( less.le_than( - *self.signature.terms[t].idx_ref(), - *self.signature.terms[t + 1].idx_ref() + *self.terms.terms[t].idx_ref(), + *self.terms.terms[t + 1].idx_ref() ) != core::cmp::Ordering::Greater ); } let mut reduction_ranges: alloc::vec::Vec> = alloc::vec![]; let mut i = 0; - while i < self.signature.terms.len() { - let outer_term_signature = &self.signature.terms[i]; + while i < self.terms.terms.len() { + let outer_term = &self.terms.terms[i]; let outer_term_idx = i; - while i < self.signature.terms.len() + while i < self.terms.terms.len() && less.free_vars_equal( - outer_term_signature.idx_ref(), - self.signature.terms[i].idx_ref(), + outer_term.idx_ref(), + self.terms.terms[i].idx_ref(), ) { i += 1; @@ -392,7 +328,7 @@ where reduction_ranges.push(outer_term_idx..i); } - self.signature.reduction_ranges = Some(reduction_ranges); + self.terms.reduction_ranges = Some(reduction_ranges); } fn robust_kernel(&self) -> Option { diff --git a/crates/sophus_opt/src/term.rs b/crates/sophus_opt/src/quadratic_cost/evaluated_term.rs similarity index 88% rename from crates/sophus_opt/src/term.rs rename to crates/sophus_opt/src/quadratic_cost/evaluated_term.rs index 2972bd9..566efb2 100644 --- a/crates/sophus_opt/src/term.rs +++ b/crates/sophus_opt/src/quadratic_cost/evaluated_term.rs @@ -1,5 +1,5 @@ -use crate::block::BlockMatrix; -use crate::block::BlockVector; +use crate::block::block_gradient::BlockGradient; +use crate::block::block_hessian::BlockHessian; use crate::prelude::*; use crate::robust_kernel; use crate::variables::VarKind; @@ -10,11 +10,11 @@ extern crate alloc; /// Evaluated cost term #[derive(Debug, Clone)] -pub struct Term { +pub struct EvaluatedCostTerm { /// Hessian - pub hessian: BlockMatrix, + pub hessian: BlockHessian, /// Gradient - pub gradient: BlockVector, + pub gradient: BlockGradient, /// cost: 0.5 * residual^T * precision_mat * residual pub cost: f64, /// indices of the variable families @@ -23,8 +23,8 @@ pub struct Term { pub num_sub_terms: usize, } -impl Term { - pub(crate) fn reduce(&mut self, other: Term) { +impl EvaluatedCostTerm { + pub(crate) fn reduce(&mut self, other: EvaluatedCostTerm) { self.hessian.mat += other.hessian.mat; self.gradient.vec += other.gradient.vec; self.cost += other.cost; @@ -38,7 +38,7 @@ trait RowLoop { idx: usize, i: usize, j: usize, - hessian: &mut BlockMatrix, + hessian: &mut BlockHessian, precision_mat: Option>, ); @@ -47,8 +47,8 @@ trait RowLoop { idx: usize, i: usize, lambda_res: &VecF64, - gradient: &mut BlockVector, - hessian: &mut BlockMatrix, + gradient: &mut BlockGradient, + hessian: &mut BlockHessian, precision_mat: Option>, ); } @@ -59,7 +59,7 @@ impl RowLoop, + _hessian: &mut BlockHessian, _precision_mat: Option>, ) { } @@ -69,8 +69,8 @@ impl RowLoop, - _gradient: &mut BlockVector, - _hessian: &mut BlockMatrix, + _gradient: &mut BlockGradient, + _hessian: &mut BlockHessian, _precision_mat: Option>, ) { } @@ -89,7 +89,7 @@ impl< idx: usize, i: usize, j: usize, - hessian: &mut BlockMatrix, + hessian: &mut BlockHessian, precision_mat: Option>, ) { if idx == i { @@ -107,8 +107,8 @@ impl< idx: usize, i: usize, lambda_res: &VecF64, - gradient: &mut BlockVector, - hessian: &mut BlockMatrix, + gradient: &mut BlockGradient, + hessian: &mut BlockHessian, precision_mat: Option>, ) { if idx == i { @@ -137,7 +137,7 @@ trait ColLoop, + hessian: &mut BlockHessian, lhs: MatF64, precision_mat: Option>, ); @@ -151,7 +151,7 @@ impl _idx: usize, _i: usize, _j: usize, - _hessian: &mut BlockMatrix, + _hessian: &mut BlockHessian, _lhs: MatF64, _precision_mat: Option>, ) { @@ -172,7 +172,7 @@ impl< idx: usize, i: usize, j: usize, - hessian: &mut BlockMatrix, + hessian: &mut BlockHessian, lhs: MatF64, precision_mat: Option>, ) { @@ -192,7 +192,7 @@ impl< } /// Trait for making n-ary terms -pub trait MakeTerm { +pub trait MakeEvaluatedCostTerm { /// make a term from a residual value, and derivatives (=self) /// /// In more detail, this function computes the Hessian, gradient and least-squares cost of the @@ -207,14 +207,14 @@ pub trait MakeTerm { /// - `precision_mat`: Precision matrix - i.e. inverse of the covariance matrix - to compute the /// least-squares cost: `0.5 * residual^T * precision_mat * residual`. /// If `None`, the identity matrix is used: `0.5 * residual^T * residual`. - fn make_term( + fn make( self, idx: [usize; N], var_kinds: [VarKind; N], residual: VecF64, robust_kernel: Option, precision_mat: Option>, - ) -> Term; + ) -> EvaluatedCostTerm; } macro_rules! nested_option_tuple { @@ -242,24 +242,24 @@ macro_rules! nested_option_tuple { }; } -macro_rules! impl_make_term_for_tuples { +macro_rules! impl_make_evaluated_cost_term_for_tuples { ($($N:literal: ($($F:ident, $D:ident, $idx:tt),+),)+) => { $( #[allow(non_snake_case)] - impl MakeTerm for ($($F,)+) + impl MakeEvaluatedCostTerm for ($($F,)+) where $( $F: FnOnce() -> MatF64, )+ { - fn make_term( + fn make( self, idx: [usize; $N], var_kinds: [VarKind; $N], residual: VecF64, robust_kernel: Option, precision_mat: Option>, - ) -> Term { + ) -> EvaluatedCostTerm { let residual = match robust_kernel { Some(rk) => rk.weight(residual.norm()) * residual, None => residual, @@ -270,8 +270,8 @@ macro_rules! impl_make_term_for_tuples { let maybe_dx = nested_option_tuple!(var_kinds; $($F, $D, $idx),+); let dims = [$($D,)+]; - let mut hessian = BlockMatrix::new(&dims); - let mut gradient = BlockVector::new(&dims); + let mut hessian = BlockHessian::new(&dims); + let mut gradient = BlockGradient::new(&dims); let lambda_res = match precision_mat { Some(pm) => pm * residual, @@ -286,7 +286,7 @@ macro_rules! impl_make_term_for_tuples { } } - Term { + EvaluatedCostTerm { hessian, gradient, cost: (residual.transpose() * lambda_res)[0], @@ -299,8 +299,8 @@ macro_rules! impl_make_term_for_tuples { } } -// implement MakeTerm for up to 25 arguments. -impl_make_term_for_tuples! { +// implement MakeEvaluatedCostTerm for up to 25 arguments. +impl_make_evaluated_cost_term_for_tuples! { 1: ( F0, D0, 0), 2: ( F0, D0, 0, F1 ,D1, 1), 3: ( F0, D0, 0, F1 ,D1, 1, F2, D2, 2), diff --git a/crates/sophus_opt/src/quadratic_cost/residual_fn.rs b/crates/sophus_opt/src/quadratic_cost/residual_fn.rs new file mode 100644 index 0000000..8e71aff --- /dev/null +++ b/crates/sophus_opt/src/quadratic_cost/residual_fn.rs @@ -0,0 +1,28 @@ +use super::evaluated_term::EvaluatedCostTerm; +use crate::{ + robust_kernel::RobustKernel, + variables::{IsVarTuple, VarKind}, +}; + +/// Residual function +/// +/// This trait is the main customization point for the user. +pub trait IsResidualFn< + const NUM: usize, + const NUM_ARGS: usize, + GlobalConstants: 'static + Send + Sync, + Args: IsVarTuple, + Constants, +>: Copy + Send + Sync + 'static +{ + /// evaluate the residual function which shall be defined by the user + fn eval( + &self, + global_constants: &GlobalConstants, + idx: [usize; NUM_ARGS], + args: Args, + derivatives: [VarKind; NUM_ARGS], + robust_kernel: Option, + constants: &Constants, + ) -> EvaluatedCostTerm; +} diff --git a/crates/sophus_opt/src/solvers.rs b/crates/sophus_opt/src/solvers.rs index 9c39c3b..356a948 100644 --- a/crates/sophus_opt/src/solvers.rs +++ b/crates/sophus_opt/src/solvers.rs @@ -1,69 +1,5 @@ -use crate::cost::IsCost; -use crate::ldlt::SparseLdlt; -use crate::ldlt::SymmetricTripletMatrix; -use crate::variables::VarKind; -use crate::variables::VarPool; -extern crate alloc; - -/// Normal equation -pub struct SparseNormalEquation { - sparse_hessian: SymmetricTripletMatrix, - neg_gradient: nalgebra::DVector, -} - -impl SparseNormalEquation { - fn from_families_and_cost( - variables: &VarPool, - costs: alloc::vec::Vec>, - nu: f64, - ) -> SparseNormalEquation { - assert!(variables.num_of_kind(VarKind::Marginalized) == 0); - assert!(variables.num_of_kind(VarKind::Free) >= 1); - - // Note let's first focus on these special cases, before attempting a - // general version covering all cases holistically. Also, it might not be trivial - // to implement VarKind::Marginalized > 1. - // - Example, the the arrow-head sparsity uses a recursive application of the Schur-Complement. - let num_var_params = variables.num_free_params(); - let mut upper_hessian_triplet = alloc::vec::Vec::new(); - let mut neg_grad = nalgebra::DVector::::zeros(num_var_params); - - for cost in costs.iter() { - cost.populate_normal_equation(variables, nu, &mut upper_hessian_triplet, &mut neg_grad); - } - - Self { - sparse_hessian: SymmetricTripletMatrix { - upper_triplets: upper_hessian_triplet, - size: num_var_params, - }, - - neg_gradient: neg_grad, - } - } - - fn solve(&mut self) -> nalgebra::DVector { - // TODO: perform permutation to reduce fill-in and symbolic factorization only once - SparseLdlt::from_triplets(&self.sparse_hessian).solve(&self.neg_gradient) - } -} - -/// Solve the normal equation -pub fn solve( - variables: &VarPool, - costs: alloc::vec::Vec>, - nu: f64, -) -> VarPool { - assert!(variables.num_of_kind(VarKind::Marginalized) <= 1); - assert!(variables.num_of_kind(VarKind::Free) >= 1); - - if variables.num_of_kind(VarKind::Marginalized) == 0 { - let mut sne = SparseNormalEquation::from_families_and_cost(variables, costs, nu); - sne.solve(); - let delta = sne.solve(); - variables.update(delta) - } else { - todo!() - } -} +/// Sparse LDL^{t} solver +pub mod sparse_ldlt; +/// normal equation +pub mod normal_equation; diff --git a/crates/sophus_opt/src/solvers/normal_equation.rs b/crates/sophus_opt/src/solvers/normal_equation.rs new file mode 100644 index 0000000..ec9677d --- /dev/null +++ b/crates/sophus_opt/src/solvers/normal_equation.rs @@ -0,0 +1,246 @@ +use super::sparse_ldlt::SimplicalSparseLdlt; +use crate::nlls::LinearSolverType; +use crate::quadratic_cost::evaluated_cost::IsEvaluatedCost; +use crate::variables::VarKind; +use crate::variables::VarPool; +use faer::prelude::SpSolver; +use snafu::Snafu; + +extern crate alloc; + +#[derive(Copy, Clone, Debug, PartialEq)] +/// Triplet type +pub enum TripletType { + /// Upper diagonal matrix + Upper, + /// Full matrix + Full, +} + +#[derive(Clone, Debug)] +/// A matrix in triplet format +pub struct SymmetricTripletMatrix { + /// diagonal triplets (either upper diagonal or full matrix, depending on the linear solver type) + pub triplets: alloc::vec::Vec<(usize, usize, f64)>, + /// trplet type + pub triplet_type: TripletType, + /// row count (== column count) + pub size: usize, +} + +impl SymmetricTripletMatrix { + /// Create an example matrix + pub fn upper_diagonal_example() -> Self { + Self { + triplets: alloc::vec![ + (0, 0, 3.05631771), + (1, 1, 60.05631771), + (2, 2, 6.05631771), + (3, 3, 5.05631771), + (4, 4, 8.05631771), + (5, 5, 5.05631771), + (6, 6, 0.05631771), + (7, 7, 10.005631771), + (0, 1, 2.41883573), + (0, 3, 2.41883573), + (0, 5, 1.88585946), + (1, 3, 1.73897015), + (1, 5, 2.12387697), + (1, 7, 1.47609157), + (2, 7, 1.4541327), + (3, 4, 2.35666066), + (3, 7, 0.94642903), + ], + triplet_type: TripletType::Upper, + size: 8, + } + } + + /// Convert to dense matrix + pub fn to_dense(&self) -> nalgebra::DMatrix { + let mut full_matrix = nalgebra::DMatrix::from_element(self.size, self.size, 0.0); + for &(row, col, value) in self.triplets.iter() { + full_matrix[(row, col)] += value; + if self.triplet_type == TripletType::Upper && row != col { + full_matrix[(col, row)] += value; + } + } + full_matrix + } + + /// Returns true if the matrix is semi-positive definite + /// + /// TODO: This is a very inefficient implementation + pub fn is_semi_positive_definite(&self) -> bool { + let full_matrix = self.to_dense(); + let eigen_comp = full_matrix.symmetric_eigen(); + eigen_comp.eigenvalues.iter().all(|&x| x >= 0.0) + } +} + +/// Normal equation +pub struct NormalEquation { + sparse_hessian: SymmetricTripletMatrix, + neg_gradient: nalgebra::DVector, + linear_solver: LinearSolverType, +} + +/// Linear solver error +#[derive(Snafu, Debug)] +pub enum SolveError { + /// Sparse LDLt error + #[snafu(display("sparse LDLt error {}", details))] + SparseLdltError { + /// details + details: faer::sparse::FaerError, + }, + /// Sparse LU error + #[snafu(display("sparse LU error {}", details))] + SparseLuError { + /// details + details: faer::sparse::LuError, + }, + /// Sparse QR error + #[snafu(display("sparse QR error {}", details))] + SparseQrError { + /// details + details: faer::sparse::FaerError, + }, + /// Dense cholesky error + #[snafu(display("dense cholesky factorization failed"))] + DenseCholeskyError, + /// Dense LU error + #[snafu(display("dense LU solve failed"))] + DenseLuError, +} + +impl NormalEquation { + fn from_families_and_cost( + variables: &VarPool, + costs: alloc::vec::Vec>, + nu: f64, + linear_solver: LinearSolverType, + ) -> NormalEquation { + assert!(variables.num_of_kind(VarKind::Marginalized) == 0); + assert!(variables.num_of_kind(VarKind::Free) >= 1); + + // Note let's first focus on these special cases, before attempting a + // general version covering all cases holistically. Also, it might not be trivial + // to implement VarKind::Marginalized > 1. + // - Example, the the arrow-head sparsity uses a recursive application of the Schur-Complement. + let num_var_params = variables.num_free_params(); + let mut hessian_triplet = alloc::vec::Vec::new(); + let mut neg_grad = nalgebra::DVector::::zeros(num_var_params); + let triplet_type = linear_solver.triplet_type(); + + for cost in costs.iter() { + cost.populate_normal_equation( + variables, + nu, + &mut hessian_triplet, + &mut neg_grad, + triplet_type == TripletType::Upper, + ); + } + + Self { + sparse_hessian: SymmetricTripletMatrix { + triplets: hessian_triplet, + triplet_type, + size: num_var_params, + }, + neg_gradient: neg_grad, + linear_solver, + } + } + + fn solve(&mut self) -> Result, SolveError> { + match self.linear_solver { + LinearSolverType::FaerSparseLdlt(ldlt_params) => Ok( + match SimplicalSparseLdlt::from_triplets(&self.sparse_hessian, ldlt_params) + .solve(&self.neg_gradient) + { + Ok(x) => x, + Err(e) => return Err(SolveError::SparseLdltError { details: e }), + }, + ), + LinearSolverType::FaerSparsePartialPivLu => { + let csr = faer::sparse::SparseColMat::try_new_from_triplets( + self.sparse_hessian.size, + self.sparse_hessian.size, + &self.sparse_hessian.triplets, + ) + .unwrap(); + let x = self.neg_gradient.clone(); + let x_slice_mut = self.neg_gradient.as_mut_slice(); + let mut x_ref = faer::mat::from_column_major_slice_mut( + x_slice_mut, + self.sparse_hessian.size, + 1, + ); + match csr.sp_lu() { + Ok(lu) => { + lu.solve_in_place(faer::reborrow::ReborrowMut::rb_mut(&mut x_ref)); + Ok(x) + } + Err(e) => Err(SolveError::SparseLuError { details: e }), + } + } + LinearSolverType::FaerSparseQR => { + let csr = faer::sparse::SparseColMat::try_new_from_triplets( + self.sparse_hessian.size, + self.sparse_hessian.size, + &self.sparse_hessian.triplets, + ) + .unwrap(); + let x = self.neg_gradient.clone(); + let x_slice_mut = self.neg_gradient.as_mut_slice(); + let mut x_ref = faer::mat::from_column_major_slice_mut( + x_slice_mut, + self.sparse_hessian.size, + 1, + ); + match csr.sp_qr() { + Ok(qr) => { + qr.solve_in_place(faer::reborrow::ReborrowMut::rb_mut(&mut x_ref)); + Ok(x) + } + Err(e) => Err(SolveError::SparseQrError { details: e }), + } + } + LinearSolverType::NalgebraDenseCholesky => { + let dense_matrix = self.sparse_hessian.to_dense(); + match dense_matrix.cholesky() { + Some(cholesky) => Ok(cholesky.solve(&self.neg_gradient)), + None => Err(SolveError::DenseCholeskyError), + } + } + LinearSolverType::NalgebraDenseFullPiVLu => { + let dense_matrix = self.sparse_hessian.to_dense(); + match dense_matrix.full_piv_lu().solve(&self.neg_gradient) { + Some(x) => Ok(x), + None => Err(SolveError::DenseLuError), + } + } + } + } +} + +/// Solve the normal equation +pub fn solve( + variables: &VarPool, + costs: alloc::vec::Vec>, + nu: f64, + linear_solver: LinearSolverType, +) -> Result { + assert!(variables.num_of_kind(VarKind::Marginalized) <= 1); + assert!(variables.num_of_kind(VarKind::Free) >= 1); + + if variables.num_of_kind(VarKind::Marginalized) == 0 { + let mut sne = NormalEquation::from_families_and_cost(variables, costs, nu, linear_solver); + let delta = sne.solve(); + Ok(variables.update(delta?)) + } else { + todo!() + } +} diff --git a/crates/sophus_opt/src/ldlt.rs b/crates/sophus_opt/src/solvers/sparse_ldlt.rs similarity index 69% rename from crates/sophus_opt/src/ldlt.rs rename to crates/sophus_opt/src/solvers/sparse_ldlt.rs index f20b611..9263296 100644 --- a/crates/sophus_opt/src/ldlt.rs +++ b/crates/sophus_opt/src/solvers/sparse_ldlt.rs @@ -1,67 +1,29 @@ +use faer::sparse::FaerError; use nalgebra::DMatrix; -extern crate alloc; +use super::normal_equation::SymmetricTripletMatrix; -/// A matrix in triplet format for sparse LDLT factorization / using faer crate -pub struct SparseLdlt { - upper_ccs: faer::sparse::SparseColMat, -} +extern crate alloc; -/// A matrix in triplet format -pub struct SymmetricTripletMatrix { - /// upper diagonal triplets - pub upper_triplets: alloc::vec::Vec<(usize, usize, f64)>, - /// row count (== column count) - pub size: usize, +/// sparse LDLT factorization parameters +#[derive(Copy, Clone, Debug)] +pub struct SparseLdltParams { + /// Regularization for LDLT factorization + pub regularization_eps: f64, } -impl SymmetricTripletMatrix { - /// Create an example matrix - pub fn example() -> Self { +impl Default for SparseLdltParams { + fn default() -> Self { Self { - upper_triplets: alloc::vec![ - (0, 0, 3.05631771), - (1, 1, 60.05631771), - (2, 2, 6.05631771), - (3, 3, 5.05631771), - (4, 4, 8.05631771), - (5, 5, 5.05631771), - (6, 6, 0.05631771), - (7, 7, 10.005631771), - (0, 1, 2.41883573), - (0, 3, 2.41883573), - (0, 5, 1.88585946), - (1, 3, 1.73897015), - (1, 5, 2.12387697), - (1, 7, 1.47609157), - (2, 7, 1.4541327), - (3, 4, 2.35666066), - (3, 7, 0.94642903), - ], - size: 8, + regularization_eps: 1e-6, } } +} - /// Convert to dense matrix - pub fn to_dense(&self) -> nalgebra::DMatrix { - let mut full_matrix = nalgebra::DMatrix::from_element(self.size, self.size, 0.0); - for &(row, col, value) in self.upper_triplets.iter() { - full_matrix[(row, col)] = value; - if row != col { - full_matrix[(col, row)] = value; - } - } - full_matrix - } - - /// Returns true if the matrix is semi-positive definite - /// - /// TODO: This is a very inefficient implementation - pub fn is_semi_positive_definite(&self) -> bool { - let full_matrix = self.to_dense(); - let eigen_comp = full_matrix.symmetric_eigen(); - eigen_comp.eigenvalues.iter().all(|&x| x >= 0.0) - } +/// A matrix in triplet format for sparse LDLT factorization / using faer crate +pub struct SimplicalSparseLdlt { + upper_ccs: faer::sparse::SparseColMat, + params: SparseLdltParams, } struct SparseLdltPerm { @@ -126,8 +88,10 @@ struct SparseLdltPermSymb { symbolic: faer::sparse::linalg::cholesky::simplicial::SymbolicSimplicialCholesky, } -impl SparseLdlt { - fn symbolic_and_perm(&self) -> SparseLdltPermSymb { +impl SimplicalSparseLdlt { + fn symbolic_and_perm(&self) -> Result { + // Following low-level example from faer crate: + // https://github.com/sarah-quinones/faer-rs/blob/main/src/sparse/linalg/cholesky.rs#L11 let dim = self.upper_ccs.ncols(); let nnz = self.upper_ccs.compute_nnz(); @@ -135,23 +99,22 @@ impl SparseLdlt { let (perm, perm_inv) = { let mut perm = alloc::vec::Vec::new(); let mut perm_inv = alloc::vec::Vec::new(); - perm.try_reserve_exact(dim).unwrap(); - perm_inv.try_reserve_exact(dim).unwrap(); + perm.try_reserve_exact(dim)?; + perm_inv.try_reserve_exact(dim)?; perm.resize(dim, 0usize); perm_inv.resize(dim, 0usize); - let mut mem = faer::dyn_stack::GlobalPodBuffer::try_new( - faer::sparse::linalg::amd::order_req::(dim, nnz).unwrap(), - ) - .unwrap(); + let mut mem = + faer::dyn_stack::GlobalPodBuffer::try_new(faer::sparse::linalg::amd::order_req::< + usize, + >(dim, nnz)?)?; faer::sparse::linalg::amd::order( &mut perm, &mut perm_inv, self.upper_ccs.symbolic(), faer::sparse::linalg::amd::Control::default(), faer::dyn_stack::PodStack::new(&mut mem), - ) - .unwrap(); + )?; (perm, perm_inv) }; @@ -164,19 +127,19 @@ impl SparseLdlt { faer::dyn_stack::StackReq::try_any_of([ faer::sparse::linalg::cholesky::simplicial::prefactorize_symbolic_cholesky_req::< usize, - >(dim, self.upper_ccs.compute_nnz()).unwrap(), + >(dim, self.upper_ccs.compute_nnz())?, faer::sparse::linalg::cholesky::simplicial::factorize_simplicial_symbolic_req::< usize, - >(dim).unwrap(), - ]).unwrap(), - ).unwrap(); + >(dim)?, + ])?, + )?; let mut stack = faer::dyn_stack::PodStack::new(&mut mem); let mut etree = alloc::vec::Vec::new(); let mut col_counts = alloc::vec::Vec::new(); - etree.try_reserve_exact(dim).unwrap(); + etree.try_reserve_exact(dim)?; etree.resize(dim, 0isize); - col_counts.try_reserve_exact(dim).unwrap(); + col_counts.try_reserve_exact(dim)?; col_counts.resize(dim, 0usize); faer::sparse::linalg::cholesky::simplicial::prefactorize_symbolic_cholesky( @@ -196,20 +159,19 @@ impl SparseLdlt { }, &col_counts, faer::reborrow::ReborrowMut::rb_mut(&mut stack), - ) - .unwrap() + )? }; - SparseLdltPermSymb { + Ok(SparseLdltPermSymb { permutation, symbolic, - } + }) } fn solve_from_symbolic( &self, b: &nalgebra::DVector, symbolic_perm: &SparseLdltPermSymb, - ) -> nalgebra::DVector { + ) -> Result, FaerError> { let dim = self.upper_ccs.ncols(); let perm_ref = unsafe { faer::perm::PermRef::new_unchecked( @@ -220,22 +182,30 @@ impl SparseLdlt { }; let symbolic = &symbolic_perm.symbolic; - let mut mem = faer::dyn_stack::GlobalPodBuffer::try_new(faer::dyn_stack::StackReq::try_all_of([ - faer::sparse::linalg::cholesky::simplicial::factorize_simplicial_numeric_ldlt_req::(dim).unwrap(), - faer::perm::permute_rows_in_place_req::(dim, 1).unwrap(), - symbolic.solve_in_place_req::(dim).unwrap(), - ]).unwrap()).unwrap(); + let mut mem = + faer::dyn_stack::GlobalPodBuffer::try_new(faer::dyn_stack::StackReq::try_all_of([ + faer::sparse::linalg::cholesky::simplicial::factorize_simplicial_numeric_ldlt_req::< + usize, + f64, + >(dim)?, + faer::perm::permute_rows_in_place_req::(dim, 1)?, + symbolic.solve_in_place_req::(dim)?, + ])?)?; let mut stack = faer::dyn_stack::PodStack::new(&mut mem); let mut l_values = alloc::vec::Vec::new(); - l_values.try_reserve_exact(symbolic.len_values()).unwrap(); + l_values.try_reserve_exact(symbolic.len_values())?; l_values.resize(symbolic.len_values(), 0.0f64); let ccs_perm_upper = symbolic_perm.permutation.perm_upper_ccs(&self.upper_ccs); faer::sparse::linalg::cholesky::simplicial::factorize_simplicial_numeric_ldlt::( &mut l_values, ccs_perm_upper.as_ref(), - faer::sparse::linalg::cholesky::LdltRegularization::default(), + faer::sparse::linalg::cholesky::LdltRegularization { + dynamic_regularization_signs: Some(&vec![1; dim]), + dynamic_regularization_delta: self.params.regularization_eps, + dynamic_regularization_epsilon: self.params.regularization_eps, + }, symbolic, faer::reborrow::ReborrowMut::rb_mut(&mut stack), ); @@ -265,18 +235,19 @@ impl SparseLdlt { faer::reborrow::ReborrowMut::rb_mut(&mut stack), ); - x + Ok(x) } /// Create a sparse LDLT factorization from a symmetric triplet matrix - pub fn from_triplets(sym_tri_mat: &SymmetricTripletMatrix) -> Self { - SparseLdlt { + pub fn from_triplets(sym_tri_mat: &SymmetricTripletMatrix, params: SparseLdltParams) -> Self { + SimplicalSparseLdlt { upper_ccs: faer::sparse::SparseColMat::try_new_from_triplets( sym_tri_mat.size, sym_tri_mat.size, - &sym_tri_mat.upper_triplets, + &sym_tri_mat.triplets, ) .unwrap(), + params, } } @@ -298,8 +269,8 @@ impl SparseLdlt { /// /// TODO: Consider a more efficient API where the symbolic factorization is /// computed once and then reused for multiple solves. - pub fn solve(&self, b: &nalgebra::DVector) -> nalgebra::DVector { - let symbolic_perm = self.symbolic_and_perm(); + pub fn solve(&self, b: &nalgebra::DVector) -> Result, FaerError> { + let symbolic_perm = self.symbolic_and_perm()?; self.solve_from_symbolic(b, &symbolic_perm) } } @@ -312,15 +283,15 @@ mod tests { #[test] fn ldlt() { - let sym_tri_mat = SymmetricTripletMatrix::example(); + let sym_tri_mat = SymmetricTripletMatrix::upper_diagonal_example(); assert!(sym_tri_mat.is_semi_positive_definite()); let dense_mat = sym_tri_mat.to_dense(); - let sparse_ldlt = SparseLdlt::from_triplets(&sym_tri_mat); + let sparse_ldlt = SimplicalSparseLdlt::from_triplets(&sym_tri_mat, Default::default()); assert_eq!(sparse_ldlt.to_dense(), dense_mat); let b = DVector::from_element(8, 1.0); - let x = sparse_ldlt.solve(&b); + let x = sparse_ldlt.solve(&b).unwrap(); approx::assert_abs_diff_eq!(dense_mat * x, b); }