1use crate::core::VarKey;
97use crate::core::variable::ManifoldVariable;
98use crate::error::ErrorLogging;
99use crate::linalg::regularization::invert_with_retry_dyn;
100use crate::linalg::schur::jacobian_ops::{column_dot, diag_jt_j};
101use crate::linalg::schur::{
102 BlockSpan, EliminatedBlocks, PcgParams, SchurOrdering, SchurPartition, SchurPreconditioner,
103 effective_landmark_keys, jt_j_vec_product, jt_vec, solve_pcg,
104};
105use crate::linalg::sparse::pattern;
106use crate::linalg::{Damping, LinAlgError, LinAlgResult, LinearSolver, SparseMode, StructureAware};
107use faer::Mat;
108use faer::sparse::SparseColMat;
109use nalgebra::DMatrix;
110use rayon::prelude::*;
111use slotmap::{SecondaryMap, SlotMap};
112use tracing::debug;
113
114#[derive(Debug, Clone, Default)]
117struct StructureCache {
118 kept_cols: Vec<usize>,
125 eliminated_cols: Vec<usize>,
126 eliminated_rows: Vec<Vec<usize>>,
128 visibility: Vec<Vec<usize>>,
130 chunk_rows: usize,
132 kept_spans: Vec<Vec<ColumnSpan>>,
140 eliminated_spans: Vec<Vec<ColumnSpan>>,
142 eliminated_block_of: Vec<u32>,
144 fingerprint: Option<pattern::PatternFingerprint>,
146}
147
148#[derive(Debug, Clone, Copy)]
154struct ColumnSpan {
155 lo: u32,
156 hi: u32,
157 local: u32,
158}
159
160fn row_chunk_len(nrows: usize) -> usize {
163 const MIN_CHUNK_ROWS: usize = 4096;
164 let chunks = rayon::current_num_threads().max(1) * 8;
165 nrows.div_ceil(chunks).max(MIN_CHUNK_ROWS)
166}
167
168struct RowSpans {
170 kept: Vec<Vec<ColumnSpan>>,
171 eliminated: Vec<Vec<ColumnSpan>>,
172 block_of: Vec<u32>,
173}
174
175impl RowSpans {
176 fn build(
179 jacobian: &SparseColMat<usize, f64>,
180 partition: &SchurPartition,
181 chunk_rows: usize,
182 ) -> LinAlgResult<Self> {
183 let chunks = jacobian.nrows().div_ceil(chunk_rows);
184 let mut kept = vec![Vec::new(); chunks];
185 for block in partition.kept_blocks() {
186 for offset in 0..block.dof {
187 let col = block.col_start + offset;
188 if let Some(local) = partition.kept_local(col) {
189 push_column_spans(jacobian, col, local, chunk_rows, &mut kept)?;
190 }
191 }
192 }
193 let mut eliminated = vec![Vec::new(); chunks];
194 let mut block_of = vec![0u32; partition.eliminated_dof()];
195 for (block_idx, block) in partition.eliminated_blocks().iter().enumerate() {
196 let owner = u32::try_from(block_idx).map_err(|_| {
197 LinAlgError::InvalidInput(format!(
198 "{block_idx} eliminated blocks exceed the implicit Schur operator's u32 range"
199 ))
200 .log()
201 })?;
202 let base = partition.eliminated_offset(block_idx);
203 for offset in 0..block.dof {
204 block_of[base + offset] = owner;
205 push_column_spans(
206 jacobian,
207 block.col_start + offset,
208 base + offset,
209 chunk_rows,
210 &mut eliminated,
211 )?;
212 }
213 }
214 Ok(Self {
215 kept,
216 eliminated,
217 block_of,
218 })
219 }
220}
221
222fn push_column_spans(
224 jacobian: &SparseColMat<usize, f64>,
225 col: usize,
226 local: usize,
227 chunk_rows: usize,
228 spans: &mut [Vec<ColumnSpan>],
229) -> LinAlgResult<()> {
230 let symbolic = jacobian.symbolic();
231 let start = symbolic.col_range(col).start;
232 let rows = symbolic.row_idx_of_col_raw(col);
233 let narrow = |x: usize| {
234 u32::try_from(x).map_err(|_| {
235 LinAlgError::InvalidInput(format!(
236 "Jacobian index {x} exceeds the implicit Schur operator's u32 range"
237 ))
238 .log()
239 })
240 };
241 let local = narrow(local)?;
242 let mut k = 0;
243 while let Some(&first_row) = rows.get(k) {
244 let chunk = first_row / chunk_rows;
245 let chunk_end = (chunk + 1) * chunk_rows;
246 let hi = k + rows[k..].partition_point(|&r| r < chunk_end);
247 let slot = spans.get_mut(chunk).ok_or_else(|| {
248 LinAlgError::InvalidInput(format!("row {first_row} outside the Jacobian")).log()
249 })?;
250 slot.push(ColumnSpan {
251 lo: narrow(start + k)?,
252 hi: narrow(start + hi)?,
253 local,
254 });
255 k = hi;
256 }
257 Ok(())
258}
259
260#[derive(Debug, Clone)]
265pub struct ImplicitSparseSchur {
266 partition: Option<SchurPartition>,
267 eliminated: EliminatedBlocks,
269 max_eliminated_dof: usize,
271 ordering: SchurOrdering,
274
275 max_cg_iterations: usize,
277 cg_tolerance: f64,
278 cg_q_tolerance: f64,
280
281 preconditioner_type: SchurPreconditioner,
283
284 jacobian: Option<SparseColMat<usize, f64>>,
291 jacobian_pattern: Option<pattern::PatternFingerprint>,
294 gradient: Option<Mat<f64>>,
296
297 workspace_rows: Vec<f64>, workspace_lm: Vec<f64>, workspace_cam: Vec<f64>, block_scratch: Vec<f64>, structure: StructureCache,
304}
305
306impl ImplicitSparseSchur {
307 pub fn new() -> Self {
310 Self {
311 partition: None,
312 eliminated: EliminatedBlocks::default(),
313 max_eliminated_dof: 1,
314 ordering: SchurOrdering::default(),
315 max_cg_iterations: 500,
316 cg_tolerance: 1e-9,
317 cg_q_tolerance: crate::linalg::schur::DEFAULT_ETA,
318 preconditioner_type: SchurPreconditioner::default(),
319 jacobian: None,
320 jacobian_pattern: None,
321 gradient: None,
322 workspace_rows: Vec::new(),
323 workspace_lm: Vec::new(),
324 workspace_cam: Vec::new(),
325 block_scratch: Vec::new(),
326 structure: StructureCache::default(),
327 }
328 }
329
330 pub fn with_cg_config(mut self, max_iterations: usize, tolerance: f64) -> Self {
332 self.max_cg_iterations = max_iterations;
333 self.cg_tolerance = tolerance;
334 self
335 }
336
337 pub fn with_preconditioner(mut self, preconditioner: SchurPreconditioner) -> Self {
343 self.preconditioner_type = preconditioner;
344 self
345 }
346
347 pub fn with_cg_q_tolerance(mut self, q_tolerance: f64) -> Self {
350 self.cg_q_tolerance = q_tolerance;
351 self
352 }
353
354 pub fn with_ordering(mut self, ordering: SchurOrdering) -> Self {
356 self.ordering = ordering;
357 self
358 }
359
360 pub fn with_cg_params(max_iterations: usize, tolerance: f64) -> Self {
363 Self::new().with_cg_config(max_iterations, tolerance)
364 }
365
366 pub fn with_config(
368 max_iterations: usize,
369 tolerance: f64,
370 preconditioner: SchurPreconditioner,
371 ) -> Self {
372 Self::new()
373 .with_cg_config(max_iterations, tolerance)
374 .with_preconditioner(preconditioner)
375 }
376
377 pub fn partition(&self) -> Option<&SchurPartition> {
379 self.partition.as_ref()
380 }
381
382 fn require_partition(&self) -> LinAlgResult<&SchurPartition> {
383 self.partition.as_ref().ok_or_else(|| {
384 LinAlgError::InvalidInput(
385 "Schur solver used before initialize_structure; the partition is unknown".into(),
386 )
387 .log()
388 })
389 }
390
391 fn ensure_structure(&mut self, jacobian: &SparseColMat<usize, f64>) -> LinAlgResult<()> {
399 let fingerprint = pattern::PatternFingerprint::of(jacobian);
400 if self.structure.fingerprint == Some(fingerprint) {
401 return Ok(());
402 }
403
404 let partition = self.require_partition()?;
405 let symbolic = jacobian.symbolic();
406
407 let mut row_owner = vec![usize::MAX; jacobian.nrows()];
409 let mut eliminated_rows = vec![Vec::new(); partition.eliminated_blocks().len()];
410 for (block_idx, block) in partition.eliminated_blocks().iter().enumerate() {
411 for offset in 0..block.dof {
412 for &row in symbolic.row_idx_of_col_raw(block.col_start + offset) {
413 match row_owner[row] {
414 usize::MAX => {
415 row_owner[row] = block_idx;
416 eliminated_rows[block_idx].push(row);
417 }
418 owned if owned == block_idx => {}
419 owned => {
420 let other = partition.eliminated_blocks()[owned];
421 return Err(LinAlgError::InvalidInput(format!(
422 "variables {:?} and {:?} are both marked for elimination but \
423 share residual row {row}, so EᵀE is not block-diagonal and \
424 the elimination would give a wrong step; eliminate only \
425 mutually unconnected variables",
426 block.key, other.key
427 ))
428 .log());
429 }
430 }
431 }
432 }
433 }
434 for rows in &mut eliminated_rows {
435 rows.sort_unstable();
436 }
437
438 let visibility: Vec<Vec<usize>> = partition
440 .kept_blocks()
441 .par_iter()
442 .map(|block| {
443 let mut seen = Vec::new();
444 for offset in 0..block.dof {
445 for &row in symbolic.row_idx_of_col_raw(block.col_start + offset) {
446 let owner = row_owner[row];
447 if owner != usize::MAX {
448 seen.push(owner);
449 }
450 }
451 }
452 seen.sort_unstable();
453 seen.dedup();
454 seen
455 })
456 .collect();
457
458 let mut kept_cols = vec![0usize; partition.kept_dof()];
459 for block in partition.kept_blocks() {
460 for offset in 0..block.dof {
461 let col = block.col_start + offset;
462 if let Some(local) = partition.kept_local(col) {
463 kept_cols[local] = col;
464 }
465 }
466 }
467 let mut eliminated_cols = vec![0usize; partition.eliminated_dof()];
468 for (block_idx, block) in partition.eliminated_blocks().iter().enumerate() {
469 let base = partition.eliminated_offset(block_idx);
470 for offset in 0..block.dof {
471 eliminated_cols[base + offset] = block.col_start + offset;
472 }
473 }
474
475 let chunk_rows = row_chunk_len(jacobian.nrows());
476 let spans = RowSpans::build(jacobian, partition, chunk_rows)?;
477
478 self.structure = StructureCache {
479 kept_cols,
480 eliminated_cols,
481 eliminated_rows,
482 visibility,
483 chunk_rows,
484 kept_spans: spans.kept,
485 eliminated_spans: spans.eliminated,
486 eliminated_block_of: spans.block_of,
487 fingerprint: Some(fingerprint),
488 };
489 Ok(())
490 }
491
492 #[allow(clippy::too_many_arguments)]
498 fn apply_schur_operator(
499 &self,
500 partition: &SchurPartition,
501 jacobian: &SparseColMat<usize, f64>,
502 damp_kept: &[f64],
503 v: &Mat<f64>,
504 out: &mut Mat<f64>,
505 out_buf: &mut [f64],
506 rows: &mut [f64],
507 temp_lm: &mut [f64],
508 scratch: &mut [f64],
509 ) {
510 let symbolic = jacobian.symbolic();
511 let row_idx = symbolic.row_idx();
512 let values = jacobian.val();
513 let chunk_rows = self.structure.chunk_rows.max(1);
514
515 rows.par_chunks_mut(chunk_rows)
517 .zip(self.structure.kept_spans.par_iter())
518 .enumerate()
519 .for_each(|(chunk, (y, spans))| {
520 y.fill(0.0);
521 let first = chunk * chunk_rows;
522 for span in spans {
523 let x = v[(span.local as usize, 0)];
524 if x == 0.0 {
525 continue;
526 }
527 for k in span.lo as usize..span.hi as usize {
528 y[row_idx[k] - first] += values[k] * x;
529 }
530 }
531 });
532
533 let rows_ref: &[f64] = rows;
535 temp_lm
536 .par_iter_mut()
537 .zip(self.structure.eliminated_cols.par_iter())
538 .for_each(|(t, &col)| {
539 let idx = symbolic.row_idx_of_col_raw(col);
540 let vals = jacobian.val_of_col(col);
541 *t = idx
542 .iter()
543 .zip(vals)
544 .map(|(&row, val)| val * rows_ref[row])
545 .sum();
546 });
547
548 let t: &[f64] = temp_lm;
551 let eliminated = &self.eliminated;
552 let block_of = &self.structure.eliminated_block_of;
553 scratch.par_iter_mut().enumerate().for_each(|(i, u)| {
554 let block_idx = block_of[i] as usize;
555 let dof = eliminated.dof(block_idx);
556 let base = partition.eliminated_offset(block_idx);
557 let inv = eliminated.block(block_idx);
558 let r = i - base;
559 let mut acc = 0.0;
560 for c in 0..dof {
561 acc += inv[c * dof + r] * t[base + c];
562 }
563 *u = acc;
564 });
565
566 let u: &[f64] = scratch;
568 rows.par_chunks_mut(chunk_rows)
569 .zip(self.structure.eliminated_spans.par_iter())
570 .enumerate()
571 .for_each(|(chunk, (y, spans))| {
572 let first = chunk * chunk_rows;
573 for span in spans {
574 let u = u[span.local as usize];
575 if u == 0.0 {
576 continue;
577 }
578 for k in span.lo as usize..span.hi as usize {
579 y[row_idx[k] - first] -= values[k] * u;
580 }
581 }
582 });
583
584 let rows_ref: &[f64] = rows;
588 out_buf
589 .par_iter_mut()
590 .enumerate()
591 .zip(self.structure.kept_cols.par_iter())
592 .for_each(|((local, o), &col)| {
593 let idx = symbolic.row_idx_of_col_raw(col);
594 let vals = jacobian.val_of_col(col);
595 let mut acc: f64 = idx
596 .iter()
597 .zip(vals)
598 .map(|(&row, val)| val * rows_ref[row])
599 .sum();
600 if let Some(d) = damp_kept.get(local) {
601 acc += d * v[(local, 0)];
602 }
603 *o = acc;
604 });
605 for (local, value) in out_buf.iter().enumerate() {
606 out[(local, 0)] = *value;
607 }
608 }
609
610 fn gather_strip(
616 jacobian: &SparseColMat<usize, f64>,
617 block: &BlockSpan,
618 rows: &[usize],
619 ) -> DMatrix<f64> {
620 let symbolic = jacobian.symbolic();
621 let mut strip = DMatrix::zeros(rows.len(), block.dof);
622 for offset in 0..block.dof {
623 let col = block.col_start + offset;
624 let idx = symbolic.row_idx_of_col_raw(col);
625 let vals = jacobian.val_of_col(col);
626 for (local_row, &row) in rows.iter().enumerate() {
627 if let Ok(k) = idx.binary_search(&row) {
628 strip[(local_row, offset)] = vals[k];
629 }
630 }
631 }
632 strip
633 }
634
635 fn build_preconditioner(
643 &self,
644 partition: &SchurPartition,
645 jacobian: &SparseColMat<usize, f64>,
646 damp_kept: &[f64],
647 ) -> Option<Vec<DMatrix<f64>>> {
648 if self.preconditioner_type == SchurPreconditioner::None {
649 return None;
650 }
651
652 let blocks = partition
653 .kept_blocks()
654 .par_iter()
655 .enumerate()
656 .map(|(kept_idx, block)| {
657 let dof = block.dof;
658 let base = partition.kept_offset(kept_idx);
659
660 let mut s_ii = DMatrix::zeros(dof, dof);
662 for c in 0..dof {
663 for r in 0..dof {
664 s_ii[(r, c)] =
665 column_dot(jacobian, block.col_start + r, block.col_start + c);
666 }
667 if let Some(d) = damp_kept.get(base + c) {
668 s_ii[(c, c)] += d;
669 }
670 }
671
672 if self.preconditioner_type == SchurPreconditioner::SchurJacobi {
676 let visible = self
677 .structure
678 .visibility
679 .get(kept_idx)
680 .map_or(&[][..], Vec::as_slice);
681 for &elim_idx in visible {
682 let elim_block = partition.eliminated_blocks()[elim_idx];
683 let rows = &self.structure.eliminated_rows[elim_idx];
684 let f_strip = Self::gather_strip(jacobian, block, rows);
685 let e_strip = Self::gather_strip(jacobian, &elim_block, rows);
686
687 let m = f_strip.transpose() * e_strip;
689 let elim_dof = elim_block.dof;
690 let inv = DMatrix::from_column_slice(
691 elim_dof,
692 elim_dof,
693 self.eliminated.block(elim_idx),
694 );
695 s_ii -= &m * inv * m.transpose();
696 }
697 }
698
699 invert_with_retry_dyn(&s_ii).unwrap_or_else(|| DMatrix::identity(dof, dof))
700 })
701 .collect();
702 Some(blocks)
703 }
704
705 fn apply_preconditioner(blocks: Option<&[DMatrix<f64>]>, r: &Mat<f64>, z: &mut Mat<f64>) {
707 let Some(blocks) = blocks else {
708 faer::zip!(z, r).for_each(|faer::unzip!(z, r)| *z = *r);
709 return;
710 };
711 let mut offset = 0usize;
712 for inv in blocks {
713 let dof = inv.nrows();
714 for row in 0..dof {
715 let mut acc = 0.0;
716 for col in 0..dof {
717 acc += inv[(row, col)] * r[(offset + col, 0)];
718 }
719 z[(offset + row, 0)] = acc;
720 }
721 offset += dof;
722 }
723 }
724
725 fn solve_from_jacobian(
729 &mut self,
730 residuals: &Mat<f64>,
731 jacobian: &SparseColMat<usize, f64>,
732 damping: Option<&Damping>,
733 ) -> LinAlgResult<Mat<f64>> {
734 self.ensure_structure(jacobian)?;
735 let partition = self.require_partition()?.clone();
736
737 if jacobian.ncols() != partition.total_dof() {
738 return Err(LinAlgError::InvalidInput(format!(
739 "Jacobian has {} columns but the partition covers {}",
740 jacobian.ncols(),
741 partition.total_dof()
742 ))
743 .log());
744 }
745
746 let gradient = jt_vec(jacobian, residuals);
749 let kept_dof = partition.kept_dof();
750 let eliminated_dof = partition.eliminated_dof();
751
752 let damp_kept: Vec<f64> = match damping {
754 None => Vec::new(),
755 Some(d) => {
756 let diag = diag_jt_j(jacobian);
757 let mut per_kept = vec![0.0; kept_dof];
758 for block in partition.kept_blocks() {
759 for offset in 0..block.dof {
760 let col = block.col_start + offset;
761 if let Some(local) = partition.kept_local(col) {
762 per_kept[local] = d.diagonal_term(diag[col]);
763 }
764 }
765 }
766 per_kept
767 }
768 };
769
770 let mut eliminated = std::mem::take(&mut self.eliminated);
772 eliminated.gather_from_jacobian(jacobian, &partition);
773 if let Some(d) = damping {
774 eliminated.damp(d);
775 }
776 let inversion = eliminated.invert_in_place(&partition);
777 self.eliminated = eliminated;
778 inversion?;
779
780 let mut g_k = Mat::<f64>::zeros(kept_dof, 1);
784 for block in partition.kept_blocks() {
785 for offset in 0..block.dof {
786 let col = block.col_start + offset;
787 if let Some(local) = partition.kept_local(col) {
788 g_k[(local, 0)] = -gradient[(col, 0)];
789 }
790 }
791 }
792 let mut g_e = Mat::<f64>::zeros(eliminated_dof, 1);
793 for (block_idx, block) in partition.eliminated_blocks().iter().enumerate() {
794 let base = partition.eliminated_offset(block_idx);
795 for offset in 0..block.dof {
796 g_e[(base + offset, 0)] = -gradient[(block.col_start + offset, 0)];
797 }
798 }
799
800 let mut hee_ge = Mat::<f64>::zeros(eliminated_dof, 1);
802 self.apply_eliminated_inverse(&partition, &g_e, &mut hee_ge);
803 let correction = self.coupling_times(&partition, jacobian, &hee_ge);
804
805 let g_reduced = Mat::from_fn(kept_dof, 1, |i, _| g_k[(i, 0)] - correction[(i, 0)]);
806
807 let precond = self.build_preconditioner(&partition, jacobian, &damp_kept);
808 let delta_k = self.solve_pcg_block(
809 &partition,
810 jacobian,
811 &damp_kept,
812 &g_reduced,
813 precond.as_deref(),
814 );
815
816 let coupling_t = self.coupling_transpose_times(&partition, jacobian, &delta_k);
818 let rhs_e = Mat::from_fn(eliminated_dof, 1, |i, _| g_e[(i, 0)] - coupling_t[(i, 0)]);
819 let mut delta_e = Mat::<f64>::zeros(eliminated_dof, 1);
820 self.apply_eliminated_inverse(&partition, &rhs_e, &mut delta_e);
821
822 let delta = self.combine_updates(&partition, &delta_k, &delta_e);
823
824 self.gradient = Some(gradient);
825 self.retain_jacobian(jacobian);
826 Ok(delta)
827 }
828
829 fn retain_jacobian(&mut self, jacobian: &SparseColMat<usize, f64>) {
837 let fingerprint = self.structure.fingerprint;
838 let reusable = self.jacobian_pattern == fingerprint
839 && self
840 .jacobian
841 .as_ref()
842 .is_some_and(|held| held.val().len() == jacobian.val().len());
843
844 match self.jacobian.as_mut() {
845 Some(held) if reusable => held.val_mut().copy_from_slice(jacobian.val()),
846 _ => {
847 self.jacobian = Some(jacobian.clone());
848 self.jacobian_pattern = fingerprint;
849 }
850 }
851 }
852
853 fn apply_eliminated_inverse(&self, partition: &SchurPartition, b: &Mat<f64>, x: &mut Mat<f64>) {
855 for block_idx in 0..partition.eliminated_blocks().len() {
856 let dof = self.eliminated.dof(block_idx);
857 let base = partition.eliminated_offset(block_idx);
858 let inv = self.eliminated.block(block_idx);
859 for r in 0..dof {
860 let mut acc = 0.0;
861 for c in 0..dof {
862 acc += inv[c * dof + r] * b[(base + c, 0)];
863 }
864 x[(base + r, 0)] = acc;
865 }
866 }
867 }
868
869 fn coupling_times(
871 &self,
872 partition: &SchurPartition,
873 jacobian: &SparseColMat<usize, f64>,
874 u: &Mat<f64>,
875 ) -> Mat<f64> {
876 let symbolic = jacobian.symbolic();
877 let mut rows = vec![0.0; jacobian.nrows()];
878 for (block_idx, block) in partition.eliminated_blocks().iter().enumerate() {
879 let base = partition.eliminated_offset(block_idx);
880 for offset in 0..block.dof {
881 let x = u[(base + offset, 0)];
882 if x == 0.0 {
883 continue;
884 }
885 let col = block.col_start + offset;
886 let idx = symbolic.row_idx_of_col_raw(col);
887 let vals = jacobian.val_of_col(col);
888 for (k, &row) in idx.iter().enumerate() {
889 rows[row] += vals[k] * x;
890 }
891 }
892 }
893
894 let mut out = Mat::<f64>::zeros(partition.kept_dof(), 1);
895 for block in partition.kept_blocks() {
896 for offset in 0..block.dof {
897 let col = block.col_start + offset;
898 let Some(local) = partition.kept_local(col) else {
899 continue;
900 };
901 let idx = symbolic.row_idx_of_col_raw(col);
902 let vals = jacobian.val_of_col(col);
903 let mut acc = 0.0;
904 for (k, &row) in idx.iter().enumerate() {
905 acc += vals[k] * rows[row];
906 }
907 out[(local, 0)] = acc;
908 }
909 }
910 out
911 }
912
913 fn coupling_transpose_times(
915 &self,
916 partition: &SchurPartition,
917 jacobian: &SparseColMat<usize, f64>,
918 v: &Mat<f64>,
919 ) -> Mat<f64> {
920 let symbolic = jacobian.symbolic();
921 let mut rows = vec![0.0; jacobian.nrows()];
922 for block in partition.kept_blocks() {
923 for offset in 0..block.dof {
924 let col = block.col_start + offset;
925 let Some(local) = partition.kept_local(col) else {
926 continue;
927 };
928 let x = v[(local, 0)];
929 if x == 0.0 {
930 continue;
931 }
932 let idx = symbolic.row_idx_of_col_raw(col);
933 let vals = jacobian.val_of_col(col);
934 for (k, &row) in idx.iter().enumerate() {
935 rows[row] += vals[k] * x;
936 }
937 }
938 }
939
940 let mut out = Mat::<f64>::zeros(partition.eliminated_dof(), 1);
941 for (block_idx, block) in partition.eliminated_blocks().iter().enumerate() {
942 let base = partition.eliminated_offset(block_idx);
943 for offset in 0..block.dof {
944 let col = block.col_start + offset;
945 let idx = symbolic.row_idx_of_col_raw(col);
946 let vals = jacobian.val_of_col(col);
947 let mut acc = 0.0;
948 for (k, &row) in idx.iter().enumerate() {
949 acc += vals[k] * rows[row];
950 }
951 out[(base + offset, 0)] = acc;
952 }
953 }
954 out
955 }
956
957 fn solve_pcg_block(
959 &mut self,
960 partition: &SchurPartition,
961 jacobian: &SparseColMat<usize, f64>,
962 damp_kept: &[f64],
963 b: &Mat<f64>,
964 precond: Option<&[DMatrix<f64>]>,
965 ) -> Mat<f64> {
966 let mut rows = std::mem::take(&mut self.workspace_rows);
969 let mut temp_lm = std::mem::take(&mut self.workspace_lm);
970 let mut out_buf = std::mem::take(&mut self.workspace_cam);
971 let mut scratch = std::mem::take(&mut self.block_scratch);
972 out_buf.clear();
973 out_buf.resize(partition.kept_dof(), 0.0);
974 rows.clear();
975 rows.resize(jacobian.nrows(), 0.0);
976 temp_lm.clear();
977 temp_lm.resize(partition.eliminated_dof(), 0.0);
978 scratch.clear();
979 scratch.resize(partition.eliminated_dof(), 0.0);
980
981 let result = solve_pcg(
982 b,
983 &PcgParams::new(self.max_cg_iterations, self.cg_tolerance)
984 .with_q_tolerance(self.cg_q_tolerance),
985 |p, ap| {
986 self.apply_schur_operator(
987 partition,
988 jacobian,
989 damp_kept,
990 p,
991 ap,
992 &mut out_buf,
993 &mut rows,
994 &mut temp_lm,
995 &mut scratch,
996 );
997 },
998 |r, z| Self::apply_preconditioner(precond, r, z),
999 );
1000
1001 self.workspace_rows = rows;
1002 self.workspace_lm = temp_lm;
1003 self.workspace_cam = out_buf;
1004 self.block_scratch = scratch;
1005
1006 debug!(
1010 "PCG: {} iterations, residual {:.3e}, {:?} ({:?})",
1011 result.iterations, result.final_residual, result.termination, self.preconditioner_type
1012 );
1013 result.x
1014 }
1015
1016 fn combine_updates(
1018 &self,
1019 partition: &SchurPartition,
1020 delta_k: &Mat<f64>,
1021 delta_e: &Mat<f64>,
1022 ) -> Mat<f64> {
1023 let mut delta = Mat::zeros(partition.total_dof(), 1);
1024 for block in partition.kept_blocks() {
1025 for offset in 0..block.dof {
1026 let global = block.col_start + offset;
1027 if let Some(local) = partition.kept_local(global) {
1028 delta[(global, 0)] = delta_k[(local, 0)];
1029 }
1030 }
1031 }
1032 for (block_idx, block) in partition.eliminated_blocks().iter().enumerate() {
1033 let base = partition.eliminated_offset(block_idx);
1034 for offset in 0..block.dof {
1035 delta[(block.col_start + offset, 0)] = delta_e[(base + offset, 0)];
1036 }
1037 }
1038 delta
1039 }
1040}
1041
1042impl Default for ImplicitSparseSchur {
1043 fn default() -> Self {
1044 Self::new()
1045 }
1046}
1047
1048impl LinearSolver<SparseMode> for ImplicitSparseSchur {
1049 fn solve_normal_equation(
1050 &mut self,
1051 residuals: &Mat<f64>,
1052 jacobian: &SparseColMat<usize, f64>,
1053 ) -> LinAlgResult<Mat<f64>> {
1054 self.solve_from_jacobian(residuals, jacobian, None)
1055 }
1056
1057 fn solve_augmented_equation(
1058 &mut self,
1059 residuals: &Mat<f64>,
1060 jacobian: &SparseColMat<usize, f64>,
1061 damping: &Damping,
1062 ) -> LinAlgResult<Mat<f64>> {
1063 self.solve_from_jacobian(residuals, jacobian, Some(damping))
1064 }
1065
1066 fn hessian_vec_product(&self, v: &Mat<f64>) -> Option<Mat<f64>> {
1067 Some(jt_j_vec_product(self.jacobian.as_ref()?, v))
1069 }
1070
1071 fn get_hessian(&self) -> Option<&SparseColMat<usize, f64>> {
1077 None
1078 }
1079
1080 fn get_gradient(&self) -> Option<&Mat<f64>> {
1081 self.gradient.as_ref()
1082 }
1083}
1084
1085impl StructureAware for ImplicitSparseSchur {
1086 fn initialize_structure(
1087 &mut self,
1088 variables: &SlotMap<VarKey, Box<dyn ManifoldVariable>>,
1089 variable_index_map: &SecondaryMap<VarKey, usize>,
1090 schur_landmark_keys: &std::collections::HashSet<VarKey>,
1091 ) -> LinAlgResult<()> {
1092 let effective_keys =
1093 effective_landmark_keys(variables, schur_landmark_keys, &self.ordering);
1094
1095 let mut kept = Vec::new();
1096 let mut eliminated = Vec::new();
1097
1098 for (key, variable) in variables {
1099 let col_start = *variable_index_map.get(key).ok_or_else(|| {
1100 LinAlgError::InvalidInput(format!("VarKey {:?} not in index map", key))
1101 })?;
1102 let span = BlockSpan {
1103 key,
1104 col_start,
1105 dof: variable.free_dof(),
1108 };
1109 if effective_keys.contains(&key) {
1110 eliminated.push(span);
1111 } else {
1112 kept.push(span);
1113 }
1114 }
1115
1116 let partition = SchurPartition::new(kept, eliminated)?;
1117 self.max_eliminated_dof = partition
1118 .eliminated_blocks()
1119 .iter()
1120 .map(|b| b.dof)
1121 .max()
1122 .unwrap_or(0)
1123 .max(1);
1124 self.eliminated = EliminatedBlocks::new(&partition);
1125 self.workspace_lm = vec![0.0; partition.eliminated_dof()];
1126 self.workspace_cam = vec![0.0; partition.kept_dof()];
1127 self.block_scratch = vec![0.0; self.max_eliminated_dof];
1128 self.workspace_rows.clear();
1129 self.structure = StructureCache::default();
1130 self.partition = Some(partition);
1131 Ok(())
1132 }
1133}
1134#[cfg(test)]
1135mod tests {
1136 use super::*;
1137 use crate::core::VarKey;
1138 use crate::core::variable::Variable;
1139 use apex_manifolds::{LieGroup, rn, se3};
1140 use faer::sparse::Triplet;
1141 use nalgebra::DVector;
1142 use slotmap::{SecondaryMap, SlotMap};
1143
1144 type TestResult = Result<(), Box<dyn std::error::Error>>;
1145
1146 type TestSetup = (
1147 SlotMap<VarKey, Box<dyn ManifoldVariable>>,
1148 SecondaryMap<VarKey, usize>,
1149 SparseColMat<usize, f64>,
1150 faer::Mat<f64>,
1151 std::collections::HashSet<VarKey>,
1152 );
1153
1154 fn create_schur_test_setup() -> Result<TestSetup, Box<dyn std::error::Error>> {
1157 let se3_id = DVector::from_vec(vec![1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]);
1158 let pt_zero = DVector::from_vec(vec![0.0, 0.0, 0.0]);
1159
1160 let mut variables: SlotMap<VarKey, Box<dyn ManifoldVariable>> = SlotMap::with_key();
1161 let cam0 = variables.insert(Box::new(Variable::new(se3::SE3::from_param_slice(
1162 se3_id.as_slice(),
1163 ))));
1164 let cam1 = variables.insert(Box::new(Variable::new(se3::SE3::from_param_slice(
1165 se3_id.as_slice(),
1166 ))));
1167 let pt0 = variables.insert(Box::new(Variable::new(rn::Rn::new(pt_zero.clone()))));
1168 let pt1 = variables.insert(Box::new(Variable::new(rn::Rn::new(pt_zero.clone()))));
1169 let pt2 = variables.insert(Box::new(Variable::new(rn::Rn::new(pt_zero.clone()))));
1170
1171 let mut variable_index_map: SecondaryMap<VarKey, usize> = SecondaryMap::new();
1172 variable_index_map.insert(cam0, 0);
1173 variable_index_map.insert(cam1, 6);
1174 variable_index_map.insert(pt0, 12);
1175 variable_index_map.insert(pt1, 15);
1176 variable_index_map.insert(pt2, 18);
1177
1178 let cam_cols = [0usize, 6];
1179 let lm_cols = [12usize, 15, 18];
1180 let mut triplets: Vec<Triplet<usize, usize, f64>> = Vec::new();
1181 for (ci, &cam_col) in cam_cols.iter().enumerate() {
1182 for (li, &lm_col) in lm_cols.iter().enumerate() {
1183 let row_base = (ci * 3 + li) * 6;
1184 for k in 0..6 {
1185 triplets.push(Triplet::new(row_base + k, cam_col + k, 1.0));
1186 triplets.push(Triplet::new(row_base + k, lm_col + (k % 3), 1.0));
1187 }
1188 }
1189 }
1190
1191 let jacobian = SparseColMat::try_new_from_triplets(36, 21, &triplets)?;
1192 let residuals = faer::Mat::from_fn(36, 1, |i, _| (i % 5) as f64 * 0.1);
1193
1194 let mut landmark_keys = std::collections::HashSet::new();
1195 landmark_keys.insert(pt0);
1196 landmark_keys.insert(pt1);
1197 landmark_keys.insert(pt2);
1198
1199 Ok((
1200 variables,
1201 variable_index_map,
1202 jacobian,
1203 residuals,
1204 landmark_keys,
1205 ))
1206 }
1207
1208 #[test]
1209 fn test_implicit_schur_creation() {
1210 let solver = ImplicitSparseSchur::new();
1211 assert_eq!(solver.max_cg_iterations, 500);
1212 assert_eq!(solver.cg_tolerance, 1e-9);
1213 assert_eq!(solver.preconditioner_type, SchurPreconditioner::SchurJacobi);
1214 }
1215
1216 #[test]
1217 fn test_with_custom_params() {
1218 let solver = ImplicitSparseSchur::with_cg_params(100, 1e-8);
1219 assert_eq!(solver.max_cg_iterations, 100);
1220 assert_eq!(solver.cg_tolerance, 1e-8);
1221 assert_eq!(solver.preconditioner_type, SchurPreconditioner::SchurJacobi);
1222 }
1223
1224 #[test]
1225 fn test_with_full_config() {
1226 let solver =
1227 ImplicitSparseSchur::with_config(200, 1e-10, SchurPreconditioner::BlockDiagonal);
1228 assert_eq!(solver.max_cg_iterations, 200);
1229 assert_eq!(solver.cg_tolerance, 1e-10);
1230 assert_eq!(
1231 solver.preconditioner_type,
1232 SchurPreconditioner::BlockDiagonal
1233 );
1234 }
1235
1236 #[test]
1237 fn test_implicit_schur_default() {
1238 let solver = ImplicitSparseSchur::default();
1239 assert_eq!(solver.max_cg_iterations, 500);
1240 assert_eq!(solver.cg_tolerance, 1e-9);
1241 assert!(solver.partition.is_none());
1242 assert!(solver.jacobian.is_none());
1243 }
1244
1245 #[test]
1246 fn test_implicit_schur_initialize_structure() -> TestResult {
1247 let (variables, variable_index_map, _, _, landmark_keys) = create_schur_test_setup()?;
1248 let mut solver = ImplicitSparseSchur::new();
1249 solver.initialize_structure(&variables, &variable_index_map, &landmark_keys)?;
1250
1251 let partition = solver.partition.as_ref().ok_or("partition is None")?;
1252 assert_eq!(partition.kept_blocks().len(), 2);
1253 assert_eq!(partition.eliminated_blocks().len(), 3);
1254 assert_eq!(partition.kept_dof(), 12);
1255 assert_eq!(partition.eliminated_dof(), 9);
1256 Ok(())
1257 }
1258
1259 #[test]
1260 fn test_implicit_schur_solve_normal_equation() -> TestResult {
1261 let (variables, variable_index_map, jacobian, residuals, landmark_keys) =
1262 create_schur_test_setup()?;
1263 let mut solver = ImplicitSparseSchur::with_cg_params(500, 1e-6);
1264 solver.initialize_structure(&variables, &variable_index_map, &landmark_keys)?;
1265
1266 let delta =
1267 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
1268 assert_eq!(delta.nrows(), 21);
1269 assert_eq!(delta.ncols(), 1);
1270 Ok(())
1271 }
1272
1273 #[test]
1274 fn test_implicit_schur_solve_augmented_equation() -> TestResult {
1275 let (variables, variable_index_map, jacobian, residuals, landmark_keys) =
1276 create_schur_test_setup()?;
1277 let mut solver = ImplicitSparseSchur::with_cg_params(500, 1e-6);
1278 solver.initialize_structure(&variables, &variable_index_map, &landmark_keys)?;
1279
1280 let delta = LinearSolver::<SparseMode>::solve_augmented_equation(
1281 &mut solver,
1282 &residuals,
1283 &jacobian,
1284 &Damping::identity(0.1),
1285 )?;
1286 assert_eq!(delta.nrows(), 21);
1287 Ok(())
1288 }
1289
1290 #[test]
1291 fn test_implicit_schur_get_hessian_gradient() -> TestResult {
1292 let (variables, variable_index_map, jacobian, residuals, landmark_keys) =
1293 create_schur_test_setup()?;
1294 let mut solver = ImplicitSparseSchur::with_cg_params(500, 1e-6);
1295 solver.initialize_structure(&variables, &variable_index_map, &landmark_keys)?;
1296
1297 assert!(LinearSolver::<SparseMode>::get_hessian(&solver).is_none());
1298 assert!(LinearSolver::<SparseMode>::get_gradient(&solver).is_none());
1299
1300 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
1301
1302 assert!(
1305 LinearSolver::<SparseMode>::get_hessian(&solver).is_none(),
1306 "matrix-free solver must not publish a Hessian"
1307 );
1308 let g = LinearSolver::<SparseMode>::get_gradient(&solver).ok_or("gradient is None")?;
1309 assert_eq!(g.nrows(), 21);
1310
1311 let v = faer::Mat::from_fn(21, 1, |i, _| ((i % 3) as f64) - 1.0);
1312 let hv = LinearSolver::<SparseMode>::hessian_vec_product(&solver, &v)
1313 .ok_or("hessian_vec_product is None")?;
1314 assert_eq!(hv.nrows(), 21);
1315 Ok(())
1316 }
1317
1318 #[test]
1324 fn every_preconditioner_reaches_the_same_step() -> TestResult {
1325 let mut steps = Vec::new();
1326 for preconditioner in [
1327 SchurPreconditioner::None,
1328 SchurPreconditioner::BlockDiagonal,
1329 SchurPreconditioner::SchurJacobi,
1330 ] {
1331 let (variables, variable_index_map, jacobian, residuals, landmark_keys) =
1332 create_schur_test_setup()?;
1333 let mut solver = ImplicitSparseSchur::new()
1334 .with_cg_config(1000, 1e-12)
1335 .with_preconditioner(preconditioner);
1336 solver.initialize_structure(&variables, &variable_index_map, &landmark_keys)?;
1337 steps.push(LinearSolver::<SparseMode>::solve_normal_equation(
1338 &mut solver,
1339 &residuals,
1340 &jacobian,
1341 )?);
1342 }
1343
1344 let reference = &steps[0];
1345 for (idx, step) in steps.iter().enumerate().skip(1) {
1346 for row in 0..reference.nrows() {
1347 let scale = reference[(row, 0)].abs().max(1.0);
1348 assert!(
1349 (reference[(row, 0)] - step[(row, 0)]).abs() / scale < 1e-6,
1350 "preconditioner {idx} diverged at row {row}: {} vs {}",
1351 step[(row, 0)],
1352 reference[(row, 0)]
1353 );
1354 }
1355 }
1356 Ok(())
1357 }
1358
1359 #[test]
1360 fn test_implicit_schur_block_diagonal_preconditioner() -> TestResult {
1361 let (variables, variable_index_map, jacobian, residuals, landmark_keys) =
1362 create_schur_test_setup()?;
1363 let mut solver =
1364 ImplicitSparseSchur::with_config(500, 1e-6, SchurPreconditioner::BlockDiagonal);
1365 solver.initialize_structure(&variables, &variable_index_map, &landmark_keys)?;
1366
1367 let delta =
1368 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
1369 assert_eq!(delta.nrows(), 21);
1370 Ok(())
1371 }
1372
1373 #[test]
1374 fn test_implicit_schur_schur_jacobi_preconditioner() -> TestResult {
1375 let (variables, variable_index_map, jacobian, residuals, landmark_keys) =
1376 create_schur_test_setup()?;
1377 let mut solver =
1378 ImplicitSparseSchur::with_config(500, 1e-6, SchurPreconditioner::SchurJacobi);
1379 solver.initialize_structure(&variables, &variable_index_map, &landmark_keys)?;
1380
1381 let delta =
1382 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
1383 assert_eq!(delta.nrows(), 21);
1384 Ok(())
1385 }
1386
1387 #[test]
1388 fn test_implicit_schur_no_preconditioner() -> TestResult {
1389 let (variables, variable_index_map, jacobian, residuals, landmark_keys) =
1390 create_schur_test_setup()?;
1391 let mut solver = ImplicitSparseSchur::with_config(500, 1e-6, SchurPreconditioner::None);
1392 solver.initialize_structure(&variables, &variable_index_map, &landmark_keys)?;
1393
1394 let delta =
1395 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
1396 assert_eq!(delta.nrows(), 21);
1397 Ok(())
1398 }
1399
1400 #[test]
1401 fn test_implicit_schur_augmented_lambda_effect() -> TestResult {
1402 let (variables, variable_index_map, jacobian, residuals, landmark_keys) =
1403 create_schur_test_setup()?;
1404
1405 let mut s1 = ImplicitSparseSchur::with_cg_params(500, 1e-6);
1406 s1.initialize_structure(&variables, &variable_index_map, &landmark_keys)?;
1407 let d1 = LinearSolver::<SparseMode>::solve_augmented_equation(
1408 &mut s1,
1409 &residuals,
1410 &jacobian,
1411 &Damping::identity(0.001),
1412 )?;
1413
1414 let mut s2 = ImplicitSparseSchur::with_cg_params(500, 1e-6);
1415 s2.initialize_structure(&variables, &variable_index_map, &landmark_keys)?;
1416 let d2 = LinearSolver::<SparseMode>::solve_augmented_equation(
1417 &mut s2,
1418 &residuals,
1419 &jacobian,
1420 &Damping::identity(100.0),
1421 )?;
1422
1423 let norm_diff: f64 = (0..21).map(|i| (d1[(i, 0)] - d2[(i, 0)]).powi(2)).sum();
1424 assert!(
1425 norm_diff > 1e-10,
1426 "Different λ should yield different updates"
1427 );
1428 Ok(())
1429 }
1430
1431 #[test]
1432 fn test_implicit_schur_solve_without_init_returns_error() -> TestResult {
1433 let triplets: Vec<Triplet<usize, usize, f64>> = vec![Triplet::new(0, 0, 1.0)];
1434 let jacobian =
1435 SparseColMat::try_new_from_triplets(1, 1, &triplets).map_err(|e| format!("{e:?}"))?;
1436 let residuals = faer::Mat::from_fn(1, 1, |_, _| 1.0);
1437 let mut solver = ImplicitSparseSchur::new();
1438
1439 let result =
1440 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian);
1441 assert!(result.is_err());
1442 Ok(())
1443 }
1444
1445 #[test]
1446 fn test_implicit_initialize_structure_no_cameras_returns_error() {
1447 use crate::core::variable::Variable;
1448 use apex_manifolds::rn;
1449
1450 let mut variables: SlotMap<VarKey, Box<dyn ManifoldVariable>> = SlotMap::with_key();
1451 let k = variables.insert(Box::new(Variable::new(rn::Rn::new(DVector::zeros(3)))));
1452 let mut variable_index_map: SecondaryMap<VarKey, usize> = SecondaryMap::new();
1453 variable_index_map.insert(k, 0);
1454 let mut landmark_keys = std::collections::HashSet::new();
1455 landmark_keys.insert(k);
1456
1457 let mut solver = ImplicitSparseSchur::new();
1458 let result = solver.initialize_structure(&variables, &variable_index_map, &landmark_keys);
1459 assert!(
1460 result.is_err(),
1461 "Expected Err when no retained variables present"
1462 );
1463 }
1464
1465 #[test]
1466 fn test_implicit_initialize_structure_no_landmarks_returns_error() {
1467 use crate::core::variable::Variable;
1468 use apex_manifolds::se3;
1469
1470 let mut variables: SlotMap<VarKey, Box<dyn ManifoldVariable>> = SlotMap::with_key();
1471 let k = variables.insert(Box::new(Variable::new(se3::SE3::from_param_slice(&[
1472 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
1473 ]))));
1474 let mut variable_index_map: SecondaryMap<VarKey, usize> = SecondaryMap::new();
1475 variable_index_map.insert(k, 0);
1476 let landmark_keys = std::collections::HashSet::<VarKey>::new();
1477
1478 let mut solver = ImplicitSparseSchur::new();
1479 let result = solver.initialize_structure(&variables, &variable_index_map, &landmark_keys);
1480 assert!(
1481 result.is_err(),
1482 "Expected Err when no eliminated variables present"
1483 );
1484 }
1485
1486 #[test]
1490 fn test_implicit_schur_supports_non_contiguous_mixed_dof_partition() -> TestResult {
1491 use crate::linalg::schur::BlockSpan;
1492
1493 let a = Variable::new(se3::SE3::from_param_slice(&[
1496 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
1497 ]));
1498 let b = Variable::new(se3::SE3::from_param_slice(&[
1499 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
1500 ]));
1501 let p = Variable::new(rn::Rn::new(DVector::from_vec(vec![0.0])));
1502 let q = Variable::new(rn::Rn::new(DVector::from_vec(vec![0.0])));
1503
1504 let mut variables: SlotMap<VarKey, Box<dyn ManifoldVariable>> = SlotMap::with_key();
1505 let key_a = variables.insert(Box::new(a));
1506 let key_p = variables.insert(Box::new(p));
1507 let key_b = variables.insert(Box::new(b));
1508 let key_q = variables.insert(Box::new(q));
1509
1510 let mut variable_index_map: SecondaryMap<VarKey, usize> = SecondaryMap::new();
1511 variable_index_map.insert(key_a, 0);
1512 variable_index_map.insert(key_p, 6);
1513 variable_index_map.insert(key_b, 7);
1514 variable_index_map.insert(key_q, 13);
1515
1516 let mut landmark_keys = std::collections::HashSet::new();
1517 landmark_keys.insert(key_p);
1518 landmark_keys.insert(key_q);
1519
1520 let mut triplets: Vec<Triplet<usize, usize, f64>> = Vec::new();
1523 let mut row = 0usize;
1524 for &cam_col in &[0usize, 7] {
1525 for &lm_col in &[6usize, 13] {
1526 for k in 0..1 {
1527 triplets.push(Triplet::new(row + k, cam_col + k, 1.0));
1528 triplets.push(Triplet::new(row + k, lm_col, 1.0));
1529 }
1530 row += 1;
1531 }
1532 }
1533 for col in 0..14 {
1534 triplets.push(Triplet::new(row, col, 0.5));
1535 row += 1;
1536 }
1537 let jacobian = SparseColMat::try_new_from_triplets(row, 14, &triplets)?;
1538 let residuals = faer::Mat::from_fn(row, 1, |i, _| 0.1 * (i as f64 + 1.0));
1539
1540 let mut solver = ImplicitSparseSchur::with_cg_params(500, 1e-10);
1541 solver.initialize_structure(&variables, &variable_index_map, &landmark_keys)?;
1542 let delta =
1543 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
1544 assert_eq!(delta.nrows(), 14);
1545 for i in 0..14 {
1546 assert!(delta[(i, 0)].is_finite());
1547 }
1548
1549 let _ = BlockSpan {
1553 key: key_a,
1554 col_start: 0,
1555 dof: 6,
1556 };
1557 Ok(())
1558 }
1559
1560 fn column_scatter_reference(
1563 solver: &ImplicitSparseSchur,
1564 partition: &SchurPartition,
1565 jacobian: &SparseColMat<usize, f64>,
1566 v: &Mat<f64>,
1567 ) -> Vec<f64> {
1568 let symbolic = jacobian.symbolic();
1569 let mut rows = vec![0.0; jacobian.nrows()];
1570 for block in partition.kept_blocks() {
1571 for offset in 0..block.dof {
1572 let col = block.col_start + offset;
1573 let Some(local) = partition.kept_local(col) else {
1574 continue;
1575 };
1576 let x = v[(local, 0)];
1577 if x == 0.0 {
1578 continue;
1579 }
1580 let idx = symbolic.row_idx_of_col_raw(col);
1581 let vals = jacobian.val_of_col(col);
1582 for (k, &row) in idx.iter().enumerate() {
1583 rows[row] += vals[k] * x;
1584 }
1585 }
1586 }
1587 let mut t = vec![0.0; partition.eliminated_dof()];
1588 for (slot, &col) in t.iter_mut().zip(&solver.structure.eliminated_cols) {
1589 let idx = symbolic.row_idx_of_col_raw(col);
1590 let vals = jacobian.val_of_col(col);
1591 *slot = idx
1592 .iter()
1593 .zip(vals)
1594 .map(|(&row, val)| val * rows[row])
1595 .sum();
1596 }
1597 for (block_idx, block) in partition.eliminated_blocks().iter().enumerate() {
1598 let dof = solver.eliminated.dof(block_idx);
1599 let base = partition.eliminated_offset(block_idx);
1600 let inv = solver.eliminated.block(block_idx);
1601 let mut u = vec![0.0; dof];
1602 for (r, out) in u.iter_mut().enumerate() {
1603 let mut acc = 0.0;
1604 for c in 0..dof {
1605 acc += inv[c * dof + r] * t[base + c];
1606 }
1607 *out = acc;
1608 }
1609 for (offset, &u) in u.iter().enumerate() {
1610 if u == 0.0 {
1611 continue;
1612 }
1613 let col = block.col_start + offset;
1614 let idx = symbolic.row_idx_of_col_raw(col);
1615 let vals = jacobian.val_of_col(col);
1616 for (k, &row) in idx.iter().enumerate() {
1617 rows[row] -= vals[k] * u;
1618 }
1619 }
1620 }
1621 solver
1622 .structure
1623 .kept_cols
1624 .iter()
1625 .map(|&col| {
1626 let idx = symbolic.row_idx_of_col_raw(col);
1627 let vals = jacobian.val_of_col(col);
1628 idx.iter()
1629 .zip(vals)
1630 .map(|(&row, val)| val * rows[row])
1631 .sum()
1632 })
1633 .collect()
1634 }
1635
1636 #[test]
1640 fn row_chunked_operator_is_bit_identical_to_column_scatter() -> TestResult {
1641 let cam = || -> Box<dyn ManifoldVariable> {
1642 Box::new(Variable::new(se3::SE3::from_param_slice(&[
1643 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
1644 ])))
1645 };
1646 let point = || -> Box<dyn ManifoldVariable> {
1647 Box::new(Variable::new(rn::Rn::new(DVector::from_vec(vec![0.0; 3]))))
1648 };
1649 let mut variables: SlotMap<VarKey, Box<dyn ManifoldVariable>> = SlotMap::with_key();
1650 let cams: Vec<VarKey> = (0..3).map(|_| variables.insert(cam())).collect();
1651 let points: Vec<VarKey> = (0..4).map(|_| variables.insert(point())).collect();
1652 let mut variable_index_map: SecondaryMap<VarKey, usize> = SecondaryMap::new();
1653 for (i, &k) in cams.iter().enumerate() {
1654 variable_index_map.insert(k, i * 6);
1655 }
1656 for (i, &k) in points.iter().enumerate() {
1657 variable_index_map.insert(k, 18 + i * 3);
1658 }
1659 let landmark_keys: std::collections::HashSet<VarKey> = points.iter().copied().collect();
1660
1661 let value = |r: usize, c: usize| ((r * 31 + c * 17) as f64 * 0.618_033_988_7).sin() + 1.5;
1665 let mut triplets: Vec<Triplet<usize, usize, f64>> = Vec::new();
1666 let mut row = 0usize;
1667 for p in 0..4 {
1668 for c in 0..3 {
1669 for _ in 0..2 {
1670 for k in 0..6 {
1671 triplets.push(Triplet::new(row, c * 6 + k, value(row, c * 6 + k)));
1672 if row % 2 == 1 {
1673 let other = ((c + 1) % 3) * 6 + k;
1674 triplets.push(Triplet::new(row, other, value(row, other)));
1675 }
1676 }
1677 for k in 0..3 {
1678 let col = 18 + p * 3 + k;
1679 triplets.push(Triplet::new(row, col, value(row, col)));
1680 }
1681 row += 1;
1682 }
1683 }
1684 }
1685 let jacobian = SparseColMat::try_new_from_triplets(row, 30, &triplets)?;
1686
1687 let mut solver = ImplicitSparseSchur::new();
1688 solver.initialize_structure(&variables, &variable_index_map, &landmark_keys)?;
1689 solver.ensure_structure(&jacobian)?;
1690 let partition = solver.require_partition()?.clone();
1691 solver
1692 .eliminated
1693 .gather_from_jacobian(&jacobian, &partition);
1694 solver.eliminated.invert_in_place(&partition)?;
1695
1696 let v = Mat::from_fn(partition.kept_dof(), 1, |i, _| {
1697 (i as f64 * 0.37).cos() - 0.2
1698 });
1699 let want = column_scatter_reference(&solver, &partition, &jacobian, &v);
1700
1701 for chunk_rows in [1usize, 2, 3, 5, 7, row, row + 4] {
1702 let spans = RowSpans::build(&jacobian, &partition, chunk_rows)?;
1703 solver.structure.chunk_rows = chunk_rows;
1704 solver.structure.kept_spans = spans.kept;
1705 solver.structure.eliminated_spans = spans.eliminated;
1706 solver.structure.eliminated_block_of = spans.block_of;
1707
1708 let mut out = Mat::zeros(partition.kept_dof(), 1);
1709 let mut out_buf = vec![0.0; partition.kept_dof()];
1710 let mut rows = vec![0.0; jacobian.nrows()];
1711 let mut temp_lm = vec![0.0; partition.eliminated_dof()];
1712 let mut scratch = vec![0.0; partition.eliminated_dof()];
1713 solver.apply_schur_operator(
1714 &partition,
1715 &jacobian,
1716 &[],
1717 &v,
1718 &mut out,
1719 &mut out_buf,
1720 &mut rows,
1721 &mut temp_lm,
1722 &mut scratch,
1723 );
1724 for (i, w) in want.iter().enumerate() {
1725 assert_eq!(
1726 out[(i, 0)].to_bits(),
1727 w.to_bits(),
1728 "chunk_rows={chunk_rows}, entry {i}: {} vs {w}",
1729 out[(i, 0)]
1730 );
1731 }
1732 }
1733 Ok(())
1734 }
1735}