From 7b6859ea6f995cf3441ae98943a98b62308cd141 Mon Sep 17 00:00:00 2001 From: vera-rubin <322284620+vera-rubin@users.noreply.github.com> Date: Wed, 14 Oct 2026 18:57:33 +0000 Subject: [PATCH] fix: preserve MultivariateNormal density and entropy in high dimensions --- src/distribution/multivariate_normal.rs | 61 +++++++++++---------- src/distribution/multivariate_students_t.rs | 12 ++-- 2 files changed, 37 insertions(+), 36 deletions(-) diff --git a/src/distribution/multivariate_normal.rs b/src/distribution/multivariate_normal.rs index b0434eef..8a6670b0 100644 --- a/src/distribution/multivariate_normal.rs +++ b/src/distribution/multivariate_normal.rs @@ -1,14 +1,14 @@ use crate::distribution::Continuous; use crate::statistics::{Max, MeanN, Min, Mode, VarianceN}; use alloc::vec::Vec; -use core::f64::consts::{E, PI}; +use core::f64::consts::PI; use nalgebra::{Cholesky, Const, DMatrix, DVector, Dim, DimMin, Dyn, OMatrix, OVector}; #[cfg(not(feature = "std"))] use num_traits::Float as _; -/// Computes both the normalization and exponential argument in the normal +/// Computes both the log normalization and exponential argument in the normal /// distribution, returning `None` on dimension mismatch. -pub(super) fn density_normalization_and_exponential( +pub(super) fn density_ln_normalization_and_exponential( mu: &OVector, cov: &OMatrix, precision: &OMatrix, @@ -20,7 +20,7 @@ where nalgebra::allocator::Allocator + nalgebra::allocator::Allocator, { Some(( - density_distribution_pdf_const(mu, cov)?, + density_distribution_ln_pdf_const(mu, cov)?, density_distribution_exponential(mu, precision, x)?, )) } @@ -50,10 +50,13 @@ where Some(exp_term) } -/// Computes the argument of the normalization term in the normal distribution, +/// Computes the log normalization term in the normal distribution, /// returning `None` on dimension mismatch. #[inline] -fn density_distribution_pdf_const(mu: &OVector, cov: &OMatrix) -> Option +fn density_distribution_ln_pdf_const( + mu: &OVector, + cov: &OMatrix, +) -> Option where D: DimMin, nalgebra::DefaultAllocator: nalgebra::allocator::Allocator @@ -63,12 +66,8 @@ where if cov.shape_generic().0 != mu.shape_generic().0 || !cov.is_square() { return None; } - let cov_det = cov.determinant(); - Some( - ((2. * PI).powi(mu.nrows() as i32) * cov_det.abs()) - .recip() - .sqrt(), - ) + let cov_ln_det = cov.clone().cholesky()?.ln_determinant(); + Some(-0.5 * (mu.nrows() as f64 * (2. * PI).ln() + cov_ln_det)) } /// Implements the [Multivariate Normal](https://en.wikipedia.org/wiki/Multivariate_normal_distribution) @@ -97,7 +96,7 @@ where mu: OVector, cov: OMatrix, precision: OMatrix, - pdf_const: f64, + ln_pdf_const: f64, } /// Represents the errors that can occur when creating a [`MultivariateNormal`]. @@ -297,7 +296,7 @@ where let precision = cholesky.inverse(); MultivariateNormal { - pdf_const: density_distribution_pdf_const(&mean, &cov).unwrap(), + ln_pdf_const: -0.5 * (mean.nrows() as f64 * (2. * PI).ln() + cholesky.ln_determinant()), cov_chol_decomp: cholesky.unpack(), mu: mean, cov, @@ -315,14 +314,7 @@ where /// /// where `Σ` is the covariance matrix and `det` is the determinant pub fn entropy(&self) -> Option { - Some( - 0.5 * self - .variance() - .unwrap() - .scale(2. * PI * E) - .determinant() - .ln(), - ) + Some(0.5 * self.mu.nrows() as f64 - self.ln_pdf_const) } /// Returns the Cholesky decomposition of the covariance matrix @@ -515,17 +507,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.ln_pdf(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.ln_pdf_const + density_distribution_exponential(&self.mu, &self.precision, x).unwrap() } } @@ -534,7 +522,7 @@ where mod tests { use core::fmt::Debug; use crate::prec; - use nalgebra::{dmatrix, dvector, matrix, vector, DimMin, OMatrix, OVector}; + use nalgebra::{dmatrix, dvector, matrix, vector, DMatrix, DVector, DimMin, OMatrix, OVector}; use crate::{ distribution::{Continuous, MultivariateNormal}, @@ -660,10 +648,11 @@ mod tests { 2.8378770664093453, entropy, ); - test_case( + test_almost( dvector![0., 0.], dmatrix![1., 0.5; 0.5, 1.], 2.694036030183455, + 1e-15, entropy, ); test_case( @@ -674,6 +663,18 @@ mod tests { ); } + #[test] + fn test_high_dimension_density_and_entropy() { + let mean = DVector::zeros(400); + let mvn = try_create(mean.clone(), DMatrix::identity(400, 400)); + prec::assert_relative_eq!( + mvn.pdf(&mean), 2.312204959487328e-160, + epsilon = 0.0, max_relative = 1e-13 + ); + prec::assert_abs_diff_eq!(mvn.ln_pdf(&mean), -367.5754132818691, epsilon = 1e-12); + prec::assert_abs_diff_eq!(mvn.entropy().unwrap(), 567.5754132818691, epsilon = 1e-12); + } + #[test] fn test_mode() { let mode = |x: MultivariateNormal<_>| x.mode(); diff --git a/src/distribution/multivariate_students_t.rs b/src/distribution/multivariate_students_t.rs index 2423b8c2..bdeba28d 100644 --- a/src/distribution/multivariate_students_t.rs +++ b/src/distribution/multivariate_students_t.rs @@ -353,15 +353,15 @@ where /// - `k` is the dimension of the distribution. fn pdf(&self, x: &OVector) -> f64 { if self.freedom.is_infinite() { - use super::multivariate_normal::density_normalization_and_exponential; - let (pdf_const, exp_arg) = density_normalization_and_exponential( + use super::multivariate_normal::density_ln_normalization_and_exponential; + let (ln_pdf_const, exp_arg) = density_ln_normalization_and_exponential( &self.location, &self.scale, &self.precision, x, ) .unwrap(); - return pdf_const * exp_arg.exp(); + return (ln_pdf_const + exp_arg).exp(); } let dv = x - &self.location; @@ -374,15 +374,15 @@ where /// student distribution at `x`. Equivalent to pdf(x).ln(). fn ln_pdf(&self, x: &OVector) -> f64 { if self.freedom.is_infinite() { - use super::multivariate_normal::density_normalization_and_exponential; - let (pdf_const, exp_arg) = density_normalization_and_exponential( + use super::multivariate_normal::density_ln_normalization_and_exponential; + let (ln_pdf_const, exp_arg) = density_ln_normalization_and_exponential( &self.location, &self.scale, &self.precision, x, ) .unwrap(); - return pdf_const.ln() + exp_arg; + return ln_pdf_const + exp_arg; } let dv = x - &self.location;