bayes_estimate 0.20.0

//!
//! Sampling Importance Resampleing estimation.
//!
//! Also known as a weighted Booststrap.
//!
//! # References
//! * \[1] "Novel approach to nonlinear-non-Gaussian Bayesian state estimation"
//!   NJ Gordon, DJ Salmond, AFM Smith IEE Proceeding-F Vol.140 No.2 April 1993
//! * \[2] Building Robust Simulation-based Filter for Evolving Data Sets"
//!   J Carpenter, P Clifford, P Fearnhead Technical Report University of Oxford
//!
//! A variety of resampling algorithms can be used for SIR.
//! There are implementations for two algorithms:
//!   standard_resample: Standard resample algorithm from \[1]
//!   systematic_resample: A Simple stratified re-sampler from \[2]
//!
//! # Numerics
//! SIR algorithms are sensitive to the PRNG properties.
//! In particular, we require that the uniform random number range be [0..1) NOT [0..1]. This implementation uses [f32]
//! to represent the random numbers.
//!
//! Resampling requires comparisons of normalised likelihoods. This will become ill-conditioned if the likelihoods have
//! a large range.
//!
//! The quantisation of the generated random number must not approach the sample size. This will result in quantisation
//! of the resampling. For example if random identically equal to 0 becomes highly probable due to quantisation
//! this will result in the first sample being selectively draw whatever its likelihood.
use nalgebra::allocator::Allocator;
use nalgebra::{Const, DefaultAllocator, Dim, OMatrix, OVector, RealField, U1};
use alloc::vec::Vec;
use alloc::vec;
use num_traits::real::Real;
use num_traits::{FromPrimitive, Pow, ToPrimitive};
use rand_core::RngCore;
use rand_distr::{Distribution, StandardNormal, Uniform};

use crate::models::{Estimator, KalmanEstimator, KalmanState};
use crate::noise::{CorrelatedNoise, CoupledNoise};

/// Sample state.
///
/// State distribution is represented as state samples and their likelihood.
pub struct SampleState<N: RealField, D: Dim>
where
    DefaultAllocator: Allocator<D>,
{
    /// State samples
    pub s: Samples<N, D>,
    /// and their likelihoods (bootstrap weights)
    pub w: Likelihoods,
}

/// Sample state and a PRNG to draw from.
pub struct SampleStateDraw<'a, N: RealField, D: Dim>
    where
        DefaultAllocator: Allocator<D>,
{
    /// Sample State
    pub state: SampleState<N, D>,
    /// A PRNG used to draw random samples from
    pub rng: &'a mut dyn RngCore,
}

/// State samples.
pub type Samples<N, D> = Vec<OVector<N, D>>;

/// Likelihoods.
pub type Likelihoods = Vec<f32>;

/// Resample count.
pub type Resamples = Vec<u32>;

/// A resampling function.
///
/// Returning a [Result] containing:
///   the vector of [Resample] counts,
///   number of unique samples,
///    smallest normalised likelihood, to determine numerical conditioning of likelihoods
/// # Errors
///  The [Result] returns the error message if there is an error in the resampling.
pub type Resampler =
    dyn FnMut(&mut Likelihoods, &mut dyn RngCore) -> Result<(Resamples, u32, f32), &'static str>;

/// A roughening function.
///
/// The samples are roughened (modified) with one or more draws from the PRNG per sample.
pub type Roughener<N, D> = dyn FnMut(&mut Samples<N, D>, &mut dyn RngCore);


impl<N: RealField, D: Dim> SampleState<N, D>
    where
        DefaultAllocator: Allocator<D, D> + Allocator<D>,
{
    /// Creates a [SampleState] with equal likelihood weights.
    pub fn equal_likelihood_samples(s: Samples<N, D>) -> SampleState<N, D> {
        let nsamples = s.len();
        SampleState {
            s,
            w: vec![1f32; nsamples],
        }
    }

    /// Resample using likelihoods and roughen the sample state.
    /// Error returns:
    ///   When the re-sampler fails due to numeric problems with the likelihoods
    /// Returns:
    ///  number of unique samples,
    ///  smallest normalised likelihood, to determine numerical conditioning of likelihoods
    pub fn update_resample(
        &mut self,
        resampler: &mut Resampler,
        roughener: &mut Roughener<N, D>,
        rng: &mut dyn RngCore,
    ) -> Result<(u32, f32), &'static str> {
        // Resample based on likelihoods
        let (resamples, unique_samples, lcond) = resampler(&mut self.w, rng)?;

        // Select live sample
        live_samples(&mut self.s, &resamples);
        // Resampling results in uniform likelihoods
        self.w.fill(1.);

        // Roughen samples
        roughener(&mut self.s, rng);

        Ok((unique_samples, lcond))
    }

    /// Predict sample state using a state prediction function 'f'.
    pub fn predict(&mut self, f: fn(&OVector<N, D>) -> OVector<N, D>) {
        self.s.iter_mut().for_each(|el| el.copy_from(&f(el)));
    }

    /// Predict sample state using a sampled state prediction function 'f'.
    /// The sampling function should predict the state and sample any noise.
    pub fn predict_sampled(
        &mut self,
        f: impl Fn(&OVector<N, D>, &mut dyn RngCore) -> OVector<N, D>,
        rng: &mut dyn RngCore,
    ) {
        // Predict particles s using supplied prediction function
        for si in 0..self.s.len() {
            let ps = f(&self.s[si], rng);
            self.s[si].copy_from(&ps)
        }
    }

    /// Observe sample likelihoods using a likelihood function 'l'.
    /// The sample likelihoods are multiplied by the observed likelihoods.
    pub fn observe<LikelihoodFn>(&mut self, l: LikelihoodFn)
        where
            LikelihoodFn: Fn(&OVector<N, D>) -> f32,
    {
        let mut wi = self.w.iter_mut();
        for si in self.s.iter() {
            let w = wi.next().unwrap();
            *w *= l(si);
        }
    }

    /// Observe sample likelihoods directly.
    /// The sample likelihoods are multiplied by these likelihoods.
    pub fn observe_likelihood(&mut self, l: Likelihoods) {
        assert_eq!(self.w.len(), l.len());
        let mut li = l.iter();
        for wi in self.w.iter_mut() {
            *wi *= li.next().unwrap();
        }
    }
}

impl<'a, N: RealField, D: Dim> SampleStateDraw<'a, N, D>
where
    DefaultAllocator: Allocator<D, D> + Allocator<D>,
{
    /// Predict sample state using a sampled state prediction function 'f'.
    /// The sampling function should predict the state and sample any noise.
    pub fn predict_sampled(
        &mut self,
        f: impl Fn(&OVector<N, D>, &mut dyn RngCore) -> OVector<N, D>,
    ) {
        self.state.predict_sampled(f, self.rng)
    }

    /// Resample using likelihoods and roughen the sample state.
    /// Error returns:
    ///   When the re-sampler fails due to numeric problems with the likelihoods
    /// Returns:
    ///  number of unique samples,
    ///  smallest normalised likelihood, to determine numerical conditioning of likelihoods
    pub fn update_resample(
        &mut self,
        resampler: &mut Resampler,
        roughener: &mut Roughener<N, D>,
    ) -> Result<(u32, f32), &'static str> {
        self.state.update_resample(resampler, roughener, self.rng)
    }
}

impl<N: FromPrimitive + RealField, D: Dim> Estimator<N, D> for SampleState<N, D>
where
    DefaultAllocator: Allocator<D>,
{
    fn state<'e>(&self) -> Result<OVector<N, D>, &'e str> {
        // Mean of distribution: mean of particles
        let s_shape = self.s[0].shape_generic();
        let mut x = OVector::zeros_generic(s_shape.0, s_shape.1);
        for s in self.s.iter() {
            x += s;
        }
        x /= N::from_usize(self.s.len()).unwrap();

        Ok(x)
    }
}

/// Standard re-sampler from \[1].
///
/// # Algorithm
/// A sample is chosen once for each time its cumulative likelihood intersects with a uniform random draw.
/// Complexity is that of Vec::sort, O(n * log(n)) worst-case.
/// This complexity is required to sort the uniform random draws made,
/// this allows comparing of the two ordered lists w (cumulative) and ur (the sorted random draws).
///
/// # Returns
/// Tuple:
/// * number of times each particle should be resampled
/// * number of unique particles (number of non zeros in resamples)
/// * conditioning of the likelihoods (min likelihood / sum likelihoods)
///
/// # Side effects
/// * 'l' becomes a normalised cumulative sum.
/// * Draws are made from 'rng' for each likelihood.
pub fn standard_resampler(
    l: &mut Likelihoods,
    rng: &mut dyn RngCore,
) -> Result<(Resamples, u32, f32), &'static str> {
    let (lmin, lcum) = cumulative_likelihood(l)?;

    // Sorted uniform random distribution [0..1) for each resample
    let uniform01: Uniform<f32> = Uniform::new(0f32, 1f32).unwrap();
    let mut ur: Vec<f32> = uniform01.sample_iter(rng).take(l.len()).collect();
    ur.sort_by(|a, b| a.partial_cmp(b).unwrap());
    assert!(*ur.first().unwrap() >= 0. && *ur.last().unwrap() < 1.); // very bad if random is incorrect

    // Scale ur to cumulative sum
    ur.iter_mut().for_each(|el| *el *= lcum);

    // Resamples based on cumulative likelihood from sorted resample random values
    let mut uri = ur.iter();
    let mut urn = uri.next();
    let mut unique: u32 = 0;
    let mut resamples = Resamples::with_capacity(l.len());

    for wi in l.iter() {
        let mut res: u32 = 0; // assume not resampled until find out otherwise
        if (urn.is_some()) && *urn.unwrap() < *wi {
            unique += 1;
            loop {
                // count resamples
                res += 1;
                urn = uri.next();
                if urn.is_none() {
                    break;
                }
                if !(*urn.unwrap() < *wi) {
                    break;
                }
            }
        }
        resamples.push(res);
    }

    if uri.peekable().peek().is_some() {
        // resample failed due no non-numeric likelihoods
        return Err("likelihoods are not numeric and cannot be resampled");
    }

    Ok((resamples, unique, lmin / lcum))
}

/// Systematic resample algorithm from \[2].
///
/// # Algorithm:
/// * A particle is chosen once for each time its cumulative likelihood intersects with an equidistant grid.
/// * A uniform random draw is chosen to position the grid within the cumulative likelihoods.
///
/// Complexity O(n)
///
/// # Returns:
/// * number of times each particle should be resampled,
/// * number of unique particles (number of non zeros in resamples),
/// * conditioning of the likelihoods (min likelihood / sum likelihoods)
///
/// # Side effects:
/// * 'l' becomes a normalised cumulative sum,
/// * Draws are made from 'rng' for each likelihood
pub fn systematic_resampler(
    l: &mut Likelihoods,
    rng: &mut dyn RngCore,
) -> Result<(Resamples, u32, f32), &'static str> {
    let (lmin, lcum) = cumulative_likelihood(l)?;

    let uniform01: Uniform<f32> = Uniform::new(0f32, 1f32).unwrap();

    // Setup grid
    let glen = l.len();
    let gstep = lcum / glen as f32;
    let goffset = uniform01.sample(rng) * gstep; // random offset

    // Resamples based on cumulative likelihoods
    let mut resamples = Resamples::with_capacity(glen);
    let mut unique: u32 = 0;

    let mut gi: u32 = 0;
    for li in l.iter() {
        let mut res: u32 = 0; // assume not resampled until find out otherwise
        if (goffset + lcum * gi as f32 / glen as f32) < *li {
            unique += 1;
            loop {
                // count resamples
                res += 1;
                gi += 1;
                if !((goffset + lcum * gi as f32 / glen as f32) < *li) {
                    break;
                }
            }
        }
        resamples.push(res);
    }
    // Numerically it is possible that not all resamples as chosen
    if (gi as usize) < glen {
        let lastres = resamples.pop().unwrap() + (glen - gi as usize) as u32;
        resamples.push(lastres);
    }

    Ok((resamples, unique, lmin / lcum))
}

/// Normalised cumulative sum of likelihoods (Kahan algorithm), and find the smallest likelihood.
fn cumulative_likelihood(l: &mut Likelihoods) -> Result<(f32, f32), &'static str> {
    let mut lmin = <f32 as Real>::max_value();
    let mut lcum = 0.;
    {
        let mut c = 0.;
        for li in l.iter_mut() {
            if *li < lmin {
                lmin = *li;
            }
            let y = *li - c;
            let t = lcum + y;
            c = t - lcum - y;
            lcum = t;
            *li = t;
        }
    }
    if lmin < 0. {
        // bad likelihoods
        return Err("negative likelihood");
    }
    if lcum <= 0. {
        // bad cumulative likelihood (previous check should actually prevent -ve
        return Err("zero cumulative likelihood sum");
    }
    // Any numerical failure should cascade into cumulative sum
    if lcum != lcum {
        // inequality due to NaN
        return Err("NaN cumulative likelihood sum");
    }
    Ok((lmin, lcum))
}

/// Update the sample 's' by selectively copying the 'resamples'.
///
/// Uses an in-place copying algorithm:
/// * First copy the live samples (those resampled) to end of 's'.
/// * Replicate live sample in-place starting at the beginning of 's'.
pub fn live_samples<N: RealField, D: Dim>(s: &mut Samples<N, D>, resamples: &Resamples)
where
    DefaultAllocator: Allocator<D>,
{
    // reverse_copy_if live
    let mut si = s.len();
    let mut livei = si;
    for pr in resamples.iter().rev() {
        si -= 1;
        if *pr > 0 {
            livei -= 1;
            s[livei] = s[si].clone();
        }
    }
    assert_eq!(si, 0);

    // Replicate live samples
    si = 0;
    for pr in resamples {
        let mut res = *pr;
        if res > 0 {
            loop {
                s[si] = s[livei].clone();
                si += 1;
                res -= 1;
                if res == 0 {
                    break;
                }
            }
            livei += 1;
        }
    }
    assert_eq!(si, s.len());
    assert_eq!(livei, s.len());
}

/// Roughen sample state using min max roughening of the samples.
///
/// Uses algorithm from Ref \[1] which applies when the states have a single locality but the distribution is unknown.
/// max-min in each state dimension in the samples determines the amount of normally distributed noise added to that
/// dimension for each sample.
///
/// 'k' is scaling factor for the normally distributed noise
///
/// # Numerics:
///  If there are very few unique samples the roughening will collapse as it is not representative of the true state distribution.
pub fn roughen_minmax<N: Copy + FromPrimitive + RealField, D: Dim>(
    s: &mut Samples<N, D>,
    k: f32,
    rng: &mut dyn RngCore,
) where
    DefaultAllocator: Allocator<D, D> + Allocator<D>,
{
    let x_dim = s[0].shape_generic().0;
    let x_size = x_dim.value();

    // Find min and max state dimension in all states
    let mut xmin = s[0].clone();
    let mut xmax = xmin.clone();
    for si in s.iter() {
        // Loop includes 0 to simplify code
        let mut mini = xmin.iter_mut();
        let mut maxi = xmax.iter_mut();

        for xd in si.iter() {
            let minx = mini.next().unwrap();
            let maxx = maxi.next().unwrap();

            if *xd < *minx {
                *minx = *xd;
            }
            if *xd > *maxx {
                *maxx = *xd;
            }
        }
    }
    // Roughening st.dev from scaled max-min and state dimensions
    let sigma_scale = k * f32::pow(s.len() as f32, -1f32 / (x_size as f32));
    let sigma = (xmax - xmin) * N::from_f32(sigma_scale).unwrap();
    let noise = normal_noise_sampler(sigma);

    for si in s.iter_mut() {
        *si += noise(&mut *rng);
    }
}

/// Coupled noise roughening.
///
/// The roughening is scaled by the deviation of the noise.
///
/// 'k' is scaling factor for the normally distributed noise
pub fn roughen_noise<N: Copy + FromPrimitive + RealField, D: Dim, ND: Dim>(
    s: &mut Samples<N, D>,
    noise: &CoupledNoise<N, D, ND>,
    k: f32,
    rng: &mut dyn RngCore,
) where
    DefaultAllocator: Allocator<D, D>
        + Allocator<D, ND>
        + Allocator<ND, ND>
        + Allocator<D>
        + Allocator<ND>,
{
    // Roughening st.dev from scaled  and state dimensions
    let sigma_scale =
        N::from_f32(k * f32::pow(s.len() as f32, -1f32 / (noise.G.nrows() as f32))).unwrap();
    let mut coupling = noise.G.clone();
    for (i, mut c) in coupling.column_iter_mut().enumerate() {
        c *= sigma_scale * noise.q[i];
    }
    let noise = normal_noise_sampler_coupled(coupling);

    for si in s.iter_mut() {
        *si += noise(&mut *rng);
    }
}

impl<N: Copy + FromPrimitive + RealField, D: Dim> SampleState<N, D>
    where
        DefaultAllocator: Allocator<D, D> + Allocator<U1, D> + Allocator<D>,
{
    pub fn from_kalman(state: KalmanState<N, D>, samples: usize, rng: &mut dyn RngCore) -> Result<SampleState<N, D>, &'static str> {
        let coupled_noise = CoupledNoise::from_correlated(&CorrelatedNoise { Q: state.X })?;
        let sampler = normal_noise_sampler_coupled(coupled_noise.G);

        let mut s = Vec::with_capacity(samples);
        for _si in 0..samples {
            let ps = sampler(rng) + &state.x;
            s.push(ps);
        }

        Ok(SampleState::equal_likelihood_samples(s))
    }
}

impl<N: Copy + FromPrimitive + RealField, D: Dim> KalmanEstimator<N, D> for SampleState<N, D>
where
    DefaultAllocator: Allocator<D, D> + Allocator<U1, D> + Allocator<D>,
{
    fn kalman_state<'e>(&self) -> Result<KalmanState<N, D>, &'e str> {
        let x = self.state()?;
        // Covariance of distribution: covariance of particles
        let s_shape = self.s[0].shape_generic();
        let mut xx = OMatrix::zeros_generic(s_shape.0, s_shape.0);

        for s in self.s.iter() {
            let sx = s - &x;
            let sxt = sx.transpose();
            xx += sx * sxt;
        }
        xx /= N::from_usize(self.s.len()).unwrap();

        Ok(KalmanState { x, X: xx })
    }
}

/// Generate as sampling function for normally distributed noise.
///
/// 'std_dev' standard deviation of normally distributed noise.
pub fn normal_noise_sampler<N: Copy + FromPrimitive + RealField, D: Dim>(
    std_dev: OVector<N, D>,
) -> impl Fn(&mut dyn RngCore) -> OVector<N, D>
where
    DefaultAllocator: Allocator<D>,
{
    // Sample with the coupled noise deviation
    move |rng: &mut dyn RngCore| -> OVector<N, D> {
        let rnormal = StandardNormal
            .sample_iter(rng)
            .enumerate()
            .map(|(i, n): (usize, f32)| N::from_f32(n).unwrap() * std_dev[i])
            .take(std_dev.nrows());
        OVector::from_iterator_generic(std_dev.shape_generic().0, Const::<1>, rnormal)
    }
}

/// Generate as sampling function for normally distributed coupled noise.
pub fn normal_noise_sampler_coupled<N: FromPrimitive + RealField, D: Dim, ND: Dim>(
    coupleing: OMatrix<N, D, ND>,
) -> impl Fn(&mut dyn RngCore) -> OVector<N, D>
where
    DefaultAllocator: Allocator<D, ND> + Allocator<D> + Allocator<ND>,
{
    // Sample with the coupled noise deviation
    move |rng: &mut dyn RngCore| -> OVector<N, D> {
        let rnormal = StandardNormal
            .sample_iter(rng)
            .map(|n: f32| N::from_f32(n).unwrap())
            .take(coupleing.ncols());
        let n = OVector::from_iterator_generic(coupleing.shape_generic().1, Const::<1>, rnormal);
        &coupleing * n
    }
}

pub fn gaussian_observation_likelihood<'r, N: Copy + RealField + ToPrimitive, D: Dim, ZD: Dim>(
    z: &'r OVector<N, ZD>,
    h: fn(&OVector<N, D>) -> OVector<N, ZD>,
    noise: &CorrelatedNoise<N, ZD>,
) -> impl Fn(&OVector<N, D>) -> f32 + 'r
where
    DefaultAllocator:
        Allocator<D, D> + Allocator<ZD, ZD> + Allocator<D> + Allocator<ZD>,
{
    // Observation Likelihood for correlated Gaussian noise
    let cholesky = noise.Q.clone().cholesky().unwrap();
    let zinv = cholesky.inverse();
    let zinv_diagonal = cholesky.l_dirty().iter().step_by(noise.Q.nrows() + 1);
    let determinate_zinv = zinv_diagonal
        .fold(N::one(), |prod: N, n: &N| prod * *n)
        .to_f32()
        .unwrap();

    move |x: &OVector<N, D>| -> f32 {
        let innov = z - h(x);
        let logl = innov.dot(&(&zinv * &innov)).to_f32().unwrap();
        (-0.5 * (logl + determinate_zinv.ln())).exp()
    }
}