bayes_estimate 0.20.0

#![allow(non_snake_case)]

//! Square-root covariance with unscented transform state estimation.
//! NOTE Experimental, the implementation has not yet been verified to produce correct results.
//!
//! A discrete Bayesian estimator that uses state with a square-root state covariance representation of the system.
//!
//! The Julier-Uhlmann 'unscented' transform is used for non-linear state predictions and observations.
//! The transform evaluates the non-linear predict and observe functions at 'sigma' points about the mean to provide an
//! estimate of the distribution. The transforms can be optimised for particular functions using parameterized weights.
//! See: Van Der Merwe, R. and Wan, E.A. The Square-Root Unscented Kalman Filter for State and Parameter-Estimation. 2001.

use nalgebra::{allocator::Allocator, stack, DefaultAllocator, Dim, Dyn, Matrix, OMatrix, OVector, RealField, StorageMut, Vector, U1};
use alloc::vec::Vec;

use crate::models::{Estimator, FunctionalObserver, FunctionalPredictor, KalmanEstimator, KalmanState};
use crate::noise::{CorrelatedNoise};
use crate::estimators::unscented::UnscentedWeights;
use num_traits::FromPrimitive;

/// 'unscented' square root state estimation.
pub struct UnscentedRootState<N: RealField, D: Dim>
where
    DefaultAllocator: Allocator<D, D> + Allocator<D>,
{
    /// State vector
    pub x: OVector<N, D>,
    /// Square root state covariance matrix (symmetric positive semi-definite)
    pub S: OMatrix<N, D, D>,

    pub w : UnscentedWeights<N>
}

impl<N: Copy + FromPrimitive + RealField, D: Dim> Estimator<N, D> for UnscentedRootState<N, D>
where
    DefaultAllocator: Allocator<D, D> + Allocator<D> + Allocator<U1, D>,
{
    fn state<'e>(&self) -> Result<OVector<N, D>, &'e str> {
        Ok(self.x.clone())
    }
}

impl<N: Copy + FromPrimitive + RealField, D: Dim> KalmanEstimator<N, D> for UnscentedRootState<N, D>
where
    DefaultAllocator: Allocator<D, D> + Allocator<D> + Allocator<U1, D>,
{
    fn kalman_state<'e>(&self) -> Result<KalmanState<N, D>, &'e str> {
        Ok(KalmanState {
            x : self.x.clone(),
            X : &self.S * &self.S.transpose()
        })
    }
}

impl<N: Copy + RealField, D: Dim> UnscentedRootState<N, D>
where
    DefaultAllocator: Allocator<D, D> + Allocator<D>,
{
    pub fn from_kalman_state(state: KalmanState<N, D>, alpha: N, beta: N) -> Result<Self, &'static str> {
        let xx = state.x.nrows();
        let mut state_chol = state.X.clone_owned();
        cholesky_psd(&mut state_chol).ok_or("state covariance not PSD") ?;

        Ok(UnscentedRootState {
            x: state.x,
            S: state_chol,
            w: UnscentedWeights::new_weights(xx, alpha, beta)
        })
    }

    /// Generates sigma points from the current state and square root covariance
    fn sigma_points(&self) -> Vec<OVector<N, D>> {
        let d = self.x.nrows();
        let mut sigma_points = Vec::with_capacity(2 * d + 1);

        // Add the mean
        sigma_points.push(self.x.clone());

        // Add scaled columns of the square root covariance
        for c in 0..d {
            let sigmaCol = self.S.column(c) * self.w.nu;
            sigma_points.push(&self.x + &sigmaCol);
            sigma_points.push(&self.x - &sigmaCol);
        }
        sigma_points
    }

}

impl<N: Copy + FromPrimitive + RealField, D: Dim> UnscentedRootState<N, D>
where
    DefaultAllocator: Allocator<D, D> + Allocator<D> + Allocator<U1, D> + Allocator<D, Dyn>
{
    /// Unscented state prediction with a functional prediction model and additive noise as a Cholesky factor.
    fn predict<QD: Dim>(
        &mut self,
        f: impl Fn(&OVector<N, D>) -> OVector<N, D>,
        noise_chol: &OMatrix<N, D, D>,
    ) -> Result<(), &'static str>
    where
        DefaultAllocator: Allocator<QD> + Allocator<D, QD>
    {
        // Generate sigma points
        let sigma_points = self.sigma_points();

        // Predict sigma points through the predict model
        let predict_points: Vec<OVector<N, D>> = sigma_points.iter()
            .map(|point| f(point))
            .collect();

        // Compute predicted state mean
        self.x = self.w.mean(&predict_points);

        // Compute predicted square root covariance
        let point_diff: Vec<OVector<N, D>> = predict_points.iter()
            .map(|point| point - &self.x)
            .collect();

        let augmented_point_matrix = stack![ OMatrix::<N, D, Dyn>::from_columns(&point_diff[1..]) * self.w.W_.sqrt(),  noise_chol ];
        let R = augmented_point_matrix.transpose().qr().unpack_r();
        let Sx: Matrix<N, D,  D, _> = R.rows_generic(0, self.x.shape_generic().0).transpose();

        self.S.copy_from(&Sx);
        cholesky_rank_one_update(&mut self.S, &mut point_diff[0].clone_owned(), self.w.Wc0);

        Ok(())
    }
}

impl<N: Copy + FromPrimitive + RealField, D: Dim> FunctionalPredictor<N, D>
    for UnscentedRootState<N, D>
where
    DefaultAllocator: Allocator<D, D> + Allocator<D> + Allocator<U1, D> + Allocator<D, Dyn>
{
    /// Unscented state prediction with a functional prediction model and additive correlated noise.
    fn predict(
        &mut self,
        f: impl Fn(&OVector<N, D>) -> OVector<N, D>,
        noise: &CorrelatedNoise<N, D>,
    ) -> Result<(), &'static str> {
        let mut noise_chol = noise.Q.clone_owned();
        cholesky_psd(&mut noise_chol).ok_or("predict noise not PSD") ?;

        self.predict(f, &noise_chol)
    }
}

impl<N: Copy + FromPrimitive + RealField, D: Dim> UnscentedRootState<N, D>
where
    DefaultAllocator: Allocator<D, D>
        + Allocator<D>
        + Allocator<U1, D>,
{
    /// Unscented state observation with a functional observation model and additive noise as a Cholesky factor.
    fn observe<ZD: Dim>(
        &mut self,
        z: &OVector<N, ZD>,
        h: impl Fn(&OVector<N, D>) -> OVector<N, ZD>,
        noise_chol: &OMatrix<N, ZD, ZD>,
    ) -> Result<(), &'static str>
    where
        DefaultAllocator:  Allocator<ZD> + Allocator<U1, ZD>
        + Allocator<ZD, D>
        + Allocator<ZD, Dyn>
        + Allocator<D, Dyn>
        + Allocator<D, ZD>
        + Allocator<ZD, ZD>
    {
        // Generate sigma points
        let sigma_points = self.sigma_points();

        // Propagate sigma points through the observe model
        let predicted_obs: Vec<OVector<N, ZD>> = sigma_points.iter()
            .map(|point| h(point))
            .collect();

        // Compute predicted observe mean
        let predicted_obs_mean = self.w.mean(&predicted_obs);

        // Compute cross-covariance
        let obs_diff: Vec<OVector<N, ZD>> = predicted_obs.iter()
            .map(|point| point - &predicted_obs_mean)
            .collect();

        let state_differences: Vec<OVector<N, D>> = sigma_points.iter()
            .map(|point| point - &self.x)
            .collect();
        let Pxy = self.w.cross_covariance(&state_differences, &obs_diff);

        // Compute square root innovation covariance
        let augmented_obs_matrix = stack![ OMatrix::<N, ZD, Dyn>::from_columns(&obs_diff[1..]) * self.w.W_.sqrt(),  noise_chol ];
        let R = augmented_obs_matrix.transpose().qr().unpack_r();
        let mut Sy = R.rows_generic(0, z.shape_generic().0).transpose();
        cholesky_rank_one_update(&mut Sy, &mut obs_diff[0].clone_owned(), self.w.Wc0);

        // Compute Kalman gain
        let PxySyI = Sy.solve_lower_triangular(&Pxy.transpose()).unwrap();
        let K = Sy.transpose().solve_upper_triangular(&PxySyI).unwrap().transpose();

        // Update state
        self.x += &K * (z - predicted_obs_mean);

        // Update square root covariance using rank 1 cholesky updates
        (K * Sy).column_iter_mut().for_each(|mut ki| {
            cholesky_rank_one_update(&mut self.S, &mut ki, N::one().neg())
        });

        Ok(())
    }
}

impl<N: Copy + FromPrimitive + RealField, D: Dim, ZD: Dim> FunctionalObserver<N, D, ZD>
for UnscentedRootState<N, D>
where
    DefaultAllocator: Allocator<D, D>
    + Allocator<D>
    + Allocator<ZD>
    + Allocator<U1, ZD>
    + Allocator<ZD, D>
    + Allocator<ZD, Dyn>
    + Allocator<D, Dyn>
    + Allocator<D, ZD>
    + Allocator<ZD, ZD>
    + Allocator<U1, D>,
{
    /// Unscented state observation with a functional observation model and additive correlate noise.
    fn observe(
        &mut self,
        z: &OVector<N, ZD>,
        h: impl Fn(&OVector<N, D>) -> OVector<N, ZD>,
        noise: &CorrelatedNoise<N, ZD>,
    ) -> Result<(), &'static str> {
        let mut noise_chol = noise.Q.clone_owned();
        cholesky_psd(&mut noise_chol).ok_or("observe noise not PSD") ?;

        self.observe(z, h, &noise_chol)
    }
}

fn cholesky_psd<N: Copy + FromPrimitive + RealField, D: Dim>(matrix: &mut OMatrix<N, D, D>) -> Option<N>
where DefaultAllocator: Allocator<D, D> {
    assert!(matrix.is_square(), "The input matrix must be square.");

    let n = matrix.nrows();

    for j in 0..n {
        for k in 0..j {
            let factor = unsafe { -matrix.get_unchecked((j, k)).clone() };

            let (mut col_j, col_k) = matrix.columns_range_pair_mut(j, k);
            let mut col_j = col_j.rows_range_mut(j..);
            let col_k = col_k.rows_range(j..);

            col_j.axpy(factor.conjugate(), &col_k, N::one());
        }

        let diag = unsafe { matrix.get_unchecked((j, j)).clone() };

        if let Some(denom) = diag.try_sqrt()
        {
            unsafe {
                *matrix.get_unchecked_mut((j, j)) = denom.clone();
            }

            let mut col = matrix.view_range_mut(j + 1.., j);
            if denom.is_zero() {
                // Possibly semi-definite, check not negative
                if col.iter().any(|e| !e.is_zero()) {
                    return None
                }
            }
            else {
                col /= denom;
            }
            continue;
        }

        // The diagonal element is either zero or its square root could not
        // be taken (e.g. for negative real numbers).
        return None;
    }

    matrix.fill_upper_triangle(N::zero(), 1);
    Some(N::one())
}

/// Given the Cholesky decomposition of a matrix `M`, a scalar `sigma` and a vector `x`,
/// performs a rank one update such that we end up with the decomposition of `M + sigma * (x * x.adjoint())`.
///
/// This helper method is called by `rank_one_update` but also `insert_column` and `remove_column`
/// where it is used on a square view of the decomposition
fn cholesky_rank_one_update<T: Copy + FromPrimitive + RealField, Dm, Sm, Rx, Sx>(
    chol: &mut Matrix<T, Dm, Dm, Sm>,
    x: &mut Vector<T, Rx, Sx>,
    sigma: T::RealField,
) where
    Dm: Dim,
    Rx: Dim,
    Sm: StorageMut<T, Dm, Dm>,
    Sx: StorageMut<T, Rx, U1>,
{
    // heavily inspired by Eigen's `llt_rank_update_lower` implementation https://eigen.tuxfamily.org/dox/LLT_8h_source.html
    let n = x.nrows();
    assert_eq!(
        n,
        chol.nrows(),
        "The input vector must be of the same size as the factorized matrix."
    );

    let mut beta = T::one();

    for j in 0..n {
        // updates the diagonal
        let diag = T::real(unsafe { chol.get_unchecked((j, j)).clone() });
        let diag2 = diag.clone() * diag.clone();
        let xj = unsafe { x.get_unchecked(j).clone() };
        let sigma_xj2 = sigma.clone() * T::modulus_squared(xj.clone());
        let gamma = diag2.clone() * beta.clone() + sigma_xj2.clone();
        let new_diag = (diag2.clone() + sigma_xj2.clone() / beta.clone()).sqrt();
        unsafe { *chol.get_unchecked_mut((j, j)) = T::from_real(new_diag.clone()) };
        beta += sigma_xj2 / diag2;
        // updates the terms of L
        let mut xjplus = x.rows_range_mut(j + 1..);
        let mut col_j = chol.view_range_mut(j + 1.., j);
        xjplus.axpy(-xj.clone() / T::from_real(diag.clone()), &col_j, T::one());
        if !gamma.is_zero() {
            // col_j = T::from_real(nljj / diag) * col_j  + (T::from_real(nljj * sigma / gamma) * T::conjugate(wj)) * temp_jplus;
            col_j.axpy(
                T::from_real(new_diag.clone() * sigma.clone() / gamma) * T::conjugate(xj),
                &xjplus,
                T::from_real(new_diag / diag),
            );
        }
    }
}