1use faer::{
2 Mat,
3 linalg::solvers::Solve,
4 sparse::linalg::solvers::{Qr, SymbolicQr},
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 SparseQRSolver {
14 factorizer: Option<Qr<usize, f64>>,
15
16 symbolic_factorization: Option<SymbolicQr<usize>>,
23
24 hessian: Option<SparseColMat<usize, f64>>,
28
29 gradient: Option<Mat<f64>>,
33
34 covariance_matrix: Option<Mat<f64>>,
38 standard_errors: Option<Mat<f64>>,
44}
45
46impl SparseQRSolver {
47 pub fn new() -> Self {
48 SparseQRSolver {
49 factorizer: None,
50 symbolic_factorization: None,
51 hessian: None,
52 gradient: None,
53 covariance_matrix: None,
54 standard_errors: None,
55 }
56 }
57
58 pub fn hessian(&self) -> Option<&SparseColMat<usize, f64>> {
59 self.hessian.as_ref()
60 }
61
62 pub fn gradient(&self) -> Option<&Mat<f64>> {
63 self.gradient.as_ref()
64 }
65
66 pub fn compute_standard_errors(&mut self) -> Option<&Mat<f64>> {
67 if self.covariance_matrix.is_none() {
69 LinearSolver::<SparseMode>::compute_covariance_matrix(self);
70 }
71
72 let hessian = self.hessian.as_ref()?;
74 let n = hessian.ncols();
75 if let Some(cov) = &self.covariance_matrix {
77 let mut std_errors = Mat::zeros(n, 1);
78 for i in 0..n {
79 let diag_val = cov[(i, i)];
80 if diag_val >= 0.0 {
81 std_errors[(i, 0)] = diag_val.sqrt();
82 } else {
83 return None;
85 }
86 }
87 self.standard_errors = Some(std_errors);
88 }
89 self.standard_errors.as_ref()
90 }
91
92 pub fn reset_covariance(&mut self) {
94 self.covariance_matrix = None;
95 self.standard_errors = None;
96 }
97}
98
99impl Default for SparseQRSolver {
100 fn default() -> Self {
101 Self::new()
102 }
103}
104
105impl LinearSolver<SparseMode> for SparseQRSolver {
106 fn solve_normal_equation(
107 &mut self,
108 residuals: &Mat<f64>,
109 jacobians: &SparseColMat<usize, f64>,
110 ) -> LinAlgResult<Mat<f64>> {
111 let jt = jacobians.as_ref().transpose();
113 let hessian = jt
114 .to_col_major()
115 .map_err(|e| {
116 LinAlgError::MatrixConversion(
117 "Failed to convert transposed Jacobian to column-major format".to_string(),
118 )
119 .log_with_source(e)
120 })?
121 .mul(jacobians.as_ref());
122
123 let gradient = jacobians.as_ref().transpose().mul(residuals);
125
126 let sym = if let Some(ref cached_sym) = self.symbolic_factorization {
129 cached_sym.clone()
134 } else {
135 let new_sym = SymbolicQr::try_new(hessian.symbolic()).map_err(|e| {
137 LinAlgError::FactorizationFailed("Symbolic QR decomposition failed".to_string())
138 .log_with_source(e)
139 })?;
140 self.symbolic_factorization = Some(new_sym.clone());
142 new_sym
143 };
144
145 let qr = Qr::try_new_with_symbolic(sym, hessian.as_ref()).map_err(|e| {
147 LinAlgError::SingularMatrix(
148 "QR factorization failed (matrix may be singular)".to_string(),
149 )
150 .log_with_source(e)
151 })?;
152
153 let dx = qr.solve(-&gradient);
155 self.hessian = Some(hessian);
156 self.gradient = Some(gradient);
157 self.factorizer = Some(qr);
158
159 Ok(dx)
160 }
161
162 fn solve_augmented_equation(
163 &mut self,
164 residuals: &Mat<f64>,
165 jacobians: &SparseColMat<usize, f64>,
166 lambda: f64,
167 ) -> LinAlgResult<Mat<f64>> {
168 let n = jacobians.ncols();
169
170 let jt = jacobians.as_ref().transpose();
172 let hessian = jt
173 .to_col_major()
174 .map_err(|e| {
175 LinAlgError::MatrixConversion(
176 "Failed to convert transposed Jacobian to column-major format".to_string(),
177 )
178 .log_with_source(e)
179 })?
180 .mul(jacobians.as_ref());
181
182 let gradient = jacobians.as_ref().transpose().mul(residuals);
184
185 let mut lambda_i_triplets = Vec::with_capacity(n);
187 for i in 0..n {
188 lambda_i_triplets.push(Triplet::new(i, i, lambda));
189 }
190 let lambda_i =
191 SparseColMat::try_new_from_triplets(n, n, &lambda_i_triplets).map_err(|e| {
192 LinAlgError::SparseMatrixCreation("Failed to create lambda*I matrix".to_string())
193 .log_with_source(e)
194 })?;
195
196 let augmented_hessian = hessian.as_ref() + lambda_i;
197
198 let sym = if let Some(ref cached_sym) = self.symbolic_factorization {
203 cached_sym.clone()
204 } else {
205 let new_sym = SymbolicQr::try_new(augmented_hessian.symbolic()).map_err(|e| {
207 LinAlgError::FactorizationFailed(
208 "Symbolic QR decomposition failed for augmented system".to_string(),
209 )
210 .log_with_source(e)
211 })?;
212 self.symbolic_factorization = Some(new_sym.clone());
214 new_sym
215 };
216
217 let qr = Qr::try_new_with_symbolic(sym, augmented_hessian.as_ref()).map_err(|e| {
219 LinAlgError::SingularMatrix(
220 "QR factorization failed (matrix may be singular)".to_string(),
221 )
222 .log_with_source(e)
223 })?;
224
225 let dx = qr.solve(-&gradient);
226 self.hessian = Some(hessian);
227 self.gradient = Some(gradient);
228 self.factorizer = Some(qr);
229
230 Ok(dx)
231 }
232
233 fn get_hessian(&self) -> Option<&SparseColMat<usize, f64>> {
234 self.hessian.as_ref()
235 }
236
237 fn get_gradient(&self) -> Option<&Mat<f64>> {
238 self.gradient.as_ref()
239 }
240
241 fn compute_covariance_matrix(&mut self) -> Option<&Mat<f64>> {
242 if self.factorizer.is_some()
244 && self.hessian.is_some()
245 && self.covariance_matrix.is_none()
246 && let (Some(factorizer), Some(hessian)) = (&self.factorizer, &self.hessian)
247 {
248 let n = hessian.ncols();
249 let identity = Mat::identity(n, n);
251
252 let cov_matrix = factorizer.solve(&identity);
254 self.covariance_matrix = Some(cov_matrix);
255 }
256 self.covariance_matrix.as_ref()
257 }
258
259 fn get_covariance_matrix(&self) -> Option<&Mat<f64>> {
260 self.covariance_matrix.as_ref()
261 }
262}
263
264#[cfg(test)]
265mod tests {
266 use super::*;
267
268 const TOLERANCE: f64 = 1e-10;
269
270 type TestResult = Result<(), Box<dyn std::error::Error>>;
271
272 fn create_test_data()
274 -> Result<(SparseColMat<usize, f64>, Mat<f64>), faer::sparse::CreationError> {
275 let triplets = vec![
277 Triplet::new(0, 0, 1.0),
278 Triplet::new(0, 1, 0.0),
279 Triplet::new(0, 2, 1.0),
280 Triplet::new(1, 0, 0.0),
281 Triplet::new(1, 1, 1.0),
282 Triplet::new(1, 2, 1.0),
283 Triplet::new(2, 0, 1.0),
284 Triplet::new(2, 1, 1.0),
285 Triplet::new(2, 2, 0.0),
286 Triplet::new(3, 0, 1.0),
287 Triplet::new(3, 1, 0.0),
288 Triplet::new(3, 2, 0.0),
289 ];
290 let jacobian = SparseColMat::try_new_from_triplets(4, 3, &triplets)?;
291
292 let residuals = Mat::from_fn(4, 1, |i, _| (i + 1) as f64);
293
294 Ok((jacobian, residuals))
295 }
296
297 #[test]
299 fn test_qr_solver_creation() {
300 let solver = SparseQRSolver::new();
301 assert!(solver.factorizer.is_none());
302
303 let default_solver = SparseQRSolver::default();
304 assert!(default_solver.factorizer.is_none());
305 }
306
307 #[test]
309 fn test_qr_solve_normal_equation() -> TestResult {
310 let mut solver = SparseQRSolver::new();
311 let (jacobian, residuals) = create_test_data()?;
312
313 let solution =
314 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
315 assert_eq!(solution.nrows(), 3); assert_eq!(solution.ncols(), 1);
317
318 assert!(solver.factorizer.is_some());
320 Ok(())
321 }
322
323 #[test]
325 fn test_qr_factorizer_caching() -> TestResult {
326 let mut solver = SparseQRSolver::new();
327 let (jacobian, residuals) = create_test_data()?;
328
329 let sol1 =
331 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
332 assert!(solver.factorizer.is_some());
333
334 let sol2 =
336 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
337
338 for i in 0..sol1.nrows() {
340 assert!((sol1[(i, 0)] - sol2[(i, 0)]).abs() < TOLERANCE);
341 }
342 Ok(())
343 }
344
345 #[test]
347 fn test_qr_solve_augmented_equation() -> TestResult {
348 let mut solver = SparseQRSolver::new();
349 let (jacobian, residuals) = create_test_data()?;
350 let lambda = 0.1;
351
352 let solution = LinearSolver::<SparseMode>::solve_augmented_equation(
353 &mut solver,
354 &residuals,
355 &jacobian,
356 lambda,
357 )?;
358 assert_eq!(solution.nrows(), 3); assert_eq!(solution.ncols(), 1);
360 Ok(())
361 }
362
363 #[test]
365 fn test_qr_augmented_different_lambdas() -> TestResult {
366 let mut solver = SparseQRSolver::new();
367 let (jacobian, residuals) = create_test_data()?;
368
369 let lambda1 = 0.01;
370 let lambda2 = 1.0;
371
372 let sol1 = LinearSolver::<SparseMode>::solve_augmented_equation(
373 &mut solver,
374 &residuals,
375 &jacobian,
376 lambda1,
377 )?;
378 let sol2 = LinearSolver::<SparseMode>::solve_augmented_equation(
379 &mut solver,
380 &residuals,
381 &jacobian,
382 lambda2,
383 )?;
384
385 let mut different = false;
387 for i in 0..sol1.nrows() {
388 if (sol1[(i, 0)] - sol2[(i, 0)]).abs() > TOLERANCE {
389 different = true;
390 break;
391 }
392 }
393 assert!(
394 different,
395 "Solutions should differ with different lambda values"
396 );
397 Ok(())
398 }
399
400 #[test]
402 fn test_qr_rank_deficient_matrix() -> TestResult {
403 let mut solver = SparseQRSolver::new();
404
405 let triplets = vec![
407 Triplet::new(0, 0, 1.0),
408 Triplet::new(0, 1, 2.0),
409 Triplet::new(0, 2, 3.0),
410 Triplet::new(1, 0, 2.0),
411 Triplet::new(1, 1, 4.0),
412 Triplet::new(1, 2, 6.0), Triplet::new(2, 0, 0.0),
414 Triplet::new(2, 1, 0.0),
415 Triplet::new(2, 2, 1.0),
416 ];
417 let jacobian = SparseColMat::try_new_from_triplets(3, 3, &triplets)?;
418 let residuals = Mat::from_fn(3, 1, |i, _| i as f64);
419
420 let result =
422 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian);
423 assert!(result.is_ok());
424 Ok(())
425 }
426
427 #[test]
429 fn test_qr_augmented_system_structure() -> TestResult {
430 let mut solver = SparseQRSolver::new();
431
432 let triplets = vec![
434 Triplet::new(0, 0, 1.0),
435 Triplet::new(0, 1, 0.0),
436 Triplet::new(1, 0, 0.0),
437 Triplet::new(1, 1, 1.0),
438 ];
439 let jacobian = SparseColMat::try_new_from_triplets(2, 2, &triplets)?;
440 let residuals = Mat::from_fn(2, 1, |i, _| (i + 1) as f64);
441 let lambda = 0.5;
442
443 let solution = LinearSolver::<SparseMode>::solve_augmented_equation(
444 &mut solver,
445 &residuals,
446 &jacobian,
447 lambda,
448 )?;
449 assert_eq!(solution.nrows(), 2); assert_eq!(solution.ncols(), 1);
451 Ok(())
452 }
453
454 #[test]
456 fn test_qr_numerical_accuracy() -> TestResult {
457 let mut solver = SparseQRSolver::new();
458
459 let triplets = vec![
461 Triplet::new(0, 0, 1.0),
462 Triplet::new(1, 1, 1.0),
463 Triplet::new(2, 2, 1.0),
464 ];
465 let jacobian = SparseColMat::try_new_from_triplets(3, 3, &triplets)?;
466
467 let residuals = Mat::from_fn(3, 1, |i, _| -((i + 1) as f64)); let solution =
470 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
471 for i in 0..3 {
473 let expected = (i + 1) as f64;
474 assert!(
475 (solution[(i, 0)] - expected).abs() < TOLERANCE,
476 "Expected {}, got {}",
477 expected,
478 solution[(i, 0)]
479 );
480 }
481 Ok(())
482 }
483
484 #[test]
486 fn test_qr_solver_clone() {
487 let solver1 = SparseQRSolver::new();
488 let solver2 = solver1.clone();
489
490 assert!(solver1.factorizer.is_none());
491 assert!(solver2.factorizer.is_none());
492 }
493
494 #[test]
496 fn test_qr_zero_lambda_augmented() -> TestResult {
497 let mut solver = SparseQRSolver::new();
498 let (jacobian, residuals) = create_test_data()?;
499
500 let normal_sol =
501 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
502 let augmented_sol = LinearSolver::<SparseMode>::solve_augmented_equation(
503 &mut solver,
504 &residuals,
505 &jacobian,
506 0.0,
507 )?;
508
509 for i in 0..normal_sol.nrows() {
511 assert!(
512 (normal_sol[(i, 0)] - augmented_sol[(i, 0)]).abs() < 1e-8,
513 "Zero lambda augmented should match normal equation"
514 );
515 }
516 Ok(())
517 }
518
519 #[test]
521 fn test_qr_covariance_computation() -> TestResult {
522 let mut solver = SparseQRSolver::new();
523 let (jacobian, residuals) = create_test_data()?;
524
525 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
527
528 let cov_matrix = LinearSolver::<SparseMode>::compute_covariance_matrix(&mut solver);
530 assert!(cov_matrix.is_some());
531
532 if let Some(cov) = cov_matrix {
533 assert_eq!(cov.nrows(), 3); assert_eq!(cov.ncols(), 3);
535
536 for i in 0..3 {
538 for j in 0..3 {
539 assert!(
540 (cov[(i, j)] - cov[(j, i)]).abs() < TOLERANCE,
541 "Covariance matrix should be symmetric"
542 );
543 }
544 }
545
546 for i in 0..3 {
548 assert!(
549 cov[(i, i)] > 0.0,
550 "Diagonal elements (variances) should be positive"
551 );
552 }
553 }
554 Ok(())
555 }
556
557 #[test]
559 fn test_qr_standard_errors_computation() -> TestResult {
560 let mut solver = SparseQRSolver::new();
561 let (jacobian, residuals) = create_test_data()?;
562
563 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
565
566 solver.compute_standard_errors();
568
569 assert!(solver.covariance_matrix.is_some());
571 assert!(solver.standard_errors.is_some());
572
573 if let (Some(cov), Some(errors)) = (&solver.covariance_matrix, &solver.standard_errors) {
574 assert_eq!(errors.nrows(), 3); assert_eq!(errors.ncols(), 1);
576
577 for i in 0..3 {
579 assert!(errors[(i, 0)] > 0.0, "Standard errors should be positive");
580 }
581
582 for i in 0..3 {
584 let expected_std_error = cov[(i, i)].sqrt();
585 assert!(
586 (errors[(i, 0)] - expected_std_error).abs() < TOLERANCE,
587 "Standard error should equal sqrt of covariance diagonal"
588 );
589 }
590 }
591 Ok(())
592 }
593
594 #[test]
596 fn test_qr_covariance_well_conditioned() -> TestResult {
597 let mut solver = SparseQRSolver::new();
598
599 let triplets = vec![
601 Triplet::new(0, 0, 2.0),
602 Triplet::new(0, 1, 0.0),
603 Triplet::new(1, 0, 0.0),
604 Triplet::new(1, 1, 3.0),
605 ];
606 let jacobian = SparseColMat::try_new_from_triplets(2, 2, &triplets)?;
607 let residuals = Mat::from_fn(2, 1, |i, _| (i + 1) as f64);
608
609 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
610
611 let cov_matrix = LinearSolver::<SparseMode>::compute_covariance_matrix(&mut solver);
612 assert!(cov_matrix.is_some());
613
614 if let Some(cov) = cov_matrix {
615 assert!((cov[(0, 0)] - 0.25).abs() < TOLERANCE);
618 assert!((cov[(1, 1)] - 1.0 / 9.0).abs() < TOLERANCE);
619 assert!(cov[(0, 1)].abs() < TOLERANCE);
620 assert!(cov[(1, 0)].abs() < TOLERANCE);
621 }
622 Ok(())
623 }
624
625 #[test]
627 fn test_qr_covariance_caching() -> TestResult {
628 let mut solver = SparseQRSolver::new();
629 let (jacobian, residuals) = create_test_data()?;
630
631 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
633
634 LinearSolver::<SparseMode>::compute_covariance_matrix(&mut solver);
636 assert!(solver.covariance_matrix.is_some());
637
638 if let Some(cov1) = &solver.covariance_matrix {
640 let cov1_ptr = cov1.as_ptr();
641
642 LinearSolver::<SparseMode>::compute_covariance_matrix(&mut solver);
644 assert!(solver.covariance_matrix.is_some());
645
646 if let Some(cov2) = &solver.covariance_matrix {
648 let cov2_ptr = cov2.as_ptr();
649
650 assert_eq!(cov1_ptr, cov2_ptr, "Covariance matrix should be cached");
652 }
653 }
654 Ok(())
655 }
656
657 #[test]
659 fn test_qr_covariance_singular_system() -> TestResult {
660 let mut solver = SparseQRSolver::new();
661
662 let triplets = vec![
664 Triplet::new(0, 0, 1.0),
665 Triplet::new(0, 1, 2.0),
666 Triplet::new(1, 0, 2.0),
667 Triplet::new(1, 1, 4.0), ];
669 let jacobian = SparseColMat::try_new_from_triplets(2, 2, &triplets)?;
670 let residuals = Mat::from_fn(2, 1, |i, _| i as f64);
671
672 let result =
674 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian);
675 if result.is_ok() {
676 let cov_matrix = LinearSolver::<SparseMode>::compute_covariance_matrix(&mut solver);
678 if let Some(cov) = cov_matrix {
680 assert!(cov.nrows() == 2);
682 assert!(cov.ncols() == 2);
683 }
684 }
685 Ok(())
686 }
687
688 #[test]
690 fn test_qr_hessian_getter() -> TestResult {
691 let mut solver = SparseQRSolver::new();
692 assert!(solver.hessian().is_none());
693
694 let (jacobian, residuals) = create_test_data()?;
695 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
696
697 assert!(solver.hessian().is_some());
698 Ok(())
699 }
700
701 #[test]
703 fn test_qr_gradient_getter() -> TestResult {
704 let mut solver = SparseQRSolver::new();
705 assert!(solver.gradient().is_none());
706
707 let (jacobian, residuals) = create_test_data()?;
708 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
709
710 assert!(solver.gradient().is_some());
711 Ok(())
712 }
713
714 #[test]
716 fn test_qr_reset_covariance() -> TestResult {
717 let mut solver = SparseQRSolver::new();
718 let (jacobian, residuals) = create_test_data()?;
719
720 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
721 LinearSolver::<SparseMode>::compute_covariance_matrix(&mut solver);
722 assert!(solver.covariance_matrix.is_some());
723
724 solver.reset_covariance();
725 assert!(solver.covariance_matrix.is_none());
726 assert!(solver.standard_errors.is_none());
727 Ok(())
728 }
729
730 #[test]
732 fn test_qr_get_hessian_trait() -> TestResult {
733 let mut solver = SparseQRSolver::new();
734 assert!(LinearSolver::<SparseMode>::get_hessian(&solver).is_none());
735
736 let (jacobian, residuals) = create_test_data()?;
737 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
738
739 assert!(LinearSolver::<SparseMode>::get_hessian(&solver).is_some());
740 Ok(())
741 }
742
743 #[test]
745 fn test_qr_get_gradient_trait() -> TestResult {
746 let mut solver = SparseQRSolver::new();
747 assert!(LinearSolver::<SparseMode>::get_gradient(&solver).is_none());
748
749 let (jacobian, residuals) = create_test_data()?;
750 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
751
752 assert!(LinearSolver::<SparseMode>::get_gradient(&solver).is_some());
753 Ok(())
754 }
755
756 #[test]
758 fn test_qr_get_covariance_matrix_getter() -> TestResult {
759 let mut solver = SparseQRSolver::new();
760 let (jacobian, residuals) = create_test_data()?;
761
762 LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
763 LinearSolver::<SparseMode>::compute_covariance_matrix(&mut solver);
764
765 let via_getter = LinearSolver::<SparseMode>::get_covariance_matrix(&solver);
766 assert!(via_getter.is_some());
767
768 if let Some(cov) = via_getter {
769 assert_eq!(cov.nrows(), 3);
770 assert_eq!(cov.ncols(), 3);
771 }
772 Ok(())
773 }
774}