diff --git a/src/statistics/dist.rs b/src/statistics/dist.rs index 39afbfc..7e35b8d 100644 --- a/src/statistics/dist.rs +++ b/src/statistics/dist.rs @@ -241,6 +241,7 @@ //! * Mean: $e^{\mu + \frac{\sigma^2}{2}}$ //! * Var: $(e^{\sigma^2} - 1)e^{2\mu + \sigma^2}$ //! * To generate log-normal random samples, Peroxide uses the `rand_distr::LogNormal` distribution from the `rand_distr` crate. +//! //! ### `MVRNG` trait //! //! * `MVRNG` trait is composed of four fields @@ -270,8 +271,8 @@ //! ### Dirichlet Distribution //! //! * Definition -//! $$ \text{Dir}(\mathbf{x} | \boldsymbol{\alpha}) = \frac{1}{\text{B}(\boldsymbol{\alpha})} \prod_{i=1}^K x_i^{\alpha_i - 1} $$ -//! where $\text{B}(\boldsymbol{\alpha}) = \frac{\prod_{i=1}^K \Gamma(\alpha_i)}{\Gamma(\sum_{i=1}^K \alpha_i)}$ +//! $$ \text{Dir}(\mathbf{x} | \boldsymbol{\alpha}) = \frac{1}{\text{B}(\boldsymbol{\alpha})} \prod_{i=1}^K x_i^{\alpha_i - 1} $$ +//! where $\text{B}(\boldsymbol{\alpha}) = \frac{\prod_{i=1}^K \Gamma(\alpha_i)}{\Gamma(\sum_{i=1}^K \alpha_i)}$ //! * Representative value //! * Mean: $\frac{\alpha_i}{\alpha_0}$ //! * Var : $\frac{\alpha_i(\alpha_0 - \alpha_i)}{\alpha_0^2(\alpha_0 + 1)}$ @@ -292,6 +293,68 @@ //! a.cov().print(); // Covariance matrix //! } //! ``` +//! +//! ### `MatRNG` trait +//! +//! * `MatRNG` trait is composed of four fields for Matrix-Variate distributions +//! * `sample`: Extract samples +//! * `sample_with_rng`: Extract samples with specific rng +//! * `pdf` : Calculate pdf value at specific point +//! * `ln_pdf` : Calculate log-pdf value at specific point +//! ```no_run +//! use rand::Rng; +//! use peroxide::fuga::*; +//! +//! pub trait MatRNG { +//! /// Extract samples of matrix-variate distributions +//! fn sample(&self, n: usize) -> Vec; +//! +//! /// Extract samples of distributions with specific rng +//! fn sample_with_rng(&self, rng: &mut R, n: usize) -> Vec; +//! +//! /// Probability Density Function +//! fn pdf(&self, x: &Matrix) -> f64; +//! +//! /// Log Probability Density Function +//! fn ln_pdf(&self, x: &Matrix) -> f64; +//! } +//! ``` +//! +//! ### Wishart Distribution +//! +//! * Definition +//! $$ \mathcal{W}_p(\nu, \mathbf{V}) = \frac{|\mathbf{X}|^{(\nu - p - 1)/2} \exp\left(-\frac{1}{2}\text{tr}(\mathbf{V}^{-1}\mathbf{X})\right)}{2^{\nu p / 2} |\mathbf{V}|^{\nu/2} \Gamma_p\left(\frac{\nu}{2}\right)} $$ +//! where $\nu$ is the degrees of freedom (`df`), $\mathbf{V}$ is the $p \times p$ scale matrix, and $\Gamma_p$ is the multivariate Gamma function. +//! * Representative value +//! * Mean: $\nu \mathbf{V}$ +//! * Variance: $\nu(V_{ij}^2 + V_{ii}V_{jj})$ (Element-wise) +//! * Covariance: $\text{Cov}(X_{ij}, X_{kl}) = \nu(V_{ik}V_{jl} + V_{il}V_{jk})$ (Vectorized to $p^2 \times p^2$) +//! * To generate Wishart random samples, Peroxide uses the Bartlett Decomposition: $\mathbf{X} = \mathbf{L} \mathbf{A} \mathbf{A}^T \mathbf{L}^T$. +//! * **Caution**: `MatDist` utilizes the existing `Statistics` trait. Covariance and correlation are returned as flattened $p^2 \times p^2$ matrices. +//! +//! ```rust +//! use peroxide::fuga::*; +//! +//! fn main() { +//! // Ensure O3 feature is enabled for LAPACK operations +//! #[cfg(feature = "O3")] +//! { +//! let mut rng = smallrng_from_seed(42); +//! let scale = matrix(vec![1.0, 0.5, 0.5, 1.0], 2, 2, Row); +//! let w = MatDist::Wishart(5.0, scale); // Wishart(df, V) +//! +//! w.sample(10)[0].print(); // Print first sample +//! w.sample_with_rng(&mut rng, 10)[0].print(); // Generate with specific rng +//! +//! let test_x = matrix(vec![2.0, 0.0, 0.0, 2.0], 2, 2, Row); +//! w.pdf(&test_x).print(); // Probability density +//! +//! w.mean().print(); // Mean matrix +//! w.var().print(); // Element-wise variance matrix +//! w.cov().print(); // Vectorized p^2 x p^2 covariance +//! } +//! } +//! ``` extern crate rand; extern crate rand_distr; @@ -303,16 +366,19 @@ pub use self::MVDist::*; pub use self::OPDist::*; pub use self::TPDist::*; use crate::special::function::*; -use crate::traits::fp::FPVector; +use crate::traits::fp::{FPMatrix, FPVector}; //use statistics::rand::ziggurat; use self::WeightedUniformError::*; use crate::statistics::{ops::C, stat::Statistics}; use crate::structure::matrix::{matrix, Matrix, Row}; -use crate::util::non_macro::{linspace, seq}; +use crate::util::non_macro::{linspace, seq, zeros}; use crate::util::useful::{auto_zip, find_interval}; use anyhow::{bail, Result}; use std::f64::consts::E; +#[cfg(feature = "O3")] +use crate::traits::matrix::{LinearAlgebra, MatrixTrait, UPLO}; + /// One parameter distribution /// /// # Distributions @@ -347,6 +413,15 @@ pub enum MVDist> { Dirichlet(Vec), } +/// Matrix-Variate Distribution +/// +/// # Distributions +/// * `Wishart(df, scale)`: Wishart distribution +#[derive(Debug, Clone)] +pub enum MatDist { + Wishart(f64, Matrix), +} + pub struct WeightedUniform> { weights: Vec, sum: T, @@ -1215,3 +1290,204 @@ impl> MVRNG for MVDist { } } } + +/// Matrix-Variate Random Number Generator Trait +pub trait MatRNG { + /// Extract samples of matrix-variate distributions + fn sample(&self, n: usize) -> Vec { + let mut rng = rand::rng(); + self.sample_with_rng(&mut rng, n) + } + + /// Extract samples of distributions with specific rng + fn sample_with_rng(&self, rng: &mut R, n: usize) -> Vec; + + /// Probability Density Function + fn pdf(&self, x: &Matrix) -> f64 { + self.ln_pdf(x).exp() + } + + /// Log Probability Density Function + fn ln_pdf(&self, x: &Matrix) -> f64; +} + +impl Statistics for MatDist { + type Array = Matrix; + type Value = Matrix; + + fn mean(&self) -> Self::Value { + match self { + MatDist::Wishart(df, scale) => scale.clone() * *df, + } + } + + fn var(&self) -> Self::Value { + match self { + MatDist::Wishart(df, scale) => { + let p = scale.row; + let mut v = matrix(vec![0f64; p * p], p, p, Row); + for i in 0..p { + for j in 0..p { + v[(i, j)] = *df * (scale[(i, j)].powi(2) + scale[(i, i)] * scale[(j, j)]); + } + } + v + } + } + } + + fn sd(&self) -> Self::Value { + self.var().fmap(|x| x.sqrt()) + } + + fn cov(&self) -> Self::Array { + match self { + MatDist::Wishart(df, scale) => { + let p = scale.row; + let p_sq = p * p; + + // Allocate a flattened p^2 x p^2 covariance matrix + let mut cov_mat = zeros(p_sq, p_sq); + + for i in 0..p { + for j in 0..p { + for k in 0..p { + for l in 0..p { + let row_idx = i * p + j; + let col_idx = k * p + l; + + cov_mat[(row_idx, col_idx)] = *df + * (scale[(i, k)] * scale[(j, l)] + + scale[(i, l)] * scale[(j, k)]); + } + } + } + } + cov_mat + } + } + } + + fn cor(&self) -> Self::Array { + let cov_mat = self.cov(); + let sd_mat = self.sd(); + let p = sd_mat.row; + let p_sq = p * p; + + let mut cor_mat = zeros(p_sq, p_sq); + + for i in 0..p { + for j in 0..p { + for k in 0..p { + for l in 0..p { + let row_idx = i * p + j; + let col_idx = k * p + l; + + // Scale the covariance element by the standard deviations + cor_mat[(row_idx, col_idx)] = + cov_mat[(row_idx, col_idx)] / (sd_mat[(i, j)] * sd_mat[(k, l)]); + } + } + } + } + cor_mat + } +} + +impl MatRNG for MatDist { + #[cfg_attr(not(feature = "O3"), allow(unused))] + fn sample_with_rng(&self, rng: &mut R, n: usize) -> Vec { + match self { + MatDist::Wishart(df, scale) => { + let p = scale.row; + + #[cfg(feature = "O3")] + { + // Bartlett Decomposition Requires Cholesky on the Scale Matrix + let l_mat = scale.cholesky(UPLO::Lower); + let mut samples = Vec::with_capacity(n); + let normal = rand_distr::StandardNormal; + + for _ in 0..n { + let mut a_mat = matrix(vec![0f64; p * p], p, p, Row); + + // Populate the Bartlett Matrix A + for i in 0..p { + for j in 0..=i { + if i == j { + let chi = rand_distr::ChiSquared::new(*df - i as f64).unwrap(); + a_mat[(i, i)] = chi.sample(rng).sqrt(); + } else { + a_mat[(i, j)] = rng.sample(normal); + } + } + } + + // Compute X = L * A * A^T * L^T + let la = &l_mat * &a_mat; + let x = &la * &la.t(); + + samples.push(x); + } + + samples + } + #[cfg(not(feature = "O3"))] + { + unimplemented!( + "Wishart sampling requires the 'O3' feature for Cholesky decomposition." + ); + } + } + } + } + + #[cfg_attr(not(feature = "O3"), allow(unused))] + fn ln_pdf(&self, x: &Matrix) -> f64 { + match self { + MatDist::Wishart(df, scale) => { + #[cfg(feature = "O3")] + { + let p_usize = scale.row; + let p = p_usize as f64; + + let det_x = x.det(); + let det_v = scale.det(); + let inv_v = scale.inv(); + + // Trace of (V^-1 * X) + let v_inv_x = &inv_v * x; + let tr_v_inv_x = v_inv_x.diag().iter().sum::(); + + let term1 = (*df - p - 1.0) / 2.0 * det_x.ln(); + let term2 = -0.5 * tr_v_inv_x; + let term3 = -(*df * p) / 2.0 * 2f64.ln(); + let term4 = -(*df / 2.0) * det_v.ln(); + let term5 = -mv_ln_gamma(p_usize, *df / 2.0); + + term1 + term2 + term3 + term4 + term5 + } + #[cfg(not(feature = "O3"))] + { + unimplemented!("Wishart PDF calculation requires the 'O3' feature for matrix inversion and determinants."); + } + } + } + } +} + +#[cfg_attr(not(feature = "O3"), allow(unused))] +fn mv_ln_gamma(p: usize, a: f64) -> f64 { + let p_f64 = p as f64; + + if a <= (p_f64 - 1.0) / 2.0 { + return f64::NAN; + } + + let mut sum = p_f64 * (p_f64 - 1.0) / 4.0 * std::f64::consts::PI.ln(); + for j in 1..=p { + sum += ln_gamma(a + (1.0 - j as f64) / 2.0); + } + + sum +} diff --git a/src/util/print.rs b/src/util/print.rs index f9041e9..3205c93 100644 --- a/src/util/print.rs +++ b/src/util/print.rs @@ -326,6 +326,62 @@ impl Printable for Matrix { } } +impl Printable for Vec { + fn print(&self) { + if self.is_empty() { + println!("[]"); + return; + } + + println!("["); + + for i in 0..self.len() { + let mat_str = self[i].to_string(); + let indented = mat_str + .lines() + .map(|line| format!(" {}", line)) + .collect::>() + .join("\n"); + + if i != self.len() - 1 { + println!("{},\n", indented); + } else { + println!("{}", indented); + } + } + + println!("]"); + } +} + +impl Printable for &Vec { + fn print(&self) { + if self.is_empty() { + println!("[]"); + return; + } + + println!("["); + + for i in 0..self.len() { + let mat_str = self[i].to_string(); + let indented = mat_str + .lines() + .map(|line| format!(" {}", line)) + .collect::>() + .join("\n"); + + if i != self.len() - 1 { + println!("{},\n", indented); + } else { + println!("{}", indented); + } + } + + println!("]"); + } +} + impl Printable for Polynomial { fn print(&self) { println!("{}", self); @@ -380,6 +436,13 @@ impl> Printable for MVD } } +#[cfg(feature = "rand")] +impl Printable for MatDist { + fn print(&self) { + println!("{:?}", self); + } +} + //impl Printable for Number { // fn print(&self) { // println!("{:?}", self) diff --git a/tests/dist.rs b/tests/dist.rs index 2549894..3f39c23 100644 --- a/tests/dist.rs +++ b/tests/dist.rs @@ -1,6 +1,6 @@ #![cfg(feature = "rand")] -extern crate peroxide; +use peroxide::c; use peroxide::fuga::*; #[test] @@ -29,3 +29,101 @@ fn test_dirichlet() { let pdf_val = dir.pdf(&[0.33333, 0.33333, 0.33333]); assert!(nearly_eq(pdf_val, 2.222155556222205)); } + +#[test] +fn test_wishart() { + #[cfg(feature = "O3")] + { + // Define Scale Matrix V: + // [ 1.0 0.5 ] + // [ 0.5 1.0 ] + let scale = matrix(c!(1.0, 0.5, 0.5, 1.0), 2, 2, Row); + + // Wishart with 5 degrees of freedom + let w = MatDist::Wishart(5.0, scale); + + // Test Sampling + let samples = w.sample(10); + samples.print(); + assert_eq!(samples.len(), 10); + assert_eq!(samples[0].nrow(), 2); + assert_eq!(samples[0].ncol(), 2); + + // Test Mean (E[X] = df * V) + let m = w.mean(); + assert!(nearly_eq(m[(0, 0)], 5.0)); + assert!(nearly_eq(m[(0, 1)], 2.5)); + assert!(nearly_eq(m[(1, 0)], 2.5)); + assert!(nearly_eq(m[(1, 1)], 5.0)); + + // Test Variance (Element-wise: Var(X_ij) = df * (V_ij^2 + V_ii * V_jj)) + let v = w.var(); + assert!(nearly_eq(v[(0, 0)], 10.0)); // 5 * (1^2 + 1*1) + assert!(nearly_eq(v[(0, 1)], 6.25)); // 5 * (0.5^2 + 1*1) = 5 * 1.25 + assert!(nearly_eq(v[(1, 0)], 6.25)); + assert!(nearly_eq(v[(1, 1)], 10.0)); + + // Test Covariance (Vectorized 4x4 matrix) + // Cov(X_ij, X_kl) = df * (V_ik * V_jl + V_il * V_jk) + let cov = w.cov(); + assert_eq!(cov.nrow(), 4); + assert_eq!(cov.ncol(), 4); + + // Cov(X_00, X_00) -> mapped to index (0, 0) + assert!(nearly_eq(cov[(0, 0)], 10.0)); + + // Cov(X_00, X_11) -> mapped to index (0, 3) + // 5 * (V_01 * V_01 + V_01 * V_01) = 5 * (0.25 + 0.25) = 2.5 + assert!(nearly_eq(cov[(0, 3)], 2.5)); + + // Cov(X_01, X_01) -> mapped to index (1, 1) + // 5 * (V_00 * V_11 + V_01 * V_10) = 5 * (1 + 0.25) = 6.25 + assert!(nearly_eq(cov[(1, 1)], 6.25)); + + // Cov(X_01, X_10) -> mapped to index (1, 2) + // 5 * (V_01 * V_10 + V_00 * V_11) = 5 * (0.25 + 1.0) = 6.25 + assert!(nearly_eq(cov[(1, 2)], 6.25)); + + // 5. Test PDF against an exact closed-form calculation + // Let test_x = [ 2.0 0.0 ] + // [ 0.0 2.0 ] + let test_x = matrix(c!(2.0, 0.0, 0.0, 2.0), 2, 2, Row); + let pdf_val = w.pdf(&test_x); + + // Exact closed-form math for this specific test case: + // Numerator: |X|^1 * exp(-0.5 * tr(V^-1 * X)) = 4 * exp(-8/3) + // Denominator: 2^5 * |V|^2.5 * Gamma_2(2.5) = 32 * (0.75)^2.5 * (0.75 * pi) + // Which simplifies cleanly to: 16 * exp(-8/3) / (27 * sqrt(3) * pi) + let expected_pdf = + 16.0 * (-8.0f64 / 3.0).exp() / (27.0 * 3.0f64.sqrt() * std::f64::consts::PI); + + assert!(nearly_eq(pdf_val, expected_pdf)); + } +} + +#[test] +fn test_wishart_1d_chi_square() { + #[cfg(feature = "O3")] + { + let df = 5.0; + let scale = matrix(c!(1.0), 1, 1, Row); + let w = MatDist::Wishart(df, scale); + + // A Chi-Square(df) distribution has Mean = df, Variance = 2*df + let m = w.mean(); + let v = w.var(); + + assert!(nearly_eq(m[(0, 0)], df)); + assert!(nearly_eq(v[(0, 0)], 2.0 * df)); + + // Check PDF Convergence + let x_val = 3.0; + let test_x = matrix(c!(x_val), 1, 1, Row); + let wishart_pdf = w.pdf(&test_x); + + let chi_pdf = (x_val.powf(df / 2.0 - 1.0) * (-x_val / 2.0).exp()) + / (2.0f64.powf(df / 2.0) * gamma(df / 2.0)); + + assert!(nearly_eq(wishart_pdf, chi_pdf)); + } +}