Skip to main content

apex_solver/linalg/sparse/schur/
implicit.rs

1//! # Implicit Sparse Schur Complement Solver
2//!
3//! Matrix-free Schur complement solved with Preconditioned Conjugate
4//! Gradients — Ceres's `ITERATIVE_SCHUR`, and matrix-free in the same sense
5//! Ceres means it: **neither `S` nor `JᵀJ` is ever formed**.
6//!
7//! ## What "implicit" has to mean
8//!
9//! Writing `J = [E | F]` for the eliminated and retained column sets, the
10//! reduced system is
11//!
12//! ```text
13//! S = FᵀF − FᵀE·(EᵀE)⁻¹·EᵀF
14//! ```
15//!
16//! and PCG only ever needs its action on a vector. Expanding that action so it
17//! reads `J` and nothing else gives the whole solver:
18//!
19//! ```text
20//! y = F·v                       scatter over the retained columns
21//! t = Eᵀ·y                      gather over the eliminated columns
22//! u = (EᵀE + λD_e)⁻¹·t          block-diagonal, one small solve per block
23//! y ← y − E·u                   scatter over the eliminated columns
24//! S·v = Fᵀ·y + λD_k·v           gather over the retained columns
25//! ```
26//!
27//! Four passes over `J`'s nonzeros, no intermediate matrix. Ceres's own
28//! documentation states the same cost model: "the cost of this evaluation
29//! scales with the number of non-zeros in the Jacobian".
30//!
31//! That distinction is the reason this solver exists. `JᵀJ` for a bundle
32//! adjustment problem carries a dense block for **every pair of cameras
33//! sharing a landmark** — fill-in `J` does not have — so forming it costs far
34//! more memory than `J`, and walking it costs far more time than walking `J`.
35//! A solver that built `JᵀJ` and only skipped `S` would avoid the smaller of
36//! the two costs while paying the larger one.
37//!
38//! ## When to use it
39//!
40//! When `JᵀJ` or `S` will not fit, or barely fits. The explicit solvers
41//! factorize an exact `S` and are faster whenever they fit in memory, so this
42//! is the large-problem fallback rather than the default —
43//! `ExplicitSparseSchur` remains what
44//! [`LevenbergMarquardtConfig::for_bundle_adjustment`](crate::optimizer::levenberg_marquardt::LevenbergMarquardtConfig::for_bundle_adjustment)
45//! selects.
46//!
47//! ## Preconditioning
48//!
49//! [`SchurPreconditioner`] selects between the three Ceres offers for
50//! `ITERATIVE_SCHUR`, all built from `J`:
51//!
52//! | This crate | Ceres | Built from |
53//! |---|---|---|
54//! | [`SchurPreconditioner::None`] | `IDENTITY` | — |
55//! | [`SchurPreconditioner::BlockDiagonal`] | `JACOBI` | diagonal blocks of `FᵀF` |
56//! | [`SchurPreconditioner::SchurJacobi`] | `SCHUR_JACOBI` | diagonal blocks of `S` |
57//!
58//! `SchurJacobi` is the default and generally the best of the three; it costs
59//! one pass over the observations to build.
60//!
61//! ## Automatic group recognition
62//!
63//! Like every other Schur solver, the eliminated ("group 0") and retained
64//! ("group 1") variable sets come from [`SchurPartition`] — the same
65//! structure [`ExplicitSparseSchur`](super::explicit::ExplicitSparseSchur) and
66//! [`ExplicitDenseSchur`](crate::linalg::dense::schur::ExplicitDenseSchur)
67//! use. This solver places no restriction on the eliminated variables' DOF or
68//! column layout: 3-D points, 1-DOF inverse-depth landmarks, or a mix of both
69//! in one problem all work identically, and the retained and eliminated
70//! columns may interleave arbitrarily.
71//!
72//! ## Usage Example
73//!
74//! ```no_run
75//! # use apex_solver::linalg::{ImplicitSparseSchur, SchurPreconditioner};
76//! # use apex_solver::linalg::StructureAware;
77//! # use apex_solver::core::VarKey;
78//! # use apex_solver::core::variable::ManifoldVariable;
79//! # use slotmap::{SlotMap, SecondaryMap};
80//! # use std::collections::HashSet;
81//! # fn example() -> Result<(), Box<dyn std::error::Error>> {
82//! # let variables: SlotMap<VarKey, Box<dyn ManifoldVariable>> = SlotMap::with_key();
83//! # let variable_index_map: SecondaryMap<VarKey, usize> = SecondaryMap::new();
84//! # let landmark_keys: HashSet<VarKey> = HashSet::new();
85//! use apex_solver::linalg::{ImplicitSparseSchur, SchurPreconditioner};
86//! use apex_solver::linalg::StructureAware;
87//!
88//! let mut solver = ImplicitSparseSchur::new()
89//!     .with_preconditioner(SchurPreconditioner::SchurJacobi)
90//!     .with_cg_config(500, 1e-9);
91//! solver.initialize_structure(&variables, &variable_index_map, &landmark_keys)?;
92//! # Ok(())
93//! # }
94//! ```
95
96use 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/// Structural facts about `J` that depend only on its sparsity, so they are
115/// rebuilt when the pattern changes and reused across every solve in between.
116#[derive(Debug, Clone, Default)]
117struct StructureCache {
118    /// Global column of each retained local index, and of each eliminated
119    /// local index.
120    ///
121    /// The operator's two gather passes write one output entry per column, so
122    /// with this mapping they parallelize over the output slice directly —
123    /// every entry written by exactly one task, no aliasing to reason about.
124    kept_cols: Vec<usize>,
125    eliminated_cols: Vec<usize>,
126    /// Sorted rows touched by each eliminated block.
127    eliminated_rows: Vec<Vec<usize>>,
128    /// Eliminated blocks visible from each kept block — they share a row.
129    visibility: Vec<Vec<usize>>,
130    /// Residual rows per chunk of the two row-space passes, `F·v` and `E·u`.
131    chunk_rows: usize,
132    /// Per row chunk, every retained column's nonzeros that fall inside it.
133    ///
134    /// The two row-space passes are scatters when run column by column, so
135    /// they cannot split over columns without two tasks writing one row.
136    /// Split over row chunks instead: each chunk owns its slice of `y`, and
137    /// replays the columns in the column-wise scatter's order, so every row
138    /// receives the same additions in the same order — bit-identical results.
139    kept_spans: Vec<Vec<ColumnSpan>>,
140    /// Same, over the eliminated columns; `local` is the eliminated index.
141    eliminated_spans: Vec<Vec<ColumnSpan>>,
142    /// Eliminated block owning each eliminated local index.
143    eliminated_block_of: Vec<u32>,
144    /// Pattern the cache was built from.
145    fingerprint: Option<pattern::PatternFingerprint>,
146}
147
148/// One column's nonzeros inside one row chunk: `J`'s value range `lo..hi`,
149/// and the column's local index in the vector the pass reads.
150///
151/// `u32` keeps the per-chunk lists cache-friendly; `ensure_structure` rejects
152/// a Jacobian whose nonzero count or dimensions do not fit.
153#[derive(Debug, Clone, Copy)]
154struct ColumnSpan {
155    lo: u32,
156    hi: u32,
157    local: u32,
158}
159
160/// Rows per chunk: enough chunks for load balance across the pool, few
161/// enough that each carries real work.
162fn 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
168/// The row-chunk span lists of [`StructureCache`] for one chunk length.
169struct RowSpans {
170    kept: Vec<Vec<ColumnSpan>>,
171    eliminated: Vec<Vec<ColumnSpan>>,
172    block_of: Vec<u32>,
173}
174
175impl RowSpans {
176    /// Push every column in exactly the order the column-wise scatter visits
177    /// them: retained blocks then offsets, eliminated blocks then offsets.
178    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
222/// Append `col`'s nonzeros, split at chunk boundaries, to the chunk lists.
223fn 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/// Implicit (matrix-free) Schur complement solver using Preconditioned
261/// Conjugate Gradients. Equivalent to Ceres's `ITERATIVE_SCHUR`.
262///
263/// Forms neither `S` nor `JᵀJ`; see the module documentation for the operator.
264#[derive(Debug, Clone)]
265pub struct ImplicitSparseSchur {
266    partition: Option<SchurPartition>,
267    /// `(EᵀE + λD_e)⁻¹` per eliminated block, rebuilt once per solve.
268    eliminated: EliminatedBlocks,
269    /// Largest eliminated block's DOF, for sizing the block-apply scratch.
270    max_eliminated_dof: usize,
271    /// Automatic eliminated/retained ("group 0"/"group 1") classification,
272    /// combined with manual marks in `initialize_structure`.
273    ordering: SchurOrdering,
274
275    // CG parameters
276    max_cg_iterations: usize,
277    cg_tolerance: f64,
278    /// Forcing-sequence parameter η; `0.0` disables the quadratic-model rule.
279    cg_q_tolerance: f64,
280
281    // Preconditioner type
282    preconditioner_type: SchurPreconditioner,
283
284    /// `J` of the last successful solve.
285    ///
286    /// Held so [`LinearSolver::hessian_vec_product`] can evaluate `Jᵀ(J·v)`
287    /// for the optimizers' quadratic model — the same arrangement
288    /// `ExplicitSparseSchur`'s chunked variant uses. One copy of `J`, against
289    /// the `JᵀJ` this solver exists to avoid.
290    jacobian: Option<SparseColMat<usize, f64>>,
291    /// Pattern the retained `jacobian` was built from, so its storage can be
292    /// reused across solves instead of reallocated.
293    jacobian_pattern: Option<pattern::PatternFingerprint>,
294    /// `+Jᵀr`, published through [`LinearSolver::get_gradient`].
295    gradient: Option<Mat<f64>>,
296
297    // Workspace buffers for the Schur operator (avoid repeated allocations).
298    workspace_rows: Vec<f64>, // residual-row sized buffer
299    workspace_lm: Vec<f64>,   // eliminated-DOF sized buffer
300    workspace_cam: Vec<f64>,  // kept-DOF sized buffer (parallel gather output)
301    block_scratch: Vec<f64>,  // eliminated-DOF sized: u = (EᵀE+λD_e)⁻¹·t
302
303    structure: StructureCache,
304}
305
306impl ImplicitSparseSchur {
307    /// Default: Schur-Jacobi preconditioner, 500 max iterations, 1e-9 relative
308    /// tolerance — matching Ceres's `ITERATIVE_SCHUR` defaults.
309    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    /// Set the PCG iteration cap and relative residual tolerance.
331    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    /// Select the preconditioner — see the table in the module documentation.
338    ///
339    /// Mirrors
340    /// [`ExplicitSparseSchur::with_preconditioner`](super::explicit::ExplicitSparseSchur::with_preconditioner),
341    /// so the two PCG-based solvers are configured the same way.
342    pub fn with_preconditioner(mut self, preconditioner: SchurPreconditioner) -> Self {
343        self.preconditioner_type = preconditioner;
344        self
345    }
346
347    /// Set the PCG forcing-sequence parameter η (`0.0` disables the
348    /// quadratic-model stopping rule).
349    pub fn with_cg_q_tolerance(mut self, q_tolerance: f64) -> Self {
350        self.cg_q_tolerance = q_tolerance;
351        self
352    }
353
354    /// Set how eliminated variables are recognized.
355    pub fn with_ordering(mut self, ordering: SchurOrdering) -> Self {
356        self.ordering = ordering;
357        self
358    }
359
360    /// Construct with explicit CG parameters, keeping the default
361    /// preconditioner.
362    pub fn with_cg_params(max_iterations: usize, tolerance: f64) -> Self {
363        Self::new().with_cg_config(max_iterations, tolerance)
364    }
365
366    /// Construct with explicit CG parameters and preconditioner.
367    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    /// The partition, once [`StructureAware::initialize_structure`] has run.
378    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    /// Rebuild the structural caches if `J`'s sparsity changed.
392    ///
393    /// Also enforces the elimination precondition against `J` directly: a
394    /// residual row touching two eliminated variables means `EᵀE` is not
395    /// block-diagonal, so inverting it blockwise would silently produce a wrong
396    /// step. `SchurPartition::verify_block_diagonal` makes the same check
397    /// against `JᵀJ`, which this solver never forms.
398    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        // Row -> owning eliminated block, and each block's sorted row set.
408        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        // Eliminated blocks visible from each kept block, via shared rows.
439        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    /// `S·v = Fᵀ(F·v − E·(EᵀE+λD_e)⁻¹·Eᵀ·F·v) + λD_k·v`, reading only `J`.
493    ///
494    /// `damp_kept` carries `λ·D_k` per retained local column, already clamped;
495    /// it is empty for an undamped solve. The eliminated side's damping is
496    /// baked into the inverted blocks before this is ever called.
497    #[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        // y = F·v — per row chunk, replaying the column-wise scatter's order.
516        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        // t = Eᵀ·y — a gather, so it parallelizes over the output entries.
534        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        // u = (EᵀE + λD_e)⁻¹·t — one output entry per task: entry `base + r`
549        // of block `b` is row `r` of that block's inverse times its slice of t.
550        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        // y ← y − E·u, again per row chunk in the scatter's order.
567        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        // S·v = Fᵀ·y + λD_k·v — also a gather, also parallel over outputs.
585        // The result lands in `out_buf` and is copied into `out` afterwards:
586        // `Mat` does not hand out a mutable slice rayon can split.
587        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    /// Dense `|rows| × dof` strip of `J` over `rows`, for the columns of
611    /// `block`.
612    ///
613    /// `rows` is sorted and each CSC column's row indices are sorted, so each
614    /// entry is one binary search. Rows the column does not touch stay zero.
615    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    /// Build the PCG preconditioner blocks from `J`.
636    ///
637    /// See the table in the module documentation for what each option is. All
638    /// of them produce one dense `dof × dof` inverse per retained variable,
639    /// or `None` for [`SchurPreconditioner::None`]. A block that will not
640    /// invert falls back to identity: a preconditioner is a convergence aid,
641    /// so a singular block must not fail the solve.
642    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                // (FᵀF + λD_k) restricted to this block.
661                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                // Schur-Jacobi additionally subtracts the elimination
673                // correction, which is what makes it a preconditioner for `S`
674                // rather than for the unreduced retained block.
675                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                        // m = F_iᵀ·E_j, the (i, j) coupling block of `FᵀE`.
688                        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    /// `z = M⁻¹·r` for the block-diagonal preconditioner, or `z = r` without one.
706    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    /// The whole solve, from `J` and `r` to the full-length update.
726    ///
727    /// `damping` is `None` for the plain normal equations.
728    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        // g = Jᵀr; the reduced system is built from −g, as in every other
747        // Schur solver here.
748        let gradient = jt_vec(jacobian, residuals);
749        let kept_dof = partition.kept_dof();
750        let eliminated_dof = partition.eliminated_dof();
751
752        // λ·D per column, from diag(JᵀJ) — the only part of `JᵀJ` needed.
753        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        // (EᵀE + λD_e)⁻¹, blockwise.
771        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        // g_reduced = −g_k + H_ke·H_ee⁻¹·g_e, assembled with `neg_gradient`
781        // playing the role of `g` (so the reduced right-hand side matches the
782        // explicit solvers' `g_k − H_ke·H_ee⁻¹·g_e` convention).
783        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        // correction = H_ke·H_ee⁻¹·g_e = Fᵀ·E·H_ee⁻¹·g_e, again through `J`.
801        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        // δ_e = H_ee⁻¹·(g_e − H_keᵀ·δ_k) = H_ee⁻¹·(g_e − Eᵀ·F·δ_k)
817        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    /// Keep `J` for [`LinearSolver::hessian_vec_product`], reusing the existing
830    /// allocation when the sparsity has not changed.
831    ///
832    /// Only the values move in that case — on the largest BAL problem a fresh
833    /// clone would allocate and free ~2 GB per optimizer iteration. The
834    /// fingerprint guards the reuse: identical nonzero counts alone would not
835    /// prove the patterns match.
836    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    /// `x = H_ee⁻¹·b`, blockwise.
854    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    /// `H_ke·u = Fᵀ·(E·u)`, from `J`.
870    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    /// `H_keᵀ·v = Eᵀ·(F·v)`, from `J`.
914    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    /// Solve `S·x = b` with the shared PCG primitive and the matrix-free operator.
958    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        // Workspaces come out of `self` so the operator closure can borrow
967        // `self` immutably alongside them.
968        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        // The PCG iteration count is the whole cost of this solver — each one
1007        // is four passes over `nnz(J)` — so it is the first number to look at
1008        // when tuning `cg_tolerance` or comparing preconditioners.
1009        debug!(
1010            "PCG: {} iterations, residual {:.3e}, {:?} ({:?})",
1011            result.iterations, result.final_residual, result.termination, self.preconditioner_type
1012        );
1013        result.x
1014    }
1015
1016    /// Scatter the two solution halves back into one full-length update.
1017    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        // `JᵀJ` was never formed, so evaluate its action from `J`.
1068        Some(jt_j_vec_product(self.jacobian.as_ref()?, v))
1069    }
1070
1071    /// Always `None`: this solver never materializes `JᵀJ`.
1072    ///
1073    /// That is the documented degradation for a matrix-free backend — callers
1074    /// use [`LinearSolver::hessian_vec_product`] instead, which is served
1075    /// exactly from `J`.
1076    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                // Free columns only — fixed tangent coordinates own no
1106                // column in the linear solve (ISSUE-0003).
1107                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    /// Build the same 2-camera + 3-landmark test setup used in explicit tests.
1155    /// Jacobian: 36 rows × 21 cols  (H_cc = 3·I₁₂, H_pp = 4·I₃ — positive definite)
1156    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        // `JᵀJ` is never formed, so `get_hessian` stays `None` by design and
1303        // the quadratic model is served through `hessian_vec_product`.
1304        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    /// A preconditioner changes how fast PCG converges, never what it
1319    /// converges to, so all three must produce the same step on the same
1320    /// system. This is what makes each option's implementation testable: a
1321    /// preconditioner that was silently ignored, or one built wrongly, shows
1322    /// up here as a disagreement or a failure to converge.
1323    #[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    /// A mixed-DOF, non-contiguous partition (1-DOF inverse depth interleaved
1487    /// with retained poses) must work end to end — the whole point of moving
1488    /// this solver onto [`SchurPartition`].
1489    #[test]
1490    fn test_implicit_schur_supports_non_contiguous_mixed_dof_partition() -> TestResult {
1491        use crate::linalg::schur::BlockSpan;
1492
1493        // Layout: kept A [0..6), eliminated P [6..7) (1-DOF), kept B [7..13),
1494        // eliminated Q [13..14) (1-DOF). Neither side is contiguous.
1495        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        // Observations: A-P, B-P, A-Q, B-Q, plus a small prior on every
1521        // column so the system is SPD.
1522        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        // A block spanning a variable never claimed by anything must still
1550        // round-trip through BlockSpan without panicking (sanity check on the
1551        // type used to build the partition above).
1552        let _ = BlockSpan {
1553            key: key_a,
1554            col_start: 0,
1555            dof: 6,
1556        };
1557        Ok(())
1558    }
1559
1560    /// The column-by-column operator the row-chunked passes replaced, kept
1561    /// verbatim as the reference they must reproduce bit for bit.
1562    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    /// Row-chunked `F·v` and `E·u` must equal the column-wise scatter exactly,
1637    /// for every chunk length — including ones that split a column's rows
1638    /// and rows that two retained blocks share, where addition order shows.
1639    #[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        // Every (camera, point) pair observed, two rows each; odd rows also
1662        // touch the next camera, so rows carry two retained blocks. Values are
1663        // irrational-ish so a reordered sum would change low bits.
1664        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}