Precompute symbolic matrix to hold jacobian (#41)

* Pre-compute symbolic matrix for jacobians

* Remove reserve

* Formatting

* Try to fix pipeline

* Try to fix pipeline

* Post-review change

* Formatting

---------

Co-authored-by: Hossam Rabbouh <hrabbouh@Hossams-MacBook-Pro-2.local>

Author: jumpinthefire

.github/workflows/rust.yml

@@ -16,6 +16,8 @@ jobs:
 
     steps:
     - uses: actions/checkout@v4
+    - name: Dependencies
+      run: sudo apt install pkg-config libfreetype6-dev libfontconfig1-dev
     - name: Format
       run: cargo fmt --check --verbose
     - name: Linting

Cargo.toml

@@ -1,8 +1,8 @@
 [package]
 name = "tiny-solver"
-version = "0.14.2"
+version = "0.15.0"
 edition = "2021"
-authors = ["Powei Lin <poweilin1994@gmail.com>"]
+authors = ["Powei Lin <poweilin1994@gmail.com>, Hossam R. <hrabbouh@gmail.com>"]
 readme = "README.md"
 license = "MIT OR Apache-2.0"
 description = "Factor graph solver"

src/optimizer/gauss_newton_optimizer.rs

@@ -43,16 +43,29 @@ impl optimizer::Optimizer for GaussNewtonOptimizer {
             LinearSolverType::SparseQR => Box::new(linear::SparseQRSolver::new()),
         };
 
+        let symbolic_structure = problem.build_symbolic_structure(
+            &parameter_blocks,
+            total_variable_dimension,
+            &variable_name_to_col_idx_dict,
+        );
+
         let mut last_err: f64 = 1.0;
 
         for i in 0..opt_option.max_iteration {
+            let start = Instant::now();
             let (residuals, jac) = problem.compute_residual_and_jacobian(
                 &parameter_blocks,
                 &variable_name_to_col_idx_dict,
                 total_variable_dimension,
+                &symbolic_structure,
             );
             let current_error = residuals.norm_l2();
-            trace!("iter:{} total err:{}", i, current_error);
+            trace!(
+                "iter:{}, total err:{}, duration: {:?}",
+                i,
+                current_error,
+                start.elapsed()
+            );
 
             if current_error < opt_option.min_error_threshold {
                 trace!("error too low");

src/optimizer/levenberg_marquardt_optimizer.rs

@@ -67,6 +67,12 @@ impl optimizer::Optimizer for LevenbergMarquardtOptimizer {
         // With LM, rather than solving A * dx = b for dx, we solve for (A + lambda * diag(A)) dx = b.
         let mut jacobi_scaling_diagonal: Option<faer::sparse::SparseColMat<usize, f64>> = None;
 
+        let s = problem.build_symbolic_structure(
+            &parameter_blocks,
+            total_variable_dimension,
+            &variable_name_to_col_idx_dict,
+        );
+
         // Damping parameter (a.k.a lambda / Marquardt parameter)
         let mut u = 1.0 / self.initial_trust_region_radius;
 
@@ -76,6 +82,7 @@ impl optimizer::Optimizer for LevenbergMarquardtOptimizer {
                 &parameter_blocks,
                 &variable_name_to_col_idx_dict,
                 total_variable_dimension,
+                &s,
             );
 
             if i == 0 {
@@ -186,7 +193,7 @@ impl optimizer::Optimizer for LevenbergMarquardtOptimizer {
                 } else {
                     // If there's too much divergence, reduce the trust region and try again with the same parameters.
                     u *= 2.0;
-                    println!("u {}", u);
+                    trace!("u {}", u);
                 }
             } else {
                 log::debug!("solve ax=b failed");

src/problem.rs

@@ -1,7 +1,7 @@
 use std::collections::{HashMap, HashSet};
 use std::sync::{Arc, Mutex};
 
-use faer::sparse::SparseColMat;
+use faer::sparse::{SparseColMat, SymbolicSparseColMat, ValuesOrder};
 use faer_ext::IntoFaer;
 use nalgebra as na;
 use rayon::prelude::*;
@@ -24,8 +24,12 @@ impl Default for Problem {
     }
 }
 
-/// (col idx in matrix, row idx in matrix, value)
-type JacobianValue = (usize, usize, f64);
+pub struct SymbolicStructure {
+    pattern: SymbolicSparseColMat<usize>,
+    order: ValuesOrder<usize>,
+}
+
+type JacobianValue = f64;
 
 impl Problem {
     pub fn new() -> Problem {
@@ -39,6 +43,51 @@ impl Problem {
         }
     }
 
+    pub fn build_symbolic_structure(
+        &self,
+        parameter_blocks: &HashMap<String, ParameterBlock>,
+        total_variable_dimension: usize,
+        variable_name_to_col_idx_dict: &HashMap<String, usize>,
+    ) -> SymbolicStructure {
+        let mut indices = Vec::<(usize, usize)>::new();
+
+        self.residual_blocks.iter().for_each(|(_, residual_block)| {
+            let mut variable_local_idx_size_list = Vec::<(usize, usize)>::new();
+            let mut count_variable_local_idx: usize = 0;
+            for var_key in &residual_block.variable_key_list {
+                if let Some(param) = parameter_blocks.get(var_key) {
+                    variable_local_idx_size_list
+                        .push((count_variable_local_idx, param.tangent_size()));
+                    count_variable_local_idx += param.tangent_size();
+                };
+            }
+            for (i, var_key) in residual_block.variable_key_list.iter().enumerate() {
+                if let Some(variable_global_idx) = variable_name_to_col_idx_dict.get(var_key) {
+                    let (_, var_size) = variable_local_idx_size_list[i];
+                    for row_idx in 0..residual_block.dim_residual {
+                        for col_idx in 0..var_size {
+                            let global_row_idx = residual_block.residual_row_start_idx + row_idx;
+                            let global_col_idx = variable_global_idx + col_idx;
+                            indices.push((global_row_idx, global_col_idx));
+                        }
+                    }
+                }
+            }
+        });
+        let start = std::time::Instant::now();
+        let (s, o) = SymbolicSparseColMat::try_new_from_indices(
+            self.total_residual_dimension,
+            total_variable_dimension,
+            &indices,
+        )
+        .unwrap();
+        log::trace!("Built symbolic matrix: {:?}", start.elapsed());
+        SymbolicStructure {
+            pattern: s,
+            order: o,
+        }
+    }
+
     pub fn get_variable_name_to_col_idx_dict(
         &self,
         parameter_blocks: &HashMap<String, ParameterBlock>,
@@ -147,16 +196,14 @@ impl Problem {
         let total_residual = Arc::new(Mutex::new(na::DVector::<f64>::zeros(
             self.total_residual_dimension,
         )));
-        self.residual_blocks
-            .par_iter()
-            .for_each(|(_, residual_block)| {
-                self.compute_residual_impl(
-                    residual_block,
-                    parameter_blocks,
-                    &total_residual,
-                    with_loss_fn,
-                )
-            });
+        self.residual_blocks.iter().for_each(|(_, residual_block)| {
+            self.compute_residual_impl(
+                residual_block,
+                parameter_blocks,
+                &total_residual,
+                with_loss_fn,
+            )
+        });
         let total_residual = Arc::try_unwrap(total_residual)
             .unwrap()
             .into_inner()
@@ -170,42 +217,46 @@ impl Problem {
         parameter_blocks: &HashMap<String, ParameterBlock>,
         variable_name_to_col_idx_dict: &HashMap<String, usize>,
         total_variable_dimension: usize,
+        symbolic_structure: &SymbolicStructure,
     ) -> (faer::Mat<f64>, SparseColMat<usize, f64>) {
         // multi
         let total_residual = Arc::new(Mutex::new(na::DVector::<f64>::zeros(
             self.total_residual_dimension,
         )));
-        let jacobian_list = Arc::new(Mutex::new(Vec::<JacobianValue>::new()));
 
-        self.residual_blocks
+        let mut start = std::time::Instant::now();
+        let jacobian_lists: Vec<JacobianValue> = self
+            .residual_blocks
             .par_iter()
-            .for_each(|(_, residual_block)| {
+            .map(|(_, residual_block)| {
                 self.compute_residual_and_jacobian_impl(
                     residual_block,
                     parameter_blocks,
                     variable_name_to_col_idx_dict,
                     &total_residual,
-                    &jacobian_list,
                 )
-            });
+            })
+            .flatten()
+            .collect();
 
+        log::debug!("compute_residual_and_jacobian_impl: {:?}, total_variable_dim: {:?}, total_residual: {:?}", start.elapsed(), total_variable_dimension, self.total_residual_dimension);
+        start = std::time::Instant::now();
         let total_residual = Arc::try_unwrap(total_residual)
             .unwrap()
             .into_inner()
             .unwrap();
-        let jacobian_list = Arc::try_unwrap(jacobian_list)
-            .unwrap()
-            .into_inner()
-            .unwrap();
-        // end
 
         let residual_faer = total_residual.view_range(.., ..).into_faer().to_owned();
-        let jacobian_faer = SparseColMat::try_new_from_triplets(
-            self.total_residual_dimension,
-            total_variable_dimension,
-            &jacobian_list,
+        log::debug!("residual_faer: {:?}", start.elapsed());
+        start = std::time::Instant::now();
+        let jacobian_faer = SparseColMat::new_from_order_and_values(
+            symbolic_structure.pattern.clone(),
+            &symbolic_structure.order,
+            jacobian_lists.as_slice(),
         )
         .unwrap();
+        log::debug!("jacobian_faer: {:?}", start.elapsed());
+        //log::debug!("rest of the function: {:?}", start.elapsed());
         (residual_faer, jacobian_faer)
     }
 
@@ -241,8 +292,7 @@ impl Problem {
         parameter_blocks: &HashMap<String, ParameterBlock>,
         variable_name_to_col_idx_dict: &HashMap<String, usize>,
         total_residual: &Arc<Mutex<na::DVector<f64>>>,
-        jacobian_list: &Arc<Mutex<Vec<JacobianValue>>>,
-    ) {
+    ) -> Vec<JacobianValue> {
         let mut params = Vec::new();
         let mut variable_local_idx_size_list = Vec::<(usize, usize)>::new();
         let mut count_variable_local_idx: usize = 0;
@@ -264,22 +314,25 @@ impl Problem {
                 .copy_from(&res);
         }
 
+        let mut local_jacobian_list = Vec::new();
+
         for (i, var_key) in residual_block.variable_key_list.iter().enumerate() {
-            if let Some(variable_global_idx) = variable_name_to_col_idx_dict.get(var_key) {
+            if variable_name_to_col_idx_dict.contains_key(var_key) {
                 let (variable_local_idx, var_size) = variable_local_idx_size_list[i];
                 let variable_jac = jac.view((0, variable_local_idx), (jac.shape().0, var_size));
-                let mut local_jacobian_list = Vec::new();
                 for row_idx in 0..jac.shape().0 {
                     for col_idx in 0..var_size {
-                        let global_row_idx = residual_block.residual_row_start_idx + row_idx;
-                        let global_col_idx = variable_global_idx + col_idx;
-                        let value = variable_jac[(row_idx, col_idx)];
-                        local_jacobian_list.push((global_row_idx, global_col_idx, value));
+                        local_jacobian_list.push(variable_jac[(row_idx, col_idx)]);
                     }
                 }
-                let mut jacobian_list = jacobian_list.lock().unwrap();
-                jacobian_list.extend(local_jacobian_list);
+            } else {
+                panic!(
+                    "Missing key {} in variable-to-column-index mapping",
+                    var_key
+                );
             }
         }
+
+        local_jacobian_list
     }
 }

tests/test_problem.rs

@@ -120,11 +120,17 @@ mod tests {
         let variable_name_to_col_idx_dict =
             problem.get_variable_name_to_col_idx_dict(&parameter_blocks);
         let total_variable_dimension = parameter_blocks.values().map(|p| p.tangent_size()).sum();
+        let symbolic_structure = problem.build_symbolic_structure(
+            &parameter_blocks,
+            total_variable_dimension,
+            &variable_name_to_col_idx_dict,
+        );
 
         let (residuals, jac) = problem.compute_residual_and_jacobian(
             &parameter_blocks,
             &variable_name_to_col_idx_dict,
             total_variable_dimension,
+            &symbolic_structure,
         );
 
         assert_eq!(residuals.nrows(), 3);