prodef 0.1.0

A simple Rust crate for handling probability distributions.
Documentation
use crate::{
    Density,
    domain::{Domain, UDomain},
    multinormal::MultiNormalDensity,
    particle::ParticleDensity,
};
use nalgebra::{
    DVector, DefaultAllocator, Dim, Dyn, MatrixView, OMatrix, OVector, RealField, U1, VectorView,
    allocator::Allocator,
};
use rand_distr::{Distribution, StandardNormal, uniform::SampleUniform};
use rayon::prelude::*;
use std::iter::Sum;

impl<T, N, B> ParticleDensity<T, N, B, MultiNormalDensity<T, N, UDomain<N>>>
where
    T: RealField + Sum,
    N: Dim,
    B: 'static + Domain<T, N>,
    StandardNormal: Distribution<T>,
    DefaultAllocator: Allocator<N> + Allocator<U1, N> + Allocator<N, N> + Allocator<N, Dyn>,
{
    /// Create a [`ParticleDensity`] from an existing matrix, with optional domain and weights.
    pub fn from_matrix(
        matrix: OMatrix<T, N, Dyn>,
        domain: B,
        opt_weights: Option<DVector<T>>,
        opt_kernel: Option<MultiNormalDensity<T, N, UDomain<N>>>,
    ) -> Option<Self>
    where
        T: Sum,
    {
        let ndim = matrix.shape_generic().0;

        let kernel = match opt_kernel {
            Some(k) => k,
            None => MultiNormalDensity::from_view(
                &matrix.as_view(),
                UDomain::new(ndim),
                opt_weights.as_ref().map(|w| w.as_slice()),
            )?,
        };

        Some(Self {
            particles: matrix,
            opt_weights,
            domain,
            kernel,
        })
    }

    /// Create a [`ParticleDensity`] from a matrix view, with optional domain and weights.
    pub fn from_view<RStride: Dim, CStride: Dim>(
        view: &MatrixView<T, N, Dyn, RStride, CStride>,
        domain: B,
        opt_weights: Option<&[T]>,
        opt_kernel: Option<MultiNormalDensity<T, N, UDomain<N>>>,
    ) -> Option<Self>
    where
        T: Sum,
    {
        let ndim = view.shape_generic().0;

        let kernel: MultiNormalDensity<T, N, UDomain<N>> = match opt_kernel {
            Some(k) => k,
            None => {
                let mut mvnk =
                    MultiNormalDensity::from_view(view, UDomain::new(ndim), opt_weights)?;

                mvnk.set_mean(OVector::from_iterator_generic(
                    ndim,
                    U1,
                    (0..ndim.value()).map(|_| T::zero()),
                ));

                mvnk
            }
        };

        let count = view.ncols();

        Some(Self {
            particles: view.clone_owned(),
            opt_weights: opt_weights.map(|w| DVector::from_iterator(count, w.iter().cloned())),
            domain,
            kernel,
        })
    }

    /// Compute non-normalized transition weights for a new particle matrix.
    pub fn transition_weights(&self, new_particles: &OMatrix<T, N, Dyn>) -> Vec<T>
    where
        <DefaultAllocator as Allocator<N>>::Buffer<T>: Sync,
        <DefaultAllocator as Allocator<N, N>>::Buffer<T>: Sync,
        <DefaultAllocator as Allocator<N, Dyn>>::Buffer<T>: Sync,
    {
        match self.opt_weights {
            Some(ref weights) => new_particles
                .par_column_iter()
                .map(|params| {
                    T::one()
                        / self
                            .particles
                            .column_iter()
                            .zip(weights.iter())
                            .map(|(params_old, weight_old)| {
                                let delta = params.clone() - params_old;

                                (weight_old.clone().ln()
                                    - self.kernel.mahalanobis_distance_sq(&delta.as_view()))
                                .exp()
                            })
                            .sum::<T>()
                })
                .collect::<Vec<T>>(),
            None => new_particles
                .par_column_iter()
                .map(|params| {
                    T::one()
                        / self
                            .particles
                            .column_iter()
                            .map(|params_old| {
                                let delta = params.clone() - params_old;

                                (-self.kernel.mahalanobis_distance_sq(&delta.as_view())).exp()
                            })
                            .sum::<T>()
                })
                .collect::<Vec<T>>(),
        }
    }
}

impl<T, N, B> Density<T, N> for ParticleDensity<T, N, B, MultiNormalDensity<T, N, UDomain<N>>>
where
    T: RealField + SampleUniform + Sum,
    N: Dim,
    B: 'static + Domain<T, N>,
    DefaultAllocator: Allocator<N> + Allocator<U1, N> + Allocator<N, N> + Allocator<N, Dyn>,
    StandardNormal: Distribution<T>,
{
    fn density<RStride: Dim, CStride: Dim>(
        &self,
        sample: &VectorView<T, N, RStride, CStride>,
    ) -> Option<T> {
        (&self).density(sample)
    }

    fn domain(&self) -> impl Domain<T, N> + 'static {
        self.domain.clone()
    }

    fn center(&self) -> OVector<T, N> {
        (&self).center()
    }

    fn is_constant(&self) -> OVector<bool, N> {
        (&self).is_constant()
    }

    fn sample(&self, rng: &mut impl rand::Rng, max_attempts: usize) -> Option<OVector<T, N>> {
        (&self).sample(rng, max_attempts)
    }
}

impl<T, N, B> Density<T, N> for &ParticleDensity<T, N, B, MultiNormalDensity<T, N, UDomain<N>>>
where
    T: RealField + SampleUniform + Sum,
    N: Dim,
    B: 'static + Domain<T, N>,
    DefaultAllocator: Allocator<N> + Allocator<U1, N> + Allocator<N, N> + Allocator<N, Dyn>,
    StandardNormal: Distribution<T>,
{
    fn density<RStride: Dim, CStride: Dim>(
        &self,
        sample: &VectorView<T, N, RStride, CStride>,
    ) -> Option<T> {
        if !self.domain.contains(sample) {
            return None;
        }

        let bandwidth = T::from_usize(self.len())
            .unwrap()
            .powf(-T::one() / (T::from_usize(4 + self.kernel.rank()).unwrap()));

        dbg!(&bandwidth);

        let normalization = self.kernel.normalization_factor();

        Some(
            normalization
                * match &self.opt_weights {
                    Some(weights) => self
                        .particles
                        .column_iter()
                        .zip(weights.iter())
                        .map(|(col, weight)| {
                            let delta_x = sample - col;

                            weight.clone()
                                * (-self.kernel.mahalanobis_distance_sq(&delta_x.as_view())
                                    / T::from_usize(2).unwrap())
                                .exp()
                        })
                        .sum::<T>(),

                    None => self
                        .particles
                        .column_iter()
                        .map(|col| {
                            let delta_x = sample - col;

                            (-self.kernel.mahalanobis_distance_sq(&delta_x.as_view())
                                / T::from_usize(2).unwrap())
                            .exp()
                        })
                        .sum::<T>(),
                },
        )
    }

    fn domain(&self) -> impl Domain<T, N> + 'static {
        self.domain.clone()
    }

    fn center(&self) -> OVector<T, N> {
        self.particles.column_mean()
    }

    fn is_constant(&self) -> OVector<bool, N> {
        OVector::<bool, N>::from_iterator_generic(
            self.particles.shape_generic().0,
            U1,
            self.particles.row_iter().map(|row| {
                let first = row[0].clone();
                row.iter().all(|x| *x == first)
            }),
        )
    }

    fn sample(&self, rng: &mut impl rand::Rng, max_attempts: usize) -> Option<OVector<T, N>> {
        let particle = self.sample_particle(rng);
        let mut candidate =
            &particle + self.kernel.sample(rng, max_attempts)? - self.kernel.center();

        while !self.domain.contains(&candidate.as_view()) {
            candidate = &particle + self.kernel.sample(rng, max_attempts)? - self.kernel.center();
        }

        Some(candidate)
    }
}

#[cfg(test)]
mod tests {
    use crate::domain::{MDomain, SDomain};

    use super::*;
    use approx::ulps_eq;
    use nalgebra::{Matrix, SVector, U2, U3, VecStorage};
    use rand::{Rng, SeedableRng};
    use rand_xoshiro::Xoshiro256PlusPlus;

    #[test]
    fn test_particle_density() {
        let mut rng = Xoshiro256PlusPlus::seed_from_u64(1);
        let uniform = StandardNormal;

        let array_0 = Matrix::<f64, U2, Dyn, VecStorage<f64, U2, Dyn>>::from_iterator(
            10000,
            (0..20000).map(|idx| {
                if idx % 2 == 0 {
                    0.1 + rng.sample::<f64, StandardNormal>(uniform)
                } else {
                    0.25 + rng.sample::<f64, StandardNormal>(uniform)
                }
            }),
        );

        let mvpdf_0 = MultiNormalDensity::from_view::<Dyn, U2>(
            &array_0.as_view(),
            MDomain::new(SVector::from([
                SDomain::new(Some(-0.75), Some(0.75)).unwrap(),
                SDomain::new(Some(-0.75), Some(0.75)).unwrap(),
            ])),
            None,
        )
        .unwrap();

        let ptpdf_0 = ParticleDensity::from_view::<U1, U2>(
            &array_0.as_view(),
            MDomain::new(SVector::from([
                SDomain::new(Some(-0.75), Some(0.75)).unwrap(),
                SDomain::new(Some(-0.75), Some(0.75)).unwrap(),
            ])),
            None,
            None,
        )
        .unwrap();

        assert!(ulps_eq!(
            mvpdf_0
                .density::<U1, U2>(&SVector::from([0.2, 0.35]).as_view())
                .unwrap(),
            0.161284,
            epsilon = 1e-5,
            max_ulps = 5
        ));

        assert!(ulps_eq!(
            ptpdf_0
                .density::<U1, U2>(&SVector::from([0.2, 0.35]).as_view())
                .unwrap(),
            0.155147,
            epsilon = 1e-5,
            max_ulps = 5
        ));

        let mut rng = Xoshiro256PlusPlus::seed_from_u64(1);

        let array = Matrix::<f64, U3, Dyn, VecStorage<f64, U3, Dyn>>::from_iterator(
            10000,
            (0..30000).map(|idx| {
                if idx % 3 == 0 {
                    0.1 + rng.sample::<f64, StandardNormal>(uniform)
                } else if idx % 3 == 1 {
                    0.0
                } else {
                    0.25 + rng.sample::<f64, StandardNormal>(uniform)
                }
            }),
        );

        let ptpdf = ParticleDensity::from_view::<U1, U3>(
            &array.as_view(),
            MDomain::new(SVector::from([
                SDomain::new(Some(-0.75), Some(0.75)).unwrap(),
                SDomain::new(Some(-0.75), Some(0.75)).unwrap(),
                SDomain::new(Some(-0.75), Some(0.75)).unwrap(),
            ])),
            None,
            None,
        )
        .unwrap();

        assert!(
            ptpdf
                .density::<U1, U3>(&SVector::from([0.2, -0.85, 0.35]).as_view())
                .is_none()
        );

        assert!(ulps_eq!(
            ptpdf
                .density::<U1, U3>(&SVector::from([0.2, 0.0, 0.35]).as_view())
                .unwrap(),
            0.155147,
            epsilon = 1e-5,
            max_ulps = 5
        ));

        assert!(ulps_eq!(
            ptpdf.sample(&mut rng, 100).unwrap(),
            SVector::from([-0.6191990262793441, 0.0, 0.6201433837562484,]),
            epsilon = 1e-5,
            max_ulps = 5
        ));
    }
}