bayes_estimate 0.20.0

#![allow(non_snake_case)]

//! Information state estimation.
//!
//! A discrete Bayesian estimator that uses a linear information representation [`InformationState`] of the system for estimation.
//!
//! A fundamental property of the information state is that information is additive. So if there is more information
//! about the system (such as by an observation) this can simply be added to i,I of the information state.
//!
//! The linear representation can also be used for non-linear systems by using linearised models.

use nalgebra::{allocator::Allocator, DefaultAllocator, Dim, OMatrix, OVector, RealField};

use crate::linalg::rcond;
use crate::models::{
    Estimator, ExtendedLinearObserver, ExtendedLinearPredictor, InformationState, KalmanEstimator,
    KalmanState,
};
use crate::noise::{CorrelatedNoise, CoupledNoise};

impl<N: RealField, D: Dim> Estimator<N, D> for InformationState<N, D>
where
    DefaultAllocator: Allocator<D, D> + Allocator<D>,
{
    fn state<'e>(&self) -> Result<OVector<N, D>, &'e str> {
        KalmanEstimator::kalman_state(self).map(|r| r.x)
    }
}

impl<N: RealField, D: Dim> TryFrom<KalmanState<N, D>> for InformationState<N, D>
    where
        DefaultAllocator: Allocator<D, D> + Allocator<D>,
{
    type Error = &'static str;

    fn try_from(state: KalmanState<N, D>) -> Result<Self, Self::Error> {
        // Information
        let I = state.X.clone().cholesky().ok_or("X not PD")?.inverse();
        // Information state
        let i = &I * &state.x;
        Ok(InformationState { i, I })
    }
}

impl<N: RealField, D: Dim> KalmanEstimator<N, D> for InformationState<N, D>
where
    DefaultAllocator: Allocator<D, D> + Allocator<D>,
{
    fn kalman_state<'e>(&self) -> Result<KalmanState<N, D>, &'e str> {
        // Covariance
        let X = self.I.clone().cholesky().ok_or("Y not PD")?.inverse();
        // State
        let x = &X * &self.i;

        Ok(KalmanState { x, X })
    }
}

impl<N: RealField, D: Dim> ExtendedLinearPredictor<N, D> for InformationState<N, D>
where
    DefaultAllocator: Allocator<D, D> + Allocator<D>,
{
    fn predict<'e>(
        &mut self,
        x_pred: &OVector<N, D>,
        fx: &OMatrix<N, D, D>,
        noise: &CorrelatedNoise<N, D>,
    ) -> Result<(), &'e str> {
        // Covariance
        let mut X = self.I.clone().cholesky().ok_or("I not PD in predict")?.inverse();

        // Predict information matrix, and state covariance
        X.quadform_tr(N::one(), &fx, &X.clone(), N::zero());
        X += &noise.Q;

        // Information
        self.I = X.cholesky().ok_or("X not PD")?.inverse();
        // Information state
        self.i = &self.I * x_pred;

        Ok(())
    }
}

impl<N: Copy + RealField, D: Dim, ZD: Dim> ExtendedLinearObserver<N, D, ZD>
    for InformationState<N, D>
where
    DefaultAllocator: Allocator<D, D>
        + Allocator<D, ZD>
        + Allocator<ZD, D>
        + Allocator<ZD, ZD>
        + Allocator<D>
        + Allocator<ZD>,
{
    fn observe_innovation<'e>(
        &mut self,
        s: &OVector<N, ZD>,
        hx: &OMatrix<N, ZD, D>,
        noise: &CorrelatedNoise<N, ZD>,
    ) -> Result<(), &'e str> {
        let x = self.state()?;
        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);

        Ok(())
    }
}

impl<N: Copy + RealField, D: Dim> InformationState<N, D>
where
    DefaultAllocator: Allocator<D, D> + Allocator<D>,
{
    /// Linear information predict.
    ///
    /// The numerical solution requires the inversion of the 'B' matrix, zero noise or small
    /// noise 'q' makes this not invertible or ill conditioned.
    pub fn predict_linear<'e, QD: Dim>(
        &mut self,
        pred_inv: OMatrix<N, D, D>, // Inverse of linear prediction model Fx
        noise: &CoupledNoise<N, D, QD>,
    ) -> Result<N, &'e str>
    where
        DefaultAllocator:
            Allocator<QD, QD> + Allocator<D, QD> + Allocator<QD, D> + Allocator<QD>,
    {
        let I_shape = self.I.shape_generic();

        // A = invFx'*Y*invFx ,Inverse Predict covariance
        let A = (&self.I * &pred_inv).tr_mul(&pred_inv);
        // B = G'*A*G+invQ ,A in coupled additive noise space
        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];
        }

        // invert B, the additive noise
        let brcond = rcond::rcond_symmetric(&B);
        rcond::check_positive(brcond, "(G'invFx'.I.inv(Fx).G + inv(Q)) not PD")?;
        B = B.cholesky().ok_or("B not PD")?.inverse();

        // G*invB*G' ,in state space
        self.I.quadform_tr(N::one(), &noise.G, &B, N::zero());
        // I - A* G*invB*G', information gain
        let ig = OMatrix::identity_generic(I_shape.0, I_shape.1) - &A * &self.I;
        // Information
        self.I = &ig * &A;
        // Information state
        let y = pred_inv.tr_mul(&self.i);
        self.i = &ig * &y;

        Ok(brcond)
    }

    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: &OMatrix<N, ZD, D>,
        noise_inv: &OMatrix<N, ZD, ZD>, // Inverse of correlated noise model
        z: &OVector<N, ZD>,
    ) -> InformationState<N, D>
    where
        DefaultAllocator:
            Allocator<ZD, ZD> + Allocator<ZD, D> + Allocator<D, ZD> + Allocator<ZD>,
    {
        // Observation Information
        let HxTZI = hx.tr_mul(noise_inv);
        // Calculate EIF i = Hx'*ZI*z
        let ii = &HxTZI * z;
        // Calculate EIF I = Hx'*ZI*Hx
        let II = &HxTZI * hx; // use column matrix trans(HxT)

        InformationState { i: ii, I: II }
    }
}


#[cfg(test)]
mod tests {
    use super::*;
    use nalgebra::{Matrix2, Vector2};

    #[test]
    fn test_predict_linear_with_partial_zero_noise() {
        // Initial information state
        let i = Vector2::new(1.0, 2.0);
        let I = Matrix2::new(2.0, 0.0, 0.0, 2.0);
        let mut info_state = InformationState { i, I };

        // Linear prediction model inverse
        let pred_inv = Matrix2::identity();

        // Partial zero noise (one component is zero)
        let q = Vector2::new(0.1, 0.0);
        let G = Matrix2::identity();
        let noise = CoupledNoise { q, G };

        // Predict should succeed with partial zero noise
        let result = info_state.predict_linear(pred_inv, &noise);
        assert!(!result.is_ok(), "predict_linear should fail with partial zero noise");
    }

    #[test]
    fn test_predict_linear_with_zero_coupling() {
        // Initial information state
        let i = Vector2::new(1.0, 2.0);
        let I = Matrix2::new(2.0, 0.0, 0.0, 2.0);
        let mut info_state = InformationState { i, I };

        // Linear prediction model inverse
        let pred_inv = Matrix2::identity();

        // Non-zero noise but zero coupling
        let q = Vector2::new(0.1, 0.1);
        let G = Matrix2::zeros();
        let noise = CoupledNoise { q, G };

        // Predict should succeed with zero coupling
        let result = info_state.predict_linear(pred_inv, &noise);
    }
}