bayes_estimate 0.20.0

#![allow(non_snake_case)]

//! UD 'square root' state estimation.
//!
//! A discrete Bayesian estimator that uses a 'square root' factorisation of the Kalman state representation [`UDState`] of the system for estimation.
//!
//! The linear representation can also be used for non-linear system by using linearised forms of the system model.
//!
//! # References
//! \[Bierman] "Factorisation Methods for Discrete Sequential Estimation" Gerald J. Bierman ISBN 0-12-097350-2

use nalgebra::{allocator::Allocator, DefaultAllocator};
use nalgebra::{Const, Dim, DimAdd, DimSum, OMatrix, OVector, RealField, U1};

use crate::linalg::cholesky::UDU;
use crate::linalg::rcond;
use crate::models::{Estimator, KalmanEstimator, KalmanState};
use crate::noise::{CorrelatedNoise, CoupledNoise, UncorrelatedNoise};

/// UD State representation.
///
/// Linear representation as a state vector and 'square root' factorisation of the state covariance matrix.
/// Numerically the this 'square root' factorisation is advantageous as conditioning for inverting is improved by the square root.
///
/// The state covariance is represented as a U.d.U' factorisation, where U is upper triangular matrix (0 diagonal) and d is a diagonal vector.
/// U and d are packed into a single UD Matrix, the lower Triangle ist not part of state representation.
pub struct UDState<N: RealField, D: Dim>
    where
        DefaultAllocator: Allocator<D, D> + Allocator<D>,
{
    /// State vector
    pub x: OVector<N, D>,
    /// UD matrix representation of state covariance
    pub UD: OMatrix<N, D, D>,
    // UDU instance for factorisations
    udu: UDU<N>,
}

impl<N: Copy + RealField, D: Dim> UDState<N, D>
    where
        DefaultAllocator: Allocator<D, D> + Allocator<D>,
{
    /// Create a UDState for given state dimensions.
    ///
    /// D is the size of states vector and rows in UD.
    pub fn new(UD: OMatrix<N, D, D>, x: OVector<N, D>) -> Self {
        assert_eq!(x.nrows(), UD.nrows(), "x rows must be == UD rows");

        UDState {
            UD,
            x,
            udu: UDU::new(),
        }
    }

    pub fn predict<QD: Dim>(
        &mut self,
        fx: &OMatrix<N, D, D>,
        x_pred: &OVector<N, D>,
        noise: &CoupledNoise<N, D, QD>,
    ) -> Result<N, &'static str>
        where
            D: DimAdd<QD>,
            DefaultAllocator: Allocator<DimSum<D, QD>, U1> + Allocator<D, QD> + Allocator<QD>,
    {
        let mut scratch = self.new_predict_scratch(noise.q.shape_generic().0);
        self.predict_use_scratch(&mut scratch, x_pred, fx, noise)
    }

    /// Implement observe using sequential observation updates.
    ///
    /// Uncorrelated observations are applied sequentially in the order they appear in z.
    ///
    //// Return: Minimum 'rcond' of all sequential observe
    pub fn observe_innovation<ZD: Dim>(
        &mut self,
        s: &OVector<N, ZD>,
        hx: &OMatrix<N, ZD, D>,
        noise: &UncorrelatedNoise<N, ZD>,
    ) -> Result<N, &'static str>
        where
            DefaultAllocator: Allocator<ZD, D> + Allocator<ZD>,
    {
        let mut scratch = self.new_observe_scratch();
        UDState::observe_innovation_use_scratch(self, &mut scratch, s, hx, noise)
    }
}

impl<N: Copy + RealField, D: Dim> Estimator<N, D> for UDState<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: Copy + RealField, D: Dim> TryFrom<KalmanState<N, D>> for UDState<N, D>
    where
        DefaultAllocator: Allocator<D, D> + Allocator<D>,
{
    type Error = &'static str;

    /// Construct the UDState with a KalmanState.
    ///
    /// The covariance matrix X is factorised into a U.d.U' as a UD matrix.
    fn try_from(state: KalmanState<N, D>) -> Result<Self, Self::Error> {
        let udu = UDU::new();
        // Factorise X into UD
        let rows = state.x.nrows();
        let mut UD = state.X.clone();
        let rcond = udu.UdUfactor(&mut UD, rows);
        rcond::check_non_negative(rcond, "X not PSD")?;

        Ok(UDState { x: state.x, UD, udu })
    }
}

impl<N: Copy + RealField, D: Dim> KalmanEstimator<N, D> for UDState<N, D>
    where
        DefaultAllocator: Allocator<D, D> + Allocator<D>,
{
    /// Derive the KalmanState from the UDState.
    ///
    /// The covariance matrix X is recomposed from U.d.U' in the UD matrix.
    fn kalman_state<'e>(&self) -> Result<KalmanState<N, D>, &'e str> {
        // assign elements of common left block of M into X
        let x_shape = self.x.shape_generic().0;

        let mut X = self.UD.columns_generic(0, x_shape).into_owned();
        UDU::UdUrecompose(&mut X);

        Ok(KalmanState {
            x: self.x.clone(),
            X,
        })
    }
}

/// Additive noise.
///
/// Noise represented as the noise covariance as a factorised UdU' matrix.
pub struct CorrelatedFactorNoise<N: RealField, D: Dim>
    where
        DefaultAllocator: Allocator<D, D>,
{
    /// Noise covariance
    pub UD: OMatrix<N, D, D>,
}

impl<N: Copy + RealField, D: Dim> CorrelatedFactorNoise<N, D>
    where
        DefaultAllocator: Allocator<D, D>,
{
    /// Creates a CorrelatedFactorNoise from an CorrelatedNoise.
    /// The CorrelatedNoise must be PSD.
    pub fn from_correlated(correlated: &CorrelatedNoise<N, D>) -> Result<Self, &'static str> {
        let udu = UDU::new();
        let mut ud: OMatrix<N, D, D> = correlated.Q.clone_owned();
        let rcond = udu.UdUfactor(&mut ud, correlated.Q.nrows());
        rcond::check_non_negative(rcond, "Q not PSD")?;

        Ok(CorrelatedFactorNoise { UD: ud })
    }
}

impl<N: Copy + RealField, D: Dim> UDState<N, D>
    where
        DefaultAllocator: Allocator<D, D> + Allocator<D>,
{
    /// Special Linear 'hx' observe for correlated factorised noise.
    ///
    /// Observation predictions are made with the linear 'hx'. This allows the observation noise
    /// to be decorrelated. Observations can then be applied for each element in the order they appear in z.
    ///
    /// Return: Minimum 'rcond' of all sequential observations
    pub fn observe_linear_correlated<ZD: Dim>(
        &mut self,
        z: &OVector<N, ZD>,
        hx: &OMatrix<N, ZD, D>,
        h_normalize: fn(&mut OVector<N, ZD>, &OVector<N, ZD>),
        noise_factor: &CorrelatedFactorNoise<N, ZD>,
    ) -> Result<N, &'static str>
        where
            DefaultAllocator: Allocator<ZD, ZD> + Allocator<ZD, D> + Allocator<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();

        // Observation prediction and normalised observation
        let mut GIHx = hx.clone();
        {
            // Solve G* GIHx = Hx for GIHx in-place
            for j in 0..x_size {
                let mut GIHx_j = GIHx.column_mut(j);
                for i in (0..z_size).rev() {
                    let UD_i = noise_factor.UD.row(i);
                    let mut t = N::zero();
                    for k in i + 1..z_size {
                        t += UD_i[k] * GIHx_j[k];
                    }
                    GIHx_j[i] -= t;
                }
            }

            // Solve G zp~ = z, G z~ = z  for zp~,z~ in-place
            for i in (0..z_size).rev() {
                let UD_i = noise_factor.UD.row(i);
                for k in i + 1..z_size {
                    let UD_ik = UD_i[k];
                    let zpt = UD_ik * zp[k];
                    zp[i] -= zpt;
                    let zpdt = UD_ik * zpdecol[k];
                    zpdecol[i] -= zpdt;
                }
            }
        }

        // Apply observations sequential as they are decorrelated
        let mut rcondmin = N::max_value().unwrap();
        for o in 0..z_size {
            // Update UD and extract gain
            GIHx.row(o).transpose_to(&mut scratch.a);
            let rcond = UDState::observeUD(self, &mut scratch, noise_factor.UD[(o, o)])
                .unwrap_or(self.udu.minus_one);
            rcond::check_positive(rcond, "S not PD in observe")?;
            if rcond < rcondmin {
                rcondmin = rcond;
            }
            // State update using linear innovation
            let s = z[o] - zpdecol[o];
            self.x += &scratch.w * s;
        }
        Ok(rcondmin)
    }
}

/// Prediction Scratch.
///
/// Provides temporary variables for prediction calculation.
pub struct PredictScratch<N: RealField, D: Dim, QD: Dim>
    where
        D: DimAdd<QD>,
        DefaultAllocator: Allocator<D, QD> + Allocator<DimSum<D, QD>>,
{
    pub G: OMatrix<N, D, QD>,
    pub d: OVector<N, DimSum<D, QD>>,
    pub dv: OVector<N, DimSum<D, QD>>,
    pub v: OVector<N, DimSum<D, QD>>,
}

/// Observe Scratch.
///
/// Provides temporary variables for observe calculation.
pub struct ObserveScratch<N: RealField, D: Dim>
    where
        DefaultAllocator: Allocator<D>,
{
    pub w: OVector<N, D>,
    pub a: OVector<N, D>,
    pub b: OVector<N, D>,
    pub S: N,
}

impl<N: Copy + RealField, D: Dim> UDState<N, D>
    where
        DefaultAllocator: Allocator<D, D> + Allocator<D>,
{
    pub fn new_predict_scratch<QD: Dim>(&self, qd: QD) -> PredictScratch<N, D, QD>
        where
            D: DimAdd<QD>,
            DefaultAllocator: Allocator<D, QD> + Allocator<DimSum<D, QD>, U1>,
    {
        // x + qd rows
        let xqd_size = self.UD.shape_generic().1.add(qd);
        PredictScratch {
            G: OMatrix::zeros_generic(self.UD.shape_generic().0, qd),
            d: OVector::zeros_generic(xqd_size, Const::<1>),
            dv: OVector::zeros_generic(xqd_size, Const::<1>),
            v: OVector::zeros_generic(xqd_size, Const::<1>),
        }
    }

    pub fn new_observe_scratch(&self) -> ObserveScratch<N, D> {
        let x_size = self.x.shape_generic().0;
        ObserveScratch {
            w: OVector::zeros_generic(x_size, Const::<1>),
            a: OVector::zeros_generic(x_size, Const::<1>),
            b: OVector::zeros_generic(x_size, Const::<1>),
            S: self.udu.zero,
        }
    }

    pub fn predict_use_scratch<QD: Dim>(
        &mut self,
        scratch: &mut PredictScratch<N, D, QD>,
        x_pred: &OVector<N, D>,
        fx: &OMatrix<N, D, D>,
        noise: &CoupledNoise<N, D, QD>,
    ) -> Result<N, &'static str>
        where
            D: DimAdd<QD>,
            DefaultAllocator: Allocator<DimSum<D, QD>> + Allocator<D, QD> + Allocator<QD>,
    {
        self.x = x_pred.clone();

        // Predict UD from model
        let rcond = UDState::predictGq(self, scratch, &fx, &noise.G, &noise.q)
            .unwrap_or(self.udu.minus_one);
        rcond::check_non_negative(rcond, "X not PSD").map(|()| rcond)
    }

    pub fn observe_innovation_use_scratch<ZD: Dim>(
        &mut self,
        scratch: &mut ObserveScratch<N, D>,
        s: &OVector<N, ZD>,
        hx: &OMatrix<N, ZD, D>,
        noise: &UncorrelatedNoise<N, ZD>,
    ) -> Result<N, &'static str>
        where
            DefaultAllocator: Allocator<ZD, D> + Allocator<ZD>,
    {
        let z_size = s.nrows();

        // Apply observations sequentially as they are uncorrelated
        let mut rcondmin = N::max_value().unwrap();

        for o in 0..z_size {
            // Check noise precondition
            if noise.q[o] < self.udu.zero {
                return Err("Zv not PSD in observe");
            }
            // Update UD and extract gain
            hx.row(o).transpose_to(&mut scratch.a);
            let rcond = UDState::observeUD(self, scratch, noise.q[o])
                .unwrap_or(self.udu.minus_one);
            if rcond < rcondmin {
                rcondmin = rcond;
            }
            // State update using normalised non-linear innovation
            self.x += &scratch.w * s[o];
        }
        Ok(rcondmin)
    }

    /// MWG-S prediction from \[Bierman\] p.132.
    ///
    /// # Return
    /// reciprocal condition number, Err if factorisation is negative
    fn predictGq<QD: Dim>(
        &mut self,
        scratch: &mut PredictScratch<N, D, QD>,
        Fx: &OMatrix<N, D, D>,
        G: &OMatrix<N, D, QD>,
        q: &OVector<N, QD>,
    ) -> Result<N, ()>
        where
            D: DimAdd<QD>,
            DefaultAllocator: Allocator<DimSum<D, QD>> + Allocator<D, QD> + Allocator<QD>,
    {
        let nx_ = self.x.shape_generic().0;
        let nx = nx_.value();
        let nq_ = q.shape_generic().0;

        // Augment d with q, UD with G
        scratch.d
            .rows_generic_mut(nx, nq_)
            .copy_from(q);
        scratch.G.copy_from(G);

        // U=Fx*U and diagonals retrieved
        for j in (1..nx).rev() {
            // nx-1..1
            let mut UD_j = self.UD.column_mut(j);
            // Prepare d as temporary
            scratch.d
                .rows_range_mut(0..j + 1)
                .copy_from(&UD_j.rows_range(0..j + 1));

            // Lower triangle of UD is implicitly empty
            let Fx_j = Fx.column(j);
            let ds = &scratch.d.rows_range(0..j);
            for i in 0..nx {
                UD_j[i] = Fx_j[i] + Fx.row(i).columns_range(0..j).tr_dot(ds);
            }
        }
        if nx > 0 {
            scratch.d[0] = self.UD[(0, 0)];
        }

        // Complete U = Fx*U
        self.UD.column_mut(0).copy_from(&Fx.column(0));

        // The MWG-S algorithm on UD transpose
        for j in (0..nx).rev() {
            // n-1..0
            let mut e = self.udu.zero;

            let mut vi = scratch.v.iter_mut();
            let mut di = scratch.d.iter();

            let UD_j = self.UD.row(j);
            let G_j = scratch.G.row(j);
            let mut udgi = UD_j.iter().chain(G_j.iter());
            for dv in &mut scratch.dv {
                let v = vi.next().unwrap();
                let d = *di.next().unwrap();
                *v = *udgi.next().unwrap();
                *dv = d * *v;
                e += *v * *dv;
            }

            // Check diagonal element
            if e > self.udu.zero {
                // Positive definite
                self.UD[(j, j)] = e;
                let diaginv = self.udu.one / e;

                for k in 0..j {
                    e = self.calce(&scratch, nx_, nq_, k);
                    e *= diaginv;
                    self.UD[(j, k)] = e;

                    let mut vi = scratch.v.iter();

                    let mut UD_k = self.UD.row_mut(k);
                    let mut G_k = scratch.G.row_mut(k);
                    for udg in UD_k.iter_mut().chain(G_k.iter_mut()) {
                        *udg -= e * *vi.next().unwrap();
                    }
                }
            } else if e == self.udu.zero {
                // Semi-definite, check not negative
                self.UD[(j, j)] = e;
                // diaginv = 1 / e is infinite

                for k in 0..j {
                    e = self.calce(&scratch, nx_, nq_, k);
                    if e != self.udu.zero {
                        return Err(());
                    }
                    // UD and G are unchanged as e == 0
                }
            } else {
                // Negative
                return Err(());
            }
        } // MWG-S loop

        // Transpose and Zero lower triangle
        self.UD.fill_upper_triangle_with_lower_triangle();
        self.UD.fill_lower_triangle(N::zero(), 1);

        // Estimate the reciprocal condition number from upper triangular part
        Ok(UDU::UdUrcond(&self.UD))
    }

    // MWG-S calculation of 'e' from for a UD row k
    fn calce<QD: Dim>(&self, scratch: & PredictScratch<N, D, QD>, nx_: D, nq_: QD, k: usize) -> N
    where
        D: DimAdd<QD>,
        DefaultAllocator: Allocator<DimSum<D, QD>> + Allocator<D, QD> + Allocator<QD>,
    {
        self.UD.row(k)
            .columns_generic(0, nx_)
            .tr_dot(&scratch.dv.rows_generic(0, nx_))
            + scratch.G.row(k)
            .columns_generic(0, nq_)
            .tr_dot(&scratch.dv.rows_generic(nx_.value(), nq_))
    }

    /// Linear UD factorisation update from \[Bierman\] p.100.
    ///
    /// # Input
    ///  scratch.a observation coefficients
    ///  R observation variance
    /// # Output
    ///  scratch.w observation Kalman gain
    ///  scratch.S observation innovation variance
    /// # Precondition
    ///  R is PSD (not checked)
    /// # Return
    ///  reciprocal condition number of UD, Err if alpha singular (negative or zero)
    fn observeUD(&mut self, scratch: &mut ObserveScratch<N, D>, R: N) -> Result<N, ()> {
        let n = self.UD.nrows();
        // a(n) is U'a
        // b(n) is Unweighted Kalman gain

        // Compute b = DU'h, a = U'h
        for j in (1..n).rev() { // n-1..1
            let t = self
                .UD
                .column(j)
                .rows_range(0..j)
                .dot(&scratch.a.rows_range(0..j));
            scratch.a[j] += t;
            scratch.b[j] = self.UD[(j, j)] * scratch.a[j];
        }
        scratch.b[0] = self.UD[(0, 0)] * scratch.a[0];

        // Update UD(0,0), d(0) modification
        scratch.S = R + scratch.b[0] * scratch.a[0];
        if scratch.S <= self.udu.zero {
            return Err(());
        }
        let mut gamma = self.udu.one / scratch.S;
        self.UD[(0, 0)] *= R * gamma;
        // Update rest of UD and gain b
        for j in 1..n {
            let mut UD_j = self.UD.column_mut(j);
            // d modification
            let alpha_jm1 = scratch.S; // alpha at j-1
            scratch.S += scratch.b[j] * scratch.a[j];
            let lamda = -scratch.a[j] * gamma;
            if scratch.S <= self.udu.zero {
                return Err(());
            }
            gamma = self.udu.one / scratch.S;
            UD_j[j] *= alpha_jm1 * gamma;
            // U modification
            for i in 0..j {
                let UD_jm1 = UD_j[i];
                UD_j[i] = UD_jm1 + lamda * scratch.b[i];
                let t = scratch.b[j] * UD_jm1;
                scratch.b[i] += t;
            }
        }
        // Update gain from b
        scratch.w.copy_from(&(&scratch.b * gamma));

        // Estimate the reciprocal condition number from upper triangular part
        Ok(UDU::UdUrcond(&self.UD))
    }
}