use crate::core::matrix::{matmul, matvec, outer, Matrix};
use crate::core::scalar::ControlScalar;
pub struct Ukf<S: ControlScalar, const N: usize, const M: usize, const I: usize> {
pub q: Matrix<S, N, N>,
pub r: Matrix<S, M, M>,
f: fn(&[S; N], &[S; I]) -> [S; N],
h: fn(&[S; N]) -> [S; M],
pub alpha: S,
pub beta: S,
pub kappa: S,
x: [S; N],
p: Matrix<S, N, N>,
}
impl<S: ControlScalar, const N: usize, const M: usize, const I: usize> Ukf<S, N, M, I> {
#[allow(clippy::too_many_arguments)]
pub fn new(
q: Matrix<S, N, N>,
r: Matrix<S, M, M>,
f: fn(&[S; N], &[S; I]) -> [S; N],
h: fn(&[S; N]) -> [S; M],
x0: [S; N],
p0: Matrix<S, N, N>,
alpha: S,
beta: S,
kappa: S,
) -> Self {
Self {
q,
r,
f,
h,
alpha,
beta,
kappa,
x: x0,
p: p0,
}
}
pub fn standard(
q: Matrix<S, N, N>,
r: Matrix<S, M, M>,
f: fn(&[S; N], &[S; I]) -> [S; N],
h: fn(&[S; N]) -> [S; M],
x0: [S; N],
p0: Matrix<S, N, N>,
) -> Self {
Self::new(
q,
r,
f,
h,
x0,
p0,
S::from_f64(0.001),
S::from_f64(2.0),
S::ZERO,
)
}
fn weights(&self) -> (S, S, S) {
let n = S::from_f64(N as f64);
let lambda = self.alpha * self.alpha * (n + self.kappa) - n;
let c = n + lambda;
let w_m0 = lambda / c;
let w_c0 = w_m0 + S::ONE - self.alpha * self.alpha + self.beta;
let w_i = S::HALF / c;
(w_m0, w_c0, w_i)
}
fn scale_factor(&self) -> S {
let n = S::from_f64(N as f64);
let lambda = self.alpha * self.alpha * (n + self.kappa) - n;
n + lambda
}
pub fn predict(&mut self, u: &[S; I]) -> bool {
let c = self.scale_factor();
let p_scaled = self.p.scale(c);
let l = match p_scaled.cholesky() {
Some(l) => l,
None => return false,
};
let (w_m0, w_c0, w_i) = self.weights();
let y0 = (self.f)(&self.x, u);
let mut y_plus: [[S; N]; N] = core::array::from_fn(|_| [S::ZERO; N]);
let mut y_minus: [[S; N]; N] = core::array::from_fn(|_| [S::ZERO; N]);
for j in 0..N {
let mut xi_plus = self.x;
let mut xi_minus = self.x;
for row in 0..N {
xi_plus[row] += l.data[row][j];
xi_minus[row] -= l.data[row][j];
}
y_plus[j] = (self.f)(&xi_plus, u);
y_minus[j] = (self.f)(&xi_minus, u);
}
let mut x_pred = [S::ZERO; N];
for k in 0..N {
x_pred[k] = w_m0 * y0[k];
for j in 0..N {
x_pred[k] += w_i * (y_plus[j][k] + y_minus[j][k]);
}
}
let d0: [S; N] = core::array::from_fn(|k| y0[k] - x_pred[k]);
let mut p_pred = self.q.add_mat(&outer(&d0, &d0).scale(w_c0));
for j in 0..N {
let dp: [S; N] = core::array::from_fn(|k| y_plus[j][k] - x_pred[k]);
let dm: [S; N] = core::array::from_fn(|k| y_minus[j][k] - x_pred[k]);
p_pred = p_pred
.add_mat(&outer(&dp, &dp).scale(w_i))
.add_mat(&outer(&dm, &dm).scale(w_i));
}
self.x = x_pred;
self.p = p_pred;
true
}
pub fn update(&mut self, z_meas: &[S; M]) -> Option<[S; M]> {
let c = self.scale_factor();
let p_scaled = self.p.scale(c);
let l = p_scaled.cholesky()?;
let (w_m0, w_c0, w_i) = self.weights();
let z0 = (self.h)(&self.x);
let mut z_plus: [[S; M]; N] = core::array::from_fn(|_| [S::ZERO; M]);
let mut z_minus: [[S; M]; N] = core::array::from_fn(|_| [S::ZERO; M]);
let mut xi_plus_arr: [[S; N]; N] = core::array::from_fn(|_| [S::ZERO; N]);
let mut xi_minus_arr: [[S; N]; N] = core::array::from_fn(|_| [S::ZERO; N]);
for j in 0..N {
let mut xi_plus = self.x;
let mut xi_minus = self.x;
for row in 0..N {
xi_plus[row] += l.data[row][j];
xi_minus[row] -= l.data[row][j];
}
xi_plus_arr[j] = xi_plus;
xi_minus_arr[j] = xi_minus;
z_plus[j] = (self.h)(&xi_plus);
z_minus[j] = (self.h)(&xi_minus);
}
let mut z_pred = [S::ZERO; M];
for k in 0..M {
z_pred[k] = w_m0 * z0[k];
for j in 0..N {
z_pred[k] += w_i * (z_plus[j][k] + z_minus[j][k]);
}
}
let dz0: [S; M] = core::array::from_fn(|k| z0[k] - z_pred[k]);
let mut s_mat = self.r.add_mat(&outer(&dz0, &dz0).scale(w_c0));
let mut c_xz = Matrix::<S, N, M>::zeros();
for j in 0..N {
let dz_p: [S; M] = core::array::from_fn(|k| z_plus[j][k] - z_pred[k]);
let dz_m: [S; M] = core::array::from_fn(|k| z_minus[j][k] - z_pred[k]);
let dy_p: [S; N] = core::array::from_fn(|k| xi_plus_arr[j][k] - self.x[k]);
let dy_m: [S; N] = core::array::from_fn(|k| xi_minus_arr[j][k] - self.x[k]);
s_mat = s_mat
.add_mat(&outer(&dz_p, &dz_p).scale(w_i))
.add_mat(&outer(&dz_m, &dz_m).scale(w_i));
for row in 0..N {
for col in 0..M {
c_xz.data[row][col] += w_i * (dy_p[row] * dz_p[col] + dy_m[row] * dz_m[col]);
}
}
}
let s_inv = s_mat.inv()?;
let k = matmul(&c_xz, &s_inv);
let innovation: [S; M] = core::array::from_fn(|k| z_meas[k] - z_pred[k]);
let k_innov = matvec(&k, &innovation);
for (i, &ki) in k_innov.iter().enumerate().take(N) {
self.x[i] += ki;
}
let k_s = matmul(&k, &s_mat);
let k_s_kt = matmul(&k_s, &k.transpose());
self.p = self.p.sub_mat(&k_s_kt);
Some(innovation)
}
pub fn state(&self) -> &[S; N] {
&self.x
}
pub fn covariance(&self) -> &Matrix<S, N, N> {
&self.p
}
pub fn set_state(&mut self, x: [S; N]) {
self.x = x;
}
pub fn reset(&mut self, x0: [S; N], p0: Matrix<S, N, N>) {
self.x = x0;
self.p = p0;
}
}
#[cfg(test)]
mod tests {
use super::*;
fn f_linear(x: &[f64; 2], u: &[f64; 1]) -> [f64; 2] {
[x[0] + 0.01 * x[1], x[1] + 0.01 * u[0]]
}
fn h_pos(x: &[f64; 2]) -> [f64; 1] {
[x[0]]
}
fn build_ukf() -> Ukf<f64, 2, 1, 1> {
let q = Matrix::<f64, 2, 2>::identity().scale(1e-4);
let r = Matrix::<f64, 1, 1>::identity().scale(0.1);
let p0 = Matrix::<f64, 2, 2>::identity();
Ukf::standard(q, r, f_linear, h_pos, [0.0_f64; 2], p0)
}
#[test]
fn predict_runs() {
let mut ukf = build_ukf();
assert!(ukf.predict(&[0.0]));
}
#[test]
fn update_returns_innovation() {
let mut ukf = build_ukf();
ukf.predict(&[0.0]);
let innov = ukf.update(&[0.1]);
assert!(innov.is_some());
}
#[test]
fn tracks_constant_position() {
let mut ukf = build_ukf();
let true_pos = 5.0_f64;
for _ in 0..200 {
ukf.predict(&[0.0]);
ukf.update(&[true_pos]);
}
assert!((ukf.state()[0] - true_pos).abs() < 0.5);
}
#[test]
fn cholesky_consistency() {
let mut ukf = build_ukf();
for _ in 0..10 {
ukf.predict(&[0.0]);
ukf.update(&[1.0]);
}
assert!(ukf.p.cholesky().is_some());
}
}