Refactor UKF Kalman gain calculation and add robust SPD solver

Author: jbrodovsky

core/src/filter.rs

@@ -15,7 +15,7 @@
 //! the UKF and particle filter. This model is used to update the state based on position
 //! measurements in the local level frame (i.e. a GPS fix).
 use crate::earth::METERS_TO_DEGREES;
-use crate::linalg::matrix_square_root;
+use crate::linalg::{matrix_square_root, symmetrize, robust_spd_solve};
 use crate::{IMUData, StrapdownState, forward, wrap_to_2pi};
 
 use std::fmt::Debug;
@@ -505,30 +505,34 @@ impl UKF {
             let state_diff = sigma_points.column(i) - &self.mean_state;
             cross_covariance += self.weights_cov[i] * state_diff * measurement_diff.transpose();
         }
-        // Calculate the Kalman gain
-        let s_inv = match s.clone().try_inverse() {
-            Some(inv) => inv,
-            None => panic!("Innovation matrix is singular"),
-        };
-        let k = &cross_covariance * &s_inv;
-        // check that the kalman gain and measurement diff are compatible to multiply
-        if k.ncols() != measurement.get_dimension() {
-            panic!("Kalman gain and measurement differential are not compatible");
-        }
+        // // Calculate the Kalman gain
+        // let s_inv = match s.clone().try_inverse() {
+        //     Some(inv) => inv,
+        //     None => panic!("Innovation matrix is singular"),
+        // };
+        // let k = &cross_covariance * &s_inv;
+        // // check that the kalman gain and measurement diff are compatible to multiply
+        // if k.ncols() != measurement.get_dimension() {
+        //     panic!("Kalman gain and measurement differential are not compatible");
+        // }
+        // K = P_xz * S^{-1} without forming S^{-1}
+        let k = self.robust_kalman_gain(&cross_covariance, &s);
+        // Update the mean and covariance
         self.mean_state += &k * (measurement.get_vector() - &z_hat);
         // wrap attitude angles to 2pi
         // TODO: #30 Refactor attitude angles to use a more robust representation
         self.mean_state[6] = wrap_to_2pi(self.mean_state[6]);
         self.mean_state[7] = wrap_to_2pi(self.mean_state[7]);
         self.mean_state[8] = wrap_to_2pi(self.mean_state[8]);
-        // Switch to Joseph Form update here
-        // P = (I - K * H) P (I - K H)^T + K R K^T
-        // UKF form:
-        // P -= K * S * K^T + K R K^T
-        self.covariance -= &k * s * &k.transpose();
+        self.covariance -= &k * &s * &k.transpose();
         // Re-symmetrize to fight round-off
         self.covariance = 0.5 * (&self.covariance + self.covariance.transpose());
-        // self.covariance -= &k * s * &k.transpose();
+    }
+    fn robust_kalman_gain(&mut self, cross_covariance: &DMatrix<f64>, s: &DMatrix<f64>) -> DMatrix<f64> {
+
+        // Solve S Kᵀ = P_xzᵀ  => K = (S^{-1} P_xz)ᵀ
+        let kt = robust_spd_solve(&symmetrize(s), &cross_covariance.transpose());
+        kt.transpose()
     }
 }
 #[derive(Clone, Debug, Default)]

core/src/linalg.rs

@@ -68,7 +68,8 @@ pub fn matrix_square_root(matrix: &DMatrix<f64>) -> DMatrix<f64> {
 /// 
 /// # Returns
 /// A symmetrized version of the input matrix.
-fn symmetrize(m: &DMatrix<f64>) -> DMatrix<f64> {
+#[inline]
+pub fn symmetrize(m: &DMatrix<f64>) -> DMatrix<f64> {
     0.5 * (m + m.transpose())
 }
 /// Plain Cholesky square root
@@ -77,7 +78,7 @@ fn symmetrize(m: &DMatrix<f64>) -> DMatrix<f64> {
 /// This is a quick way to initially attempt to calculate a matrix square root.
 /// 
 /// # Arguments
-/// * `p` - the matrix to factor
+/// * ``p` - the matrix to factor
 /// 
 /// # Returns
 /// A lower triangular matrix L such that P ≈ L Lᵀ, or None if it fails.
@@ -131,6 +132,64 @@ fn evd_symmetric_sqrt_with_floor(p: &DMatrix<f64>, floor: f64) -> DMatrix<f64> {
 
 }
 
+#[derive(Debug, Clone, Copy)]
+pub struct SolveOptions {
+    pub initial_jitter: f64,   // e.g., 1e-12
+    pub max_jitter: f64,       // e.g., 1e-6
+    pub max_tries: usize,      // e.g., 6
+}
+
+impl Default for SolveOptions {
+    fn default() -> Self {
+        Self { initial_jitter: 1e-12, max_jitter: 1e-6, max_tries: 6 }
+    }
+}
+/// Solve A X = B for SPD-ish A via Cholesky, with jitter retries.
+/// Returns None if all attempts fail.
+pub fn chol_solve_spd(
+    a: &DMatrix<f64>,
+    b: &DMatrix<f64>,
+    opt: SolveOptions,
+) -> Option<DMatrix<f64>> {
+    assert!(a.is_square(), "chol_solve_spd: A must be square");
+    assert_eq!(a.nrows(), b.nrows(), "chol_solve_spd: A and B incompatible");
+
+    // Symmetrize first (SPD drift is common).
+    let a_sym = symmetrize(a);
+
+    // Try plain Cholesky
+    if let Some(ch) = Cholesky::new(a_sym.clone()) {
+        return Some(ch.solve(b));
+    }
+
+    // Jitter ramp
+    let n = a_sym.nrows();
+    let mut jitter = opt.initial_jitter;
+    for _ in 0..opt.max_tries {
+        let mut a_j = a_sym.clone();
+        for i in 0..n { a_j[(i, i)] += jitter; }
+        if let Some(ch) = Cholesky::new(a_j) {
+            return Some(ch.solve(b));
+        }
+        jitter *= 10.0;
+        if jitter > opt.max_jitter { break; }
+    }
+    None
+}
+
+/// Robust SPD solve with sane defaults:
+/// - Cholesky + jitter (preferred)
+/// - Last resort: explicit inverse
+pub fn robust_spd_solve(a: &DMatrix<f64>, b: &DMatrix<f64>) -> DMatrix<f64> {
+    if let Some(x) = chol_solve_spd(a, b, SolveOptions::default()) {
+        x
+    } else if let Some(inv) = symmetrize(a).try_inverse() {
+        &inv * b
+    } else {
+        panic!("robust_spd_solve: A is not invertible (even after jitter).");
+    }
+}
+
 /* =============================== Tests ==================================== */
 
 #[cfg(test)]