bayes_estimate 0.20.0

#![allow(non_snake_case)]

//! Covariance with unscented transform state estimation.
//!
//! A discrete Bayesian estimator that uses the [`KalmanState`] linear 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, Cholesky, DefaultAllocator, Dim, OMatrix, OVector, RealField, U1};
use alloc::vec::Vec;
use crate::models::{Estimator, FunctionalObserver, FunctionalPredictor, KalmanEstimator, KalmanState};
use crate::noise::CorrelatedNoise;
use num_traits::FromPrimitive;

/// 'unscented' state estimation.
///
/// Simply the Kalman state with the weights required for the 'unscented' transform.
pub struct UnscentedState<N: RealField, D: Dim>
where
    DefaultAllocator: Allocator<D, D> + Allocator<D>,
{
    pub kalman: KalmanState<N, D>,
    pub w: UnscentedWeights<N>,
}

pub struct UnscentedWeights<N: RealField> {
    /// Sigma point generation scaling
    pub nu: N,
    /// Sigma point 0 weight for mean
    pub Wm0: N,
    /// Sigma point 0 weight for covariance
    pub Wc0: N,
    /// Sigma point (other than 0) weight for mean and covariance
    pub W_: N
}

impl<N: Copy + RealField> UnscentedWeights<N> {
    /// Calculate [UnscentedWeights] from an alpha, beta parameterization
    /// # Arguments
    /// * `dd` - state dimension
    /// * `alpha` - parameters which spreads the unscented points
    /// * `beta` - parameter which adjusts the unscented points
    pub fn new_weights(dd: usize, alpha: N, beta: N) -> Self {
        let d = N::from_usize(dd).unwrap();
        //  Calculate scaling factor for all off-center points
        let lamda = d * alpha * alpha - d;
        let nu = (d + lamda).sqrt();
        let Wm0 = lamda / (d + lamda);
        let Wc0 = Wm0 + (N::one() - alpha * alpha + beta);
        let two = N::from_usize(2).unwrap();
        let W_ = (two * (d + lamda)).recip();

        UnscentedWeights { nu, Wm0, Wc0, W_ }
    }
}

impl<N: Copy + RealField> UnscentedWeights<N>
{
    /// Computes the mean of a set of sigma points
    pub fn mean<PD: Dim>(&self, points: &[OVector<N, PD>]) -> OVector<N, PD>
    where DefaultAllocator: Allocator<PD> {
        let mut pi = points.iter();
        let mut mean = pi.next().unwrap() * self.Wm0;
        pi.for_each(|point| mean += point * self.W_);
        mean
    }

    /// Computes the cross-covariance of sigma points
    pub fn cross_covariance<AD: Dim, BD: Dim>(
        &self,
        a: &[OVector<N, AD>],
        b: &[OVector<N, BD>],
    ) -> OMatrix<N, AD, BD> where DefaultAllocator: Allocator<AD> + Allocator<BD> + Allocator<U1, BD> + Allocator<AD, BD> {
        let mut zipi = a.iter().zip(b);
        let point0 = zipi.next().unwrap();
        let mut cc = point0.0 * point0.1.transpose() * self.Wc0;
        zipi.for_each(|point| cc += point.0 * point.1.transpose() * self.W_);
        cc
    }

}


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

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

impl<N: Copy + RealField, D: Dim> UnscentedState<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 d = state.x.nrows();

        Ok(UnscentedState {
            kalman: state,
            w: UnscentedWeights::new_weights(d, alpha, beta)
        })
    }

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

        let sigma = self.kalman.X.clone()
            .cholesky()
            .ok_or("sigma_point, X not PSD")?
            .l();

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

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

}

impl<N: Copy + FromPrimitive + RealField, D: Dim> FunctionalPredictor<N, D>
    for UnscentedState<N, D>
where
    DefaultAllocator: Allocator<D, D> + Allocator<D> + Allocator<U1, D>,
{
    /// 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> {
        // 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.kalman.x = self.w.mean(&predict_points);

        // Compute predicted covariance
        let point_diff: Vec<OVector<N, D>> = predict_points.iter()
            .map(|point| point - &self.kalman.x)
            .collect();
        // State covariance with additive noise
        self.kalman.X = self.w.cross_covariance(&point_diff, &point_diff) + &noise.Q;

        Ok(())
    }
}

impl<N: Copy + FromPrimitive + RealField, D: Dim, ZD: Dim> FunctionalObserver<N, D, ZD>
    for UnscentedState<N, D>
where
    DefaultAllocator: Allocator<D, D>
        + Allocator<D>
        + Allocator<ZD>
        + Allocator<U1, ZD>
        + Allocator<ZD, D>
        + 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> {
        // 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.kalman.x)
            .collect();

        // Innovation covariance
        let S = self.w.cross_covariance(&obs_diff, &obs_diff) + &noise.Q;

        // Inverse innovation covariance
        let SI = S.clone().cholesky().ok_or("S not PD in observe")?.inverse();

        // Kalman gain, X*Hx'*SI
        let W = self.w.cross_covariance(&state_differences, &obs_diff) * &SI;

        // State update
        self.kalman.x += &W * (z - predicted_obs_mean);
        // X -= W.S.W'
        self.kalman.X.quadform_tr(N::one().neg(), &W, &S, N::one());

        Ok(())
    }

}

/// Result of an observation innovation calculation.
///
/// Contains the innovation, innovation covariance, and its factorisation and inverse.
/// Useful to compute likelihood mearuese of the measurement given the state.
#[derive(Clone, Debug)]
pub struct ObserveInnovation<N: RealField, ZD: Dim>
where
    DefaultAllocator: Allocator<ZD> + Allocator<ZD, ZD> + Allocator<U1, ZD>,
{
    /// Innovation (residual): z - z_predicted
    pub s: OVector<N, ZD>,
    /// Innovation covariance S
    pub S: OMatrix<N, ZD, ZD>,
    /// Innovation covariance S Cholseky factorisation
    pub S_chol: Cholesky<N, ZD>,
    /// Innovation covariance S inverse
    pub SI: OMatrix<N, ZD, ZD>,
}

impl<N: Copy + FromPrimitive + RealField, D: Dim> UnscentedState<N, D>
where
    DefaultAllocator: Allocator<D, D>
    + Allocator<D>
    + Allocator<U1, D>,
{
    /// Unscented state observation with a functional observation model and additive correlated noise.
    /// Returns the observation innovation calculations.
    pub fn observe_with_innovation<ZD: Dim>(
        &mut self,
        z: &OVector<N, ZD>,
        h: impl Fn(&OVector<N, D>) -> OVector<N, ZD>,
        noise: &CorrelatedNoise<N, ZD>,
    ) -> Result<ObserveInnovation<N, ZD>, &'static str>
    where
        DefaultAllocator: Allocator<ZD>
        + Allocator<U1, ZD>
        + Allocator<ZD, D>
        + Allocator<D, ZD>
        + Allocator<ZD, ZD>
        + Allocator<U1, D>,
    {
        // 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.kalman.x).collect();

        // Innovation covariance
        let S = self.w.cross_covariance(&obs_diff, &obs_diff) + &noise.Q;

        // Compute Cholesky decomposition for innovation covariance
        let S_chol = S.clone().cholesky().ok_or("S not PD in observe")?;

        // Inverse innovation covariance
        let SI = S_chol.inverse();

        // Innovation (residual)
        let s = z - &predicted_obs_mean;

        // Kalman gain, X*Hx'*SI
        let W = self.w.cross_covariance(&state_differences, &obs_diff) * &SI;

        // State update
        self.kalman.x += &W * &s;
        // X -= W.S.W'
        self.kalman.X.quadform_tr(N::one().neg(), &W, &S, N::one());

        Ok(ObserveInnovation { s, S, S_chol, SI })
    }
}