diff --git a/Cargo.toml b/Cargo.toml index 96a52797..d9aace77 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,3 +19,4 @@ path = "src/lib.rs" [dependencies] rand = "0.7" +nalgebra = "0.19" \ No newline at end of file diff --git a/src/distribution/mod.rs b/src/distribution/mod.rs index 8eb7acc9..0c30279e 100644 --- a/src/distribution/mod.rs +++ b/src/distribution/mod.rs @@ -20,6 +20,7 @@ pub use self::hypergeometric::Hypergeometric; pub use self::inverse_gamma::InverseGamma; pub use self::log_normal::LogNormal; pub use self::multinomial::Multinomial; +pub use self::multivariate_normal::MultivariateNormal; pub use self::normal::Normal; pub use self::pareto::Pareto; pub use self::poisson::Poisson; @@ -48,6 +49,7 @@ mod internal; mod inverse_gamma; mod log_normal; mod multinomial; +mod multivariate_normal; mod normal; mod pareto; mod poisson; diff --git a/src/distribution/multivariate_normal.rs b/src/distribution/multivariate_normal.rs new file mode 100644 index 00000000..80b4e812 --- /dev/null +++ b/src/distribution/multivariate_normal.rs @@ -0,0 +1,411 @@ +use crate::distribution::Continuous; +use crate::distribution::Normal; +use crate::statistics::{Covariance, Entropy, Max, Mean, Min, Mode}; +use crate::{Result, StatsError}; +use nalgebra::{ + base::allocator::Allocator, + base::{dimension::DimName, MatrixN, VectorN}, + Cholesky, DefaultAllocator, Dim, DimMin, LU, U1, +}; +use rand::distributions::Distribution; +use rand::Rng; +use std::f64; +use std::f64::consts::{E, PI}; + +/// Implements the [Multivariate Normal](https://en.wikipedia.org/wiki/Multivariate_normal_distribution) +/// distribution using the "nalgebra" crate for matrix operations +/// +/// # Examples +/// +/// ``` +/// use statrs::distribution::{MultivariateNormal, Continuous}; +/// use nalgebra::base::dimension::U2; +/// use nalgebra::{Vector2, Matrix2}; +/// use statrs::statistics::{Mean, Covariance}; +/// +/// let mvn = MultivariateNormal::::new(&Vector2::zeros(), &Matrix2::identity()).unwrap(); +/// assert_eq!(mvn.mean(), Vector2::new(0., 0.)); +/// assert_eq!(mvn.variance(), Matrix2::new(1., 0., 0., 1.)); +/// assert_eq!(mvn.pdf(Vector2::new(1., 1.)), 0.05854983152431917); +/// ``` +#[derive(Debug, Clone)] +pub struct MultivariateNormal +where + N: Dim + DimMin + DimName, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator<(usize, usize), >::Output>, +{ + cov_chol_decomp: MatrixN, + mu: VectorN, + cov: MatrixN, + precision: MatrixN, + pdf_const: f64, +} + +impl MultivariateNormal +where + N: Dim + DimMin + DimName, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator<(usize, usize), >::Output>, +{ + /// Constructs a new multivariate normal distribution with a mean of `mean` + /// and covariance matrix `cov` + /// + /// # Errors + /// + /// Returns an error if the given covariance matrix is not + /// symmetric or positive-definite + pub fn new(mean: &VectorN, cov: &MatrixN) -> Result { + // Check that the provided covariance matrix is symmetric + // Check that mean and covariance do not contain NaN + if cov.lower_triangle() != cov.upper_triangle().transpose() + || mean.iter().any(|f| f.is_nan()) + || cov.iter().any(|f| f.is_nan()) + { + return Err(StatsError::BadParams); + } + let cov_det = LU::new(cov.clone()).determinant(); + let pdf_const = ((2. * PI).powi(mean.nrows() as i32) * cov_det.abs()) + .recip() + .sqrt(); + // Store the Cholesky decomposition of the covariance matrix + // for sampling + match Cholesky::new(cov.clone()) { + None => Err(StatsError::BadParams), + Some(cholesky_decomp) => Ok(MultivariateNormal { + cov_chol_decomp: cholesky_decomp.clone().unpack(), + mu: mean.clone(), + cov: cov.clone(), + precision: cholesky_decomp.inverse(), + pdf_const: pdf_const, + }), + } + } +} + +impl Distribution> for MultivariateNormal +where + N: Dim + DimMin + DimName, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator<(usize, usize), >::Output>, +{ + /// Samples from the multivariate normal distribution + /// + /// # Formula + /// L * Z + μ + /// + /// where `L` is the Cholesky decomposition of the covariance matrix, + /// `Z` is a vector of normally distributed random variables, and + /// `μ` is the mean vector + + fn sample(&self, rng: &mut R) -> VectorN { + let d = Normal::new(0., 1.).unwrap(); + let z = VectorN::::from_distribution(&d, rng); + (self.cov_chol_decomp.clone() * z) + self.mu.clone() + } +} + +impl Min> for MultivariateNormal +where + N: Dim + DimMin + DimName, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator<(usize, usize), >::Output>, +{ + /// Returns the minimum value in the domain of the + /// multivariate normal distribution represented by a real vector + fn min(&self) -> VectorN { + VectorN::::repeat(f64::NEG_INFINITY) + } +} + +impl Max> for MultivariateNormal +where + N: Dim + DimMin + DimName, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator<(usize, usize), >::Output>, +{ + /// Returns the maximum value in the domain of the + /// multivariate normal distribution represented by a real vector + fn max(&self) -> VectorN { + VectorN::::repeat(f64::INFINITY) + } +} + +impl Mean> for MultivariateNormal +where + N: Dim + DimMin + DimName, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator<(usize, usize), >::Output>, +{ + /// Returns the mean of the normal distribution + /// + /// # Remarks + /// + /// This is the same mean used to construct the distribution + fn mean(&self) -> VectorN { + self.mu.clone() + } +} + +impl Covariance> for MultivariateNormal +where + N: Dim + DimMin + DimName, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator<(usize, usize), >::Output>, +{ + /// Returns the covariance matrix of the multivariate normal distribution + fn variance(&self) -> MatrixN { + self.cov.clone() + } +} + +impl Entropy for MultivariateNormal +where + N: Dim + DimMin + DimName, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator<(usize, usize), >::Output>, +{ + /// Returns the entropy of the multivariate normal distribution + /// + /// # Formula + /// + /// ```ignore + /// (1 / 2) * ln(det(2 * π * e * Σ)) + /// ``` + /// + /// where `Σ` is the covariance matrix and `det` is the determinant + fn entropy(&self) -> f64 { + 0.5 * LU::new(self.variance().clone().scale(2. * PI * E)) + .determinant() + .ln() + } +} + +impl Mode> for MultivariateNormal +where + N: Dim + DimMin + DimName, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator<(usize, usize), >::Output>, +{ + /// Returns the mode of the multivariate normal distribution + /// + /// # Formula + /// + /// ```ignore + /// μ + /// ``` + /// + /// where `μ` is the mean + fn mode(&self) -> VectorN { + self.mu.clone() + } +} + +impl Continuous, f64> for MultivariateNormal +where + N: Dim + DimMin + DimName, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator<(usize, usize), >::Output>, +{ + /// Calculates the probability density function for the multivariate + /// normal distribution at `x` + /// + /// # Formula + /// + /// ```ignore + /// (2 * π) ^ (-k / 2) * det(Σ) ^ (1 / 2) * e ^ ( -(1 / 2) * transpose(x - μ) * inv(Σ) * (x - μ)) + /// ``` + /// + /// where `μ` is the mean, `inv(Σ)` is the precision matrix, `det(Σ)` is the determinant + /// of the covariance matrix, and `k` is the dimension of the distribution + fn pdf(&self, x: VectorN) -> f64 { + let dv = x - &self.mu; + let exp_term = -0.5 + * *(&dv.transpose() * &self.precision * &dv) + .get((0, 0)) + .unwrap(); + self.pdf_const * exp_term.exp() + } + /// Calculates the log probability density function for the multivariate + /// normal distribution at `x`. Equivalent to pdf(x).ln(). + fn ln_pdf(&self, x: VectorN) -> f64 { + let dv = x - &self.mu; + let exp_term = -0.5 + * *(&dv.transpose() * &self.precision * &dv) + .get((0, 0)) + .unwrap(); + self.pdf_const.ln() + exp_term + } +} + +#[cfg_attr(rustfmt, rustfmt_skip)] +#[cfg(test)] +mod test { + use std::f64; + use crate::statistics::*; + use crate::distribution::{MultivariateNormal, Continuous}; + use nalgebra::{Matrix2, Vector2, Matrix3, Vector3, VectorN, MatrixN, Dim, DimMin, DimName, DefaultAllocator, U1}; + use nalgebra::base::allocator::Allocator; + use core::fmt::Debug; + + fn try_create(mean: VectorN, covariance: MatrixN) -> MultivariateNormal + where + N: Dim + DimMin + DimName, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator<(usize, usize), >::Output>, + { + let mvn = MultivariateNormal::new(&mean, &covariance); + assert!(mvn.is_ok()); + mvn.unwrap() + } + + fn create_case(mean: VectorN, covariance: MatrixN) + where + N: Dim + DimMin + DimName, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator<(usize, usize), >::Output>, + { + let mvn = try_create(mean.clone(), covariance.clone()); + assert_eq!(mean, mvn.mean()); + assert_eq!(covariance, mvn.variance()); + } + + fn bad_create_case(mean: VectorN, covariance: MatrixN) + where + N: Dim + DimMin + DimName, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator<(usize, usize), >::Output>, + { + let mvn = MultivariateNormal::new(&mean, &covariance); + assert!(mvn.is_err()); + } + + fn test_case(mean: VectorN, covariance: MatrixN, expected: T, eval: F) + where + T: Debug + PartialEq, + F: Fn(MultivariateNormal) -> T, + N: Dim + DimMin + DimName, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator<(usize, usize), >::Output>, + { + let mvn = try_create(mean, covariance); + let x = eval(mvn); + assert_eq!(expected, x); + } + + fn test_almost(mean: VectorN, covariance: MatrixN, expected: f64, acc: f64, eval: F) + where + F: Fn(MultivariateNormal) -> f64, + N: Dim + DimMin + DimName, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator, + DefaultAllocator: Allocator<(usize, usize), >::Output>, + { + let mvn = try_create(mean, covariance); + let x = eval(mvn); + assert_almost_eq!(expected, x, acc); + } + + #[test] + fn test_create() { + create_case(Vector2::new(0., 0.), Matrix2::new(1., 0., 0., 1.)); + create_case(Vector2::new(10., 5.), Matrix2::new(2., 1., 1., 2.)); + create_case(Vector3::new(4., 5., 6.), Matrix3::new(2., 1., 0., 1., 2., 1., 0., 1., 2.)); + create_case(Vector2::new(0., f64::INFINITY), Matrix2::identity()); + create_case(Vector2::zeros(), Matrix2::new(f64::INFINITY, 0., 0., f64::INFINITY)); + } + + #[test] + fn test_bad_create() { + // Covariance not symmetric + bad_create_case(Vector2::zeros(), Matrix2::new(1., 1., 0., 1.)); + // Covariance not positive-definite + bad_create_case(Vector2::zeros(), Matrix2::new(1., 2., 2., 1.)); + // NaN in mean + bad_create_case(Vector2::new(0., f64::NAN), Matrix2::identity()); + // NaN in Covariance Matrix + bad_create_case(Vector2::zeros(), Matrix2::new(1., 0., 0., f64::NAN)); + } + + #[test] + fn test_variance() { + test_case(Vector2::zeros(), Matrix2::identity(), Matrix2::new(1., 0., 0., 1.), |x| x.variance()); + test_case(Vector2::zeros(), Matrix2::new(f64::INFINITY, 0., 0., f64::INFINITY), Matrix2::new(f64::INFINITY, 0., 0., f64::INFINITY), |x| x.variance()); + } + + #[test] + fn test_entropy() { + test_case(Vector2::zeros(), Matrix2::identity(), 2.8378770664093453, |x| x.entropy()); + test_case(Vector2::zeros(), Matrix2::new(1., 0.5, 0.5, 1.), 2.694036030183455, |x| x.entropy()); + test_case(Vector2::zeros(), Matrix2::new(f64::INFINITY, 0., 0., f64::INFINITY), f64::INFINITY, |x| x.entropy()); + } + + #[test] + fn test_mode() { + test_case(Vector2::zeros(), Matrix2::identity(), Vector2::new(0., 0.), |x| x.mode()); + test_case(Vector2::::repeat(f64::INFINITY), Matrix2::identity(), Vector2::new(f64::INFINITY, f64::INFINITY), |x| x.mode()); + } + + #[test] + fn test_min_max() { + test_case(Vector2::zeros(), Matrix2::identity(), Vector2::new(f64::NEG_INFINITY, f64::NEG_INFINITY), |x| x.min()); + test_case(Vector2::zeros(), Matrix2::identity(), Vector2::new(f64::INFINITY, f64::INFINITY), |x| x.max()); + test_case(Vector2::new(10., 1.), Matrix2::identity(), Vector2::new(f64::NEG_INFINITY, f64::NEG_INFINITY), |x| x.min()); + test_case(Vector2::new(-3., 5.), Matrix2::identity(), Vector2::new(f64::INFINITY, f64::INFINITY), |x| x.max()); + } + + #[test] + fn test_pdf() { + test_case(Vector2::zeros(), Matrix2::identity(), 0.05854983152431917, |x| x.pdf(Vector2::new(1., 1.))); + test_almost(Vector2::zeros(), Matrix2::identity(), 0.013064233284684921, 1e-15, |x| x.pdf(Vector2::new(1., 2.))); + test_almost(Vector2::zeros(), Matrix2::identity(), 1.8618676045881531e-23, 1e-35, |x| x.pdf(Vector2::new(1., 10.))); + test_almost(Vector2::zeros(), Matrix2::identity(), 5.920684802611216e-45, 1e-58, |x| x.pdf(Vector2::new(10., 10.))); + test_almost(Vector2::zeros(), Matrix2::new(1., 0.9, 0.9, 1.), 1.6576716577547003e-05, 1e-18, |x| x.pdf(Vector2::new(1., -1.))); + test_almost(Vector2::zeros(), Matrix2::new(1., 0.99, 0.99, 1.), 4.1970621773477824e-44, 1e-54, |x| x.pdf(Vector2::new(1., -1.))); + test_almost(Vector2::new(0.5, -0.2), Matrix2::new(2.0, 0.3, 0.3, 0.5), 0.0013075203140666656, 1e-15, |x| x.pdf(Vector2::new(2., 2.))); + test_case(Vector2::zeros(), Matrix2::new(f64::INFINITY, 0., 0., f64::INFINITY), 0.0, |x| x.pdf(Vector2::new(10., 10.))); + test_case(Vector2::zeros(), Matrix2::new(f64::INFINITY, 0., 0., f64::INFINITY), 0.0, |x| x.pdf(Vector2::new(100., 100.))); + } + + #[test] + fn test_ln_pdf() { + test_case(Vector2::zeros(), Matrix2::identity(), (0.05854983152431917f64).ln(), |x| x.ln_pdf(Vector2::new(1., 1.))); + test_almost(Vector2::zeros(), Matrix2::identity(), (0.013064233284684921f64).ln(), 1e-15, |x| x.ln_pdf(Vector2::new(1., 2.))); + test_almost(Vector2::zeros(), Matrix2::identity(), (1.8618676045881531e-23f64).ln(), 1e-15, |x| x.ln_pdf(Vector2::new(1., 10.))); + test_almost(Vector2::zeros(), Matrix2::identity(), (5.920684802611216e-45f64).ln(), 1e-15, |x| x.ln_pdf(Vector2::new(10., 10.))); + test_almost(Vector2::zeros(), Matrix2::new(1., 0.9, 0.9, 1.), (1.6576716577547003e-05f64).ln(), 1e-14, |x| x.ln_pdf(Vector2::new(1., -1.))); + test_almost(Vector2::zeros(), Matrix2::new(1., 0.99, 0.99, 1.), (4.1970621773477824e-44f64).ln(), 1e-12, |x| x.ln_pdf(Vector2::new(1., -1.))); + test_almost(Vector2::new(0.5, -0.2), Matrix2::new(2.0, 0.3, 0.3, 0.5), (0.0013075203140666656f64).ln(), 1e-15, |x| x.ln_pdf(Vector2::new(2., 2.))); + test_case(Vector2::zeros(), Matrix2::new(f64::INFINITY, 0., 0., f64::INFINITY), f64::NEG_INFINITY, |x| x.ln_pdf(Vector2::new(10., 10.))); + test_case(Vector2::zeros(), Matrix2::new(f64::INFINITY, 0., 0., f64::INFINITY), f64::NEG_INFINITY, |x| x.ln_pdf(Vector2::new(100., 100.))); + } +} diff --git a/src/statistics/traits.rs b/src/statistics/traits.rs index 4e3b3adf..f068548f 100644 --- a/src/statistics/traits.rs +++ b/src/statistics/traits.rs @@ -104,6 +104,10 @@ pub trait Variance: Mean { fn std_dev(&self) -> T; } +pub trait Covariance { + fn variance(&self) -> T; +} + pub trait CheckedVariance: CheckedMean { /// Returns the variance. /// # Examples