1use faer::{
4 Mat,
5 sparse::{SparseColMat, SymbolicSparseColMat},
6};
7use rayon::prelude::*;
8use slotmap::{SecondaryMap, SlotMap};
9
10use crate::core::VarKey;
11use crate::error::ErrorLogging;
12use crate::linearizer::{
13 AssemblyWorkspace, LinearizerError, LinearizerResult, compute_block_into,
14 split_by_row_offsets_mut,
15};
16
17use crate::core::problem::Problem;
18use crate::core::variable::ManifoldVariable;
19
20pub struct SymbolicStructure {
22 pub pattern: SymbolicSparseColMat<usize>,
23 pub scatter_ops: Vec<ScatterOp>,
30 pub scatter_has_duplicates: bool,
33}
34
35#[derive(Debug, Clone, Copy)]
37pub struct ScatterOp {
38 pub dest: u32,
39 pub src: u32,
40 pub accumulate: bool,
41}
42
43fn free_local_columns(
48 variable: &dyn ManifoldVariable,
49 var_size: usize,
50) -> impl Iterator<Item = (usize, usize)> + '_ {
51 (0..var_size).filter_map(move |col| variable.local_free_offset(col).map(|free| (col, free)))
52}
53
54pub fn build_symbolic_structure(
66 problem: &Problem,
67 variables: &SlotMap<VarKey, Box<dyn ManifoldVariable>>,
68 variable_index_map: &SecondaryMap<VarKey, usize>,
69 total_dof: usize,
70) -> LinearizerResult<SymbolicStructure> {
71 let mut triples = Vec::<(usize, usize, u32)>::new();
76 let mut total_jac_len = 0usize;
77
78 let mut blocks: Vec<_> = problem.residual_blocks().iter().collect();
79 blocks.sort_by_key(|(_, block)| block.residual_row_start_idx);
80 for (_, block) in &blocks {
81 let residual_dim = block.factor.residual_dim();
82 let (jac_rows, jac_cols) = block.factor.jacobian_shape();
83 debug_assert_eq!(jac_rows, residual_dim);
84 let jac_base = total_jac_len;
85 total_jac_len += jac_rows * jac_cols;
86
87 let mut var_local_sizes = Vec::<(usize, usize)>::new();
88 let mut local_offset = 0;
89
90 for &var_key in &block.variable_keys {
91 if let Some(variable) = variables.get(var_key) {
92 var_local_sizes.push((local_offset, variable.dof()));
93 local_offset += variable.dof();
94 }
95 }
96
97 for (i, &var_key) in block.variable_keys.iter().enumerate() {
98 if let Some(variable) = variables.get(var_key) {
99 let Some(&global_col) = variable_index_map.get(var_key) else {
100 return Err(LinearizerError::Variable(format!(
101 "VarKey {var_key:?} missing in variable-to-column-index mapping"
102 ))
103 .log());
104 };
105 let Some(&(local_col, var_size)) = var_local_sizes.get(i) else {
106 continue;
107 };
108 let free_cols: smallvec::SmallVec<[(usize, usize); 16]> =
109 free_local_columns(variable.as_ref(), var_size).collect();
110 triples.reserve(residual_dim * free_cols.len());
111 for row in 0..residual_dim {
112 for &(col, free_col) in &free_cols {
113 triples.push((
116 global_col + free_col,
117 block.residual_row_start_idx + row,
118 (jac_base + (local_col + col) * residual_dim + row) as u32,
119 ));
120 }
121 }
122 }
123 }
124 }
125
126 let plan = csc_from_triples(triples, total_dof)?;
127 let pattern = SymbolicSparseColMat::new_checked(
128 problem.total_residual_dimension,
129 total_dof,
130 plan.col_ptr,
131 None,
132 plan.row_idx,
133 );
134
135 Ok(SymbolicStructure {
136 pattern,
137 scatter_ops: plan.scatter_ops,
138 scatter_has_duplicates: plan.has_duplicates,
139 })
140}
141
142struct CscPlan {
145 col_ptr: Vec<usize>,
146 row_idx: Vec<usize>,
147 scatter_ops: Vec<ScatterOp>,
148 has_duplicates: bool,
149}
150
151fn csc_from_triples(
152 triples: Vec<(usize, usize, u32)>,
153 total_dof: usize,
154) -> LinearizerResult<CscPlan> {
155 let mut bucket = vec![0usize; total_dof + 1];
163 for &(c, _, _) in &triples {
164 let Some(count) = bucket.get_mut(c + 1) else {
165 return Err(LinearizerError::SymbolicStructure(format!(
166 "Jacobian column {c} outside the {total_dof} solve columns"
167 ))
168 .log());
169 };
170 *count += 1;
171 }
172 for c in 0..total_dof {
173 bucket[c + 1] += bucket[c];
174 }
175 let mut next = bucket.clone();
176 let mut by_column = vec![(0usize, 0u32); triples.len()];
177 for &(c, r, src) in &triples {
178 by_column[next[c]] = (r, src);
179 next[c] += 1;
180 }
181 drop(triples);
182
183 let mut col_ptr = vec![0usize; total_dof + 1];
184 let mut row_idx = Vec::with_capacity(by_column.len());
185 let mut scatter_ops = Vec::<ScatterOp>::with_capacity(by_column.len());
186 let mut has_duplicates = false;
187 let mut last_dest = 0u32;
188 for c in 0..total_dof {
189 let entries = &mut by_column[bucket[c]..bucket[c + 1]];
190 if !entries.is_sorted_by_key(|&(r, _)| r) {
191 entries.sort_by_key(|&(r, _)| r);
192 }
193 let mut prev: Option<usize> = None;
194 for &(r, src) in entries.iter() {
195 if prev == Some(r) {
196 has_duplicates = true;
199 scatter_ops.push(ScatterOp {
200 dest: last_dest,
201 src,
202 accumulate: true,
203 });
204 } else {
205 row_idx.push(r);
206 last_dest = row_idx.len() as u32 - 1;
207 scatter_ops.push(ScatterOp {
208 dest: last_dest,
209 src,
210 accumulate: false,
211 });
212 }
213 prev = Some(r);
214 }
215 col_ptr[c + 1] = row_idx.len();
216 }
217
218 Ok(CscPlan {
219 col_ptr,
220 row_idx,
221 scatter_ops,
222 has_duplicates,
223 })
224}
225
226pub fn assemble_sparse(
231 problem: &Problem,
232 variables: &SlotMap<VarKey, Box<dyn ManifoldVariable>>,
233 _variable_index_map: &SecondaryMap<VarKey, usize>,
234 symbolic_structure: &SymbolicStructure,
235 workspace: &mut AssemblyWorkspace,
236) -> LinearizerResult<(Mat<f64>, SparseColMat<usize, f64>)> {
237 let total_nnz = symbolic_structure.pattern.compute_nnz();
238
239 workspace.residual_buf.fill(0.0);
242 let residual_slices =
243 split_by_row_offsets_mut(&mut workspace.residual_buf, &workspace.offsets_lens);
244 let jac_slices = split_by_row_offsets_mut(&mut workspace.jac_arena, &workspace.jac_offsets);
245 let residual_blocks = problem.residual_blocks();
246
247 residual_slices
250 .into_par_iter()
251 .zip(jac_slices)
252 .zip(workspace.block_order.par_iter())
253 .map(|((res_slice, jac_buf), key)| {
254 let block = &residual_blocks[*key];
255 jac_buf.fill(0.0);
256 compute_block_into(block, variables, res_slice, Some(jac_buf)).map(|_| ())
257 })
258 .collect::<LinearizerResult<()>>()?;
259
260 let mut csc_values = std::mem::take(&mut workspace.jacobian_values);
263 csc_values.resize(total_nnz, 0.0);
264 if symbolic_structure.scatter_has_duplicates {
265 for op in &symbolic_structure.scatter_ops {
266 let v = workspace.jac_arena[op.src as usize];
267 if op.accumulate {
268 csc_values[op.dest as usize] += v;
269 } else {
270 csc_values[op.dest as usize] = v;
271 }
272 }
273 } else {
274 csc_values
275 .par_iter_mut()
276 .zip(symbolic_structure.scatter_ops.par_iter())
277 .for_each(|(v, op)| *v = workspace.jac_arena[op.src as usize]);
278 }
279 let jacobian_sparse = SparseColMat::new(symbolic_structure.pattern.clone(), csc_values);
280
281 workspace.jacobian_values = Vec::with_capacity(total_nnz);
284
285 let n = problem.total_residual_dimension;
287 let residual_faer = faer::Mat::from_fn(n, 1, |i, _| workspace.residual_buf[i]);
288
289 Ok((residual_faer, jacobian_sparse))
290}
291
292#[cfg(test)]
293mod tests {
294 use super::*;
295 use crate::{core::problem::Problem, factors, linalg::JacobianMode};
296 use apex_manifolds::ManifoldType;
297 use faer::prelude::ReborrowMut;
298 use nalgebra::dvector;
299
300 type TestResult = Result<(), Box<dyn std::error::Error>>;
301
302 struct LinearFactor {
303 target: f64,
304 }
305
306 impl factors::Factor for LinearFactor {
307 fn linearize(
308 &self,
309 params: &[&[f64]],
310 residual: &mut [f64],
311 jacobian: Option<faer::mat::MatMut<'_, f64>>,
312 ) {
313 residual[0] = params[0][0] - self.target;
314 if let Some(mut jac) = jacobian {
315 *jac.rb_mut().get_mut(0, 0) = 1.0;
316 }
317 }
318 fn residual_dim(&self) -> usize {
319 1
320 }
321 fn jacobian_shape(&self) -> (usize, usize) {
322 (1, 1)
323 }
324 }
325
326 fn one_var_problem() -> (Problem, VarKey) {
327 let mut problem = Problem::new(JacobianMode::Sparse);
328 let k = problem.add_variable(ManifoldType::RN, dvector![5.0]);
329 problem.add_residual_block(&[k], Box::new(LinearFactor { target: 0.0 }), None);
330 (problem, k)
331 }
332
333 struct ScaledLinearFactor {
336 target: f64,
337 gain: f64,
338 }
339
340 impl factors::Factor for ScaledLinearFactor {
341 fn linearize(
342 &self,
343 params: &[&[f64]],
344 residual: &mut [f64],
345 jacobian: Option<faer::mat::MatMut<'_, f64>>,
346 ) {
347 residual[0] = params[0][0] - self.target;
348 if let Some(mut jac) = jacobian {
349 *jac.rb_mut().get_mut(0, 0) = self.gain;
350 }
351 }
352 fn residual_dim(&self) -> usize {
353 1
354 }
355 fn jacobian_shape(&self) -> (usize, usize) {
356 (1, 1)
357 }
358 }
359
360 fn build_index_map(problem: &Problem) -> (SecondaryMap<VarKey, usize>, usize) {
361 let mut map = SecondaryMap::new();
362 let mut offset = 0;
363 for (k, v) in &problem.variables {
364 map.insert(k, offset);
365 offset += v.free_dof();
366 }
367 (map, offset)
368 }
369
370 #[test]
371 fn test_build_symbolic_structure_nnz() -> TestResult {
372 let (problem, _) = one_var_problem();
373 let (index_map, total_dof) = build_index_map(&problem);
374 let sym = build_symbolic_structure(&problem, &problem.variables, &index_map, total_dof)?;
375 assert_eq!(sym.pattern.compute_nnz(), 1);
376 Ok(())
377 }
378
379 #[test]
380 fn test_build_symbolic_structure_dimensions() -> TestResult {
381 let (problem, _) = one_var_problem();
382 let (index_map, total_dof) = build_index_map(&problem);
383 let sym = build_symbolic_structure(&problem, &problem.variables, &index_map, total_dof)?;
384 assert_eq!(sym.pattern.nrows(), 1);
385 assert_eq!(sym.pattern.ncols(), 1);
386 Ok(())
387 }
388
389 #[test]
390 fn test_build_symbolic_structure_two_factors() -> TestResult {
391 let mut problem = Problem::new(JacobianMode::Sparse);
392 let k = problem.add_variable(ManifoldType::RN, dvector![5.0]);
393 problem.add_residual_block(&[k], Box::new(LinearFactor { target: 0.0 }), None);
394 problem.add_residual_block(&[k], Box::new(LinearFactor { target: 1.0 }), None);
395 let (index_map, total_dof) = build_index_map(&problem);
396 let sym = build_symbolic_structure(&problem, &problem.variables, &index_map, total_dof)?;
397 assert_eq!(sym.pattern.compute_nnz(), 2);
398 Ok(())
399 }
400
401 #[test]
413 fn test_grouped_sparse_assembly_matches_dense() -> TestResult {
414 let mut problem = Problem::new(JacobianMode::Sparse);
415 let c = problem.add_variable(ManifoldType::RN, dvector![3.0]);
416 let b = problem.add_variable(ManifoldType::RN, dvector![2.0]);
417 let a = problem.add_variable(ManifoldType::RN, dvector![1.0]);
418 for (key, target, gain) in [(c, 30.0, 3.0), (b, 20.0, 2.0), (a, 10.0, 1.0)] {
420 problem.add_residual_block(&[key], Box::new(ScaledLinearFactor { target, gain }), None);
421 }
422 problem.mark_for_elimination(c);
423 assert!(
424 problem.group_rows_for_elimination(),
425 "insertion order must differ from grouped order for this test"
426 );
427
428 let (index_map, total_dof) = build_index_map(&problem);
429 let sym = build_symbolic_structure(&problem, &problem.variables, &index_map, total_dof)?;
430 let mut workspace = AssemblyWorkspace::build(&problem);
431 let (r_sparse, j_sparse) = assemble_sparse(
432 &problem,
433 &problem.variables,
434 &index_map,
435 &sym,
436 &mut workspace,
437 )?;
438 let (r_dense, j_dense) = crate::linearizer::cpu::dense::assemble_dense(
439 &problem,
440 &problem.variables,
441 &index_map,
442 total_dof,
443 &mut workspace,
444 )?;
445
446 assert_eq!(r_sparse.nrows(), 3);
447 for i in 0..3 {
448 assert!(
449 (r_sparse[(i, 0)] - r_dense[(i, 0)]).abs() < 1e-12,
450 "residual row {i} differs"
451 );
452 }
453 assert_eq!(j_sparse.as_ref().compute_nnz(), 3);
457 for col in 0..total_dof {
458 let rows = j_sparse.symbolic().row_idx_of_col_raw(col);
459 let vals = j_sparse.val_of_col(col);
460 assert_eq!(rows.len(), 1, "col {col} must hold exactly one entry");
461 assert!(
462 (vals[0] - j_dense[(rows[0], col)]).abs() < 1e-12,
463 "sparse value {} disagrees with dense {} at ({}, {col})",
464 vals[0],
465 j_dense[(rows[0], col)],
466 rows[0]
467 );
468 for row in 0..3 {
469 if row != rows[0] {
470 assert!(
471 j_dense[(row, col)].abs() < 1e-12,
472 "dense has an unexpected entry at ({row}, {col})"
473 );
474 }
475 }
476 }
477 Ok(())
478 }
479
480 #[test]
481 fn test_assemble_sparse_basic() -> TestResult {
482 let (problem, _) = one_var_problem();
483 let (index_map, total_dof) = build_index_map(&problem);
484 let sym = build_symbolic_structure(&problem, &problem.variables, &index_map, total_dof)?;
485 let (residual, _) = assemble_sparse(
486 &problem,
487 &problem.variables,
488 &index_map,
489 &sym,
490 &mut AssemblyWorkspace::build(&problem),
491 )?;
492 assert!((residual[(0, 0)] - 5.0).abs() < 1e-12);
493 Ok(())
494 }
495
496 #[test]
497 fn test_assemble_sparse_jacobian_value() -> TestResult {
498 let (problem, _) = one_var_problem();
499 let (index_map, total_dof) = build_index_map(&problem);
500 let sym = build_symbolic_structure(&problem, &problem.variables, &index_map, total_dof)?;
501 let (_, jacobian) = assemble_sparse(
502 &problem,
503 &problem.variables,
504 &index_map,
505 &sym,
506 &mut AssemblyWorkspace::build(&problem),
507 )?;
508 let val = jacobian.as_ref().val_of_col(0)[0];
509 assert!((val - 1.0).abs() < 1e-12);
510 Ok(())
511 }
512
513 #[test]
514 fn test_assemble_sparse_zero_residual() -> TestResult {
515 let mut problem = Problem::new(JacobianMode::Sparse);
516 let k = problem.add_variable(ManifoldType::RN, dvector![3.0]);
517 problem.add_residual_block(&[k], Box::new(LinearFactor { target: 3.0 }), None);
518 let (index_map, total_dof) = build_index_map(&problem);
519 let sym = build_symbolic_structure(&problem, &problem.variables, &index_map, total_dof)?;
520 let (residual, _) = assemble_sparse(
521 &problem,
522 &problem.variables,
523 &index_map,
524 &sym,
525 &mut AssemblyWorkspace::build(&problem),
526 )?;
527 assert!(residual[(0, 0)].abs() < 1e-12);
528 Ok(())
529 }
530
531 #[test]
532 fn test_assemble_sparse_dimensions() -> TestResult {
533 let (problem, _) = one_var_problem();
534 let (index_map, total_dof) = build_index_map(&problem);
535 let sym = build_symbolic_structure(&problem, &problem.variables, &index_map, total_dof)?;
536 let (residual, jacobian) = assemble_sparse(
537 &problem,
538 &problem.variables,
539 &index_map,
540 &sym,
541 &mut AssemblyWorkspace::build(&problem),
542 )?;
543 assert_eq!(residual.nrows(), 1);
544 assert_eq!(residual.ncols(), 1);
545 assert_eq!(jacobian.nrows(), 1);
546 assert_eq!(jacobian.ncols(), 1);
547 Ok(())
548 }
549
550 #[test]
551 fn test_assemble_sparse_two_variables() -> TestResult {
552 let mut problem = Problem::new(JacobianMode::Sparse);
553 let kx = problem.add_variable(ManifoldType::RN, dvector![2.0]);
554 let ky = problem.add_variable(ManifoldType::RN, dvector![7.0]);
555 problem.add_residual_block(&[kx], Box::new(LinearFactor { target: 0.0 }), None);
556 problem.add_residual_block(&[ky], Box::new(LinearFactor { target: 0.0 }), None);
557 let (index_map, total_dof) = build_index_map(&problem);
558 let sym = build_symbolic_structure(&problem, &problem.variables, &index_map, total_dof)?;
559 let (residual, _) = assemble_sparse(
560 &problem,
561 &problem.variables,
562 &index_map,
563 &sym,
564 &mut AssemblyWorkspace::build(&problem),
565 )?;
566 assert_eq!(residual.nrows(), 2);
567 let rsum = residual[(0, 0)].abs() + residual[(1, 0)].abs();
568 assert!((rsum - 9.0).abs() < 1e-12);
569 Ok(())
570 }
571
572 #[test]
573 fn test_assemble_sparse_missing_variable_key_returns_error() -> TestResult {
574 let (problem, _) = one_var_problem();
575 let (_, total_dof) = build_index_map(&problem);
576 let empty: SecondaryMap<VarKey, usize> = SecondaryMap::new();
580 let result = build_symbolic_structure(&problem, &problem.variables, &empty, total_dof);
581 assert!(result.is_err(), "expected Err for missing variable key");
582 Ok(())
583 }
584
585 type PlanParts = (Vec<usize>, Vec<usize>, Vec<(u32, u32, bool)>, bool);
587
588 fn csc_by_stable_sort(mut triples: Vec<(usize, usize, u32)>, total_dof: usize) -> PlanParts {
591 triples.sort_by_key(|&(c, r, _)| (c, r));
592 let mut col_ptr = vec![0usize; total_dof + 1];
593 let mut row_idx = Vec::new();
594 let mut ops = Vec::new();
595 let mut dup = false;
596 let mut prev = None;
597 let mut last = 0u32;
598 for &(c, r, src) in &triples {
599 if prev == Some((c, r)) {
600 dup = true;
601 ops.push((last, src, true));
602 } else {
603 col_ptr[c + 1] += 1;
604 row_idx.push(r);
605 last = row_idx.len() as u32 - 1;
606 ops.push((last, src, false));
607 }
608 prev = Some((c, r));
609 }
610 for c in 0..total_dof {
611 col_ptr[c + 1] += col_ptr[c];
612 }
613 (col_ptr, row_idx, ops, dup)
614 }
615
616 fn plan_parts(plan: CscPlan) -> PlanParts {
617 let ops = plan
618 .scatter_ops
619 .iter()
620 .map(|op| (op.dest, op.src, op.accumulate))
621 .collect();
622 (plan.col_ptr, plan.row_idx, ops, plan.has_duplicates)
623 }
624
625 #[test]
630 fn counting_sort_plan_matches_stable_comparison_sort() -> TestResult {
631 let total_dof = 7;
632 let mut ordered = Vec::new();
633 let mut src = 0u32;
634 for row in 0..40usize {
635 for col in [row % 7, (row * 3 + 1) % 7, (row * 5 + 2) % 7] {
636 ordered.push((col, row, src));
637 src += 1;
638 if row % 6 == 0 {
639 ordered.push((col, row, src));
641 src += 1;
642 }
643 }
644 }
645 let mut shuffled = ordered.clone();
646 let n = shuffled.len();
647 for i in 0..n {
648 shuffled.swap(i, (i * 7919 + 13) % n);
649 }
650 for triples in [ordered, shuffled] {
651 let want = csc_by_stable_sort(triples.clone(), total_dof);
652 let got = plan_parts(csc_from_triples(triples, total_dof)?);
653 assert_eq!(got, want);
654 }
655 assert!(csc_from_triples(vec![(7, 0, 0)], total_dof).is_err());
656 Ok(())
657 }
658}