#![allow(non_snake_case)]
use nalgebra as na;
use na::{allocator::Allocator, DefaultAllocator, storage::Storage, Dim, U1, MatrixMN, MatrixN, RealField, VectorN};
use crate::linalg::rcond;
use crate::matrix::{check_positive};
use crate::models::{InformationState, KalmanEstimator, KalmanState, ExtendedLinearPredictor, Estimator, ExtendedLinearObserver};
use crate::noise::{CorrelatedNoise, CoupledNoise};
impl<N: RealField, D: Dim> InformationState<N, D>
where
DefaultAllocator: Allocator<N, D, D> + Allocator<N, D>,
{
pub fn new_zero(d: D) -> InformationState<N, D> {
InformationState {
i: VectorN::zeros_generic(d, U1),
I: MatrixN::zeros_generic(d, d),
}
}
}
impl<N: RealField, D: Dim> Estimator<N, D> for InformationState<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 InformationState<N, D>
where
DefaultAllocator: Allocator<N, D, D> + Allocator<N, D>,
{
fn init(&mut self, state: &KalmanState<N, D>) -> Result<(), &'static str> {
self.I = state.X.clone().cholesky().ok_or("X not PD")?.inverse();
self.i = &self.I * &state.x;
Ok(())
}
fn kalman_state(&self) -> Result<KalmanState<N, D>, &'static str> {
let X = self.I.clone().cholesky().ok_or("Y not PD")?.inverse();
let x = &X * &self.i;
Ok(KalmanState { x, X })
}
}
impl<N: RealField, D: Dim> ExtendedLinearPredictor<N, D> for InformationState<N, D>
where
DefaultAllocator: Allocator<N, D, D> + Allocator<N, D>
{
fn predict(
&mut self,
x_pred: &VectorN<N, D>,
fx: &MatrixN<N, D>,
noise: &CorrelatedNoise<N, D>,
) -> Result<(), &'static str> {
let mut X = self.I.clone().cholesky().ok_or("I not PD in predict")?.inverse();
X.quadform_tr(N::one(), &fx, &X.clone(), N::zero());
X += &noise.Q;
self.init(&KalmanState { x: x_pred.clone_owned(), X })?;
return Ok(())
}
}
impl<N: RealField, D: Dim, ZD: Dim> ExtendedLinearObserver<N, D, ZD> for InformationState<N, D>
where
DefaultAllocator: Allocator<N, D, D> + Allocator<N, D, ZD> + Allocator<N, ZD, D> + Allocator<N, ZD, ZD> + Allocator<N, D> + Allocator<N, ZD>,
{
fn observe_innovation(
&mut self,
s: &VectorN<N, ZD>,
hx: &MatrixMN<N, ZD, D>,
noise: &CorrelatedNoise<N, ZD>,
) -> Result<(), &'static str>
{
let x = self.state().unwrap();
let noise_inv = noise.Q.clone().cholesky().ok_or("Q not PD in observe")?.inverse();
let info = self.observe_info(hx, &noise_inv, &(s + hx * x));
self.add_information(&info);
return Ok(())
}
}
impl<N: RealField, D: Dim> InformationState<N, D>
where
DefaultAllocator: Allocator<N, D, D> + Allocator<N, D>,
{
pub fn predict_linear<QD: Dim>(
&mut self,
pred_inv: MatrixN<N, D>,
noise: &CoupledNoise<N, D, QD>,
) -> Result<N, &'static str>
where
DefaultAllocator: Allocator<N, QD, QD> + Allocator<N, D, QD> + Allocator<N, QD, D> + Allocator<N, QD>,
{
let I_shape = self.I.data.shape();
let A = (&self.I * &pred_inv).tr_mul(&pred_inv);
let mut B = (&A * &noise.G).tr_mul(&noise.G);
for i in 0..noise.q.nrows() {
B[(i, i)] += N::one() / noise.q[i];
}
B = B.cholesky().ok_or("B not PD")?.inverse();
let rcond = rcond::rcond_symetric(&B);
check_positive(rcond, "(G'invFx'.I.inv(Fx).G + inv(Q)) not PD")?;
self.I.quadform_tr(N::one(), &noise.G, &B, N::zero());
let ig = MatrixMN::identity_generic(I_shape.0, I_shape.1) - &A * &self.I;
self.I = &ig * &A;
let y = pred_inv.tr_mul(&self.i);
self.i = &ig * &y;
Ok(rcond)
}
pub fn add_information(&mut self, information: &InformationState<N, D>) {
self.i += &information.i;
self.I += &information.I;
}
pub fn observe_info<ZD: Dim>(
&self,
hx: &MatrixMN<N, ZD, D>,
noise_inv: &MatrixN<N, ZD>,
z: &VectorN<N, ZD>
) -> InformationState<N, D>
where
DefaultAllocator: Allocator<N, ZD, ZD> + Allocator<N, ZD, D> + Allocator<N, D, ZD> + Allocator<N, ZD>
{
let HxTZI = hx.tr_mul(noise_inv);
let ii = &HxTZI * z;
let II = &HxTZI * hx;
InformationState { i: ii, I: II }
}
}