prodef 0.1.0

//! A module that implements multiple univariate and a joint probability density functions.

mod constant;
mod cosine;
mod loguniform;
mod normal;
mod uniform;

pub use constant::*;
pub use cosine::*;
pub use loguniform::*;
pub use normal::*;
pub use uniform::*;

use crate::{
    Density, Domain,
    domain::{MDomain, SDomain},
};
use derive_more::{Deref, DerefMut, IntoIterator};
use nalgebra::{
    DefaultAllocator, Dim, OVector, RealField, SVector, U1, VectorView, allocator::Allocator,
};
use rand::Rng;
use rand_distr::{Distribution, StandardNormal, uniform::SampleUniform};
use serde::{Deserialize, Serialize};
use std::{f64, fmt::Debug};

/// An algebraic data type (ADT) that contains all univariate probability density functions.
#[allow(missing_docs)]
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(tag = "type", content = "content")]
pub enum UnivariateDensity<T>
where
    T: RealField,
{
    Constant(ConstantDensity<T>),
    Cosine(CosineDensity<T>),
    Loguniform(LogUniformDensity<T>),
    Normal(NormalDensity<T>),
    Uniform(UniformDensity<T>),
}

impl<T> From<ConstantDensity<T>> for UnivariateDensity<T>
where
    T: RealField,
{
    fn from(value: ConstantDensity<T>) -> Self {
        Self::Constant(value)
    }
}

impl<T> From<CosineDensity<T>> for UnivariateDensity<T>
where
    T: RealField,
{
    fn from(value: CosineDensity<T>) -> Self {
        Self::Cosine(value)
    }
}

impl<T> From<LogUniformDensity<T>> for UnivariateDensity<T>
where
    T: RealField,
{
    fn from(value: LogUniformDensity<T>) -> Self {
        Self::Loguniform(value)
    }
}

impl<T> From<NormalDensity<T>> for UnivariateDensity<T>
where
    T: RealField,
{
    fn from(value: NormalDensity<T>) -> Self {
        Self::Normal(value)
    }
}

impl<T> From<UniformDensity<T>> for UnivariateDensity<T>
where
    T: RealField,
{
    fn from(value: UniformDensity<T>) -> Self {
        Self::Uniform(value)
    }
}

impl<T> UnivariateDensity<T>
where
    T: RealField,
{
    /// Returns a reference to the inner [`ConstantDensity`] if this is a `UnivariateDensity::Constant`, `None` otherwise.
    pub fn as_constant(&self) -> Option<&ConstantDensity<T>> {
        match self {
            UnivariateDensity::Constant(pdf) => Some(pdf),
            _ => None,
        }
    }

    /// Returns a reference to the inner [`CosineDensity`] if this is a `UnivariateDensity::Cosine`, `None` otherwise.
    pub fn as_cosine(&self) -> Option<&CosineDensity<T>> {
        match self {
            UnivariateDensity::Cosine(pdf) => Some(pdf),
            _ => None,
        }
    }

    /// Returns a reference to the inner [`LogUniformDensity`] if this is a `UnivariateDensity::Loguniform`, `None` otherwise.
    pub fn as_loguniform(&self) -> Option<&LogUniformDensity<T>> {
        match self {
            UnivariateDensity::Loguniform(pdf) => Some(pdf),
            _ => None,
        }
    }

    /// Returns a reference to the inner [`NormalDensity`] if this is a `UnivariateDensity::Normal`, `None` otherwise.
    pub fn as_normal(&self) -> Option<&NormalDensity<T>> {
        match self {
            UnivariateDensity::Normal(pdf) => Some(pdf),
            _ => None,
        }
    }

    /// Returns a reference to the inner [`UniformDensity`] if this is a `UnivariateDensity::Uniform`, `None` otherwise.
    pub fn as_uniform(&self) -> Option<&UniformDensity<T>> {
        match self {
            UnivariateDensity::Uniform(pdf) => Some(pdf),
            _ => None,
        }
    }
}

/// A joint probability density function composed of `D` univariate PDFs.
///
/// ```
/// # use nalgebra::{Const, SVector};
/// # use prodef::{SDomain, univariate::{ConstantDensity, CosineDensity, LogUniformDensity, MultivariateDensity, NormalDensity, UniformDensity}};
/// # let mvpdf = &MultivariateDensity::<f64, Const<5>>::new(SVector::from([
/// #    ConstantDensity::new(1.0).into(),
/// #    CosineDensity::new(0.1, 0.2).unwrap().into(),
/// #    LogUniformDensity::new(0.1, 0.5).unwrap().into(),
/// #    NormalDensity::new(0.1, 0.25, Some(-0.5), Some(1.5)).unwrap().into(),
/// #    UniformDensity::new(1.0, 2.0).unwrap().into(),
/// # ]));
/// ```
#[derive(Clone, Debug, Deref, DerefMut, Deserialize, IntoIterator, Serialize)]
#[serde(bound(serialize = "OVector<UnivariateDensity<T>, D>: Serialize"))]
#[serde(bound(deserialize = "OVector<UnivariateDensity<T>, D>: Deserialize<'de>"))]
pub struct MultivariateDensity<T, D>(
    #[into_iterator(owned, ref, ref_mut)] OVector<UnivariateDensity<T>, D>,
)
where
    T: RealField,
    D: Dim,
    DefaultAllocator: Allocator<D>;

impl<T, D> MultivariateDensity<T, D>
where
    T: RealField,
    D: Dim,
    DefaultAllocator: Allocator<D>,
{
    /// Create a new [`MDomain`] from a vector of [`UnivariateDensity`]s.
    pub fn new(domains: OVector<UnivariateDensity<T>, D>) -> Self {
        Self(domains)
    }
}

impl<T, D> Density<T, D> for MultivariateDensity<T, D>
where
    T: RealField + SampleUniform + Send + Sync,
    D: Dim,
    MDomain<T, D>: Domain<T, D>,
    for<'a> &'a MDomain<T, D>: Domain<T, D>,
    StandardNormal: Distribution<T>,
    DefaultAllocator: Allocator<D>,
{
    fn density<RStride: Dim, CStride: Dim>(
        &self,
        sample: &VectorView<T, D, RStride, CStride>,
    ) -> Option<T> {
        (&self).density(sample)
    }

    fn domain(&self) -> impl Domain<T, D> + 'static {
        MDomain::new(OVector::from_iterator_generic(
            self.0.shape_generic().0,
            U1,
            self.0.iter().map(|uvpdf| {
                let (a, b) = match uvpdf {
                    UnivariateDensity::Constant(pdf) => (
                        Domain::<T, U1>::minimum_values(&pdf.domain())
                            .map(|value| value[0].clone()),
                        Domain::<T, U1>::maximum_values(&pdf.domain())
                            .map(|value| value[0].clone()),
                    ),
                    UnivariateDensity::Cosine(pdf) => (
                        Domain::<T, U1>::minimum_values(&pdf.domain())
                            .map(|value| value[0].clone()),
                        Domain::<T, U1>::maximum_values(&pdf.domain())
                            .map(|value| value[0].clone()),
                    ),
                    UnivariateDensity::Loguniform(pdf) => (
                        Domain::<T, U1>::minimum_values(&pdf.domain())
                            .map(|value| value[0].clone()),
                        Domain::<T, U1>::maximum_values(&pdf.domain())
                            .map(|value| value[0].clone()),
                    ),
                    UnivariateDensity::Normal(pdf) => (
                        Domain::<T, U1>::minimum_values(&pdf.domain())
                            .map(|value| value[0].clone()),
                        Domain::<T, U1>::maximum_values(&pdf.domain())
                            .map(|value| value[0].clone()),
                    ),
                    UnivariateDensity::Uniform(pdf) => (
                        Domain::<T, U1>::minimum_values(&pdf.domain())
                            .map(|value| value[0].clone()),
                        Domain::<T, U1>::maximum_values(&pdf.domain())
                            .map(|value| value[0].clone()),
                    ),
                };

                SDomain::new(a, b).unwrap()
            }),
        ))
    }

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

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

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

impl<T, D> Density<T, D> for &MultivariateDensity<T, D>
where
    T: RealField + SampleUniform + Send + Sync,
    D: Dim,
    MDomain<T, D>: Domain<T, D>,
    for<'a> &'a MDomain<T, D>: Domain<T, D>,
    StandardNormal: Distribution<T>,
    DefaultAllocator: Allocator<D>,
{
    fn density<RStride: Dim, CStride: Dim>(
        &self,
        sample: &VectorView<T, D, RStride, CStride>,
    ) -> Option<T> {
        if !self.domain().contains(sample) {
            return None;
        }

        let mut rlh = T::one();

        self.0.iter().zip(sample.iter()).for_each(|(uvpdf, value)| {
            let vec = SVector::from([value.clone()]);

            rlh *= match uvpdf {
                UnivariateDensity::Constant(pdf) => {
                    Density::<T, U1>::density::<U1, U1>(&pdf, &vec.as_view())
                }
                UnivariateDensity::Cosine(pdf) => {
                    Density::<T, U1>::density::<U1, U1>(&pdf, &vec.as_view())
                }
                UnivariateDensity::Loguniform(pdf) => {
                    Density::<T, U1>::density::<U1, U1>(&pdf, &vec.as_view())
                }
                UnivariateDensity::Normal(pdf) => {
                    Density::<T, U1>::density::<U1, U1>(&pdf, &vec.as_view())
                }
                UnivariateDensity::Uniform(pdf) => {
                    Density::<T, U1>::density::<U1, U1>(&pdf, &vec.as_view())
                }
            }
            .unwrap_or(T::from_f64(f64::NAN).unwrap());
        });

        Some(rlh)
    }

    fn domain(&self) -> impl Domain<T, D> + 'static {
        MDomain::new(OVector::from_iterator_generic(
            self.0.shape_generic().0,
            U1,
            self.0.iter().map(|uvpdf| {
                let (a, b) = match uvpdf {
                    UnivariateDensity::Constant(pdf) => (
                        Domain::<T, U1>::minimum_values(&pdf.domain())
                            .map(|value| value[0].clone()),
                        Domain::<T, U1>::maximum_values(&pdf.domain())
                            .map(|value| value[0].clone()),
                    ),
                    UnivariateDensity::Cosine(pdf) => (
                        Domain::<T, U1>::minimum_values(&pdf.domain())
                            .map(|value| value[0].clone()),
                        Domain::<T, U1>::maximum_values(&pdf.domain())
                            .map(|value| value[0].clone()),
                    ),
                    UnivariateDensity::Loguniform(pdf) => (
                        Domain::<T, U1>::minimum_values(&pdf.domain())
                            .map(|value| value[0].clone()),
                        Domain::<T, U1>::maximum_values(&pdf.domain())
                            .map(|value| value[0].clone()),
                    ),
                    UnivariateDensity::Normal(pdf) => (
                        Domain::<T, U1>::minimum_values(&pdf.domain())
                            .map(|value| value[0].clone()),
                        Domain::<T, U1>::maximum_values(&pdf.domain())
                            .map(|value| value[0].clone()),
                    ),
                    UnivariateDensity::Uniform(pdf) => (
                        Domain::<T, U1>::minimum_values(&pdf.domain())
                            .map(|value| value[0].clone()),
                        Domain::<T, U1>::maximum_values(&pdf.domain())
                            .map(|value| value[0].clone()),
                    ),
                };

                SDomain::new(a, b).unwrap()
            }),
        ))
    }

    fn center(&self) -> OVector<T, D> {
        OVector::<T, D>::from_iterator_generic(
            self.0.shape_generic().0,
            U1,
            self.0.iter().map(|uvpdf| match uvpdf {
                UnivariateDensity::Constant(pdf) => Density::<T, U1>::center(&pdf)[0].clone(),
                UnivariateDensity::Cosine(pdf) => Density::<T, U1>::center(&pdf)[0].clone(),
                UnivariateDensity::Loguniform(pdf) => Density::<T, U1>::center(&pdf)[0].clone(),
                UnivariateDensity::Normal(pdf) => Density::<T, U1>::center(&pdf)[0].clone(),
                UnivariateDensity::Uniform(pdf) => Density::<T, U1>::center(&pdf)[0].clone(),
            }),
        )
    }

    fn is_constant(&self) -> OVector<bool, D> {
        OVector::<bool, D>::from_iterator_generic(
            self.0.shape_generic().0,
            U1,
            self.0.iter().map(|uvpdf| match uvpdf {
                UnivariateDensity::Constant(pdf) => Density::<T, U1>::is_constant(&pdf)[0],
                UnivariateDensity::Cosine(pdf) => Density::<T, U1>::is_constant(&pdf)[0],
                UnivariateDensity::Loguniform(pdf) => Density::<T, U1>::is_constant(&pdf)[0],
                UnivariateDensity::Normal(pdf) => Density::<T, U1>::is_constant(&pdf)[0],
                UnivariateDensity::Uniform(pdf) => Density::<T, U1>::is_constant(&pdf)[0],
            }),
        )
    }

    fn sample(&self, rng: &mut impl Rng, max_attempts: usize) -> Option<OVector<T, D>> {
        let mut draw = OVector::<T, D>::zeros_generic(self.0.shape_generic().0, U1);

        for i in 0..self.0.shape_generic().0.value() {
            match &self.0[i] {
                UnivariateDensity::Constant(pdf) => {
                    draw[i] = match Density::<T, U1>::sample(pdf, rng, max_attempts) {
                        Some(draw) => draw[0].clone(),
                        None => return None,
                    };
                }
                UnivariateDensity::Cosine(pdf) => {
                    draw[i] = match Density::<T, U1>::sample(pdf, rng, max_attempts) {
                        Some(draw) => draw[0].clone(),
                        None => return None,
                    };
                }
                UnivariateDensity::Loguniform(pdf) => {
                    draw[i] = match Density::<T, U1>::sample(pdf, rng, max_attempts) {
                        Some(draw) => draw[0].clone(),
                        None => return None,
                    };
                }
                UnivariateDensity::Normal(pdf) => {
                    draw[i] = match Density::<T, U1>::sample(pdf, rng, max_attempts) {
                        Some(draw) => draw[0].clone(),
                        None => return None,
                    };
                }
                UnivariateDensity::Uniform(pdf) => {
                    draw[i] = match Density::<T, U1>::sample(pdf, rng, max_attempts) {
                        Some(draw) => draw[0].clone(),
                        None => return None,
                    };
                }
            }
        }

        Some(draw)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use approx::ulps_eq;
    use nalgebra::{OVector, SVector, U5};
    use rand::SeedableRng;
    use rand_xoshiro::Xoshiro256PlusPlus;

    #[test]
    fn test_multivariate_density() {
        let const_a = ConstantDensity::new(1.0);
        let const_b = ConstantDensity::new(2.0);

        assert!(ulps_eq!(
            const_a
                .density::<U1, U1>(&OVector::from([1.0]).as_view())
                .unwrap(),
            const_b
                .density::<U1, U1>(&OVector::from([2.0]).as_view())
                .unwrap()
        ));

        assert!(
            const_a
                .density::<U1, U1>(&OVector::from([1.0]).as_view())
                .unwrap()
                == const_a
                    .density::<U1, U1>(&OVector::from([1.0]).as_view())
                    .unwrap()
        );

        let cosine = CosineDensity::new(-0.1, 0.35).unwrap();

        assert!(ulps_eq!(
            cosine
                .density::<U1, U1>(&OVector::from([0.0]).as_view())
                .unwrap(),
            1.0
        ));

        assert!(ulps_eq!(
            cosine
                .density::<U1, U1>(&OVector::from([0.3]).as_view())
                .unwrap(),
            0.955336489125606
        ));

        assert!(ulps_eq!(
            cosine
                .density::<U1, U1>(&OVector::from([0.3]).as_view())
                .unwrap(),
            0.955336489125606
        ));

        let loguniform = LogUniformDensity::new(0.5, 2.5).unwrap();

        assert!(ulps_eq!(
            loguniform
                .density::<U1, U1>(&OVector::from([1.0]).as_view())
                .unwrap(),
            2.0 * loguniform
                .density::<U1, U1>(&OVector::from([2.0]).as_view())
                .unwrap()
        ));

        let normal = NormalDensity::new(1.0, 1.0, None, None).unwrap();

        assert!(ulps_eq!(
            normal
                .density::<U1, U1>(&OVector::from([1.0]).as_view())
                .unwrap(),
            0.3989422804014327
        ));

        assert!(ulps_eq!(
            normal
                .density::<U1, U1>(&OVector::from([2.0]).as_view())
                .unwrap(),
            0.24197072451914337
        ));

        let uniform = UniformDensity::new(0.0, 1.0).unwrap();

        assert!(ulps_eq!(
            uniform
                .density::<U1, U1>(&OVector::from([0.5]).as_view())
                .unwrap(),
            1.0
        ));

        let uvpdf = &MultivariateDensity::new(SVector::from([
            ConstantDensity::new(1.0).into(),
            CosineDensity::new(0.1, 0.2).unwrap().into(),
            LogUniformDensity::new(0.1, 0.5).unwrap().into(),
            NormalDensity::new(0.1, 0.25, Some(-0.5), Some(1.5))
                .unwrap()
                .into(),
            UniformDensity::new(1.0, 2.0).unwrap().into(),
        ]));

        assert!(ulps_eq!(
            uvpdf
                .density::<U1, U5>(&SVector::from([1.0f64, 0.15, 0.15, 0.2, 1.5]).as_view())
                .unwrap(),
            6.033325,
            epsilon = 1e-5,
            max_ulps = 5
        ));

        assert!(
            uvpdf
                .density::<U1, U5>(&SVector::from([1.0, 0.05, 0.2, 0.15, 1.5]).as_view())
                .is_none()
        );

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

        assert!(ulps_eq!(
            uvpdf.sample(&mut rng, 100).unwrap(),
            OVector::from([1.0, 0.1810371, 0.33281568, -0.37896788, 1.7462168,]),
            epsilon = 1e-5,
            max_ulps = 5
        ));

        assert!(
            uvpdf
                .domain()
                .contains::<U1, U5>(&uvpdf.sample(&mut rng, 100).unwrap().as_view())
        );
    }
}