Even closer to get example "basic" up and running!

Cargo.toml

@@ -10,7 +10,7 @@ ndarray = "0.13"
 ndarray-linalg = { version = "0.12" }
 cblas = "0.2"
 lapacke = "0.2"
-statrs = "0.10"
+statrs = "0.12"
 ordered-float = "1.0"
 rand = "0.6"
 rand_xoshiro = "0.1"

src/fitter.rs

@@ -4,7 +4,10 @@ pub use recursive::RecursiveFitter;
 
 pub trait Fitter {
     fn add_sample(&mut self, t: f64) -> usize;
+
     fn allocate(&mut self);
+    fn is_allocated(&self) -> bool;
+
     fn fit(&mut self);
 
     fn vs(&self, idx: usize) -> f64;

src/fitter/recursive.rs

@@ -1,5 +1,6 @@
 use ndarray::prelude::*;
 use ndarray::stack;
+use ndarray_linalg::Inverse;
 
 use crate::kernel::Kernel;
 
@@ -15,15 +16,15 @@ pub struct RecursiveFitter {
     xs: ArrayD<f64>,
     is_fitted: bool,
     h: Array1<f64>,
-    i: ArrayD<f64>,
+    i: Array2<f64>,
     a: Vec<Array2<f64>>,
     q: Vec<Array2<f64>>,
-    m_p: ArrayD<f64>,
+    m_p: Vec<Array1<f64>>,
-    p_p: ArrayD<f64>,
+    p_p: Vec<Array2<f64>>,
-    m_f: ArrayD<f64>,
+    m_f: Vec<Array1<f64>>,
-    p_f: ArrayD<f64>,
+    p_f: Vec<Array2<f64>>,
-    m_s: ArrayD<f64>,
+    m_s: Vec<Array1<f64>>,
-    p_s: ArrayD<f64>,
+    p_s: Vec<Array2<f64>>,
 }
 
 impl RecursiveFitter {
@@ -41,15 +42,15 @@ impl RecursiveFitter {
             xs: Array::zeros(0).into_dyn(),
             is_fitted: true,
             h,
-            i: Array::eye(m).into_dyn(),
+            i: Array::eye(m),
             a: Vec::new(),
             q: Vec::new(),
-            m_p: Array::zeros((0, m)).into_dyn(),
+            m_p: Vec::new(),
-            p_p: Array::zeros((0, m, m)).into_dyn(),
+            p_p: Vec::new(),
-            m_f: Array::zeros((0, m)).into_dyn(),
+            m_f: Vec::new(),
-            p_f: Array::zeros((0, m, m)).into_dyn(),
+            p_f: Vec::new(),
-            m_s: Array::zeros((0, m)).into_dyn(),
+            m_s: Vec::new(),
-            p_s: Array::zeros((0, m, m)).into_dyn(),
+            p_s: Vec::new(),
         }
     }
 }
@@ -81,43 +82,19 @@ impl Fitter for RecursiveFitter {
         self.xs = stack![Axis(0), self.xs, zeros];
 
         // Initialize the predictive, filtering and smoothing distributions.
-        let mean = self
+        for t in &self.ts_new {
-            .ts_new
+            let mean = self.kernel.state_mean(*t);
-            .iter()
-            .flat_map(|t| self.kernel.state_mean(*t).to_vec().into_iter())
-            .collect::<Vec<f64>>();
 
-        let mean = Array::from_shape_vec((self.ts_new.len(), self.kernel.order()), mean)
+            self.m_p.push(mean.clone());
-            .expect("failed to create mean matrix")
+            self.m_f.push(mean.clone());
-            .into_dyn();
+            self.m_s.push(mean);
 
-        let cov = self
+            let cov = self.kernel.state_cov(*t);
-            .ts_new
-            .iter()
-            .flat_map(|t| {
-                self.kernel
-                    .state_cov(*t)
-                    .iter()
-                    .cloned()
-                    .collect::<Vec<f64>>()
-            })
-            .collect::<Vec<f64>>();
 
-        let cov = Array3::from_shape_vec(
+            self.p_p.push(cov.clone());
-            (self.ts_new.len(), self.kernel.order(), self.kernel.order()),
+            self.p_f.push(cov.clone());
-            cov,
+            self.p_s.push(cov);
-        )
+        }
-        .expect("failed to create cov matrix")
-        .into_dyn();
-
-        self.m_p = stack![Axis(0), self.m_p, mean];
-        self.p_p = stack![Axis(0), self.p_p, cov];
-
-        self.m_f = stack![Axis(0), self.m_f, mean];
-        self.p_f = stack![Axis(0), self.p_f, cov];
-
-        self.m_s = stack![Axis(0), self.m_s, mean];
-        self.p_s = stack![Axis(0), self.p_s, cov];
 
         // Compute the new transition and noise covariance matrices.
         for i in (self.ts.len() - n_new)..self.ts.len() {
@@ -139,8 +116,80 @@ impl Fitter for RecursiveFitter {
         self.ts_new.clear();
     }
 
+    fn is_allocated(&self) -> bool {
+        self.ts_new.is_empty()
+    }
+
     fn fit(&mut self) {
-        todo!();
+        if !self.is_allocated() {
+            // raise RuntimeError("new data since last call to `allocate()`")
+        }
+
+        if self.ts.is_empty() {
+            self.is_fitted = true;
+
+            return;
+        }
+
+        // Forward pass (Kalman filter).
+        for i in 0..self.ts.len() {
+            if i > 0 {
+                self.m_p[i] = self.a[i - 1].dot(&self.m_f[i - 1]);
+                self.p_p[i] =
+                    self.a[i - 1].dot(&self.p_f[i - 1]).dot(&self.a[i - 1].t()) + &self.q[i - 1];
+            }
+
+            // These are slightly modified equations to work with tau and nu.
+            let k = self.p_p[i].dot(&self.h)
+                / (1.0 + self.xs[i] * self.h.dot(&self.p_p[i]).dot(&self.h));
+
+            let k = Array1::from(k);
+
+            self.m_f[i] =
+                (&self.m_p[i] + &k) * (&self.ns[i] - &self.xs[i] * &self.h.dot(&self.m_p[i]));
+
+            // Covariance matrix is computed using the Joseph form.
+            let outer = (self.xs[i] * &k)
+                .iter()
+                .flat_map(|a| self.h.iter().map(move |b| a * b))
+                .collect::<Vec<f64>>();
+
+            let outer = Array::from_shape_vec((self.h.len(), self.h.len()), outer)
+                .expect("failed to create outer matrix");
+
+            let z = &self.i - &outer;
+
+            let outer = k
+                .iter()
+                .flat_map(|a| k.iter().map(move |b| a * b))
+                .collect::<Vec<f64>>();
+
+            let outer = Array::from_shape_vec((self.h.len(), self.h.len()), outer)
+                .expect("failed to create outer matrix");
+
+            self.p_f[i] = z.dot(&self.p_p[i]).dot(&z.t()) + self.xs[i] * outer;
+        }
+
+        // Backward pass (RTS smoother).
+        for i in (0..self.ts.len()).rev() {
+            if i == self.ts.len() - 1 {
+                self.m_s[i] = self.m_f[i].clone();
+                self.p_s[i] = self.p_f[i].clone();
+            } else {
+                let g = self.a[i]
+                    .dot(&self.p_f[i])
+                    .dot(&self.p_p[i + 1].inv().expect("failed to inverse matrix"));
+
+                self.m_s[i] = &self.m_f[i] + &g.dot(&(&self.m_s[i + 1] - &self.m_p[i + 1]));
+                self.p_s[i] =
+                    &self.p_f[i] + &g.dot(&(&self.p_s[i + 1] - &self.p_p[i + 1])).dot(&g.t());
+            }
+
+            self.ms[i] = self.h.dot(&self.m_s[i]);
+            self.vs[i] = self.h.dot(&self.p_s[i]).dot(&self.h);
+        }
+
+        self.is_fitted = true;
     }
 
     fn vs(&self, idx: usize) -> f64 {

src/utils.rs

@@ -38,7 +38,7 @@ const QS: [f64; 6] = [
 fn normcdf(x: f64) -> f64 {
     // If X ~ N(0,1), returns P(X < x).
     // erfc(-x / SQRT2) / 2.0
-
+    // https://docs.rs/statrs/0.12.0/statrs/function/erf/fn.erfc.html
     todo!();
 }
 
@@ -85,29 +85,4 @@ pub fn logphi(z: f64) -> (f64, f64) {
 
         (res, dres)
     }
-
-    /*
-    if z * z < 0.0492:
-        # First case: z close to zero.
-        coef = -z / SQRT2PI
-        val = 0
-        for c in CS:
-            val = coef * (c + val)
-        res = -2 * val - log(2)
-        dres = exp(-(z * z) / 2 - res) / SQRT2PI
-    elif z < -11.3137:
-        # Second case: z very small.
-        num = 0.5641895835477550741
-        for r in RS:
-            num = -z * num / SQRT2 + r
-        den = 1.0
-        for q in QS:
-            den = -z * den / SQRT2 + q
-        res = log(num / (2 * den)) - (z * z) / 2
-        dres = abs(den / num) * sqrt(2.0 / pi)
-    else:
-        res = log(normcdf(z))
-        dres = exp(-(z * z) / 2 - res) / SQRT2PI
-    return res, dres
-    */
 }