1use faer::{
2 Mat, Side,
3 linalg::solvers::Solve,
4 sparse::linalg::solvers::{Llt, SymbolicLlt},
5 sparse::{SparseColMat, Triplet},
6};
7use std::ops::Mul;
8
9use crate::error::ErrorLogging;
10use crate::linalg::{LinAlgError, LinAlgResult, LinearSolver, SparseMode};
11
12#[derive(Debug, Clone)]
13pub struct SparseCholeskySolver {
14 factorizer: Option<Llt<usize, f64>>,
15
16 symbolic_factorization: Option<SymbolicLlt<usize>>,
21
22 hessian: Option<SparseColMat<usize, f64>>,
26
27 gradient: Option<Mat<f64>>,
31
32 covariance_matrix: Option<Mat<f64>>,
36 standard_errors: Option<Mat<f64>>,
42}
43
44impl SparseCholeskySolver {
45 pub fn new() -> Self {
46 SparseCholeskySolver {
47 factorizer: None,
48 symbolic_factorization: None,
49 hessian: None,
50 gradient: None,
51 covariance_matrix: None,
52 standard_errors: None,
53 }
54 }
55
56 pub fn hessian(&self) -> Option<&SparseColMat<usize, f64>> {
57 self.hessian.as_ref()
58 }
59
60 pub fn gradient(&self) -> Option<&Mat<f64>> {
61 self.gradient.as_ref()
62 }
63
64 pub fn compute_standard_errors(&mut self) -> Option<&Mat<f64>> {
65 if self.covariance_matrix.is_none() {
67 LinearSolver::<SparseMode>::compute_covariance_matrix(self);
68 }
69
70 let hessian = self.hessian.as_ref()?;
72 let n = hessian.ncols();
73 if let Some(cov) = &self.covariance_matrix {
75 let mut std_errors = Mat::zeros(n, 1);
76 for i in 0..n {
77 let diag_val = cov[(i, i)];
78 if diag_val >= 0.0 {
79 std_errors[(i, 0)] = diag_val.sqrt();
80 } else {
81 return None;
83 }
84 }
85 self.standard_errors = Some(std_errors);
86 }
87 self.standard_errors.as_ref()
88 }
89
90 pub fn reset_covariance(&mut self) {
92 self.covariance_matrix = None;
93 self.standard_errors = None;
94 }
95}
96
97impl Default for SparseCholeskySolver {
98 fn default() -> Self {
99 Self::new()
100 }
101}
102impl LinearSolver<SparseMode> for SparseCholeskySolver {
103 fn solve_normal_equation(
104 &mut self,
105 residuals: &Mat<f64>,
106 jacobians: &SparseColMat<usize, f64>,
107 ) -> LinAlgResult<Mat<f64>> {
108 let jt = jacobians.as_ref().transpose();
110 let hessian = jt
111 .to_col_major()
112 .map_err(|e| {
113 LinAlgError::MatrixConversion(
114 "Failed to convert transposed Jacobian to column-major format".to_string(),
115 )
116 .log_with_source(e)
117 })?
118 .mul(jacobians.as_ref());
119
120 let gradient = jacobians.as_ref().transpose().mul(residuals);
122
123 let sym = if let Some(ref cached_sym) = self.symbolic_factorization {
124 cached_sym.clone()
129 } else {
130 let new_sym = SymbolicLlt::try_new(hessian.symbolic(), Side::Lower).map_err(|e| {
132 LinAlgError::FactorizationFailed(
133 "Symbolic Cholesky decomposition failed".to_string(),
134 )
135 .log_with_source(e)
136 })?;
137 self.symbolic_factorization = Some(new_sym.clone());
139 new_sym
140 };
141
142 let cholesky =
144 Llt::try_new_with_symbolic(sym, hessian.as_ref(), Side::Lower).map_err(|e| {
145 LinAlgError::SingularMatrix(
146 "Cholesky factorization failed (matrix may be singular)".to_string(),
147 )
148 .log_with_source(e)
149 })?;
150
151 let dx = cholesky.solve(-&gradient);
152 self.hessian = Some(hessian);
153 self.gradient = Some(gradient);
154 self.factorizer = Some(cholesky);
155
156 Ok(dx)
157 }
158
159 fn solve_augmented_equation(
160 &mut self,
161 residuals: &Mat<f64>,
162 jacobians: &SparseColMat<usize, f64>,
163 lambda: f64,
164 ) -> LinAlgResult<Mat<f64>> {
165 let n = jacobians.ncols();
166
167 let jt = jacobians.as_ref().transpose();
169 let hessian = jt
170 .to_col_major()
171 .map_err(|e| {
172 LinAlgError::MatrixConversion(
173 "Failed to convert transposed Jacobian to column-major format".to_string(),
174 )
175 .log_with_source(e)
176 })?
177 .mul(jacobians.as_ref());
178
179 let gradient = jacobians.as_ref().transpose().mul(residuals);
181
182 let mut lambda_i_triplets = Vec::with_capacity(n);
184 for i in 0..n {
185 lambda_i_triplets.push(Triplet::new(i, i, lambda));
186 }
187 let lambda_i =
188 SparseColMat::try_new_from_triplets(n, n, &lambda_i_triplets).map_err(|e| {
189 LinAlgError::SparseMatrixCreation("Failed to create lambda*I matrix".to_string())
190 .log_with_source(e)
191 })?;
192
193 let augmented_hessian = &hessian + lambda_i;
194
195 let sym = if let Some(ref cached_sym) = self.symbolic_factorization {
196 cached_sym.clone()
201 } else {
202 let new_sym =
204 SymbolicLlt::try_new(augmented_hessian.symbolic(), Side::Lower).map_err(|e| {
205 LinAlgError::FactorizationFailed(
206 "Symbolic Cholesky decomposition failed for augmented system".to_string(),
207 )
208 .log_with_source(e)
209 })?;
210 self.symbolic_factorization = Some(new_sym.clone());
212 new_sym
213 };
214
215 let cholesky = Llt::try_new_with_symbolic(sym, augmented_hessian.as_ref(), Side::Lower)
217 .map_err(|e| {
218 LinAlgError::SingularMatrix(
219 "Cholesky factorization failed (matrix may be singular)".to_string(),
220 )
221 .log_with_source(e)
222 })?;
223
224 let dx = cholesky.solve(-&gradient);
225 self.hessian = Some(hessian);
226 self.gradient = Some(gradient);
227 self.factorizer = Some(cholesky);
228
229 Ok(dx)
230 }
231
232 fn get_hessian(&self) -> Option<&SparseColMat<usize, f64>> {
233 self.hessian.as_ref()
234 }
235
236 fn get_gradient(&self) -> Option<&Mat<f64>> {
237 self.gradient.as_ref()
238 }
239
240 fn compute_covariance_matrix(&mut self) -> Option<&Mat<f64>> {
241 if self.factorizer.is_some()
243 && self.hessian.is_some()
244 && self.covariance_matrix.is_none()
245 && let (Some(factorizer), Some(hessian)) = (&self.factorizer, &self.hessian)
246 {
247 let n = hessian.ncols();
248 let identity = Mat::identity(n, n);
250
251 let cov_matrix = factorizer.solve(&identity);
253 self.covariance_matrix = Some(cov_matrix);
254 }
255 self.covariance_matrix.as_ref()
256 }
257
258 fn get_covariance_matrix(&self) -> Option<&Mat<f64>> {
259 self.covariance_matrix.as_ref()
260 }
261}
262
263#[cfg(test)]
264mod tests {
265 use super::*;
266
267 const TOLERANCE: f64 = 1e-10;
268
269 type TestResult = Result<(), Box<dyn std::error::Error>>;
270
271 fn create_test_data()
273 -> Result<(SparseColMat<usize, f64>, Mat<f64>), faer::sparse::CreationError> {
274 let triplets = vec![
276 Triplet::new(0, 0, 2.0),
277 Triplet::new(0, 1, 1.0),
278 Triplet::new(1, 0, 1.0),
279 Triplet::new(1, 1, 3.0),
280 Triplet::new(1, 2, 1.0),
281 Triplet::new(2, 1, 1.0),
282 Triplet::new(2, 2, 2.0),
283 Triplet::new(3, 0, 1.5), Triplet::new(3, 2, 0.5),
285 ];
286 let jacobian = SparseColMat::try_new_from_triplets(4, 3, &triplets)?;
287
288 let residuals = Mat::from_fn(4, 1, |i, _| match i {
289 0 => 1.0,
290 1 => -2.0,
291 2 => 0.5,
292 3 => 1.2,
293 _ => 0.0,
294 });
295
296 Ok((jacobian, residuals))
297 }
298
299 #[test]
301 fn test_solver_creation() {
302 let solver = SparseCholeskySolver::new();
303 assert!(solver.factorizer.is_none());
304
305 let default_solver = SparseCholeskySolver::default();
306 assert!(default_solver.factorizer.is_none());
307 }
308
309 #[test]
311 fn test_solve_normal_equation_well_conditioned() -> TestResult {
312 let mut solver = SparseCholeskySolver::new();
313 let (jacobian, residuals) = create_test_data()?;
314
315 let solution =
316 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
317 assert_eq!(solution.nrows(), 3);
318 assert_eq!(solution.ncols(), 1);
319
320 assert!(solver.factorizer.is_some());
322 Ok(())
323 }
324
325 #[test]
327 fn test_symbolic_pattern_caching() -> TestResult {
328 let mut solver = SparseCholeskySolver::new();
329 let (jacobian, residuals) = create_test_data()?;
330
331 let sol1 =
333 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
334 assert!(solver.factorizer.is_some());
335
336 let sol2 =
338 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
339
340 for i in 0..sol1.nrows() {
342 assert!((sol1[(i, 0)] - sol2[(i, 0)]).abs() < TOLERANCE);
343 }
344 Ok(())
345 }
346
347 #[test]
349 fn test_solve_augmented_equation() -> TestResult {
350 let mut solver = SparseCholeskySolver::new();
351 let (jacobian, residuals) = create_test_data()?;
352 let lambda = 0.1;
353
354 let solution = LinearSolver::<SparseMode>::solve_augmented_equation(
355 &mut solver,
356 &residuals,
357 &jacobian,
358 lambda,
359 )?;
360 assert_eq!(solution.nrows(), 3);
361 assert_eq!(solution.ncols(), 1);
362 Ok(())
363 }
364
365 #[test]
367 fn test_augmented_equation_different_lambdas() -> TestResult {
368 let mut solver = SparseCholeskySolver::new();
369 let (jacobian, residuals) = create_test_data()?;
370
371 let lambda1 = 0.01;
372 let lambda2 = 1.0;
373
374 let sol1 = LinearSolver::<SparseMode>::solve_augmented_equation(
375 &mut solver,
376 &residuals,
377 &jacobian,
378 lambda1,
379 )?;
380 let sol2 = LinearSolver::<SparseMode>::solve_augmented_equation(
381 &mut solver,
382 &residuals,
383 &jacobian,
384 lambda2,
385 )?;
386
387 let mut different = false;
389 for i in 0..sol1.nrows() {
390 if (sol1[(i, 0)] - sol2[(i, 0)]).abs() > TOLERANCE {
391 different = true;
392 break;
393 }
394 }
395 assert!(
396 different,
397 "Solutions should differ with different lambda values"
398 );
399 Ok(())
400 }
401
402 #[test]
404 fn test_singular_matrix() -> TestResult {
405 let mut solver = SparseCholeskySolver::new();
406
407 let triplets = vec![
409 Triplet::new(0, 0, 1.0),
410 Triplet::new(0, 1, 2.0),
411 Triplet::new(1, 0, 2.0),
412 Triplet::new(1, 1, 4.0), ];
414 let singular_jacobian = SparseColMat::try_new_from_triplets(2, 2, &triplets)?;
415 let residuals = Mat::from_fn(2, 1, |i, _| i as f64);
416
417 let result = LinearSolver::<SparseMode>::solve_normal_equation(
418 &mut solver,
419 &residuals,
420 &singular_jacobian,
421 );
422 assert!(result.is_err(), "Singular matrix should return Err");
424 Ok(())
425 }
426
427 #[test]
429 fn test_empty_matrix() -> TestResult {
430 let mut solver = SparseCholeskySolver::new();
431
432 let empty_jacobian = SparseColMat::try_new_from_triplets(0, 0, &[])?;
433 let empty_residuals = Mat::zeros(0, 1);
434
435 let result = LinearSolver::<SparseMode>::solve_normal_equation(
436 &mut solver,
437 &empty_residuals,
438 &empty_jacobian,
439 );
440 if let Ok(solution) = result {
441 assert_eq!(solution.nrows(), 0);
442 }
443 Ok(())
444 }
445
446 #[test]
448 fn test_numerical_accuracy() -> TestResult {
449 let mut solver = SparseCholeskySolver::new();
450
451 let triplets = vec![
453 Triplet::new(0, 0, 1.0),
454 Triplet::new(0, 1, 0.0),
455 Triplet::new(1, 0, 0.0),
456 Triplet::new(1, 1, 1.0),
457 ];
458 let jacobian = SparseColMat::try_new_from_triplets(2, 2, &triplets)?;
459 let residuals = Mat::from_fn(2, 1, |i, _| -((i + 1) as f64)); let solution =
462 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
463 assert!((solution[(0, 0)] - 1.0).abs() < TOLERANCE);
465 assert!((solution[(1, 0)] - 2.0).abs() < TOLERANCE);
466 Ok(())
467 }
468
469 #[test]
471 fn test_solver_clone() {
472 let solver1 = SparseCholeskySolver::new();
473 let solver2 = solver1.clone();
474
475 assert!(solver1.factorizer.is_none());
476 assert!(solver2.factorizer.is_none());
477 }
478
479 #[test]
481 fn test_cholesky_covariance_computation() -> TestResult {
482 let mut solver = SparseCholeskySolver::new();
483 let (jacobian, residuals) = create_test_data()?;
484
485 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
487
488 let cov_matrix = LinearSolver::<SparseMode>::compute_covariance_matrix(&mut solver);
490 assert!(cov_matrix.is_some());
491
492 if let Some(cov) = cov_matrix {
493 assert_eq!(cov.nrows(), 3); assert_eq!(cov.ncols(), 3);
495
496 for i in 0..3 {
498 for j in 0..3 {
499 assert!(
500 (cov[(i, j)] - cov[(j, i)]).abs() < TOLERANCE,
501 "Covariance matrix should be symmetric"
502 );
503 }
504 }
505
506 for i in 0..3 {
508 assert!(
509 cov[(i, i)] > 0.0,
510 "Diagonal elements (variances) should be positive"
511 );
512 }
513 }
514 Ok(())
515 }
516
517 #[test]
519 fn test_cholesky_standard_errors_computation() -> TestResult {
520 let mut solver = SparseCholeskySolver::new();
521 let (jacobian, residuals) = create_test_data()?;
522
523 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
525
526 solver.compute_standard_errors();
528
529 assert!(solver.covariance_matrix.is_some());
531 assert!(solver.standard_errors.is_some());
532
533 if let (Some(cov), Some(errors)) = (&solver.covariance_matrix, &solver.standard_errors) {
534 assert_eq!(errors.nrows(), 3); assert_eq!(errors.ncols(), 1);
536
537 for i in 0..3 {
539 assert!(errors[(i, 0)] > 0.0, "Standard errors should be positive");
540 }
541
542 for i in 0..3 {
544 let expected_std_error = cov[(i, i)].sqrt();
545 assert!(
546 (errors[(i, 0)] - expected_std_error).abs() < TOLERANCE,
547 "Standard error should equal sqrt of covariance diagonal"
548 );
549 }
550 }
551 Ok(())
552 }
553
554 #[test]
556 fn test_cholesky_covariance_positive_definite() -> TestResult {
557 let mut solver = SparseCholeskySolver::new();
558
559 let triplets = vec![
561 Triplet::new(0, 0, 3.0),
562 Triplet::new(0, 1, 1.0),
563 Triplet::new(1, 0, 1.0),
564 Triplet::new(1, 1, 2.0),
565 ];
566 let jacobian = SparseColMat::try_new_from_triplets(2, 2, &triplets)?;
567 let residuals = Mat::from_fn(2, 1, |i, _| (i + 1) as f64);
568
569 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
570
571 let cov_matrix = LinearSolver::<SparseMode>::compute_covariance_matrix(&mut solver);
572 assert!(cov_matrix.is_some());
573
574 if let Some(cov) = cov_matrix {
575 assert!((cov[(0, 0)] - 0.2).abs() < TOLERANCE);
578 assert!((cov[(1, 1)] - 0.4).abs() < TOLERANCE);
579 assert!((cov[(0, 1)] - (-0.2)).abs() < TOLERANCE);
580 assert!((cov[(1, 0)] - (-0.2)).abs() < TOLERANCE);
581 }
582 Ok(())
583 }
584
585 #[test]
587 fn test_cholesky_covariance_caching() -> TestResult {
588 let mut solver = SparseCholeskySolver::new();
589 let (jacobian, residuals) = create_test_data()?;
590
591 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
593
594 LinearSolver::<SparseMode>::compute_covariance_matrix(&mut solver);
596 assert!(solver.covariance_matrix.is_some());
597
598 if let Some(cov1) = &solver.covariance_matrix {
600 let cov1_ptr = cov1.as_ptr();
601
602 LinearSolver::<SparseMode>::compute_covariance_matrix(&mut solver);
604 assert!(solver.covariance_matrix.is_some());
605
606 if let Some(cov2) = &solver.covariance_matrix {
608 let cov2_ptr = cov2.as_ptr();
609
610 assert_eq!(cov1_ptr, cov2_ptr, "Covariance matrix should be cached");
612 }
613 }
614 Ok(())
615 }
616
617 #[test]
619 fn test_cholesky_decomposition_properties() -> TestResult {
620 let mut solver = SparseCholeskySolver::new();
621
622 let triplets = vec![Triplet::new(0, 0, 2.0), Triplet::new(1, 1, 3.0)];
624 let jacobian = SparseColMat::try_new_from_triplets(2, 2, &triplets)?;
625 let residuals = Mat::from_fn(2, 1, |i, _| (i + 1) as f64);
626
627 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
628
629 assert!(solver.factorizer.is_some());
631 assert!(solver.hessian.is_some());
632
633 if let Some(hessian) = &solver.hessian {
635 assert_eq!(hessian.nrows(), 2);
636 assert_eq!(hessian.ncols(), 2);
637 }
638 Ok(())
639 }
640
641 #[test]
643 fn test_cholesky_numerical_stability() -> TestResult {
644 let mut solver = SparseCholeskySolver::new();
645
646 let triplets = vec![
648 Triplet::new(0, 0, 1.0),
649 Triplet::new(1, 1, 1.0),
650 Triplet::new(2, 2, 1.0),
651 ];
652 let jacobian = SparseColMat::try_new_from_triplets(3, 3, &triplets)?;
653 let residuals = Mat::from_fn(3, 1, |i, _| -((i + 1) as f64)); let solution =
656 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
657 for i in 0..3 {
659 let expected = (i + 1) as f64;
660 assert!(
661 (solution[(i, 0)] - expected).abs() < TOLERANCE,
662 "Expected {}, got {}",
663 expected,
664 solution[(i, 0)]
665 );
666 }
667
668 let cov_matrix = LinearSolver::<SparseMode>::compute_covariance_matrix(&mut solver);
670 assert!(cov_matrix.is_some());
671 if let Some(cov) = cov_matrix {
672 for i in 0..3 {
673 for j in 0..3 {
674 let expected = if i == j { 1.0 } else { 0.0 };
675 assert!(
676 (cov[(i, j)] - expected).abs() < TOLERANCE,
677 "Covariance[{}, {}] expected {}, got {}",
678 i,
679 j,
680 expected,
681 cov[(i, j)]
682 );
683 }
684 }
685 }
686 Ok(())
687 }
688
689 #[test]
691 fn test_cholesky_hessian_getter() -> TestResult {
692 let mut solver = SparseCholeskySolver::new();
693 assert!(solver.hessian().is_none());
694
695 let (jacobian, residuals) = create_test_data()?;
696 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
697
698 assert!(solver.hessian().is_some());
699 Ok(())
700 }
701
702 #[test]
704 fn test_cholesky_gradient_getter() -> TestResult {
705 let mut solver = SparseCholeskySolver::new();
706 assert!(solver.gradient().is_none());
707
708 let (jacobian, residuals) = create_test_data()?;
709 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
710
711 assert!(solver.gradient().is_some());
712 Ok(())
713 }
714
715 #[test]
717 fn test_cholesky_reset_covariance() -> TestResult {
718 let mut solver = SparseCholeskySolver::new();
719 let (jacobian, residuals) = create_test_data()?;
720
721 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
722 LinearSolver::<SparseMode>::compute_covariance_matrix(&mut solver);
723 assert!(solver.covariance_matrix.is_some());
724
725 solver.reset_covariance();
726 assert!(solver.covariance_matrix.is_none());
727 assert!(solver.standard_errors.is_none());
728 Ok(())
729 }
730
731 #[test]
733 fn test_cholesky_get_hessian_trait() -> TestResult {
734 let mut solver = SparseCholeskySolver::new();
735 assert!(LinearSolver::<SparseMode>::get_hessian(&solver).is_none());
736
737 let (jacobian, residuals) = create_test_data()?;
738 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
739
740 assert!(LinearSolver::<SparseMode>::get_hessian(&solver).is_some());
741 Ok(())
742 }
743
744 #[test]
746 fn test_cholesky_get_gradient_trait() -> TestResult {
747 let mut solver = SparseCholeskySolver::new();
748 assert!(LinearSolver::<SparseMode>::get_gradient(&solver).is_none());
749
750 let (jacobian, residuals) = create_test_data()?;
751 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
752
753 assert!(LinearSolver::<SparseMode>::get_gradient(&solver).is_some());
754 Ok(())
755 }
756
757 #[test]
759 fn test_cholesky_get_covariance_matrix_getter() -> TestResult {
760 let mut solver = SparseCholeskySolver::new();
761 let (jacobian, residuals) = create_test_data()?;
762
763 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
764 LinearSolver::<SparseMode>::compute_covariance_matrix(&mut solver);
765
766 let via_getter = LinearSolver::<SparseMode>::get_covariance_matrix(&solver);
767 assert!(via_getter.is_some());
768
769 if let Some(cov) = via_getter {
770 assert_eq!(cov.nrows(), 3);
771 assert_eq!(cov.ncols(), 3);
772 }
773 Ok(())
774 }
775
776 #[test]
781 fn test_compute_standard_errors_before_solve_returns_none() {
782 let mut solver = SparseCholeskySolver::new();
783 let result = solver.compute_standard_errors();
785 assert!(
786 result.is_none(),
787 "compute_standard_errors on uninitialized solver (no hessian) should return None"
788 );
789 }
790}