#![allow(non_snake_case)]
use nalgebra as na;
use na::{allocator::Allocator, DefaultAllocator, storage::Storage};
use na::{Dim, RealField, U1, MatrixMN, MatrixN, VectorN};
use crate::linalg::cholesky::UDU;
use crate::matrix;
use crate::models::{KalmanEstimator, KalmanState, Estimator};
use crate::noise::{UncorrelatedNoise, CoupledNoise, CorrelatedFactorNoise};
use nalgebra::{DimAdd, DimSum};
pub struct UDState<N: RealField, D: Dim>
where
DefaultAllocator: Allocator<N, D, D> + Allocator<N, D>,
{
pub x: VectorN<N, D>,
pub UD: MatrixN<N, D>,
udu: UDU<N>,
}
impl<N: RealField, D: Dim> UDState<N, D>
where
DefaultAllocator: Allocator<N, D, D> + Allocator<N, D>,
{
pub fn new(UD: MatrixN<N, D>, x: VectorN<N, D>) -> Self {
assert!(x.nrows() == UD.nrows(), "x rows must be == UD rows");
UDState {
UD,
x,
udu: UDU::new(),
}
}
pub fn new_zero(d: D) -> Self {
UDState {
UD: MatrixN::<N, D>::zeros_generic(d, d),
x: VectorN::zeros_generic(d, U1),
udu: UDU::new(),
}
}
pub fn predict<QD: Dim>(
&mut self,
fx: &MatrixN<N, D>,
x_pred: &VectorN<N, D>,
noise: &CoupledNoise<N, D, QD>,
) -> Result<N, &'static str>
where
D: DimAdd<QD>,
DefaultAllocator: Allocator<N, DimSum<D, QD>, U1> + Allocator<N, D, QD> + Allocator<N, QD>
{
let mut scratch = self.new_predict_scratch(noise.q.data.shape().0);
self.predict_use_scratch(&mut scratch, x_pred, fx, noise)
}
pub fn observe_innovation<ZD: Dim>(
&mut self,
hx: &MatrixMN<N, ZD, D>,
noise: &UncorrelatedNoise<N, ZD>,
s: &VectorN<N, ZD>,
) -> Result<N, &'static str>
where
DefaultAllocator: Allocator<N, ZD, D> + Allocator<N, ZD>
{
let mut scratch = self.new_observe_scratch();
UDState::observe_innovation_use_scratch(self, &mut scratch, s, hx, noise)
}
}
impl<N: RealField, D: Dim> Estimator<N, D> for UDState<N, D>
where
DefaultAllocator: Allocator<N, D, D> + Allocator<N, D>,
{
fn state(&self) -> Result<VectorN<N, D>, &'static str> {
KalmanEstimator::kalman_state(self).map(|r| r.x)
}
}
impl<N: RealField, D: Dim> KalmanEstimator<N, D> for UDState<N, D>
where
DefaultAllocator: Allocator<N, D, D> + Allocator<N, D>,
{
fn init(&mut self, state: &KalmanState<N, D>) -> Result<(), &'static str> {
self.x = state.x.clone();
let rows = self.UD.nrows();
matrix::copy_from(&mut self.UD.columns_mut(0, rows), &state.X);
let rcond = self.udu.UdUfactor_variant2(&mut self.UD, rows);
matrix::check_non_negativ(rcond, "X not PSD")?;
Ok(())
}
fn kalman_state(&self) -> Result<KalmanState<N, D>, &'static str> {
let x_shape = self.x.data.shape().0;
let mut X = matrix::as_zeros((x_shape, x_shape));
matrix::copy_from(&mut X, &self.UD.columns(0, self.UD.nrows()));
UDU::UdUrecompose(&mut X);
Ok(KalmanState { x: self.x.clone(), X })
}
}
impl<N: RealField, D: Dim> UDState<N, D>
where
DefaultAllocator: Allocator<N, D, D> + Allocator<N, D>,
{
pub fn observe_correlated<ZD: Dim>(
&mut self,
z: &VectorN<N, ZD>,
hx: &MatrixMN<N, ZD, D>,
h_normalize: fn(&mut VectorN<N, ZD>, &VectorN<N, ZD>),
noise_factor: &CorrelatedFactorNoise<N, ZD>,
) -> Result<N, &'static str>
where
DefaultAllocator: Allocator<N, ZD, ZD> + Allocator<N, ZD, D> + Allocator<N, ZD>
{
let x_size = self.x.nrows();
let z_size = z.nrows();
let mut scratch = self.new_observe_scratch();
let mut zp = hx * &self.x;
h_normalize(&mut zp, z);
let mut zpdecol = zp.clone();
let mut GIHx = hx.clone();
{
for j in 0..x_size {
for i in (0..z_size).rev() {
for k in i + 1..z_size {
let t = noise_factor.UD[(i, k)] * GIHx[(k, j)];
GIHx[(i, j)] -= t;
}
}
}
for i in (0..z_size).rev() {
for k in i + 1..z_size {
let zpt = noise_factor.UD[(i, k)] * zp[k];
zp[i] -= zpt;
let zpdt = noise_factor.UD[(i, k)] * zpdecol[k];
zpdecol[i] -= zpdt;
}
}
}
let mut rcondmin = N::max_value();
for o in 0..z_size {
let mut S = self.udu.zero;
GIHx.row(o).transpose_to(&mut scratch.a);
let rcond = UDState::observeUD(self, &mut scratch, &mut S, noise_factor.UD[(o, o)]);
matrix::check_positive(rcond, "S not PD in observe")?;
if rcond < rcondmin {
rcondmin = rcond;
}
let s = z[o] - zpdecol[o];
self.x += &scratch.w * s;
}
Ok(rcondmin)
}
}
pub struct PredictScratch<N: RealField, D: Dim, QD: Dim>
where
D: DimAdd<QD>,
DefaultAllocator: Allocator<N, D, QD> + Allocator<N, DimSum<D, QD>>,
{
pub G: MatrixMN<N, D, QD>,
pub d: VectorN<N, DimSum<D, QD>>,
pub dv: VectorN<N, DimSum<D, QD>>,
pub v: VectorN<N, DimSum<D, QD>>,
}
pub struct ObserveScratch<N: RealField, D: Dim>
where
DefaultAllocator: Allocator<N, D>,
{
pub w: VectorN<N, D>,
pub a: VectorN<N, D>,
pub b: VectorN<N, D>,
}
impl<N: RealField, D: Dim> UDState<N, D>
where
DefaultAllocator: Allocator<N, D, D> + Allocator<N, D>,
{
pub fn new_predict_scratch<QD: Dim>(&self, qd: QD) -> PredictScratch<N, D, QD>
where
D: DimAdd<QD>,
DefaultAllocator: Allocator<N, D, QD> + Allocator<N, DimSum<D, QD>, U1>
{
let ud_col_vec_shape = (self.UD.data.shape().1.add(qd), U1);
PredictScratch {
G: matrix::as_zeros((self.UD.data.shape().0, qd)),
d: matrix::as_zeros(ud_col_vec_shape),
dv: matrix::as_zeros(ud_col_vec_shape),
v: matrix::as_zeros(ud_col_vec_shape),
}
}
pub fn new_observe_scratch(&self) -> ObserveScratch<N, D> {
let x_vec_shape = self.x.data.shape();
ObserveScratch {
w: matrix::as_zeros(x_vec_shape),
a: matrix::as_zeros(x_vec_shape),
b: matrix::as_zeros(x_vec_shape),
}
}
pub fn predict_use_scratch<QD: Dim>(
&mut self,
scratch: &mut PredictScratch<N, D, QD>,
x_pred: &VectorN<N, D>,
fx: &MatrixN<N, D>,
noise: &CoupledNoise<N, D, QD>,
) -> Result<N, &'static str>
where
D: DimAdd<QD>,
DefaultAllocator: Allocator<N, DimSum<D, QD>> + Allocator<N, D, QD> + Allocator<N, QD>,
{
self.x = x_pred.clone();
let rcond = UDState::predictGq(self, scratch, &fx, &noise.G, &noise.q);
matrix::check_non_negativ(rcond, "X not PSD")
}
pub fn observe_innovation_use_scratch<ZD: Dim>(
&mut self,
scratch: &mut ObserveScratch<N, D>,
s: &VectorN<N, ZD>,
hx: &MatrixMN<N, ZD, D>,
noise: &UncorrelatedNoise<N, ZD>,
) -> Result<N, &'static str>
where
DefaultAllocator: Allocator<N, ZD, D> + Allocator<N, ZD>
{
let z_size = s.nrows();
let mut rcondmin = N::max_value();
for o in 0..z_size {
if noise.q[o] < self.udu.zero {
return Err("Zv not PSD in observe");
}
let mut S = self.udu.zero;
hx.row(o).transpose_to(&mut scratch.a);
let rcond = UDState::observeUD(self, scratch, &mut S, noise.q[o]);
if rcond < rcondmin {
rcondmin = rcond;
}
scratch.w *= s[0];
self.x += &scratch.w;
}
Ok(rcondmin)
}
fn predictGq<QD: Dim>(
&mut self,
scratch: &mut PredictScratch<N, D, QD>,
Fx: &MatrixN<N, D>,
G: &MatrixMN<N, D, QD>,
q: &VectorN<N, QD>,
) -> N
where
D: DimAdd<QD>,
DefaultAllocator: Allocator<N, DimSum<D, QD>> + Allocator<N, D, QD> + Allocator<N, QD>,
{
let n = self.x.nrows();
let Nq = q.nrows();
let NN = n + Nq;
for i in 0..Nq {
scratch.d[i + n] = q[i];
}
for j in 0..n {
for i in 0..Nq {
scratch.G[(j, i)] = G[(j, i)];
}
}
for j in (1..n).rev() {
for i in 0..=j {
scratch.d[i] = self.UD[(i, j)];
}
for i in 0..n {
self.UD[(i, j)] = Fx[(i, j)];
for k in 0..j {
self.UD[(i, j)] += Fx[(i, k)] * scratch.d[k];
}
}
}
if n > 0 {
scratch.d[0] = self.UD[(0, 0)];
}
for j in 0..n {
self.UD[(j, 0)] = Fx[(j, 0)];
}
for j in (0..n).rev() {
let mut e = self.udu.zero;
for k in 0..n {
scratch.v[k] = self.UD[(j, k)];
scratch.dv[k] = scratch.d[k] * scratch.v[k];
e += scratch.v[k] * scratch.dv[k];
}
for k in n..NN {
scratch.v[k] = scratch.G[(j, k - n)];
scratch.dv[k] = scratch.d[k] * scratch.v[k];
e += scratch.v[k] * scratch.dv[k];
}
if e > self.udu.zero {
self.UD[(j, j)] = e;
let diaginv = self.udu.one / e;
for k in 0..j {
e = self.udu.zero;
for i in 0..n {
e += self.UD[(k, i)] * scratch.dv[i];
}
for i in n..NN {
e += scratch.G[(k, i - n)] * scratch.dv[i];
}
e *= diaginv;
self.UD[(j, k)] = e;
for i in 0..n {
self.UD[(k, i)] -= e * scratch.v[i]
}
for i in n..NN {
scratch.G[(k, i - n)] -= e * scratch.v[i]
}
}
} else if e == self.udu.zero {
self.UD[(j, j)] = e;
for k in 0..j {
for i in 0..n {
e = self.UD[(k, i)] * scratch.dv[i];
if e != self.udu.zero {
return self.udu.minus_one;
}
}
for i in n..NN {
e = scratch.G[(k, i - n)] * scratch.dv[i];
if e != self.udu.zero {
return self.udu.minus_one;
}
}
}
} else {
return self.udu.minus_one;
}
}
for j in 1..n {
for i in 0..j {
self.UD[(i, j)] = self.UD[(j, i)];
self.UD[(j, i)] = self.udu.zero;
}
}
UDU::UdUrcond(&self.UD)
}
fn observeUD(&mut self, scratch: &mut ObserveScratch<N, D>, alpha: &mut N, q: N)
-> N {
let n = self.UD.nrows();
for j in (1..n).rev() {
for k in 0..j {
let t = self.UD[(k, j)] * scratch.a[k];
scratch.a[j] += t;
}
scratch.b[j] = self.UD[(j, j)] * scratch.a[j];
}
scratch.b[0] = self.UD[(0, 0)] * scratch.a[0];
*alpha = q + scratch.b[0] * scratch.a[0];
if *alpha <= self.udu.zero {
return self.udu.minus_one;
}
let mut gamma = self.udu.one / *alpha;
self.UD[(0, 0)] *= q * gamma;
for j in 1..n {
let alpha_jm1 = *alpha;
*alpha += scratch.b[j] * scratch.a[j];
let lamda = -scratch.a[j] * gamma;
if *alpha <= self.udu.zero {
return self.udu.minus_one;
}
gamma = self.udu.one / *alpha;
self.UD[(j, j)] *= alpha_jm1 * gamma;
for i in 0..j {
let UD_jm1 = self.UD[(i, j)];
self.UD[(i, j)] = UD_jm1 + lamda * scratch.b[i];
let t = scratch.b[j] * UD_jm1;
scratch.b[i] += t;
}
}
scratch.w.copy_from(&(&scratch.b * gamma));
UDU::UdUrcond(&self.UD)
}
}