From be40c8ca6052e1410c6e88db3e0a5b3c103f7cd6 Mon Sep 17 00:00:00 2001 From: Orion Yeung <11580988+orionyeung001@users.noreply.github.com> Date: Mon, 9 Sep 2024 12:56:29 -0500 Subject: [PATCH 1/3] test: mark display implementations as ignored from coverage --- src/distribution/bernoulli.rs | 1 + src/distribution/chi_squared.rs | 1 + src/distribution/erlang.rs | 1 + 3 files changed, 3 insertions(+) diff --git a/src/distribution/bernoulli.rs b/src/distribution/bernoulli.rs index 28f9d104..4b9613ad 100644 --- a/src/distribution/bernoulli.rs +++ b/src/distribution/bernoulli.rs @@ -79,6 +79,7 @@ impl Bernoulli { } impl std::fmt::Display for Bernoulli { + #[cfg_attr(coverage_nightly, coverage(off))] fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "Bernoulli({})", self.p()) } diff --git a/src/distribution/chi_squared.rs b/src/distribution/chi_squared.rs index f61d6e19..38306f20 100644 --- a/src/distribution/chi_squared.rs +++ b/src/distribution/chi_squared.rs @@ -95,6 +95,7 @@ impl ChiSquared { } impl std::fmt::Display for ChiSquared { + #[cfg_attr(coverage_nightly, coverage(off))] fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "χ^2_{}", self.freedom) } diff --git a/src/distribution/erlang.rs b/src/distribution/erlang.rs index 2ad017f3..25df524f 100644 --- a/src/distribution/erlang.rs +++ b/src/distribution/erlang.rs @@ -77,6 +77,7 @@ impl Erlang { } impl std::fmt::Display for Erlang { + #[cfg_attr(coverage_nightly, coverage(off))] fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "E({}, {})", self.rate(), self.shape()) } From e88c40bf7944be42615b84a1e62cee2ebe018c4c Mon Sep 17 00:00:00 2001 From: Orion Yeung <11580988+orionyeung001@users.noreply.github.com> Date: Fri, 6 Sep 2024 15:14:18 -0500 Subject: [PATCH 2/3] feat: reusing buffers for Dyn multivariate takes approach of making a few fields `Option`, to use `take` since creating `nalgebra::Cholesky` requires ownership --- src/distribution/multivariate_normal.rs | 162 ++++++++++++-------- src/distribution/multivariate_students_t.rs | 22 +-- 2 files changed, 104 insertions(+), 80 deletions(-) diff --git a/src/distribution/multivariate_normal.rs b/src/distribution/multivariate_normal.rs index b336d9db..20eb3e9f 100644 --- a/src/distribution/multivariate_normal.rs +++ b/src/distribution/multivariate_normal.rs @@ -5,76 +5,39 @@ use nalgebra::{Cholesky, Const, DMatrix, DVector, Dim, DimMin, Dyn, OMatrix, OVe use std::f64; use std::f64::consts::{E, PI}; -/// computes both the normalization and exponential argument in the normal distribution -/// # Errors -/// will error on dimension mismatch -pub(super) fn density_normalization_and_exponential( - mu: &OVector, - cov: &OMatrix, - precision: &OMatrix, - x: &OVector, -) -> std::result::Result<(f64, f64), StatsError> -where - D: DimMin, - nalgebra::DefaultAllocator: nalgebra::allocator::Allocator - + nalgebra::allocator::Allocator - + nalgebra::allocator::Allocator<(usize, usize), D>, -{ - Ok(( - density_distribution_pdf_const(mu, cov)?, - density_distribution_exponential(mu, precision, x)?, - )) -} - /// computes the argument of the exponential term in the normal distribution -/// ```text -/// ``` -/// # Errors -/// will error on dimension mismatch +/// # Panics +/// will panic on dimension mismatch #[inline] -pub(super) fn density_distribution_exponential( +pub(super) fn pdf_exponent_unchecked( mu: &OVector, precision: &OMatrix, x: &OVector, -) -> std::result::Result +) -> f64 where D: Dim, nalgebra::DefaultAllocator: nalgebra::allocator::Allocator + nalgebra::allocator::Allocator, { - if x.shape_generic().0 != precision.shape_generic().0 - || x.shape_generic().0 != mu.shape_generic().0 - || !precision.is_square() - { - return Err(StatsError::ContainersMustBeSameLength); - } let dv = x - mu; - let exp_term: f64 = -0.5 * (precision * &dv).dot(&dv); - Ok(exp_term) - // TODO update to dimension mismatch error + -0.5 * (precision * &dv).dot(&dv) } /// computes the argument of the normalization term in the normal distribution -/// # Errors -/// will error on dimension mismatch +/// # Panics +/// will panic on dimension mismatch #[inline] -pub(super) fn density_distribution_pdf_const( - mu: &OVector, - cov: &OMatrix, -) -> std::result::Result +pub(super) fn pdf_const_unchecked(mu: &OVector, cov: &OMatrix) -> f64 where D: DimMin, nalgebra::DefaultAllocator: nalgebra::allocator::Allocator + nalgebra::allocator::Allocator + nalgebra::allocator::Allocator<(usize, usize), D>, { - if cov.shape_generic().0 != mu.shape_generic().0 || !cov.is_square() { - return Err(StatsError::ContainersMustBeSameLength); - } let cov_det = cov.determinant(); - Ok(((2. * PI).powi(mu.nrows() as i32) * cov_det.abs()) + ((2. * PI).powi(mu.nrows() as i32) * cov_det.abs()) .recip() - .sqrt()) + .sqrt() } /// Implements the [Multivariate Normal](https://en.wikipedia.org/wiki/Multivariate_normal_distribution) @@ -92,18 +55,18 @@ where /// assert_eq!(mvn.variance().unwrap(), matrix![1., 0.; 0., 1.]); /// assert_eq!(mvn.pdf(&vector![1., 1.]), 0.05854983152431917); /// ``` -#[derive(Clone, PartialEq, Debug)] +#[derive(Clone, Debug)] pub struct MultivariateNormal where D: Dim, nalgebra::DefaultAllocator: nalgebra::allocator::Allocator + nalgebra::allocator::Allocator, { - cov_chol_decomp: OMatrix, mu: OVector, cov: OMatrix, - precision: OMatrix, pdf_const: f64, + cov_chol: Option>, + precision: Option>, } /// Represents the errors that can occur when creating a [`MultivariateNormal`]. @@ -199,20 +162,95 @@ where // for sampling match Cholesky::new(cov.clone()) { None => Err(MultivariateNormalError::CholeskyFailed), - Some(cholesky_decomp) => { - let precision = cholesky_decomp.inverse(); - Ok(MultivariateNormal { - // .unwrap() because prerequisites are already checked above - pdf_const: density_distribution_pdf_const(&mean, &cov).unwrap(), - cov_chol_decomp: cholesky_decomp.unpack(), - mu: mean, + cov_chol @ Some(_) => Ok(MultivariateNormal { + precision: cov_chol.as_ref().map(|ll| ll.inverse()), + cov_chol, + pdf_const: pdf_const_unchecked(&mean, &cov), + mu: mean, + cov, + }), + } + } + + /// Constructs a new multivariate normal distribution with a mean of `mean` + /// and covariance matrix `cov` from an existing `MultivariateNormal` without + /// requiring reallocation. + /// + /// # Errors + /// + /// Returns an error if the given covariance matrix is not + /// symmetric or positive-definite + pub fn into_new_from_nalgebra( + self, + other_mean: &OVector, + other_cov: &OMatrix, + ) -> Result { + let Self { + mut mu, + mut cov, + pdf_const, + cov_chol, + .. + } = self; + + mu.clone_from(other_mean); + cov.clone_from(other_cov); + + // clone if storage needed, take from previous cholesky otherwise + let other_cov = cov_chol.map_or((*other_cov).clone(), |chol| chol.unpack()); + + match other_cov.cholesky() { + None => Err(StatsError::BadParams), + cov_chol @ Some(_) => { + let precision = cov_chol.as_ref().map(|chol| chol.inverse()); + + Ok(Self { + mu, cov, + pdf_const, + cov_chol, precision, }) } } } + /// updates the covariance of the distribution without requiring reallocation + /// + /// interally uses `clone_from` and [`set_cov_with`] + /// # Errors + /// if dimensions change, then this will error with dimension mismatch + #[inline] + pub fn set_cov(&mut self, new_cov: &OMatrix) -> Result<(), StatsError> { + self.set_cov_with(|old_cov: &mut OMatrix| old_cov.clone_from(new_cov)) + } + + /// updates the covariance of the distribution without requiring reallocation + /// by calling a function + /// # Errors + /// if dimensions change, then this will error with dimension mismatch + pub fn set_cov_with( + &mut self, + f: impl FnOnce(&mut OMatrix), + ) -> Result<(), StatsError> { + let old_shape = self.cov.shape_generic(); + f(&mut self.cov); + if old_shape != self.cov.shape_generic() { + Err(StatsError::ContainersMustBeSameLength) + } else { + match self.cov_chol.take() { + None => (), + Some(l_old) => { + let mut l = l_old.unpack(); + l.clone_from(&self.cov); + // ignore possible fallibility for now + self.cov_chol = Some(l.cholesky().unwrap()); + } + } + Ok(()) + } + } + /// Returns the entropy of the multivariate normal distribution /// /// # Formula @@ -266,7 +304,7 @@ where fn sample(&self, rng: &mut R) -> OVector { let d = crate::distribution::Normal::new(0., 1.).unwrap(); let z = OVector::from_distribution_generic(self.mu.shape_generic().0, Const::<1>, &d, rng); - (&self.cov_chol_decomp * z) + &self.mu + self.cov_chol.as_ref().unwrap().l() * z + &self.mu } } @@ -362,17 +400,13 @@ where /// 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: &OVector) -> f64 { - self.pdf_const - * density_distribution_exponential(&self.mu, &self.precision, x) - .unwrap() - .exp() + self.pdf_const * pdf_exponent_unchecked(&self.mu, self.precision.as_ref().unwrap(), x).exp() } /// Calculates the log probability density function for the multivariate /// normal distribution at `x`. Equivalent to pdf(x).ln(). fn ln_pdf(&self, x: &OVector) -> f64 { - self.pdf_const.ln() - + density_distribution_exponential(&self.mu, &self.precision, x).unwrap() + self.pdf_const.ln() + pdf_exponent_unchecked(&self.mu, self.precision.as_ref().unwrap(), x) } } diff --git a/src/distribution/multivariate_students_t.rs b/src/distribution/multivariate_students_t.rs index 73e8f8f2..96d3d277 100644 --- a/src/distribution/multivariate_students_t.rs +++ b/src/distribution/multivariate_students_t.rs @@ -346,14 +346,9 @@ where /// of the scale matrix, and `k` is the dimension of the distribution. fn pdf(&self, x: &'a OVector) -> f64 { if self.freedom.is_infinite() { - use super::multivariate_normal::density_normalization_and_exponential; - let (pdf_const, exp_arg) = density_normalization_and_exponential( - &self.location, - &self.scale, - &self.precision, - x, - ) - .unwrap(); + use super::multivariate_normal as mvn; + let pdf_const = mvn::pdf_const_unchecked(&self.location, &self.scale); + let exp_arg = mvn::pdf_exponent_unchecked(&self.location, &self.precision, x); return pdf_const * exp_arg.exp(); } @@ -367,14 +362,9 @@ where /// student distribution at `x`. Equivalent to pdf(x).ln(). fn ln_pdf(&self, x: &'a OVector) -> f64 { if self.freedom.is_infinite() { - use super::multivariate_normal::density_normalization_and_exponential; - let (pdf_const, exp_arg) = density_normalization_and_exponential( - &self.location, - &self.scale, - &self.precision, - x, - ) - .unwrap(); + use super::multivariate_normal as mvn; + let pdf_const = mvn::pdf_const_unchecked(&self.location, &self.scale); + let exp_arg = mvn::pdf_exponent_unchecked(&self.location, &self.precision, x); return pdf_const.ln() + exp_arg; } From 3657843f5cfc47cf7ca7133a5b8f75570c0102a7 Mon Sep 17 00:00:00 2001 From: Orion Yeung <11580988+orionyeung001@users.noreply.github.com> Date: Tue, 10 Sep 2024 10:55:59 -0500 Subject: [PATCH 3/3] style: drop `_unchecked` suffix to match conventions from std --- src/distribution/multivariate_normal.rs | 12 ++++++------ src/distribution/multivariate_students_t.rs | 8 ++++---- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/src/distribution/multivariate_normal.rs b/src/distribution/multivariate_normal.rs index 20eb3e9f..7f5923c0 100644 --- a/src/distribution/multivariate_normal.rs +++ b/src/distribution/multivariate_normal.rs @@ -9,7 +9,7 @@ use std::f64::consts::{E, PI}; /// # Panics /// will panic on dimension mismatch #[inline] -pub(super) fn pdf_exponent_unchecked( +pub(super) fn pdf_exponent( mu: &OVector, precision: &OMatrix, x: &OVector, @@ -23,11 +23,11 @@ where -0.5 * (precision * &dv).dot(&dv) } -/// computes the argument of the normalization term in the normal distribution +/// computes normalization term in the normal distribution /// # Panics /// will panic on dimension mismatch #[inline] -pub(super) fn pdf_const_unchecked(mu: &OVector, cov: &OMatrix) -> f64 +pub(super) fn pdf_const(mu: &OVector, cov: &OMatrix) -> f64 where D: DimMin, nalgebra::DefaultAllocator: nalgebra::allocator::Allocator @@ -165,7 +165,7 @@ where cov_chol @ Some(_) => Ok(MultivariateNormal { precision: cov_chol.as_ref().map(|ll| ll.inverse()), cov_chol, - pdf_const: pdf_const_unchecked(&mean, &cov), + pdf_const: pdf_const(&mean, &cov), mu: mean, cov, }), @@ -400,13 +400,13 @@ where /// 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: &OVector) -> f64 { - self.pdf_const * pdf_exponent_unchecked(&self.mu, self.precision.as_ref().unwrap(), x).exp() + self.pdf_const * pdf_exponent(&self.mu, self.precision.as_ref().unwrap(), x).exp() } /// Calculates the log probability density function for the multivariate /// normal distribution at `x`. Equivalent to pdf(x).ln(). fn ln_pdf(&self, x: &OVector) -> f64 { - self.pdf_const.ln() + pdf_exponent_unchecked(&self.mu, self.precision.as_ref().unwrap(), x) + self.pdf_const.ln() + pdf_exponent(&self.mu, self.precision.as_ref().unwrap(), x) } } diff --git a/src/distribution/multivariate_students_t.rs b/src/distribution/multivariate_students_t.rs index 96d3d277..2d290555 100644 --- a/src/distribution/multivariate_students_t.rs +++ b/src/distribution/multivariate_students_t.rs @@ -347,8 +347,8 @@ where fn pdf(&self, x: &'a OVector) -> f64 { if self.freedom.is_infinite() { use super::multivariate_normal as mvn; - let pdf_const = mvn::pdf_const_unchecked(&self.location, &self.scale); - let exp_arg = mvn::pdf_exponent_unchecked(&self.location, &self.precision, x); + let pdf_const = mvn::pdf_const(&self.location, &self.scale); + let exp_arg = mvn::pdf_exponent(&self.location, &self.precision, x); return pdf_const * exp_arg.exp(); } @@ -363,8 +363,8 @@ where fn ln_pdf(&self, x: &'a OVector) -> f64 { if self.freedom.is_infinite() { use super::multivariate_normal as mvn; - let pdf_const = mvn::pdf_const_unchecked(&self.location, &self.scale); - let exp_arg = mvn::pdf_exponent_unchecked(&self.location, &self.precision, x); + let pdf_const = mvn::pdf_const(&self.location, &self.scale); + let exp_arg = mvn::pdf_exponent(&self.location, &self.precision, x); return pdf_const.ln() + exp_arg; }