use crate::core::VarKey;
use crate::core::variable::ManifoldVariable;
use crate::error::ErrorLogging;
use crate::linalg::regularization::invert_with_retry_dyn;
use crate::linalg::schur::{BlockSpan, EliminatedBlocks, SchurOrdering, SchurPartition};
use crate::linalg::schur::{PcgParams, SchurPreconditioner, solve_pcg};
use crate::linalg::sparse::normal_eq::{LazyNormalEquations, NormalEquations};
use crate::linalg::sparse::pattern;
use crate::linalg::{Damping, LinAlgError, LinAlgResult, LinearSolver, SparseMode, StructureAware};
use faer::sparse::{SparseColMat, SymbolicSparseColMat, Triplet};
use faer::{
Accum, Mat, Side,
linalg::solvers::Solve,
sparse::linalg::solvers::{Llt, SymbolicLlt},
};
use nalgebra::DMatrix;
use rayon::prelude::*;
use slotmap::{SecondaryMap, SlotMap};
use tracing::debug;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum ExplicitSchurVariant {
#[default]
Sparse,
Iterative,
Chunked,
}
#[derive(Debug, Clone)]
pub struct ExplicitSparseSchur {
partition: Option<SchurPartition>,
eliminated: EliminatedBlocks,
chunked: Option<super::chunk_eliminator::ChunkedSchurEliminator>,
chunked_jacobian: Option<SparseColMat<usize, f64>>,
verified_pattern: Option<pattern::PatternFingerprint>,
ordering: SchurOrdering,
variant: ExplicitSchurVariant,
preconditioner: SchurPreconditioner,
cg_q_tolerance: f64,
cg_max_iterations: usize,
cg_tolerance: f64,
ne_cache: LazyNormalEquations,
extraction: Option<ExtractionCache>,
s_bitmap: Option<(pattern::PatternFingerprint, Vec<u64>)>,
hessian: Option<SparseColMat<usize, f64>>,
gradient: Option<Mat<f64>>,
s_llt_cache: Option<(u64, SymbolicLlt<usize>)>,
}
#[derive(Debug, Clone)]
struct ExtractionCache {
fingerprint: pattern::PatternFingerprint,
h_kk_symbolic: SymbolicSparseColMat<usize>,
h_ke_symbolic: SymbolicSparseColMat<usize>,
h_kk_src: Vec<u32>,
h_ke_src: Vec<u32>,
}
impl ExtractionCache {
fn build(partition: &SchurPartition, hessian: &SparseColMat<usize, f64>) -> LinAlgResult<Self> {
let symbolic = hessian.symbolic();
let mut kept_cols = vec![u32::MAX; hessian.ncols()];
for block in partition.kept_blocks() {
for offset in 0..block.dof {
let global = block.col_start + offset;
if let Some(local) = partition.kept_local(global) {
kept_cols[global] = local as u32;
}
}
}
let mut elim_cols = vec![u32::MAX; hessian.ncols()];
for (block_idx, block) in partition.eliminated_blocks().iter().enumerate() {
let base = partition.eliminated_offset(block_idx);
for offset in 0..block.dof {
elim_cols[block.col_start + offset] = (base + offset) as u32;
}
}
let mut kk_triplets: Vec<(usize, usize, u32)> = Vec::new();
let mut ke_triplets: Vec<(usize, usize, u32)> = Vec::new();
for col in 0..hessian.ncols() {
let col_start = symbolic.col_range(col).start;
let rows = symbolic.row_idx_of_col_raw(col);
let c_kept = kept_cols[col];
let c_elim = elim_cols[col];
for (k, &row) in rows.iter().enumerate() {
let src = (col_start + k) as u32;
if c_kept != u32::MAX {
if let Some(r) = partition.kept_local(row) {
kk_triplets.push((r, c_kept as usize, src));
}
} else if c_elim != u32::MAX
&& let Some(r) = partition.kept_local(row)
{
ke_triplets.push((r, c_elim as usize, src));
}
}
}
let (h_kk_symbolic, h_kk_src) =
Self::symbolic_and_src(kk_triplets, partition.kept_dof(), partition.kept_dof())?;
let (h_ke_symbolic, h_ke_src) = Self::symbolic_and_src(
ke_triplets,
partition.kept_dof(),
partition.eliminated_dof(),
)?;
Ok(Self {
fingerprint: pattern::PatternFingerprint::of(hessian),
h_kk_symbolic,
h_ke_symbolic,
h_kk_src,
h_ke_src,
})
}
fn symbolic_and_src(
mut triplets: Vec<(usize, usize, u32)>,
nrows: usize,
ncols: usize,
) -> LinAlgResult<(SymbolicSparseColMat<usize>, Vec<u32>)> {
triplets.sort_unstable_by_key(|&(r, c, _)| (c, r));
let nnz = triplets.len();
let mut col_ptr = vec![0usize; ncols + 1];
let mut row_idx = Vec::with_capacity(nnz);
let mut src = Vec::with_capacity(nnz);
for &(r, c, s) in &triplets {
col_ptr[c + 1] += 1;
row_idx.push(r);
src.push(s);
}
for c in 0..ncols {
col_ptr[c + 1] += col_ptr[c];
}
let symbolic = SymbolicSparseColMat::new_checked(nrows, ncols, col_ptr, None, row_idx);
Ok((symbolic, src))
}
fn extract(
&self,
hessian: &SparseColMat<usize, f64>,
) -> (SparseColMat<usize, f64>, SparseColMat<usize, f64>) {
let values = hessian.as_ref().val();
let kk: Vec<f64> = self
.h_kk_src
.par_iter()
.map(|&s| values[s as usize])
.collect();
let ke: Vec<f64> = self
.h_ke_src
.par_iter()
.map(|&s| values[s as usize])
.collect();
(
SparseColMat::new(self.h_kk_symbolic.clone(), kk),
SparseColMat::new(self.h_ke_symbolic.clone(), ke),
)
}
}
#[inline]
fn set_dense_bit(bits: &mut [u64], wpc: usize, row: usize, col: usize) {
bits[col * wpc + (row >> 6)] |= 1u64 << (row & 63);
}
fn build_schur_output(
kept_dof: usize,
wpc: usize,
bits: &[u64],
s_dense: &[f64],
) -> (SymbolicSparseColMat<usize>, Vec<f64>, u64) {
const FNV_OFFSET: u64 = 0xcbf29ce484222325;
const FNV_PRIME: u64 = 0x100000001b3;
let mut pattern_hash = FNV_OFFSET;
let mut col_counts = vec![0usize; kept_dof];
let mut nnz = 0usize;
for c in 0..kept_dof {
let base = c * wpc;
for w in 0..wpc {
let mut word = bits[base + w];
while word != 0 {
let bit = word.trailing_zeros() as usize;
word &= word - 1;
let r = (w << 6) | bit;
if r < kept_dof {
let avg = (s_dense[r * kept_dof + c] + s_dense[c * kept_dof + r]) * 0.5;
if avg.abs() > 1e-12 {
col_counts[c] += 1;
nnz += 1;
}
}
}
}
}
let mut col_ptr = Vec::with_capacity(kept_dof + 1);
col_ptr.push(0usize);
for c in 0..kept_dof {
col_ptr.push(col_ptr[c] + col_counts[c]);
}
let mut row_idx = vec![0usize; nnz];
let mut s_values = vec![0.0f64; nnz];
let mut offsets = col_ptr[..kept_dof].to_vec();
for c in 0..kept_dof {
let base = c * wpc;
pattern_hash = (pattern_hash ^ (c as u64)).wrapping_mul(FNV_PRIME);
pattern_hash = (pattern_hash ^ (col_counts[c] as u64)).wrapping_mul(FNV_PRIME);
for w in 0..wpc {
let mut word = bits[base + w];
while word != 0 {
let bit = word.trailing_zeros() as usize;
word &= word - 1;
let r = (w << 6) | bit;
if r < kept_dof {
let avg = (s_dense[r * kept_dof + c] + s_dense[c * kept_dof + r]) * 0.5;
if avg.abs() > 1e-12 {
let slot = offsets[c];
offsets[c] += 1;
row_idx[slot] = r;
s_values[slot] = avg;
pattern_hash = (pattern_hash ^ (r as u64)).wrapping_mul(FNV_PRIME);
}
}
}
}
}
let symbolic = SymbolicSparseColMat::new_checked(kept_dof, kept_dof, col_ptr, None, row_idx);
(symbolic, s_values, pattern_hash)
}
impl ExplicitSparseSchur {
pub fn new() -> Self {
Self {
partition: None,
eliminated: EliminatedBlocks::default(),
chunked: None,
chunked_jacobian: None,
verified_pattern: None,
ordering: SchurOrdering::default(),
variant: ExplicitSchurVariant::default(),
preconditioner: SchurPreconditioner::default(),
cg_q_tolerance: crate::linalg::schur::DEFAULT_ETA,
cg_max_iterations: 200, cg_tolerance: 1e-6, ne_cache: LazyNormalEquations::default(),
extraction: None,
s_bitmap: None,
hessian: None,
gradient: None,
s_llt_cache: None,
}
}
pub fn with_ordering(mut self, ordering: SchurOrdering) -> Self {
self.ordering = ordering;
self
}
pub fn with_variant(mut self, variant: ExplicitSchurVariant) -> Self {
self.variant = variant;
self
}
pub fn with_preconditioner(mut self, preconditioner: SchurPreconditioner) -> Self {
self.preconditioner = preconditioner;
self
}
pub fn with_cg_q_tolerance(mut self, q_tolerance: f64) -> Self {
self.cg_q_tolerance = q_tolerance;
self
}
pub fn with_cg_params(mut self, max_iter: usize, tol: f64) -> Self {
self.cg_max_iterations = max_iter;
self.cg_tolerance = tol;
self
}
pub fn partition(&self) -> Option<&SchurPartition> {
self.partition.as_ref()
}
fn ensure_block_diagonal(&mut self, hessian: &SparseColMat<usize, f64>) -> LinAlgResult<()> {
let fingerprint = pattern::PatternFingerprint::of(hessian);
if self.verified_pattern == Some(fingerprint) {
return Ok(());
}
self.require_partition()?.verify_block_diagonal(hessian)?;
self.verified_pattern = Some(fingerprint);
Ok(())
}
fn require_partition(&self) -> LinAlgResult<&SchurPartition> {
self.partition.as_ref().ok_or_else(|| {
LinAlgError::InvalidInput(
"Block structure not built. Call initialize_structure() first.".to_string(),
)
.log()
})
}
fn build_partition(
&mut self,
variables: &SlotMap<VarKey, Box<dyn ManifoldVariable>>,
variable_index_map: &SecondaryMap<VarKey, usize>,
eliminate_keys: &std::collections::HashSet<VarKey>,
) -> LinAlgResult<()> {
let mut kept = Vec::new();
let mut eliminated = Vec::new();
for (key, variable) in variables {
let col_start = *variable_index_map.get(key).ok_or_else(|| {
LinAlgError::InvalidInput(format!("VarKey {:?} not found in index map", key)).log()
})?;
let span = BlockSpan {
key,
col_start,
dof: variable.free_dof(),
};
if eliminate_keys.contains(&key) {
eliminated.push(span);
} else {
kept.push(span);
}
}
let partition = SchurPartition::new(kept, eliminated)?;
self.eliminated = EliminatedBlocks::new(&partition);
self.partition = Some(partition);
self.verified_pattern = None;
Ok(())
}
fn extract_kept_and_coupling(
&mut self,
hessian: &SparseColMat<usize, f64>,
) -> LinAlgResult<(SparseColMat<usize, f64>, SparseColMat<usize, f64>)> {
let fingerprint = pattern::PatternFingerprint::of(hessian);
let cache = match self.extraction.take() {
Some(cache) if cache.fingerprint == fingerprint => cache,
_ => ExtractionCache::build(self.require_partition()?, hessian)?,
};
let out = cache.extract(hessian);
self.extraction = Some(cache);
Ok(out)
}
#[cfg_attr(not(test), allow(dead_code))]
fn extract_kept_block(
&self,
hessian: &SparseColMat<usize, f64>,
) -> LinAlgResult<SparseColMat<usize, f64>> {
let partition = self.require_partition()?;
let kept_dof = partition.kept_dof();
let symbolic = hessian.symbolic();
let mut triplets = Vec::new();
for block in partition.kept_blocks() {
for offset in 0..block.dof {
let global_col = block.col_start + offset;
let Some(local_col) = partition.kept_local(global_col) else {
continue;
};
let rows = symbolic.row_idx_of_col_raw(global_col);
let vals = hessian.val_of_col(global_col);
for (idx, &global_row) in rows.iter().enumerate() {
if let Some(local_row) = partition.kept_local(global_row) {
triplets.push(Triplet::new(local_row, local_col, vals[idx]));
}
}
}
}
SparseColMat::try_new_from_triplets(kept_dof, kept_dof, &triplets)
.map_err(|e| LinAlgError::SparseMatrixCreation(format!("H_kk: {:?}", e)))
}
#[cfg_attr(not(test), allow(dead_code))]
fn extract_coupling_block(
&self,
hessian: &SparseColMat<usize, f64>,
) -> LinAlgResult<SparseColMat<usize, f64>> {
let partition = self.require_partition()?;
let symbolic = hessian.symbolic();
let mut triplets = Vec::new();
for (block_idx, block) in partition.eliminated_blocks().iter().enumerate() {
let col_base = partition.eliminated_offset(block_idx);
for offset in 0..block.dof {
let global_col = block.col_start + offset;
let rows = symbolic.row_idx_of_col_raw(global_col);
let vals = hessian.val_of_col(global_col);
for (idx, &global_row) in rows.iter().enumerate() {
if let Some(local_row) = partition.kept_local(global_row) {
triplets.push(Triplet::new(local_row, col_base + offset, vals[idx]));
}
}
}
}
SparseColMat::try_new_from_triplets(
partition.kept_dof(),
partition.eliminated_dof(),
&triplets,
)
.map_err(|e| LinAlgError::SparseMatrixCreation(format!("H_ke: {:?}", e)))
}
fn extract_gradient_blocks(&self, gradient: &Mat<f64>) -> LinAlgResult<(Mat<f64>, Mat<f64>)> {
let partition = self.require_partition()?;
let mut g_k = Mat::zeros(partition.kept_dof(), 1);
let mut g_e = Mat::zeros(partition.eliminated_dof(), 1);
for block in partition.kept_blocks() {
for offset in 0..block.dof {
let global = block.col_start + offset;
if let Some(local) = partition.kept_local(global) {
g_k[(local, 0)] = gradient[(global, 0)];
}
}
}
for (block_idx, block) in partition.eliminated_blocks().iter().enumerate() {
let base = partition.eliminated_offset(block_idx);
for offset in 0..block.dof {
g_e[(base + offset, 0)] = gradient[(block.col_start + offset, 0)];
}
}
Ok((g_k, g_e))
}
fn solve_with_cholesky(
&mut self,
a: &SparseColMat<usize, f64>,
b: &Mat<f64>,
pattern_hash: Option<u64>,
) -> LinAlgResult<Mat<f64>> {
if let (Some(hash), Some((cached_hash, sym))) = (pattern_hash, self.s_llt_cache.as_ref())
&& *cached_hash == hash
{
if let Ok(cholesky) = Llt::try_new_with_symbolic(sym.clone(), a.as_ref(), Side::Lower) {
return Ok(cholesky.solve(b));
}
self.s_llt_cache = None;
}
let sym = SymbolicLlt::try_new(a.symbolic(), Side::Lower).map_err(|e| {
LinAlgError::FactorizationFailed(format!("Symbolic Cholesky failed: {:?}", e))
})?;
self.s_llt_cache = pattern_hash.map(|hash| (hash, sym.clone()));
match Llt::try_new_with_symbolic(sym.clone(), a.as_ref(), Side::Lower) {
Ok(cholesky) => return Ok(cholesky.solve(b)),
Err(e) => {
debug!(
"Cholesky factorization failed: {:?}. Applying regularization.",
e
);
}
}
let n = a.nrows();
let symbolic = a.symbolic();
let mut trace = 0.0;
let mut max_diag = 0.0f64;
for col in 0..n {
let row_indices = symbolic.row_idx_of_col_raw(col);
let col_values = a.val_of_col(col);
for (idx, &row) in row_indices.iter().enumerate() {
if row == col {
trace += col_values[idx];
max_diag = max_diag.max(col_values[idx].abs());
}
}
}
let avg_diag = trace / n as f64;
let base_reg = avg_diag.max(max_diag).max(1.0);
for attempt in 0..5 {
let reg = base_reg * 10.0f64.powi(attempt - 4); debug!(
"Cholesky attempt {}: regularization = {:.2e}",
attempt + 2,
reg
);
let mut triplets = Vec::with_capacity(n * 10);
for col in 0..n {
let row_indices = symbolic.row_idx_of_col_raw(col);
let col_values = a.val_of_col(col);
for (idx, &row) in row_indices.iter().enumerate() {
triplets.push(Triplet::new(row, col, col_values[idx]));
}
}
for i in 0..n {
triplets.push(Triplet::new(i, i, reg));
}
let a_reg = match SparseColMat::try_new_from_triplets(n, n, &triplets) {
Ok(m) => m,
Err(e) => {
debug!("Failed to create regularized matrix: {:?}", e);
continue;
}
};
let sym_reg = match SymbolicLlt::try_new(a_reg.symbolic(), Side::Lower) {
Ok(s) => s,
Err(e) => {
debug!("Symbolic factorization failed: {:?}", e);
continue;
}
};
match Llt::try_new_with_symbolic(sym_reg, a_reg.as_ref(), Side::Lower) {
Ok(cholesky) => {
debug!("Cholesky succeeded with regularization {:.2e}", reg);
return Ok(cholesky.solve(b));
}
Err(e) => {
debug!("Cholesky failed with reg {:.2e}: {:?}", reg, e);
}
}
}
Err(LinAlgError::SingularMatrix(format!(
"Schur complement singular after 5 regularization attempts (max reg = {:.2e})",
base_reg
)))
}
fn solve_with_pcg(
&self,
a: &SparseColMat<usize, f64>,
b: &Mat<f64>,
h_kk: &SparseColMat<usize, f64>,
) -> LinAlgResult<Mat<f64>> {
let precond = self.build_preconditioner(a, h_kk)?;
let result = solve_pcg(
b,
&PcgParams::new(self.cg_max_iterations, self.cg_tolerance)
.with_q_tolerance(self.cg_q_tolerance),
|p, ap| {
faer::sparse::linalg::matmul::sparse_dense_matmul(
ap.as_mut(),
Accum::Replace,
a.as_ref(),
p.as_ref(),
1.0,
faer::get_global_parallelism(),
);
},
|r, z| apply_block_preconditioner(precond.as_deref(), r, z),
);
Ok(result.x)
}
fn build_preconditioner(
&self,
s: &SparseColMat<usize, f64>,
h_kk: &SparseColMat<usize, f64>,
) -> LinAlgResult<Option<Vec<DMatrix<f64>>>> {
let source = match self.preconditioner {
SchurPreconditioner::None => return Ok(None),
SchurPreconditioner::SchurJacobi => s,
SchurPreconditioner::BlockDiagonal => h_kk,
};
let partition = self.require_partition()?;
let blocks = partition
.kept_blocks()
.par_iter()
.enumerate()
.map(|(idx, block)| {
let offset = partition.kept_offset(idx);
let dense = diagonal_block(source, offset, block.dof);
invert_with_retry_dyn(&dense)
.unwrap_or_else(|| DMatrix::identity(block.dof, block.dof))
})
.collect();
Ok(Some(blocks))
}
fn compute_schur_complement(
&mut self,
h_kk: &SparseColMat<usize, f64>,
h_ke: &SparseColMat<usize, f64>,
h_ee_inv: &EliminatedBlocks,
) -> LinAlgResult<(SparseColMat<usize, f64>, Option<u64>)> {
let partition = self.require_partition()?.clone();
let kept_dof = h_kk.nrows();
let h_ke_symbolic = h_ke.symbolic();
let wpc = kept_dof.div_ceil(64);
let mut s_dense = vec![0.0f64; kept_dof * kept_dof];
let extraction_fp = self.extraction.as_ref().map(|c| c.fingerprint);
let words = kept_dof * wpc;
let mut bits = match self.s_bitmap.take() {
Some((fp, b)) if Some(fp) == extraction_fp && b.len() == words => b,
_ => vec![0u64; words],
};
let h_kk_symbolic = h_kk.symbolic();
for col in 0..h_kk.ncols() {
let rows = h_kk_symbolic.row_idx_of_col_raw(col);
let vals = h_kk.val_of_col(col);
for (idx, &row) in rows.iter().enumerate() {
s_dense[row * kept_dof + col] += vals[idx];
set_dense_bit(&mut bits, wpc, row, col);
}
}
let max_dof = partition
.eliminated_blocks()
.iter()
.map(|b| b.dof)
.max()
.unwrap_or(0)
.max(1);
let mut kept_rows: Vec<usize> = Vec::with_capacity(32);
let mut h_ke_rows: Vec<f64> = Vec::with_capacity(32 * max_dof);
let mut contrib: Vec<f64> = Vec::with_capacity(32 * max_dof);
let mut cursors: Vec<usize> = Vec::with_capacity(max_dof);
let mut col_rows: Vec<&[usize]> = Vec::with_capacity(max_dof);
let mut col_vals: Vec<&[f64]> = Vec::with_capacity(max_dof);
for block_idx in 0..h_ee_inv.len() {
let dof = h_ee_inv.dof(block_idx);
if dof == 0 {
continue;
}
let col_base = partition.eliminated_offset(block_idx);
kept_rows.clear();
h_ke_rows.clear();
cursors.clear();
cursors.resize(dof, 0);
col_rows.clear();
col_vals.clear();
for local in 0..dof {
let col = col_base + local;
col_rows.push(h_ke_symbolic.row_idx_of_col_raw(col));
col_vals.push(h_ke.val_of_col(col));
}
loop {
let mut min_row = usize::MAX;
for local in 0..dof {
if let Some(&row) = col_rows[local].get(cursors[local]) {
min_row = min_row.min(row);
}
}
if min_row == usize::MAX {
break;
}
let strip = h_ke_rows.len();
h_ke_rows.resize(strip + dof, 0.0);
for local in 0..dof {
if col_rows[local].get(cursors[local]) == Some(&min_row) {
h_ke_rows[strip + local] = col_vals[local][cursors[local]];
cursors[local] += 1;
}
}
kept_rows.push(min_row);
}
if kept_rows.is_empty() {
continue;
}
let inv = h_ee_inv.block(block_idx);
contrib.clear();
contrib.resize(kept_rows.len() * dof, 0.0);
for r in 0..kept_rows.len() {
let row = &h_ke_rows[r * dof..(r + 1) * dof];
for c in 0..dof {
let inv_col = &inv[c * dof..(c + 1) * dof];
let mut acc = 0.0;
for k in 0..dof {
acc += row[k] * inv_col[k];
}
contrib[r * dof + c] = acc;
}
}
if dof == 3 {
for (i, &row_i) in kept_rows.iter().enumerate() {
let (c0, c1, c2) = (contrib[i * 3], contrib[i * 3 + 1], contrib[i * 3 + 2]);
let base = row_i * kept_dof;
for (j, &row_j) in kept_rows.iter().enumerate() {
let dot = c0 * h_ke_rows[j * 3]
+ c1 * h_ke_rows[j * 3 + 1]
+ c2 * h_ke_rows[j * 3 + 2];
s_dense[base + row_j] -= dot;
set_dense_bit(&mut bits, wpc, row_i, row_j);
}
}
} else {
for (i, &row_i) in kept_rows.iter().enumerate() {
let contrib_i = &contrib[i * dof..(i + 1) * dof];
let base = row_i * kept_dof;
for (j, &row_j) in kept_rows.iter().enumerate() {
let h_ke_j = &h_ke_rows[j * dof..(j + 1) * dof];
let mut dot = 0.0;
for k in 0..dof {
dot += contrib_i[k] * h_ke_j[k];
}
s_dense[base + row_j] -= dot;
set_dense_bit(&mut bits, wpc, row_i, row_j);
}
}
}
}
let output = if self.extraction.is_some() {
let (symbolic, s_values, pattern_hash) =
build_schur_output(kept_dof, wpc, &bits, &s_dense);
let s = SparseColMat::new(symbolic, s_values);
return Ok((s, Some(pattern_hash)));
} else {
for i in 0..kept_dof {
for j in (i + 1)..kept_dof {
let avg = (s_dense[i * kept_dof + j] + s_dense[j * kept_dof + i]) * 0.5;
s_dense[i * kept_dof + j] = avg;
s_dense[j * kept_dof + i] = avg;
}
}
let mut s_triplets: Vec<Triplet<usize, usize, f64>> =
Vec::with_capacity(kept_dof.saturating_mul(8));
for row in 0..kept_dof {
let row_base = row * kept_dof;
for col in 0..kept_dof {
let val = s_dense[row_base + col];
if val.abs() > 1e-12 {
s_triplets.push(Triplet::new(row, col, val));
}
}
}
SparseColMat::try_new_from_triplets(kept_dof, kept_dof, &s_triplets)
.map_err(|e| LinAlgError::SparseMatrixCreation(format!("Schur S: {:?}", e)))?
};
self.s_bitmap = Some((extraction_fp.unwrap_or_default(), bits));
Ok((output, None))
}
fn compute_reduced_gradient(
&self,
g_k: &Mat<f64>,
g_e: &Mat<f64>,
h_ke: &SparseColMat<usize, f64>,
h_ee_inv: &EliminatedBlocks,
) -> LinAlgResult<Mat<f64>> {
let partition = self.require_partition()?;
let kept_dof = g_k.nrows();
let mut hee_inv_ge = Mat::zeros(g_e.nrows(), 1);
for block_idx in 0..h_ee_inv.len() {
let dof = h_ee_inv.dof(block_idx);
let base = partition.eliminated_offset(block_idx);
let inv = h_ee_inv.block(block_idx);
for r in 0..dof {
let mut acc = 0.0;
for c in 0..dof {
acc += inv[c * dof + r] * g_e[(base + c, 0)];
}
hee_inv_ge[(base + r, 0)] = acc;
}
}
let mut correction = Mat::<f64>::zeros(kept_dof, 1);
let symbolic = h_ke.symbolic();
for col in 0..h_ke.ncols() {
let rows = symbolic.row_idx_of_col_raw(col);
let vals = h_ke.val_of_col(col);
for (idx, &row) in rows.iter().enumerate() {
correction[(row, 0)] += vals[idx] * hee_inv_ge[(col, 0)];
}
}
let mut g_reduced = Mat::zeros(kept_dof, 1);
for i in 0..kept_dof {
g_reduced[(i, 0)] = g_k[(i, 0)] - correction[(i, 0)];
}
Ok(g_reduced)
}
fn back_substitute(
&self,
delta_k: &Mat<f64>,
g_e: &Mat<f64>,
h_ke: &SparseColMat<usize, f64>,
h_ee_inv: &EliminatedBlocks,
) -> LinAlgResult<Mat<f64>> {
let partition = self.require_partition()?;
let eliminated_dof = g_e.nrows();
let mut hke_t_delta = Mat::<f64>::zeros(eliminated_dof, 1);
let symbolic = h_ke.symbolic();
for col in 0..h_ke.ncols() {
let rows = symbolic.row_idx_of_col_raw(col);
let vals = h_ke.val_of_col(col);
let mut acc = 0.0;
for (idx, &row) in rows.iter().enumerate() {
acc += vals[idx] * delta_k[(row, 0)];
}
hke_t_delta[(col, 0)] = acc;
}
let mut delta_e = Mat::zeros(eliminated_dof, 1);
for block_idx in 0..h_ee_inv.len() {
let dof = h_ee_inv.dof(block_idx);
let base = partition.eliminated_offset(block_idx);
let inv = h_ee_inv.block(block_idx);
for r in 0..dof {
let mut acc = 0.0;
for c in 0..dof {
let rhs = g_e[(base + c, 0)] - hke_t_delta[(base + c, 0)];
acc += inv[c * dof + r] * rhs;
}
delta_e[(base + r, 0)] = acc;
}
}
Ok(delta_e)
}
}
impl Default for ExplicitSparseSchur {
fn default() -> Self {
Self::new()
}
}
impl StructureAware for ExplicitSparseSchur {
fn initialize_structure(
&mut self,
variables: &SlotMap<VarKey, Box<dyn ManifoldVariable>>,
variable_index_map: &SecondaryMap<VarKey, usize>,
schur_landmark_keys: &std::collections::HashSet<VarKey>,
) -> LinAlgResult<()> {
let effective_keys = crate::linalg::schur::effective_landmark_keys(
variables,
schur_landmark_keys,
&self.ordering,
);
self.build_partition(variables, variable_index_map, &effective_keys)?;
Ok(())
}
}
impl LinearSolver<SparseMode> for ExplicitSparseSchur {
fn solve_normal_equation(
&mut self,
residuals: &Mat<f64>,
jacobian: &SparseColMat<usize, f64>,
) -> LinAlgResult<Mat<f64>> {
self.require_partition()?;
if self.variant == ExplicitSchurVariant::Chunked {
return self.solve_chunked(residuals, jacobian, None);
}
let NormalEquations { hessian, gradient } = self.ne_cache.compute(residuals, jacobian)?;
let mut neg_gradient = Mat::zeros(gradient.nrows(), 1);
for i in 0..gradient.nrows() {
neg_gradient[(i, 0)] = -gradient[(i, 0)];
}
self.ensure_block_diagonal(&hessian)?;
let (h_kk, h_ke) = self.extract_kept_and_coupling(&hessian)?;
let (g_k, g_e) = self.extract_gradient_blocks(&neg_gradient)?;
let mut eliminated = std::mem::take(&mut self.eliminated);
let inversion = (|| -> LinAlgResult<()> {
let partition = self.require_partition()?;
eliminated.gather(&hessian, partition);
eliminated.invert_in_place(partition)
})();
if let Err(e) = inversion {
self.eliminated = eliminated;
return Err(e);
}
let result = self.solve_reduced_system(&h_kk, &h_ke, &g_k, &g_e, &eliminated);
self.eliminated = eliminated;
if result.is_ok() {
self.hessian = Some(hessian);
self.gradient = Some(gradient);
}
result
}
fn solve_augmented_equation(
&mut self,
residuals: &Mat<f64>,
jacobian: &SparseColMat<usize, f64>,
damping: &Damping,
) -> LinAlgResult<Mat<f64>> {
self.require_partition()?;
if self.variant == ExplicitSchurVariant::Chunked {
return self.solve_chunked(residuals, jacobian, Some(damping));
}
let NormalEquations { hessian, gradient } = self.ne_cache.compute(residuals, jacobian)?;
let mut neg_gradient = Mat::zeros(gradient.nrows(), 1);
for i in 0..gradient.nrows() {
neg_gradient[(i, 0)] = -gradient[(i, 0)];
}
self.ensure_block_diagonal(&hessian)?;
let partition = self.require_partition()?;
let kept_dof = partition.kept_dof();
let num_eliminated_blocks = partition.eliminated_blocks().len();
let (h_kk, h_ke) = self.extract_kept_and_coupling(&hessian)?;
let (g_k, g_e) = self.extract_gradient_blocks(&neg_gradient)?;
debug!(
"Iteration matrices: H {}×{}, H_kk {}×{}, H_ke {}×{}, {} eliminated blocks",
hessian.nrows(),
hessian.ncols(),
h_kk.nrows(),
h_kk.ncols(),
h_ke.nrows(),
h_ke.ncols(),
num_eliminated_blocks
);
let h_kk_damped = damp_camera_block(&h_kk, kept_dof, damping)?;
let mut eliminated = std::mem::take(&mut self.eliminated);
let prepared = (|| -> LinAlgResult<()> {
let partition = self.require_partition()?;
eliminated.gather(&hessian, partition);
eliminated.damp(damping);
eliminated.invert_in_place(partition)
})();
if let Err(e) = prepared {
self.eliminated = eliminated;
return Err(e);
}
let result = self.solve_reduced_system(&h_kk_damped, &h_ke, &g_k, &g_e, &eliminated);
self.eliminated = eliminated;
if result.is_ok() {
self.hessian = Some(hessian);
self.gradient = Some(gradient);
}
result
}
fn hessian_vec_product(&self, v: &Mat<f64>) -> Option<Mat<f64>> {
if let Some(h) = self.hessian.as_ref() {
return Some(
<SparseMode as crate::linearizer::AssemblyBackend>::hessian_vec_product(h, v),
);
}
let j = self.chunked_jacobian.as_ref()?;
let symbolic = j.symbolic();
let mut jv = Mat::<f64>::zeros(j.nrows(), 1);
for col in 0..j.ncols() {
let x = v[(col, 0)];
if x == 0.0 {
continue;
}
let rows = symbolic.row_idx_of_col_raw(col);
let vals = j.val_of_col(col);
for (i, &row) in rows.iter().enumerate() {
jv[(row, 0)] += vals[i] * x;
}
}
let mut out = Mat::<f64>::zeros(j.ncols(), 1);
for col in 0..j.ncols() {
let rows = symbolic.row_idx_of_col_raw(col);
let vals = j.val_of_col(col);
let mut acc = 0.0;
for (i, &row) in rows.iter().enumerate() {
acc += vals[i] * jv[(row, 0)];
}
out[(col, 0)] = acc;
}
Some(out)
}
fn get_hessian(&self) -> Option<&SparseColMat<usize, f64>> {
self.hessian.as_ref()
}
fn get_gradient(&self) -> Option<&Mat<f64>> {
self.gradient.as_ref()
}
}
impl ExplicitSparseSchur {
fn solve_chunked(
&mut self,
residuals: &Mat<f64>,
jacobian: &SparseColMat<usize, f64>,
damping: Option<&Damping>,
) -> LinAlgResult<Mat<f64>> {
if self.chunked.as_ref().is_none_or(|c| !c.matches(jacobian)) {
let partition = self.require_partition()?;
self.chunked = Some(super::chunk_eliminator::ChunkedSchurEliminator::new(
jacobian, partition,
)?);
}
let mut eliminator = match self.chunked.take() {
Some(e) => e,
None => {
return Err(LinAlgError::InvalidState(
"chunk eliminator not initialized".to_string(),
)
.log());
}
};
let outcome = (|| -> LinAlgResult<Mat<f64>> {
let partition = self.require_partition()?.clone();
let reduced = eliminator.eliminate(jacobian, residuals, &partition, damping)?;
let kept_dof = reduced.kept_dof;
let mut triplets: Vec<Triplet<usize, usize, f64>> =
Vec::with_capacity(kept_dof.saturating_mul(8));
for col in 0..kept_dof {
for row in 0..kept_dof {
let v = reduced.s[(row, col)];
if v.abs() > 1e-12 {
triplets.push(Triplet::new(row, col, v));
}
}
}
let s = SparseColMat::try_new_from_triplets(kept_dof, kept_dof, &triplets)
.map_err(|e| LinAlgError::SparseMatrixCreation(format!("Schur S: {e:?}")))?;
let mut rhs = Mat::<f64>::zeros(kept_dof, 1);
for i in 0..kept_dof {
rhs[(i, 0)] = -reduced.g_reduced[(i, 0)];
}
let delta_k = self.solve_with_cholesky(&s, &rhs, None)?;
let delta_e = self.back_substitute_chunked(&delta_k, &reduced, jacobian, &partition)?;
self.combine_updates(&delta_k, &delta_e)
})();
self.chunked = Some(eliminator);
if outcome.is_ok() {
self.gradient = Some(Self::full_gradient(jacobian, residuals));
self.hessian = None;
self.chunked_jacobian = Some(jacobian.clone());
}
outcome
}
fn full_gradient(jacobian: &SparseColMat<usize, f64>, residuals: &Mat<f64>) -> Mat<f64> {
let symbolic = jacobian.symbolic();
let per_col: Vec<f64> = (0..jacobian.ncols())
.into_par_iter()
.map(|col| {
let rows = symbolic.row_idx_of_col_raw(col);
let vals = jacobian.val_of_col(col);
let mut acc = 0.0;
for (i, &row) in rows.iter().enumerate() {
acc += vals[i] * residuals[(row, 0)];
}
acc
})
.collect();
Mat::from_fn(jacobian.ncols(), 1, |r, _| per_col[r])
}
fn back_substitute_chunked(
&self,
delta_k: &Mat<f64>,
reduced: &super::chunk_eliminator::ReducedSystem,
jacobian: &SparseColMat<usize, f64>,
partition: &SchurPartition,
) -> LinAlgResult<Mat<f64>> {
let symbolic = jacobian.symbolic();
let mut j_delta = Mat::<f64>::zeros(jacobian.nrows(), 1);
let mut local = 0usize;
for block in partition.kept_blocks() {
for offset in 0..block.dof {
let col = block.col_start + offset;
let rows = symbolic.row_idx_of_col_raw(col);
let vals = jacobian.val_of_col(col);
let x = delta_k[(local, 0)];
if x != 0.0 {
for (i, &row) in rows.iter().enumerate() {
j_delta[(row, 0)] += vals[i] * x;
}
}
local += 1;
}
}
let mut delta_e = Mat::<f64>::zeros(partition.eliminated_dof(), 1);
let max_dof = partition
.eliminated_blocks()
.iter()
.map(|block| block.dof)
.max()
.unwrap_or(0);
let mut etjd = vec![0.0f64; max_dof];
for (block_idx, block) in partition.eliminated_blocks().iter().enumerate() {
let dof = block.dof;
let base = partition.eliminated_offset(block_idx);
let inv = reduced.eliminated_inverse.block(block_idx);
for (a, slot) in etjd[..dof].iter_mut().enumerate() {
let col = block.col_start + a;
let rows = symbolic.row_idx_of_col_raw(col);
let vals = jacobian.val_of_col(col);
let mut acc = 0.0;
for (i, &row) in rows.iter().enumerate() {
acc += vals[i] * j_delta[(row, 0)];
}
*slot = acc;
}
for r in 0..dof {
let mut acc = -reduced.eliminated_rhs[(base + r, 0)];
for (c, &term) in etjd[..dof].iter().enumerate() {
acc -= inv[c * dof + r] * term;
}
delta_e[(base + r, 0)] = acc;
}
}
Ok(delta_e)
}
fn solve_reduced_system(
&mut self,
h_kk: &SparseColMat<usize, f64>,
h_ke: &SparseColMat<usize, f64>,
g_k: &Mat<f64>,
g_e: &Mat<f64>,
h_ee_inv: &EliminatedBlocks,
) -> LinAlgResult<Mat<f64>> {
let (s, s_pattern_hash) = self.compute_schur_complement(h_kk, h_ke, h_ee_inv)?;
let g_reduced = self.compute_reduced_gradient(g_k, g_e, h_ke, h_ee_inv)?;
let delta_k = match self.variant {
ExplicitSchurVariant::Iterative => self.solve_with_pcg(&s, &g_reduced, h_kk)?,
ExplicitSchurVariant::Sparse | ExplicitSchurVariant::Chunked => {
self.solve_with_cholesky(&s, &g_reduced, s_pattern_hash)?
}
};
let delta_e = self.back_substitute(&delta_k, g_e, h_ke, h_ee_inv)?;
self.combine_updates(&delta_k, &delta_e)
}
fn combine_updates(&self, delta_k: &Mat<f64>, delta_e: &Mat<f64>) -> LinAlgResult<Mat<f64>> {
let partition = self.require_partition()?;
let mut delta = Mat::zeros(partition.total_dof(), 1);
for block in partition.kept_blocks() {
for offset in 0..block.dof {
let global = block.col_start + offset;
if let Some(local) = partition.kept_local(global) {
delta[(global, 0)] = delta_k[(local, 0)];
}
}
}
for (block_idx, block) in partition.eliminated_blocks().iter().enumerate() {
let base = partition.eliminated_offset(block_idx);
for offset in 0..block.dof {
delta[(block.col_start + offset, 0)] = delta_e[(base + offset, 0)];
}
}
debug!(
"Update norms: delta_kept={:.6e}, delta_eliminated={:.6e}, combined={:.6e}",
delta_k.norm_l2(),
delta_e.norm_l2(),
delta.norm_l2()
);
Ok(delta)
}
}
fn diagonal_block(m: &SparseColMat<usize, f64>, offset: usize, dof: usize) -> DMatrix<f64> {
let symbolic = m.symbolic();
let mut block = DMatrix::zeros(dof, dof);
for local_col in 0..dof {
let col = offset + local_col;
let rows = symbolic.row_idx_of_col_raw(col);
let vals = m.val_of_col(col);
for (idx, &row) in rows.iter().enumerate() {
if row >= offset && row < offset + dof {
block[(row - offset, local_col)] = vals[idx];
}
}
}
block
}
fn apply_block_preconditioner(blocks: Option<&[DMatrix<f64>]>, r: &Mat<f64>, z: &mut Mat<f64>) {
let Some(blocks) = blocks else {
faer::zip!(z, r).for_each(|faer::unzip!(z, r)| *z = *r);
return;
};
let mut offset = 0usize;
for inv in blocks {
let dof = inv.nrows();
for row in 0..dof {
let mut acc = 0.0;
for col in 0..dof {
acc += inv[(row, col)] * r[(offset + col, 0)];
}
z[(offset + row, 0)] = acc;
}
offset += dof;
}
}
fn damp_camera_block(
h_cc: &SparseColMat<usize, f64>,
cam_size: usize,
damping: &Damping,
) -> LinAlgResult<SparseColMat<usize, f64>> {
let symbolic = h_cc.symbolic();
let diag_pos: Vec<Option<usize>> = (0..h_cc.ncols())
.map(|col| {
let col_start = symbolic.col_range(col).start;
symbolic
.row_idx_of_col_raw(col)
.iter()
.position(|&row| row == col)
.map(|local| col_start + local)
})
.collect();
if diag_pos.iter().all(Option::is_some) {
let owned = symbolic.to_owned().map_err(|e| {
LinAlgError::SparseMatrixCreation(format!("Damped H_cc symbolic: {e:?}"))
})?;
let mut values = h_cc.as_ref().val().to_vec();
for pos in diag_pos.iter().flatten() {
values[*pos] += damping.diagonal_term(values[*pos]);
}
return Ok(SparseColMat::new(owned, values));
}
let mut triplets = Vec::with_capacity(h_cc.compute_nnz() + cam_size);
for col in 0..h_cc.ncols() {
let row_indices = symbolic.row_idx_of_col_raw(col);
let col_values = h_cc.val_of_col(col);
for (idx, &row) in row_indices.iter().enumerate() {
let val = col_values[idx];
let val = if row == col {
val + damping.diagonal_term(val)
} else {
val
};
triplets.push(Triplet::new(row, col, val));
}
}
for (i, pos) in diag_pos.iter().enumerate().take(cam_size) {
if pos.is_none() {
triplets.push(Triplet::new(i, i, damping.diagonal_term(0.0)));
}
}
SparseColMat::try_new_from_triplets(cam_size, cam_size, &triplets)
.map_err(|e| LinAlgError::SparseMatrixCreation(format!("Damped H_cc: {e:?}")))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::core::VarKey;
use crate::core::variable::Variable;
use apex_manifolds::{LieGroup, ManifoldType, rn, se3};
use nalgebra::DVector;
use slotmap::{SecondaryMap, SlotMap};
type TestResult = Result<(), Box<dyn std::error::Error>>;
type TestSetup = (
SlotMap<VarKey, Box<dyn ManifoldVariable>>,
SecondaryMap<VarKey, usize>,
SparseColMat<usize, f64>,
Mat<f64>,
std::collections::HashSet<VarKey>,
);
fn create_schur_test_setup() -> Result<TestSetup, Box<dyn std::error::Error>> {
let se3_id = DVector::from_vec(vec![1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]);
let pt_zero = DVector::from_vec(vec![0.0, 0.0, 0.0]);
let mut variables: SlotMap<VarKey, Box<dyn ManifoldVariable>> = SlotMap::with_key();
let cam0 = variables.insert(Box::new(Variable::new(se3::SE3::from_param_slice(
se3_id.as_slice(),
))));
let cam1 = variables.insert(Box::new(Variable::new(se3::SE3::from_param_slice(
se3_id.as_slice(),
))));
let pt0 = variables.insert(Box::new(Variable::new(rn::Rn::new(pt_zero.clone()))));
let pt1 = variables.insert(Box::new(Variable::new(rn::Rn::new(pt_zero.clone()))));
let pt2 = variables.insert(Box::new(Variable::new(rn::Rn::new(pt_zero.clone()))));
let mut variable_index_map: SecondaryMap<VarKey, usize> = SecondaryMap::new();
variable_index_map.insert(cam0, 0);
variable_index_map.insert(cam1, 6);
variable_index_map.insert(pt0, 12);
variable_index_map.insert(pt1, 15);
variable_index_map.insert(pt2, 18);
let n_rows = 36;
let n_cols = 21;
let cam_cols = [0usize, 6];
let lm_cols = [12usize, 15, 18];
let mut triplets: Vec<Triplet<usize, usize, f64>> = Vec::new();
for (ci, &cam_col) in cam_cols.iter().enumerate() {
for (li, &lm_col) in lm_cols.iter().enumerate() {
let row_base = (ci * 3 + li) * 6;
for k in 0..6 {
triplets.push(Triplet::new(row_base + k, cam_col + k, 1.0));
triplets.push(Triplet::new(row_base + k, lm_col + (k % 3), 1.0));
}
}
}
let jacobian = SparseColMat::try_new_from_triplets(n_rows, n_cols, &triplets)?;
let residuals = Mat::from_fn(n_rows, 1, |i, _| (i % 5) as f64 * 0.1);
let mut landmark_keys = std::collections::HashSet::new();
landmark_keys.insert(pt0);
landmark_keys.insert(pt1);
landmark_keys.insert(pt2);
Ok((
variables,
variable_index_map,
jacobian,
residuals,
landmark_keys,
))
}
#[test]
fn cached_extraction_matches_direct_reference() -> TestResult {
let (variables, variable_index_map, jacobian, _residuals, landmark_keys) =
create_schur_test_setup()?;
let mut solver = ExplicitSparseSchur::new();
solver.initialize_structure(&variables, &variable_index_map, &landmark_keys)?;
let h_full = {
let mut cache =
crate::linalg::sparse::normal_eq::NormalEquationsCache::try_new(&jacobian)?;
let residuals = Mat::zeros(jacobian.nrows(), 1);
cache.compute(&residuals, &jacobian)?.hessian
};
let kk_ref = solver.extract_kept_block(&h_full)?;
let ke_ref = solver.extract_coupling_block(&h_full)?;
let (kk_first, ke_first) = solver.extract_kept_and_coupling(&h_full)?;
let (kk_second, ke_second) = solver.extract_kept_and_coupling(&h_full)?;
for (cached, direct, what) in [
(&kk_first, &kk_ref, "H_kk first call"),
(&kk_second, &kk_ref, "H_kk cached call"),
(&ke_first, &ke_ref, "H_ke first call"),
(&ke_second, &ke_ref, "H_ke cached call"),
] {
assert_eq!(
(cached.nrows(), cached.ncols()),
(direct.nrows(), direct.ncols()),
"{what}: shape mismatch"
);
assert_eq!(cached.compute_nnz(), direct.compute_nnz(), "{what}: nnz");
for col in 0..direct.ncols() {
assert_eq!(
cached.symbolic().row_idx_of_col_raw(col),
direct.symbolic().row_idx_of_col_raw(col),
"{what}: pattern mismatch in column {col}"
);
assert_eq!(
cached.val_of_col(col),
direct.val_of_col(col),
"{what}: value mismatch in column {col}"
);
}
}
Ok(())
}
#[test]
fn schur_complement_matches_naive_dense() -> TestResult {
let (variables, variable_index_map, jacobian, _residuals, landmark_keys) =
create_schur_test_setup()?;
let mut solver = ExplicitSparseSchur::new();
solver.initialize_structure(&variables, &variable_index_map, &landmark_keys)?;
let nrows = jacobian.nrows();
let ncols = jacobian.ncols();
let mut dense_j = nalgebra::DMatrix::zeros(nrows, ncols);
for col in 0..ncols {
let rows = jacobian.symbolic().row_idx_of_col_raw(col);
let vals = jacobian.val_of_col(col);
for (i, &row) in rows.iter().enumerate() {
dense_j[(row, col)] = vals[i];
}
}
let h = dense_j.transpose() * &dense_j;
let mut elim_cols = Vec::new();
for key in &landmark_keys {
let start = variable_index_map.get(*key).ok_or("elim key missing")?;
let dof = variables.get(*key).ok_or("elim var missing")?.dof();
elim_cols.extend(*start..*start + dof);
}
elim_cols.sort_unstable();
let kept_cols: Vec<usize> = (0..ncols).filter(|c| !elim_cols.contains(c)).collect();
let submatrix = |rows: &[usize], cols: &[usize]| -> nalgebra::DMatrix<f64> {
nalgebra::DMatrix::from_fn(rows.len(), cols.len(), |r, c| h[(rows[r], cols[c])])
};
let h_kk = submatrix(&kept_cols, &kept_cols);
let h_ke = submatrix(&kept_cols, &elim_cols);
let h_ee = submatrix(&elim_cols, &elim_cols);
let h_ee_inv = h_ee.try_inverse().ok_or("naive H_ee must invert")?;
let s_naive = h_kk - &h_ke * h_ee_inv * h_ke.transpose();
let h_full = {
let mut cache =
crate::linalg::sparse::normal_eq::NormalEquationsCache::try_new(&jacobian)?;
let residuals = Mat::zeros(nrows, 1);
cache.compute(&residuals, &jacobian)?.hessian
};
let h_kk2 = solver.extract_kept_block(&h_full)?;
let h_ke2 = solver.extract_coupling_block(&h_full)?;
let mut blocks = EliminatedBlocks::new(solver.partition().ok_or("partition missing")?);
blocks.gather(&h_full, solver.partition().ok_or("partition missing")?);
blocks.invert_in_place(solver.partition().ok_or("partition missing")?)?;
let (s, _) = solver.compute_schur_complement(&h_kk2, &h_ke2, &blocks)?;
let s_val = |i: usize, j: usize| -> f64 {
let rows = s.symbolic().row_idx_of_col_raw(j);
let vals = s.val_of_col(j);
rows.iter()
.position(|&r| r == i)
.map(|p| vals[p])
.unwrap_or(0.0)
};
let k = kept_cols.len();
assert_eq!((s.nrows(), s.ncols()), (k, k));
for i in 0..k {
for j in 0..k {
let (a, b) = (s_val(i, j), s_val(j, i));
assert!(
(a - b).abs() < 1e-12,
"S not symmetric at ({i},{j}): {a} vs {b}"
);
assert!(
(a - s_naive[(i, j)]).abs() < 1e-9,
"S[{i},{j}] = {a}, naive dense = {}",
s_naive[(i, j)]
);
}
}
Ok(())
}
#[test]
fn test_schur_ordering_rn3_eliminated() {
let ordering = SchurOrdering::default();
assert!(ordering.should_eliminate(&ManifoldType::RN, 3));
}
#[test]
fn test_schur_ordering_se3_not_eliminated() {
let ordering = SchurOrdering::default();
assert!(!ordering.should_eliminate(&ManifoldType::SE3, 6));
}
#[test]
fn test_schur_ordering_wrong_type_not_eliminated() {
let ordering = SchurOrdering::default();
assert!(!ordering.should_eliminate(&ManifoldType::SE3, 6));
}
#[test]
fn test_schur_ordering_wrong_size_not_eliminated() {
let ordering = SchurOrdering::default();
assert!(!ordering.should_eliminate(&ManifoldType::RN, 6));
assert!(!ordering.should_eliminate(&ManifoldType::RN, 2));
}
#[test]
fn test_schur_ordering_by_name_matches_type() {
let ordering = SchurOrdering::default();
assert!(ordering.should_eliminate_by_name("Rn", 3));
assert!(!ordering.should_eliminate_by_name("SE3", 6));
assert!(!ordering.should_eliminate_by_name("NotAManifold", 3));
}
#[test]
fn test_auto_elimination_is_opt_in() -> Result<(), Box<dyn std::error::Error>> {
use crate::core::variable::Variable;
use apex_manifolds::{rn, se3};
use nalgebra::DVector;
use slotmap::SlotMap;
let mut variables: SlotMap<VarKey, Box<dyn ManifoldVariable>> = SlotMap::with_key();
let se3_data = DVector::from_vec(vec![1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]);
let cam0 = variables.insert(Box::new(Variable::new(se3::SE3::from_param_slice(
se3_data.as_slice(),
))));
let cam1 = variables.insert(Box::new(Variable::new(se3::SE3::from_param_slice(
se3_data.as_slice(),
))));
let pt_data = DVector::from_vec(vec![0.0, 0.0, 0.0]);
let pt0 = variables.insert(Box::new(Variable::new(rn::Rn::new(pt_data.clone()))));
let pt1 = variables.insert(Box::new(Variable::new(rn::Rn::new(pt_data.clone()))));
let mut index_map = SecondaryMap::new();
index_map.insert(cam0, 0);
index_map.insert(cam1, 6);
index_map.insert(pt0, 12);
index_map.insert(pt1, 15);
let empty_marks = std::collections::HashSet::new();
let mut marks = std::collections::HashSet::new();
marks.insert(pt0);
let mut solver = ExplicitSparseSchur::new();
solver.initialize_structure(&variables, &index_map, &marks)?;
let structure = solver.partition().ok_or("partition missing")?;
assert_eq!(
structure.eliminated_blocks().len(),
1,
"auto-detection must be off by default"
);
assert_eq!(structure.kept_blocks().len(), 3);
assert!(structure.kept_blocks().iter().any(|b| b.key == pt1));
let ordering = SchurOrdering::default().with_auto_detect(true);
let mut solver = ExplicitSparseSchur::new().with_ordering(ordering);
solver.initialize_structure(&variables, &index_map, &empty_marks)?;
let structure = solver.partition().ok_or("partition missing")?;
assert_eq!(
structure.eliminated_blocks().len(),
2,
"Rn(3) variables must auto-eliminate"
);
assert_eq!(
structure.kept_blocks().len(),
2,
"SE3 variables must stay cameras"
);
let mut marks = std::collections::HashSet::new();
marks.insert(cam1);
let mut solver = ExplicitSparseSchur::new();
solver.initialize_structure(&variables, &index_map, &marks)?;
let structure = solver.partition().ok_or("partition missing")?;
assert!(structure.eliminated_blocks().iter().any(|b| b.key == cam1));
Ok(())
}
#[test]
fn test_partition_requires_both_sides() {
assert!(SchurPartition::new(Vec::new(), Vec::new()).is_err());
}
#[test]
fn test_solver_creation() {
let solver = ExplicitSparseSchur::new();
assert!(solver.partition().is_none());
}
#[test]
fn test_schur_variants() {
let solver = ExplicitSparseSchur::new()
.with_variant(ExplicitSchurVariant::Iterative)
.with_preconditioner(SchurPreconditioner::BlockDiagonal)
.with_cg_params(100, 1e-8);
assert_eq!(solver.cg_max_iterations, 100);
assert!((solver.cg_tolerance - 1e-8).abs() < 1e-12);
assert_eq!(solver.variant, ExplicitSchurVariant::Iterative);
assert_eq!(solver.preconditioner, SchurPreconditioner::BlockDiagonal);
}
#[test]
fn preconditioner_selection_reaches_pcg() -> Result<(), LinAlgError> {
use faer::sparse::Triplet;
let kept = vec![BlockSpan {
key: VarKey::default(),
col_start: 0,
dof: 2,
}];
let eliminated = vec![BlockSpan {
key: VarKey::default(),
col_start: 2,
dof: 1,
}];
let partition = SchurPartition::new(kept, eliminated)?;
let mut solver = ExplicitSparseSchur::new();
solver.partition = Some(partition);
let mk = |vals: [f64; 4]| -> Result<SparseColMat<usize, f64>, LinAlgError> {
SparseColMat::try_new_from_triplets(
2,
2,
&[
Triplet::new(0, 0, vals[0]),
Triplet::new(1, 0, vals[1]),
Triplet::new(0, 1, vals[2]),
Triplet::new(1, 1, vals[3]),
],
)
.map_err(|e| LinAlgError::SparseMatrixCreation(format!("{e:?}")))
};
let h_kk = mk([4.0, 1.0, 1.0, 5.0])?;
let s = mk([2.0, 0.5, 0.5, 3.0])?;
solver.preconditioner = SchurPreconditioner::None;
assert!(
solver.build_preconditioner(&s, &h_kk)?.is_none(),
"None must leave PCG unpreconditioned"
);
solver.preconditioner = SchurPreconditioner::SchurJacobi;
let blocks_missing = || LinAlgError::InvalidInput("expected blocks".into());
let schur = solver
.build_preconditioner(&s, &h_kk)?
.ok_or_else(blocks_missing)?;
solver.preconditioner = SchurPreconditioner::BlockDiagonal;
let plain = solver
.build_preconditioner(&s, &h_kk)?
.ok_or_else(blocks_missing)?;
let expect_schur = DMatrix::from_row_slice(2, 2, &[2.0, 0.5, 0.5, 3.0])
.try_inverse()
.ok_or_else(|| LinAlgError::InvalidInput("S block must invert".into()))?;
for r in 0..2 {
for c in 0..2 {
assert!(
(schur[0][(r, c)] - expect_schur[(r, c)]).abs() < 1e-12,
"SchurJacobi must invert S's diagonal block"
);
assert!(
(schur[0][(r, c)] - plain[0][(r, c)]).abs() > 1e-9,
"SchurJacobi and BlockDiagonal must not coincide at ({r},{c})"
);
}
}
Ok(())
}
fn solver_with_block(
kept_dof: usize,
dof: usize,
inverse: &[f64],
) -> Result<(ExplicitSparseSchur, EliminatedBlocks), LinAlgError> {
let kept = (0..kept_dof)
.map(|i| BlockSpan {
key: VarKey::default(),
col_start: i,
dof: 1,
})
.collect();
let eliminated = vec![BlockSpan {
key: VarKey::default(),
col_start: kept_dof,
dof,
}];
let partition = SchurPartition::new(kept, eliminated)?;
let mut blocks = EliminatedBlocks::new(&partition);
blocks.block_mut(0).copy_from_slice(inverse);
let mut solver = ExplicitSparseSchur::new();
solver.partition = Some(partition);
Ok((solver, blocks))
}
#[test]
fn test_compute_schur_complement_known_matrix() -> Result<(), LinAlgError> {
use faer::sparse::Triplet;
let inv = [0.5, 0.0, 0.0, 0.0, 0.5, 0.0, 0.0, 0.0, 0.5];
let (mut solver, blocks) = solver_with_block(2, 3, &inv)?;
let h_kk = SparseColMat::try_new_from_triplets(
2,
2,
&[Triplet::new(0, 0, 4.0), Triplet::new(1, 1, 5.0)],
)
.map_err(|e| LinAlgError::SparseMatrixCreation(format!("{e:?}")))?;
let h_ke = SparseColMat::try_new_from_triplets(
2,
3,
&[Triplet::new(0, 0, 1.0), Triplet::new(1, 1, 2.0)],
)
.map_err(|e| LinAlgError::SparseMatrixCreation(format!("{e:?}")))?;
let (s, _) = solver.compute_schur_complement(&h_kk, &h_ke, &blocks)?;
assert_eq!(s.nrows(), 2);
assert_eq!(s.ncols(), 2);
let dense = dense_of(&s);
assert!((dense[0][0] - 3.5).abs() < 1e-12, "got {}", dense[0][0]);
assert!((dense[1][1] - 3.0).abs() < 1e-12, "got {}", dense[1][1]);
Ok(())
}
#[test]
fn test_back_substitute() -> Result<(), LinAlgError> {
use faer::sparse::Triplet;
let inv = [1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0];
let (solver, blocks) = solver_with_block(2, 3, &inv)?;
let delta_c = Mat::from_fn(2, 1, |i, _| (i + 1) as f64); let g_p = Mat::from_fn(3, 1, |i, _| (i + 1) as f64);
let h_cp = SparseColMat::try_new_from_triplets(
2,
3,
&[Triplet::new(0, 0, 1.0), Triplet::new(1, 1, 1.0)],
)
.map_err(|e| LinAlgError::SparseMatrixCreation(format!("{e:?}")))?;
let delta_p = solver.back_substitute(&delta_c, &g_p, &h_cp, &blocks)?;
assert_eq!(delta_p.nrows(), 3);
assert!((delta_p[(0, 0)]).abs() < 1e-10);
assert!((delta_p[(1, 0)]).abs() < 1e-10);
assert!((delta_p[(2, 0)] - 3.0).abs() < 1e-10);
Ok(())
}
#[test]
fn test_compute_reduced_gradient() -> Result<(), LinAlgError> {
use faer::sparse::Triplet;
let inv = [2.0, 0.0, 0.0, 0.0, 2.0, 0.0, 0.0, 0.0, 2.0];
let (solver, blocks) = solver_with_block(2, 3, &inv)?;
let g_c = Mat::from_fn(2, 1, |i, _| (i + 1) as f64); let g_p = Mat::from_fn(3, 1, |i, _| (i + 1) as f64);
let h_cp = SparseColMat::try_new_from_triplets(
2,
3,
&[Triplet::new(0, 0, 1.0), Triplet::new(1, 1, 1.0)],
)
.map_err(|e| LinAlgError::SparseMatrixCreation(format!("{e:?}")))?;
let g_reduced = solver.compute_reduced_gradient(&g_c, &g_p, &h_cp, &blocks)?;
assert_eq!(g_reduced.nrows(), 2);
assert!((g_reduced[(0, 0)] + 1.0).abs() < 1e-10);
assert!((g_reduced[(1, 0)] + 2.0).abs() < 1e-10);
Ok(())
}
#[test]
fn test_solver_default() {
let solver = ExplicitSparseSchur::default();
assert!(solver.partition().is_none());
assert!(solver.hessian.is_none());
assert!(solver.gradient.is_none());
}
#[test]
fn test_partition_getter() -> TestResult {
let (variables, variable_index_map, _, _, landmark_keys) = create_schur_test_setup()?;
let mut solver = ExplicitSparseSchur::new();
assert!(solver.partition().is_none());
solver.initialize_structure(&variables, &variable_index_map, &landmark_keys)?;
assert!(solver.partition().is_some());
Ok(())
}
#[test]
fn test_with_ordering_builder() {
let ordering = SchurOrdering {
eliminate_types: vec![ManifoldType::RN],
eliminate_rn_size: Some(3),
auto_detect: false,
};
let solver = ExplicitSparseSchur::new().with_ordering(ordering);
assert_eq!(solver.ordering.eliminate_rn_size, Some(3));
}
#[test]
fn test_eliminated_block_inversion_3dof() -> TestResult {
let inv_seed = [2.0, 0.0, 0.0, 0.0, 3.0, 0.0, 0.0, 0.0, 4.0];
let kept = vec![BlockSpan {
key: VarKey::default(),
col_start: 0,
dof: 1,
}];
let eliminated = vec![BlockSpan {
key: VarKey::default(),
col_start: 1,
dof: 3,
}];
let partition = SchurPartition::new(kept, eliminated)?;
let mut blocks = EliminatedBlocks::new(&partition);
blocks.block_mut(0).copy_from_slice(&inv_seed);
blocks.invert_in_place(&partition)?;
assert!((blocks.at(0, 0, 0) - 0.5).abs() < 1e-10);
assert!((blocks.at(0, 1, 1) - 1.0 / 3.0).abs() < 1e-10);
assert!((blocks.at(0, 2, 2) - 0.25).abs() < 1e-10);
Ok(())
}
#[test]
fn test_explicit_schur_initialize_structure() -> TestResult {
let (variables, variable_index_map, _, _, landmark_keys) = create_schur_test_setup()?;
let mut solver = ExplicitSparseSchur::new();
solver.initialize_structure(&variables, &variable_index_map, &landmark_keys)?;
let bs = solver.partition().ok_or("partition is None")?;
assert_eq!(bs.kept_blocks().len(), 2);
assert_eq!(bs.eliminated_blocks().len(), 3);
assert_eq!(bs.kept_dof(), 12); assert_eq!(bs.eliminated_dof(), 9); Ok(())
}
#[test]
fn test_extract_gradient_blocks() -> TestResult {
let (variables, variable_index_map, _, _, landmark_keys) = create_schur_test_setup()?;
let mut solver = ExplicitSparseSchur::new();
solver.initialize_structure(&variables, &variable_index_map, &landmark_keys)?;
let gradient = Mat::from_fn(21, 1, |i, _| i as f64);
let (g_c, g_p) = solver.extract_gradient_blocks(&gradient)?;
assert_eq!(g_c.nrows(), 12); assert_eq!(g_p.nrows(), 9); Ok(())
}
#[test]
fn test_explicit_schur_solve_normal_equation() -> TestResult {
let (variables, variable_index_map, jacobian, residuals, landmark_keys) =
create_schur_test_setup()?;
let mut solver = ExplicitSparseSchur::new().with_variant(ExplicitSchurVariant::Sparse);
solver.initialize_structure(&variables, &variable_index_map, &landmark_keys)?;
let delta =
LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
assert_eq!(delta.nrows(), 21);
assert_eq!(delta.ncols(), 1);
Ok(())
}
#[test]
fn test_explicit_schur_solve_augmented_equation() -> TestResult {
let (variables, variable_index_map, jacobian, residuals, landmark_keys) =
create_schur_test_setup()?;
let mut solver = ExplicitSparseSchur::new().with_variant(ExplicitSchurVariant::Sparse);
solver.initialize_structure(&variables, &variable_index_map, &landmark_keys)?;
let delta = LinearSolver::<SparseMode>::solve_augmented_equation(
&mut solver,
&residuals,
&jacobian,
&Damping::identity(0.1),
)?;
assert_eq!(delta.nrows(), 21);
Ok(())
}
#[test]
fn test_explicit_schur_solve_iterative_variant() -> TestResult {
let (variables, variable_index_map, jacobian, residuals, landmark_keys) =
create_schur_test_setup()?;
let mut solver = ExplicitSparseSchur::new()
.with_variant(ExplicitSchurVariant::Iterative)
.with_cg_params(200, 1e-10);
solver.initialize_structure(&variables, &variable_index_map, &landmark_keys)?;
let delta =
LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
assert_eq!(delta.nrows(), 21);
let mut cholesky_solver =
ExplicitSparseSchur::new().with_variant(ExplicitSchurVariant::Sparse);
cholesky_solver.initialize_structure(&variables, &variable_index_map, &landmark_keys)?;
let cholesky_delta = LinearSolver::<SparseMode>::solve_normal_equation(
&mut cholesky_solver,
&residuals,
&jacobian,
)?;
for i in 0..21 {
assert!(
(delta[(i, 0)] - cholesky_delta[(i, 0)]).abs() < 1e-4,
"PCG-on-S must agree with Cholesky-on-S at {i}: {} vs {}",
delta[(i, 0)],
cholesky_delta[(i, 0)]
);
}
Ok(())
}
#[test]
fn test_explicit_schur_get_hessian_gradient() -> TestResult {
let (variables, variable_index_map, jacobian, residuals, landmark_keys) =
create_schur_test_setup()?;
let mut solver = ExplicitSparseSchur::new();
solver.initialize_structure(&variables, &variable_index_map, &landmark_keys)?;
assert!(LinearSolver::<SparseMode>::get_hessian(&solver).is_none());
assert!(LinearSolver::<SparseMode>::get_gradient(&solver).is_none());
LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
let h = LinearSolver::<SparseMode>::get_hessian(&solver);
let g = LinearSolver::<SparseMode>::get_gradient(&solver);
assert!(h.is_some());
assert!(g.is_some());
let h = h.ok_or("hessian is None")?;
let g = g.ok_or("gradient is None")?;
assert_eq!(h.nrows(), 21);
assert_eq!(g.nrows(), 21);
Ok(())
}
#[test]
fn test_explicit_schur_augmented_lambda_effect() -> TestResult {
let (variables, variable_index_map, jacobian, residuals, landmark_keys) =
create_schur_test_setup()?;
let mut solver1 = ExplicitSparseSchur::new();
solver1.initialize_structure(&variables, &variable_index_map, &landmark_keys)?;
let delta1 = LinearSolver::<SparseMode>::solve_augmented_equation(
&mut solver1,
&residuals,
&jacobian,
&Damping::identity(0.001),
)?;
let mut solver2 = ExplicitSparseSchur::new();
solver2.initialize_structure(&variables, &variable_index_map, &landmark_keys)?;
let delta2 = LinearSolver::<SparseMode>::solve_augmented_equation(
&mut solver2,
&residuals,
&jacobian,
&Damping::identity(100.0),
)?;
let norm_diff: f64 = (0..21)
.map(|i| (delta1[(i, 0)] - delta2[(i, 0)]).powi(2))
.sum();
assert!(
norm_diff > 1e-10,
"Different λ should yield different updates"
);
Ok(())
}
#[test]
fn test_combine_updates() -> TestResult {
let (variables, variable_index_map, _, _, landmark_keys) = create_schur_test_setup()?;
let mut solver = ExplicitSparseSchur::new();
solver.initialize_structure(&variables, &variable_index_map, &landmark_keys)?;
let delta_c = Mat::from_fn(12, 1, |_, _| 1.0);
let delta_p = Mat::from_fn(9, 1, |_, _| 2.0);
let combined = solver.combine_updates(&delta_c, &delta_p)?;
assert_eq!(combined.nrows(), 21);
for i in 0..12 {
assert!((combined[(i, 0)] - 1.0).abs() < 1e-10);
}
for i in 12..21 {
assert!((combined[(i, 0)] - 2.0).abs() < 1e-10);
}
Ok(())
}
#[test]
fn test_explicit_schur_solve_without_init_returns_error() -> TestResult {
let triplets: Vec<Triplet<usize, usize, f64>> = vec![Triplet::new(0, 0, 1.0)];
let jacobian =
SparseColMat::try_new_from_triplets(1, 1, &triplets).map_err(|e| format!("{e:?}"))?;
let residuals = Mat::from_fn(1, 1, |_, _| 1.0);
let mut solver = ExplicitSparseSchur::new();
let result =
LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian);
assert!(result.is_err());
Ok(())
}
#[test]
fn test_schur_ordering_new_equals_default() {
let a = SchurOrdering::new();
let b = SchurOrdering::default();
assert_eq!(a.eliminate_rn_size, b.eliminate_rn_size);
assert_eq!(a.eliminate_types.len(), b.eliminate_types.len());
}
#[test]
fn test_extract_camera_block() -> TestResult {
let (variables, variable_index_map, jacobian, residuals, landmark_keys) =
create_schur_test_setup()?;
let mut solver = ExplicitSparseSchur::new();
solver.initialize_structure(&variables, &variable_index_map, &landmark_keys)?;
LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
let hessian = solver.hessian.clone().ok_or("hessian is None")?;
let mut fresh = ExplicitSparseSchur::new();
fresh.initialize_structure(&variables, &variable_index_map, &landmark_keys)?;
let h_cc = fresh.extract_kept_block(&hessian)?;
assert_eq!(h_cc.nrows(), 12);
assert_eq!(h_cc.ncols(), 12);
Ok(())
}
#[test]
fn test_extract_coupling_block() -> TestResult {
let (variables, variable_index_map, jacobian, residuals, landmark_keys) =
create_schur_test_setup()?;
let mut solver = ExplicitSparseSchur::new();
solver.initialize_structure(&variables, &variable_index_map, &landmark_keys)?;
LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
let hessian = solver.hessian.clone().ok_or("hessian is None")?;
let mut fresh = ExplicitSparseSchur::new();
fresh.initialize_structure(&variables, &variable_index_map, &landmark_keys)?;
let h_cp = fresh.extract_coupling_block(&hessian)?;
assert_eq!(h_cp.nrows(), 12);
assert_eq!(h_cp.ncols(), 9);
Ok(())
}
#[test]
fn test_gather_eliminated_blocks() -> TestResult {
let (variables, variable_index_map, jacobian, residuals, landmark_keys) =
create_schur_test_setup()?;
let mut solver = ExplicitSparseSchur::new();
solver.initialize_structure(&variables, &variable_index_map, &landmark_keys)?;
LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
let hessian = solver.hessian.clone().ok_or("hessian is None")?;
let mut fresh = ExplicitSparseSchur::new();
fresh.initialize_structure(&variables, &variable_index_map, &landmark_keys)?;
let partition = fresh.partition().ok_or("partition is None")?;
let mut blocks = EliminatedBlocks::new(partition);
blocks.gather(&hessian, partition);
assert_eq!(blocks.len(), 3);
for i in 0..blocks.len() {
assert_eq!(blocks.dof(i), 3);
}
Ok(())
}
#[test]
fn schur_matches_cholesky_on_assembled_pose_chain() -> TestResult {
use crate::core::problem::Problem;
use crate::factors::pose::{BetweenFactor, PriorFactor};
use crate::linalg::JacobianMode;
use crate::linalg::sparse::cholesky::SparseCholeskySolver;
use apex_manifolds::{ManifoldType, se2::SE2};
use nalgebra::dvector;
let mut problem = Problem::new(JacobianMode::Sparse);
let mut keys = Vec::new();
for i in 0..5 {
keys.push(
problem.add_variable(ManifoldType::SE2, dvector![i as f64 + 0.1, 0.05, 0.02]),
);
}
for w in keys.windows(2) {
problem.add_residual_block(
&[w[0], w[1]],
Box::new(BetweenFactor::new(SE2::from_xy_angle(1.0, 0.0, 0.0))),
None,
);
}
problem.add_residual_block(
&[keys[0]],
Box::new(PriorFactor::new(SE2::from_xy_angle(0.0, 0.0, 0.0))),
None,
);
problem.mark_for_elimination(keys[1]);
problem.mark_for_elimination(keys[3]);
problem.group_rows_for_elimination();
let state = crate::optimizer::initialize_optimization_state(&mut problem)?;
let symbolic = state.symbolic_structure.as_ref().ok_or("sparse symbolic")?;
let mut workspace = crate::linearizer::AssemblyWorkspace::build(&problem);
let (residuals, jacobian) = crate::linearizer::cpu::sparse::assemble_sparse(
&problem,
&state.variables,
&state.variable_index_map,
symbolic,
&mut workspace,
)?;
let mut eliminate = std::collections::HashSet::new();
eliminate.insert(keys[1]);
eliminate.insert(keys[3]);
let mut cholesky = SparseCholeskySolver::new();
let reference =
crate::linalg::LinearSolver::<crate::linalg::SparseMode>::solve_normal_equation(
&mut cholesky,
&residuals,
&jacobian,
)?;
let mut schur = ExplicitSparseSchur::new();
schur.initialize_structure(&state.variables, &state.variable_index_map, &eliminate)?;
let first =
crate::linalg::LinearSolver::<crate::linalg::SparseMode>::solve_normal_equation(
&mut schur, &residuals, &jacobian,
)?;
let second =
crate::linalg::LinearSolver::<crate::linalg::SparseMode>::solve_normal_equation(
&mut schur, &residuals, &jacobian,
)?;
let scale = reference.norm_l2().max(1.0);
for (label, step) in [("first", &first), ("second", &second)] {
assert_eq!(reference.nrows(), step.nrows(), "{label}: length");
for i in 0..reference.nrows() {
let diff = (reference[(i, 0)] - step[(i, 0)]).abs();
assert!(
diff / scale < 1e-9,
"{label}: component {i} differs — cholesky {}, schur {} (rel {:.3e})",
reference[(i, 0)],
step[(i, 0)],
diff / scale
);
}
}
Ok(())
}
#[test]
fn schur_tracks_cholesky_across_iterations() -> TestResult {
use crate::core::problem::Problem;
use crate::factors::pose::{BetweenFactor, PriorFactor};
use crate::linalg::JacobianMode;
use crate::linalg::sparse::cholesky::SparseCholeskySolver;
use apex_manifolds::{ManifoldType, se2::SE2};
use nalgebra::dvector;
let mut problem = Problem::new(JacobianMode::Sparse);
let mut keys = Vec::new();
for i in 0..5 {
keys.push(
problem.add_variable(ManifoldType::SE2, dvector![i as f64 + 0.1, 0.05, 0.02]),
);
}
for w in keys.windows(2) {
problem.add_residual_block(
&[w[0], w[1]],
Box::new(BetweenFactor::new(SE2::from_xy_angle(1.0, 0.0, 0.0))),
None,
);
}
problem.add_residual_block(
&[keys[0]],
Box::new(PriorFactor::new(SE2::from_xy_angle(0.0, 0.0, 0.0))),
None,
);
problem.mark_for_elimination(keys[1]);
problem.mark_for_elimination(keys[3]);
problem.group_rows_for_elimination();
let mut state = crate::optimizer::initialize_optimization_state(&mut problem)?;
let mut eliminate = std::collections::HashSet::new();
eliminate.insert(keys[1]);
eliminate.insert(keys[3]);
let mut cholesky = SparseCholeskySolver::new();
let mut schur = ExplicitSparseSchur::new();
schur.initialize_structure(&state.variables, &state.variable_index_map, &eliminate)?;
for iter in 0..3 {
let symbolic = state.symbolic_structure.as_ref().ok_or("sparse symbolic")?;
let mut workspace = crate::linearizer::AssemblyWorkspace::build(&problem);
let (residuals, jacobian) = crate::linearizer::cpu::sparse::assemble_sparse(
&problem,
&state.variables,
&state.variable_index_map,
symbolic,
&mut workspace,
)?;
let reference =
crate::linalg::LinearSolver::<crate::linalg::SparseMode>::solve_normal_equation(
&mut cholesky,
&residuals,
&jacobian,
)?;
let step =
crate::linalg::LinearSolver::<crate::linalg::SparseMode>::solve_normal_equation(
&mut schur, &residuals, &jacobian,
)?;
let scale = reference.norm_l2().max(1.0);
let mut max_rel: f64 = 0.0;
for i in 0..reference.nrows() {
max_rel = max_rel.max((reference[(i, 0)] - step[(i, 0)]).abs() / scale);
}
assert!(
max_rel < 1e-9,
"iteration {iter}: Schur step diverges from Cholesky (max rel {max_rel:.3e})"
);
crate::optimizer::apply_parameter_step(
&mut state.variables,
step.as_ref(),
&state.sorted_vars,
);
}
Ok(())
}
#[test]
fn schur_tracks_cholesky_with_shared_workspace() -> TestResult {
use crate::core::problem::Problem;
use crate::factors::pose::{BetweenFactor, PriorFactor};
use crate::linalg::JacobianMode;
use crate::linalg::sparse::cholesky::SparseCholeskySolver;
use apex_manifolds::{ManifoldType, se2::SE2};
use nalgebra::dvector;
let mut problem = Problem::new(JacobianMode::Sparse);
let mut keys = Vec::new();
for i in 0..5 {
keys.push(
problem.add_variable(ManifoldType::SE2, dvector![i as f64 + 0.1, 0.05, 0.02]),
);
}
for w in keys.windows(2) {
problem.add_residual_block(
&[w[0], w[1]],
Box::new(BetweenFactor::new(SE2::from_xy_angle(1.0, 0.0, 0.0))),
None,
);
}
problem.add_residual_block(
&[keys[0]],
Box::new(PriorFactor::new(SE2::from_xy_angle(0.0, 0.0, 0.0))),
None,
);
problem.mark_for_elimination(keys[1]);
problem.mark_for_elimination(keys[3]);
problem.group_rows_for_elimination();
let mut state = crate::optimizer::initialize_optimization_state(&mut problem)?;
let mut eliminate = std::collections::HashSet::new();
eliminate.insert(keys[1]);
eliminate.insert(keys[3]);
let costs = |variables: &slotmap::SlotMap<
crate::core::VarKey,
Box<dyn crate::core::variable::ManifoldVariable>,
>|
-> Result<f64, Box<dyn std::error::Error>> {
Ok(problem.compute_residual_and_cost_sparse(variables)?.1)
};
let mut cholesky = SparseCholeskySolver::new();
let mut schur = ExplicitSparseSchur::new();
schur.initialize_structure(&state.variables, &state.variable_index_map, &eliminate)?;
for iter in 0..3 {
let symbolic = state.symbolic_structure.as_ref().ok_or("sparse symbolic")?;
let (residuals, jacobian) = crate::linearizer::cpu::sparse::assemble_sparse(
&problem,
&state.variables,
&state.variable_index_map,
symbolic,
&mut state.workspace,
)?;
let reference =
crate::linalg::LinearSolver::<crate::linalg::SparseMode>::solve_normal_equation(
&mut cholesky,
&residuals,
&jacobian,
)?;
let step =
crate::linalg::LinearSolver::<crate::linalg::SparseMode>::solve_normal_equation(
&mut schur, &residuals, &jacobian,
)?;
let scale = reference.norm_l2().max(1.0);
let mut max_rel: f64 = 0.0;
for i in 0..reference.nrows() {
max_rel = max_rel.max((reference[(i, 0)] - step[(i, 0)]).abs() / scale);
}
assert!(
max_rel < 1e-9,
"iteration {iter}: Schur step diverges from Cholesky (max rel {max_rel:.3e})"
);
crate::optimizer::apply_parameter_step(
&mut state.variables,
reference.as_ref(),
&state.sorted_vars,
);
let _ = costs(&state.variables)?;
}
Ok(())
}
#[test]
fn test_solve_with_cholesky_small_spd() -> TestResult {
let mut solver = ExplicitSparseSchur::new();
let triplets = vec![
Triplet::new(0usize, 0usize, 4.0f64),
Triplet::new(1usize, 0usize, 1.0f64),
Triplet::new(0usize, 1usize, 1.0f64),
Triplet::new(1usize, 1usize, 3.0f64),
];
let a =
SparseColMat::try_new_from_triplets(2, 2, &triplets).map_err(|e| format!("{e:?}"))?;
let b = Mat::from_fn(2, 1, |i, _| (i + 1) as f64);
let x = solver.solve_with_cholesky(&a, &b, None)?;
assert_eq!(x.nrows(), 2);
let ax0 = 4.0 * x[(0, 0)] + 1.0 * x[(1, 0)];
let ax1 = 1.0 * x[(0, 0)] + 3.0 * x[(1, 0)];
assert!((ax0 - 1.0).abs() < 1e-8, "A·x[0] = {ax0}");
assert!((ax1 - 2.0).abs() < 1e-8, "A·x[1] = {ax1}");
Ok(())
}
#[test]
fn test_solve_with_pcg_diagonal_system() -> TestResult {
let inv = [1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0];
let (solver, _) = solver_with_block(2, 3, &inv)?;
let triplets = vec![
Triplet::new(0usize, 0usize, 2.0f64),
Triplet::new(1usize, 1usize, 3.0f64),
];
let a =
SparseColMat::try_new_from_triplets(2, 2, &triplets).map_err(|e| format!("{e:?}"))?;
let b = Mat::from_fn(2, 1, |i, _| (i + 1) as f64);
let x = solver.solve_with_pcg(&a, &b, &a)?;
assert!((x[(0, 0)] - 0.5).abs() < 1e-6, "x[0] = {}", x[(0, 0)]);
assert!((x[(1, 0)] - 2.0 / 3.0).abs() < 1e-6, "x[1] = {}", x[(1, 0)]);
Ok(())
}
#[test]
fn test_initialize_structure_no_cameras_returns_error() {
use crate::core::variable::Variable;
use apex_manifolds::rn;
use nalgebra::DVector;
let mut variables: SlotMap<VarKey, Box<dyn ManifoldVariable>> = SlotMap::with_key();
let k = variables.insert(Box::new(Variable::new(rn::Rn::new(DVector::zeros(3)))));
let mut variable_index_map: SecondaryMap<VarKey, usize> = SecondaryMap::new();
variable_index_map.insert(k, 0);
let mut landmark_keys = std::collections::HashSet::new();
landmark_keys.insert(k);
let mut solver = ExplicitSparseSchur::new();
let result = solver.initialize_structure(&variables, &variable_index_map, &landmark_keys);
assert!(
result.is_err(),
"Expected Err when no camera variables present"
);
}
#[test]
fn test_initialize_structure_no_landmarks_returns_error() {
use crate::core::variable::Variable;
use apex_manifolds::se3;
let mut variables: SlotMap<VarKey, Box<dyn ManifoldVariable>> = SlotMap::with_key();
let _k = variables.insert(Box::new(Variable::new(se3::SE3::from_param_slice(&[
1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
]))));
let mut variable_index_map: SecondaryMap<VarKey, usize> = SecondaryMap::new();
variable_index_map.insert(_k, 0);
let landmark_keys = std::collections::HashSet::<VarKey>::new();
let mut solver = ExplicitSparseSchur::new();
let result = solver.initialize_structure(&variables, &variable_index_map, &landmark_keys);
assert!(
result.is_err(),
"Expected Err when no landmark variables present"
);
}
fn dense_of(m: &SparseColMat<usize, f64>) -> Vec<Vec<f64>> {
let mut out = vec![vec![0.0; m.ncols()]; m.nrows()];
for (col, out_col) in (0..m.ncols()).zip(0..m.ncols()) {
let rows = m.symbolic().row_idx_of_col_raw(col);
let vals = m.val_of_col(col);
for (i, &r) in rows.iter().enumerate() {
out[r][out_col] = vals[i];
}
}
out
}
#[test]
fn test_damp_camera_block_full_diagonal() -> TestResult {
let triplets = vec![
Triplet::new(0usize, 0usize, 4.0_f64),
Triplet::new(1, 0, 1.0),
Triplet::new(0, 1, 1.0),
Triplet::new(1, 1, 9.0),
];
let h_cc = SparseColMat::try_new_from_triplets(2, 2, &triplets)?;
let damping = Damping::new(0.5, 1e-6, 1e32)?;
let damped = dense_of(&damp_camera_block(&h_cc, 2, &damping)?);
assert!((damped[0][0] - (4.0 + 0.5 * 4.0)).abs() < 1e-12);
assert!((damped[1][1] - (9.0 + 0.5 * 9.0)).abs() < 1e-12);
assert!((damped[0][1] - 1.0).abs() < 1e-12, "off-diagonal changed");
assert!((damped[1][0] - 1.0).abs() < 1e-12, "off-diagonal changed");
Ok(())
}
#[test]
fn test_damp_camera_block_missing_diagonal() -> TestResult {
let triplets = vec![
Triplet::new(0usize, 0usize, 4.0_f64),
Triplet::new(0, 1, 2.0),
];
let h_cc = SparseColMat::try_new_from_triplets(2, 2, &triplets)?;
let damping = Damping::new(0.5, 1e-3, 1e32)?;
let damped = dense_of(&damp_camera_block(&h_cc, 2, &damping)?);
assert!((damped[0][0] - (4.0 + 0.5 * 4.0)).abs() < 1e-12);
assert!((damped[1][1] - 0.5 * 1e-3).abs() < 1e-12);
assert!((damped[0][1] - 2.0).abs() < 1e-12, "off-diagonal changed");
Ok(())
}
}