apex_solver/linalg/
mod.rs

1pub mod dense;
2pub mod sparse;
3pub mod utils;
4
5use crate::core::{VarKey, variable::ManifoldVariable};
6use faer::Mat;
7use slotmap::{SecondaryMap, SlotMap};
8use std::collections::HashSet;
9use std::fmt::{self, Debug, Display, Formatter};
10use thiserror::Error;
11
12pub use sparse::{
13    IterativeSchurSolver, SchurBlockStructure, SchurOrdering, SchurPreconditioner, SchurVariant,
14    SparseCholeskySolver, SparseQRSolver, SparseSchurComplementSolver,
15};
16
17pub use dense::{DenseCholeskySolver, DenseQRSolver};
18
19pub use crate::linearizer::cpu::{DenseMode, LinearizationMode, SparseMode};
20
21// ============================================================================
22// Jacobian mode selection
23// ============================================================================
24
25/// Controls which Jacobian assembly strategy the Problem uses.
26///
27/// Set this when constructing a [`Problem`](crate::core::problem::Problem):
28/// - `Problem::new(JacobianMode::Sparse)` — sparse (default, best for large-scale problems)
29/// - `Problem::new(JacobianMode::Dense)` — dense (best for small-to-medium problems < ~500 DOF)
30/// - `Problem::default()` — equivalent to `JacobianMode::Sparse`
31///
32/// The optimizer reads this field and dispatches to the appropriate assembly path.
33/// `LinearSolverType` selects the specific algorithm within the sparse path.
34#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
35pub enum JacobianMode {
36    /// Sparse Jacobian using symbolic structure and `SparseColMat`. Best for large problems.
37    #[default]
38    Sparse,
39    /// Dense Jacobian using `Mat<f64>`. Best for small-to-medium problems (< ~500 DOF).
40    Dense,
41}
42
43// ============================================================================
44// Linear solver type selection
45// ============================================================================
46
47#[non_exhaustive]
48#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
49pub enum LinearSolverType {
50    #[default]
51    SparseCholesky,
52    SparseQR,
53    SparseSchurComplement,
54    DenseCholesky,
55    DenseQR,
56}
57
58impl Display for LinearSolverType {
59    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
60        match self {
61            LinearSolverType::SparseCholesky => write!(f, "Sparse Cholesky"),
62            LinearSolverType::SparseQR => write!(f, "Sparse QR"),
63            LinearSolverType::SparseSchurComplement => write!(f, "Sparse Schur Complement"),
64            LinearSolverType::DenseCholesky => write!(f, "Dense Cholesky"),
65            LinearSolverType::DenseQR => write!(f, "Dense QR"),
66        }
67    }
68}
69
70// ============================================================================
71// Error types
72// ============================================================================
73
74/// Linear algebra specific error types for apex-solver
75#[derive(Debug, Clone, Error)]
76pub enum LinAlgError {
77    /// Matrix factorization failed (Cholesky, QR, etc.)
78    #[error("Matrix factorization failed: {0}")]
79    FactorizationFailed(String),
80
81    /// Singular or near-singular matrix detected
82    #[error("Singular matrix detected: {0}")]
83    SingularMatrix(String),
84
85    /// Failed to create sparse matrix from triplets
86    #[error("Failed to create sparse matrix: {0}")]
87    SparseMatrixCreation(String),
88
89    /// Matrix format conversion failed
90    #[error("Matrix conversion failed: {0}")]
91    MatrixConversion(String),
92
93    /// Invalid input provided to linear solver
94    #[error("Invalid input: {0}")]
95    InvalidInput(String),
96
97    /// Solver in invalid state (e.g., initialized incorrectly)
98    #[error("Invalid solver state: {0}")]
99    InvalidState(String),
100}
101
102/// Result type for linear algebra operations
103pub type LinAlgResult<T> = Result<T, LinAlgError>;
104
105// ============================================================================
106// StructureAware
107// ============================================================================
108
109/// For solvers that need variable structure information before solving.
110///
111/// Implemented by Schur complement solvers, which must partition variables
112/// into camera and landmark blocks before performing any linear solves.
113/// Call [`initialize_structure`](StructureAware::initialize_structure) once
114/// during solver setup, before passing the solver to an optimizer.
115pub trait StructureAware {
116    /// Initialize the solver's block structure from problem variables.
117    fn initialize_structure(
118        &mut self,
119        variables: &SlotMap<VarKey, Box<dyn ManifoldVariable>>,
120        variable_index_map: &SecondaryMap<VarKey, usize>,
121        schur_landmark_keys: &HashSet<VarKey>,
122    ) -> LinAlgResult<()>;
123}
124
125// ============================================================================
126// LinearizationMode — re-exported from linearizer/cpu where it is defined
127// ============================================================================
128
129// ============================================================================
130// LinearSolver trait (unified solver interface, generic over LinearizationMode)
131// ============================================================================
132
133/// Unified linear solver interface parameterized by [`LinearizationMode`].
134///
135/// This is the single trait implemented by all linear solvers. When `M` is
136/// a concrete type (e.g., `SparseMode`), this trait is object-safe and can
137/// be used as `dyn LinearSolver<SparseMode>` or `dyn LinearSolver<DenseMode>`.
138///
139/// - Sparse solvers (`SparseCholeskySolver`, `SparseQRSolver`, `SchurSolverAdapter`)
140///   implement `LinearSolver<SparseMode>`.
141/// - Dense solvers (`DenseCholeskySolver`, `DenseQRSolver`)
142///   implement `LinearSolver<DenseMode>`.
143pub trait LinearSolver<M: LinearizationMode> {
144    /// Solve the normal equations: (J^T · J) · dx = −J^T · r
145    fn solve_normal_equation(
146        &mut self,
147        residuals: &Mat<f64>,
148        jacobian: &M::Jacobian,
149    ) -> LinAlgResult<Mat<f64>>;
150
151    /// Solve the augmented equations: (J^T · J + λI) · dx = −J^T · r
152    fn solve_augmented_equation(
153        &mut self,
154        residuals: &Mat<f64>,
155        jacobian: &M::Jacobian,
156        lambda: f64,
157    ) -> LinAlgResult<Mat<f64>>;
158
159    /// Get the cached Hessian matrix (J^T · J) from the last solve
160    fn get_hessian(&self) -> Option<&M::Hessian>;
161
162    /// Get the cached gradient vector (J^T · r) from the last solve
163    fn get_gradient(&self) -> Option<&Mat<f64>>;
164
165    /// Compute the covariance matrix (H^{-1}) by inverting the cached Hessian.
166    ///
167    /// Returns `None` for solvers that do not support covariance estimation
168    /// (e.g., QR solvers, Schur complement solvers). Only Cholesky-based
169    /// solvers provide a real implementation.
170    fn compute_covariance_matrix(&mut self) -> Option<&Mat<f64>> {
171        None
172    }
173
174    /// Get the cached covariance matrix (H^{-1}) computed from the Hessian.
175    ///
176    /// Returns `None` if covariance has not been computed or is not supported.
177    fn get_covariance_matrix(&self) -> Option<&Mat<f64>> {
178        None
179    }
180}
181
182// ============================================================================
183// Utility functions
184// ============================================================================
185
186/// Extract per-variable covariance blocks from the full covariance matrix.
187///
188/// Given the full covariance matrix H^{-1} (inverse of information matrix),
189/// this function extracts the diagonal blocks corresponding to each individual variable.
190pub(crate) fn extract_variable_covariances(
191    full_covariance: &Mat<f64>,
192    variables: &SlotMap<VarKey, Box<dyn ManifoldVariable>>,
193    variable_index_map: &SecondaryMap<VarKey, usize>,
194) -> SecondaryMap<VarKey, Mat<f64>> {
195    let mut result = SecondaryMap::new();
196
197    for (key, var) in variables {
198        if let Some(&start_idx) = variable_index_map.get(key) {
199            let dim = var.dof();
200            let mut var_cov = Mat::zeros(dim, dim);
201            for i in 0..dim {
202                for j in 0..dim {
203                    var_cov[(i, j)] = full_covariance[(start_idx + i, start_idx + j)];
204                }
205            }
206            result.insert(key, var_cov);
207        }
208    }
209
210    result
211}
212
213#[cfg(test)]
214mod tests {
215    use super::*;
216    use crate::core::variable::{ManifoldVariable, Variable};
217    use crate::error::ErrorLogging;
218    use apex_manifolds::rn::Rn;
219    use faer::Mat;
220    use nalgebra::dvector;
221
222    // -------------------------------------------------------------------------
223    // JacobianMode
224    // -------------------------------------------------------------------------
225
226    #[test]
227    fn test_jacobian_mode_default_is_sparse() {
228        assert_eq!(JacobianMode::default(), JacobianMode::Sparse);
229    }
230
231    #[test]
232    fn test_jacobian_mode_equality() {
233        assert_eq!(JacobianMode::Sparse, JacobianMode::Sparse);
234        assert_eq!(JacobianMode::Dense, JacobianMode::Dense);
235        assert_ne!(JacobianMode::Sparse, JacobianMode::Dense);
236    }
237
238    // -------------------------------------------------------------------------
239    // LinearSolverType Display + Default
240    // -------------------------------------------------------------------------
241
242    #[test]
243    fn test_linear_solver_type_default_is_cholesky() {
244        assert_eq!(
245            LinearSolverType::default(),
246            LinearSolverType::SparseCholesky
247        );
248    }
249
250    #[test]
251    fn test_linear_solver_type_display_all_variants() {
252        assert_eq!(
253            format!("{}", LinearSolverType::SparseCholesky),
254            "Sparse Cholesky"
255        );
256        assert_eq!(format!("{}", LinearSolverType::SparseQR), "Sparse QR");
257        assert_eq!(
258            format!("{}", LinearSolverType::SparseSchurComplement),
259            "Sparse Schur Complement"
260        );
261        assert_eq!(
262            format!("{}", LinearSolverType::DenseCholesky),
263            "Dense Cholesky"
264        );
265        assert_eq!(format!("{}", LinearSolverType::DenseQR), "Dense QR");
266    }
267
268    // -------------------------------------------------------------------------
269    // LinAlgError Display — one per variant
270    // -------------------------------------------------------------------------
271
272    #[test]
273    fn test_lin_alg_error_factorization_failed_display() {
274        let e = LinAlgError::FactorizationFailed("non-positive definite".into());
275        assert!(e.to_string().contains("non-positive definite"));
276    }
277
278    #[test]
279    fn test_lin_alg_error_singular_matrix_display() {
280        let e = LinAlgError::SingularMatrix("rank deficient".into());
281        assert!(e.to_string().contains("rank deficient"));
282    }
283
284    #[test]
285    fn test_lin_alg_error_sparse_matrix_creation_display() {
286        let e = LinAlgError::SparseMatrixCreation("bad triplets".into());
287        assert!(e.to_string().contains("bad triplets"));
288    }
289
290    #[test]
291    fn test_lin_alg_error_matrix_conversion_display() {
292        let e = LinAlgError::MatrixConversion("size mismatch".into());
293        assert!(e.to_string().contains("size mismatch"));
294    }
295
296    #[test]
297    fn test_lin_alg_error_invalid_input_display() {
298        let e = LinAlgError::InvalidInput("null jacobian".into());
299        assert!(e.to_string().contains("null jacobian"));
300    }
301
302    #[test]
303    fn test_lin_alg_error_invalid_state_display() {
304        let e = LinAlgError::InvalidState("not initialized".into());
305        assert!(e.to_string().contains("not initialized"));
306    }
307
308    // -------------------------------------------------------------------------
309    // log() / log_with_source() return self
310    // -------------------------------------------------------------------------
311
312    #[test]
313    fn test_lin_alg_error_log_returns_self() {
314        let e = LinAlgError::InvalidInput("log_test".into());
315        let returned = e.log();
316        assert!(returned.to_string().contains("log_test"));
317    }
318
319    #[test]
320    fn test_lin_alg_error_log_with_source_returns_self() {
321        let e = LinAlgError::SingularMatrix("source_test".into());
322        let source = std::io::Error::other("src");
323        let returned = e.log_with_source(source);
324        assert!(returned.to_string().contains("source_test"));
325    }
326
327    // -------------------------------------------------------------------------
328    // LinAlgResult type alias
329    // -------------------------------------------------------------------------
330
331    #[test]
332    fn test_lin_alg_result_ok() {
333        let r: LinAlgResult<i32> = Ok(7);
334        assert!(matches!(r, Ok(7)));
335    }
336
337    #[test]
338    fn test_lin_alg_result_err() {
339        let r: LinAlgResult<i32> = Err(LinAlgError::InvalidInput("oops".into()));
340        assert!(r.is_err());
341    }
342
343    // -------------------------------------------------------------------------
344    // extract_variable_covariances
345    // -------------------------------------------------------------------------
346
347    fn make_rn_var(val: f64) -> Box<dyn ManifoldVariable> {
348        Box::new(Variable::new(Rn::new(dvector![val])))
349    }
350
351    #[test]
352    fn test_extract_variable_covariances_single_variable() {
353        use crate::core::VarKey;
354        use slotmap::{SecondaryMap, SlotMap};
355
356        let mut variables: SlotMap<VarKey, Box<dyn ManifoldVariable>> = SlotMap::with_key();
357        let kx = variables.insert(make_rn_var(1.0));
358        let mut index_map: SecondaryMap<VarKey, usize> = SecondaryMap::new();
359        index_map.insert(kx, 0);
360
361        let full_cov = Mat::from_fn(1, 1, |_, _| 2.5);
362        let result = extract_variable_covariances(&full_cov, &variables, &index_map);
363        assert_eq!(result.len(), 1);
364        assert!((result[kx][(0, 0)] - 2.5).abs() < 1e-12);
365    }
366
367    #[test]
368    fn test_extract_variable_covariances_two_variables() {
369        use crate::core::VarKey;
370        use slotmap::{SecondaryMap, SlotMap};
371
372        let mut variables: SlotMap<VarKey, Box<dyn ManifoldVariable>> = SlotMap::with_key();
373        let ka = variables.insert(make_rn_var(1.0));
374        let kb = variables.insert(make_rn_var(2.0));
375        let mut index_map: SecondaryMap<VarKey, usize> = SecondaryMap::new();
376        index_map.insert(ka, 0);
377        index_map.insert(kb, 1);
378
379        let full_cov = Mat::from_fn(2, 2, |i, j| if i == j { [3.0, 7.0][i] } else { 0.0 });
380        let result = extract_variable_covariances(&full_cov, &variables, &index_map);
381        assert_eq!(result.len(), 2);
382        assert!((result[ka][(0, 0)] - 3.0).abs() < 1e-12);
383        assert!((result[kb][(0, 0)] - 7.0).abs() < 1e-12);
384    }
385
386    #[test]
387    fn test_extract_variable_covariances_empty_variables() {
388        use crate::core::VarKey;
389        use slotmap::{SecondaryMap, SlotMap};
390
391        let variables: SlotMap<VarKey, Box<dyn ManifoldVariable>> = SlotMap::with_key();
392        let index_map: SecondaryMap<VarKey, usize> = SecondaryMap::new();
393        let full_cov = Mat::zeros(0, 0);
394        let result = extract_variable_covariances(&full_cov, &variables, &index_map);
395        assert!(result.is_empty());
396    }
397
398    #[test]
399    fn test_extract_variable_covariances_var_not_in_index_map() {
400        use crate::core::VarKey;
401        use slotmap::{SecondaryMap, SlotMap};
402
403        let mut variables: SlotMap<VarKey, Box<dyn ManifoldVariable>> = SlotMap::with_key();
404        variables.insert(make_rn_var(1.0));
405        let index_map: SecondaryMap<VarKey, usize> = SecondaryMap::new(); // empty
406
407        let full_cov = Mat::from_fn(1, 1, |_, _| 5.0);
408        let result = extract_variable_covariances(&full_cov, &variables, &index_map);
409        assert!(result.is_empty());
410    }
411}
apex_solver/linalg/mod.rs

apex_solver/linalg/
mod.rs
