qualia_core_db/solvers/learning/sequential/
kalman.rs

1//! Kalman filter (PRML ch 13.3) — exact inference for a linear-Gaussian state-space
2//! model. Recursively estimates the hidden state `x` and its covariance `P` from
3//! noisy linear observations. The matrix products reuse `linear_algebra::gemm` and
4//! the innovation-covariance inverse reuses `linear_algebra::cholesky` (no new
5//! solver). Kernel-class `DenseLinear`.
6//!
7//! Model: `xₜ = F xₜ₋₁ + w` (`w ~ N(0, Q)`), `zₜ = H xₜ + v` (`v ~ N(0, R)`).
8//! Predict: `x ← Fx`, `P ← FPFᵀ + Q`.
9//! Update with `z`: `S = HPHᵀ + R`, `K = PHᵀS⁻¹`, `x ← x + K(z − Hx)`,
10//! `P ← (I − KH)P`.
11
12use crate::solvers::learning::LearningError;
13use crate::solvers::linear_algebra::cholesky::{cholesky_factor, cholesky_solve};
14use crate::solvers::linear_algebra::gemm::{gemm, matvec, Transpose};
15
16/// A linear-Gaussian Kalman filter with current state estimate.
17#[derive(Debug, Clone)]
18pub struct KalmanFilter {
19    f: Vec<f64>, // n_x × n_x transition
20    h: Vec<f64>, // n_z × n_x observation
21    q: Vec<f64>, // n_x × n_x process noise
22    r: Vec<f64>, // n_z × n_z measurement noise
23    x: Vec<f64>, // n_x state estimate
24    p: Vec<f64>, // n_x × n_x state covariance
25    nx: usize,
26    nz: usize,
27}
28
29fn no(
30    m: usize,
31    n: usize,
32    k: usize,
33    a: &[f64],
34    b: &[f64],
35    out: &mut [f64],
36) -> Result<(), LearningError> {
37    gemm(Transpose::No, Transpose::No, m, n, k, 1.0, a, b, 0.0, out).map_err(Into::into)
38}
39fn no_t(
40    m: usize,
41    n: usize,
42    k: usize,
43    a: &[f64],
44    b: &[f64],
45    out: &mut [f64],
46) -> Result<(), LearningError> {
47    // out(m×n) = a(m×k) · b(n×k)ᵀ
48    gemm(Transpose::No, Transpose::Yes, m, n, k, 1.0, a, b, 0.0, out).map_err(Into::into)
49}
50
51impl KalmanFilter {
52    /// Construct with model matrices and an initial state estimate `(x0, p0)`.
53    pub fn new(
54        f: Vec<f64>,
55        h: Vec<f64>,
56        q: Vec<f64>,
57        r: Vec<f64>,
58        x0: Vec<f64>,
59        p0: Vec<f64>,
60        nx: usize,
61        nz: usize,
62    ) -> Result<Self, LearningError> {
63        if nx == 0
64            || nz == 0
65            || f.len() != nx * nx
66            || h.len() != nz * nx
67            || q.len() != nx * nx
68            || r.len() != nz * nz
69            || x0.len() != nx
70            || p0.len() != nx * nx
71        {
72            return Err(LearningError::InvalidDimension);
73        }
74        Ok(Self {
75            f,
76            h,
77            q,
78            r,
79            x: x0,
80            p: p0,
81            nx,
82            nz,
83        })
84    }
85
86    pub fn state(&self) -> &[f64] {
87        &self.x
88    }
89    pub fn covariance(&self) -> &[f64] {
90        &self.p
91    }
92
93    /// Time update (predict): advance the state and inflate the covariance.
94    pub fn predict(&mut self) -> Result<(), LearningError> {
95        let nx = self.nx;
96        // x ← F x.
97        let mut xn = vec![0.0; nx];
98        matvec(Transpose::No, nx, nx, &self.f, &self.x, &mut xn)?;
99        self.x = xn;
100        // P ← F P Fᵀ + Q.
101        let mut fp = vec![0.0; nx * nx];
102        no(nx, nx, nx, &self.f, &self.p, &mut fp)?;
103        let mut pn = vec![0.0; nx * nx];
104        no_t(nx, nx, nx, &fp, &self.f, &mut pn)?;
105        for i in 0..nx * nx {
106            pn[i] += self.q[i];
107        }
108        self.p = pn;
109        Ok(())
110    }
111
112    /// Measurement update (correct) with observation `z`.
113    pub fn update(&mut self, z: &[f64]) -> Result<(), LearningError> {
114        let (nx, nz) = (self.nx, self.nz);
115        if z.len() != nz {
116            return Err(LearningError::InvalidDimension);
117        }
118        // Innovation y = z − H x.
119        let mut hx = vec![0.0; nz];
120        matvec(Transpose::No, nz, nx, &self.h, &self.x, &mut hx)?;
121        let y: Vec<f64> = z.iter().zip(&hx).map(|(zi, hxi)| zi - hxi).collect();
122        // HP (nz×nx) and S = HP Hᵀ + R (nz×nz).
123        let mut hp = vec![0.0; nz * nx];
124        no(nz, nx, nx, &self.h, &self.p, &mut hp)?;
125        let mut s = vec![0.0; nz * nz];
126        no_t(nz, nz, nx, &hp, &self.h, &mut s)?;
127        for i in 0..nz * nz {
128            s[i] += self.r[i];
129        }
130        // S⁻¹ via Cholesky.
131        let mut l = vec![0.0; nz * nz];
132        cholesky_factor(nz, &s, &mut l).map_err(|_| LearningError::Singular)?;
133        let mut s_inv = vec![0.0; nz * nz];
134        let mut ej = vec![0.0; nz];
135        let mut cj = vec![0.0; nz];
136        for j in 0..nz {
137            ej.iter_mut().for_each(|v| *v = 0.0);
138            ej[j] = 1.0;
139            cholesky_solve(nz, &l, &ej, &mut cj)?;
140            for i in 0..nz {
141                s_inv[i * nz + j] = cj[i];
142            }
143        }
144        // P Hᵀ (nx×nz), then K = P Hᵀ S⁻¹ (nx×nz).
145        let mut pht = vec![0.0; nx * nz];
146        no_t(nx, nz, nx, &self.p, &self.h, &mut pht)?;
147        let mut kgain = vec![0.0; nx * nz];
148        no(nx, nz, nz, &pht, &s_inv, &mut kgain)?;
149        // x ← x + K y.
150        let mut ky = vec![0.0; nx];
151        matvec(Transpose::No, nx, nz, &kgain, &y, &mut ky)?;
152        for i in 0..nx {
153            self.x[i] += ky[i];
154        }
155        // P ← P − K (H P) = P − K·HP.
156        let mut khp = vec![0.0; nx * nx];
157        no(nx, nx, nz, &kgain, &hp, &mut khp)?;
158        for i in 0..nx * nx {
159            self.p[i] -= khp[i];
160        }
161        Ok(())
162    }
163
164    /// Filter a sequence of observations (row-major `t × nz`), returning the state
165    /// estimate after each (predict→update) step, row-major `t × nx`.
166    pub fn filter(&mut self, observations: &[f64], t: usize) -> Result<Vec<f64>, LearningError> {
167        if observations.len() != t * self.nz {
168            return Err(LearningError::InvalidDimension);
169        }
170        let mut out = vec![0.0; t * self.nx];
171        for step in 0..t {
172            self.predict()?;
173            self.update(&observations[step * self.nz..(step + 1) * self.nz])?;
174            out[step * self.nx..(step + 1) * self.nx].copy_from_slice(&self.x);
175        }
176        Ok(out)
177    }
178}
179
180#[cfg(test)]
181mod tests {
182    use super::*;
183
184    #[test]
185    fn tracks_a_constant_with_noisy_measurements() {
186        // 1-D random-walk model tracking a constant true value 5, noisy obs.
187        let mut kf = KalmanFilter::new(
188            vec![1.0],  // F
189            vec![1.0],  // H
190            vec![1e-4], // Q (nearly constant state)
191            vec![1.0],  // R (noisy measurement)
192            vec![0.0],  // x0
193            vec![10.0], // P0 (weak prior → the data dominates)
194            1,
195            1,
196        )
197        .unwrap();
198        // Measurements jitter around 5.
199        let obs = [5.3, 4.7, 5.1, 4.9, 5.2, 4.8, 5.0, 5.1, 4.9, 5.0];
200        let est = kf.filter(&obs, 10).unwrap();
201        // The final estimate is close to 5 and the covariance has shrunk.
202        assert!((est[9] - 5.0).abs() < 0.3, "estimate {}", est[9]);
203        assert!(kf.covariance()[0] < 1.0, "covariance should shrink");
204    }
205
206    #[test]
207    fn smooths_better_than_raw_measurements() {
208        // The filtered estimate has lower variance than the raw noisy obs.
209        let mut kf = KalmanFilter::new(
210            vec![1.0],
211            vec![1.0],
212            vec![1e-3],
213            vec![1.0],
214            vec![10.0],
215            vec![1.0],
216            1,
217            1,
218        )
219        .unwrap();
220        let obs = [10.5, 9.4, 10.6, 9.5, 10.4, 9.6, 10.5, 9.5];
221        let est = kf.filter(&obs, 8).unwrap();
222        let var = |v: &[f64]| {
223            let m = v.iter().sum::<f64>() / v.len() as f64;
224            v.iter().map(|x| (x - m) * (x - m)).sum::<f64>() / v.len() as f64
225        };
226        let obs_var = var(&obs);
227        let est_var = var(&est[2..]); // skip the warm-up
228        assert!(
229            est_var < obs_var,
230            "filter should smooth: {est_var} !< {obs_var}"
231        );
232    }
233
234    #[test]
235    fn guards() {
236        assert_eq!(
237            KalmanFilter::new(
238                vec![1.0],
239                vec![1.0],
240                vec![1.0],
241                vec![1.0],
242                vec![0.0],
243                vec![1.0, 0.0],
244                1,
245                1
246            )
247            .unwrap_err(),
248            LearningError::InvalidDimension
249        );
250    }
251}
qualia_core_db/solvers/learning/sequential/kalman.rs

qualia_core_db/solvers/learning/sequential/
kalman.rs
