Betteragain

Author: adamchalmers

kcl-ezpz/src/solver.rs

@@ -101,8 +101,8 @@ pub(crate) struct Model<'c> {
     lu_symbolic: SymbolicLu<usize>,
     jt_sym: SymbolicSparseColMat<usize>,
     jtj_symbolic: (SymbolicSparseColMat<usize>, SparseMatMulInfo),
-    a_sym: SymbolicSparseColMat<usize>,
     jt_value_indices: Vec<usize>,
+    lambda_diag_indices: Vec<usize>,
 }
 
 fn validate_variables(
@@ -219,8 +219,8 @@ impl<'c> Model<'c> {
         let lu_symbolic = Self::precompute_symbolic_lu(&jc.sym, &lambda_i)?;
         let jt_sym = jc.sym.transpose().to_col_major()?;
         let jtj_symbolic = Self::precompute_symbolic_jtj(&jt_sym, &jc.sym)?;
-        let a_sym = Self::precompute_symbolic_a(&jtj_symbolic.0, &lambda_i)?;
         let jt_value_indices = Self::precompute_jt_value_indices(&jc.sym, &jt_sym);
+        let lambda_diag_indices = Self::precompute_lambda_diag_indices(&jtj_symbolic.0);
 
         // All done.
         Ok(Self {
@@ -234,8 +234,8 @@ impl<'c> Model<'c> {
             lu_symbolic,
             jt_sym,
             jtj_symbolic,
-            a_sym,
             jt_value_indices,
+            lambda_diag_indices,
         })
     }
 
@@ -250,17 +250,6 @@ impl<'c> Model<'c> {
         Ok(jtj_sym)
     }
 
-    fn precompute_symbolic_a(
-        jtj_sym: &SymbolicSparseColMat<usize>,
-        lambda_i: &faer::sparse::SparseColMat<usize, f64>,
-    ) -> Result<SymbolicSparseColMat<usize>, NonLinearSystemError> {
-        // Any non-zero values will do; we only care about the sparsity pattern of JᵀJ + λI.
-        let ones = vec![1.0; jtj_sym.compute_nnz()];
-        let jtj = SparseColMatRef::new(jtj_sym.as_ref(), &ones);
-        let a = jtj + lambda_i;
-        Ok(a.symbolic().to_owned()?)
-    }
-
     fn precompute_jt_value_indices(
         jc_sym: &SymbolicSparseColMat<usize>,
         jt_sym: &SymbolicSparseColMat<usize>,
@@ -285,6 +274,19 @@ impl<'c> Model<'c> {
         indices
     }
 
+    fn precompute_lambda_diag_indices(jtj_sym: &SymbolicSparseColMat<usize>) -> Vec<usize> {
+        let mut diag = Vec::with_capacity(jtj_sym.ncols());
+        let row_idx = jtj_sym.row_idx();
+        for col in 0..jtj_sym.ncols() {
+            let mut col_range = jtj_sym.col_range(col);
+            let idx = col_range
+                .find(|idx| row_idx[*idx] == col)
+                .expect("diagonal must exist in J^T J");
+            diag.push(idx);
+        }
+        diag
+    }
+
     /// This is used in the core Newton solving, but it can be calculated entirely from
     /// the symbolic structure of the constraints. So let's do it here, before running
     /// the newton loop, to keep that loop fast.

kcl-ezpz/src/solver/newton.rs

@@ -6,7 +6,6 @@ use faer::{
     sparse::{
         SparseColMatMut, SparseColMatRef,
         linalg::{matmul, solvers::Lu},
-        ops,
     },
 };
 
@@ -32,7 +31,6 @@ impl Model<'_> {
         };
         let mut jtj_vals = vec![0.0; jtj_nnz];
         let mut jtj_mem = MemBuffer::new(jtj_scratch_req);
-        let mut a_vals = vec![0.0; self.a_sym.compute_nnz()];
         let mut jt_vals = vec![0.0; self.jt_sym.compute_nnz()];
 
         for this_iteration in 0..config.max_iterations {
@@ -79,26 +77,11 @@ impl Model<'_> {
                 get_global_parallelism(),
                 jtj_stack,
             );
-            let jtj = SparseColMatRef::new(jtj_sym.as_ref(), &jtj_vals);
-
-            a_vals.fill(0.0);
-            ops::binary_op_assign_into(
-                SparseColMatMut::new(self.a_sym.as_ref(), &mut a_vals),
-                jtj,
-                |dst, src| {
-                    *dst = *src.unwrap_or(&0.0);
-                },
-            );
-            ops::binary_op_assign_into(
-                SparseColMatMut::new(self.a_sym.as_ref(), &mut a_vals),
-                self.lambda_i.as_ref(),
-                |dst, src| {
-                    if let Some(val) = src {
-                        *dst += *val;
-                    }
-                },
-            );
-            let a = SparseColMatRef::new(self.a_sym.as_ref(), &a_vals);
+            let mut a_vals = jtj_vals.clone();
+            for (col, diag_idx) in self.lambda_diag_indices.iter().copied().enumerate() {
+                a_vals[diag_idx] += self.lambda_i.val_of_col(col)[0];
+            }
+            let a = SparseColMatRef::new(jtj_sym.as_ref(), &a_vals);
             let b = j.transpose() * -ColRef::from_slice(&global_residual);
 
             // Solve linear system