use nalgebra::{VectorN, MatrixMN, DefaultAllocator, Dim, RealField, linalg,
allocator::Allocator};
#[derive(Debug)]
pub struct Error {
kind: ErrorKind,
}
impl Error {
pub fn kind(&self) -> &ErrorKind {
&self.kind
}
}
impl std::error::Error for Error {}
impl std::fmt::Display for Error {
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
write!(f, "{:?}", self.kind)
}
}
#[derive(Debug,Clone)]
pub enum ErrorKind {
NotDefinitePositive,
}
#[derive(Debug,Clone)]
pub struct MultivariateNormal<Real,N>
where
Real: RealField,
N: Dim + nalgebra::DimMin<N, Output = N>,
DefaultAllocator: Allocator<Real, N>,
DefaultAllocator: Allocator<Real, N, N>,
DefaultAllocator: Allocator<Real, nalgebra::U1, N>,
DefaultAllocator: Allocator<(usize, usize), <N as nalgebra::DimMin<N>>::Output>,
{
neg_mu: nalgebra::VectorN<Real,N>,
precision: nalgebra::MatrixN<Real,N>,
fac: Real,
}
impl<Real,N> MultivariateNormal<Real,N>
where
Real: RealField,
N: Dim + nalgebra::DimMin<N, Output = N> + nalgebra::DimSub<nalgebra::Dynamic>,
DefaultAllocator: Allocator<Real, N>,
DefaultAllocator: Allocator<Real, N, N>,
DefaultAllocator: Allocator<Real, nalgebra::U1, N>,
DefaultAllocator: Allocator<(usize, usize), <N as nalgebra::DimMin<N>>::Output>,
{
pub fn from_mean_and_precision(
mu: &nalgebra::VectorN<Real,N>,
precision: &nalgebra::MatrixN<Real,N>,
) -> Self {
let precision_det = nalgebra::linalg::LU::new(precision.clone()).determinant();
let det = Real::one()/precision_det;
let ndim = mu.nrows();
let fac: Real = Real::one() / ( Real::two_pi().powi(ndim as i32)
* det.abs() ).sqrt();
Self { neg_mu: -mu, precision: precision.clone(), fac }
}
pub fn from_mean_and_covariance(
mu: &nalgebra::VectorN<Real,N>,
covariance: &nalgebra::MatrixN<Real,N>,
) -> Result<Self,Error> {
let precision = linalg::Cholesky::new(covariance.clone())
.ok_or(Error{kind: ErrorKind::NotDefinitePositive})?.inverse();
let result = Self::from_mean_and_precision(mu, &precision);
Ok(result)
}
pub fn mean(&self) -> nalgebra::VectorN<Real,N> {
-&self.neg_mu
}
pub fn precision(&self) -> &nalgebra::MatrixN<Real,N> {
&self.precision
}
fn inner_pdf<Count>(
&self,
xs_t: &nalgebra::MatrixMN<Real,Count,N>,
) -> nalgebra::VectorN<Real,Count>
where
Count: Dim,
DefaultAllocator: Allocator<Real, Count>,
DefaultAllocator: Allocator<Real, N, Count>,
DefaultAllocator: Allocator<Real, Count, N>,
DefaultAllocator: Allocator<Real, Count, Count>,
{
let dvs: nalgebra::MatrixMN<Real,Count,N> = broadcast_add(&xs_t,&self.neg_mu);
let left: nalgebra::MatrixMN<Real,Count,N> = &dvs * &self.precision;
let ny2_tmp: nalgebra::MatrixMN<Real,Count,N> = left.component_mul( &dvs );
let ones = nalgebra::MatrixMN::<Real,N,nalgebra::U1>::repeat_generic(
N::from_usize(self.neg_mu.nrows()),
nalgebra::U1,
nalgebra::convert::<f64,Real>(1.0),
);
let ny2: nalgebra::VectorN<Real,Count> = ny2_tmp * ones;
let y: nalgebra::VectorN<Real,Count> = ny2 * nalgebra::convert::<f64,Real>(-0.5);
y
}
pub fn pdf<Count>(
&self,
xs: &nalgebra::MatrixMN<Real,Count,N>,
) -> nalgebra::VectorN<Real,Count>
where
Count: Dim,
DefaultAllocator: Allocator<Real, Count>,
DefaultAllocator: Allocator<Real, N, Count>,
DefaultAllocator: Allocator<Real, Count, N>,
DefaultAllocator: Allocator<Real, Count, Count>,
{
let y = self.inner_pdf(xs);
vec_exp(&y) * self.fac
}
pub fn logpdf<Count>(
&self,
xs: &nalgebra::MatrixMN<Real,Count,N>,
) -> nalgebra::VectorN<Real,Count>
where
Count: Dim,
DefaultAllocator: Allocator<Real, Count>,
DefaultAllocator: Allocator<Real, N, Count>,
DefaultAllocator: Allocator<Real, Count, N>,
DefaultAllocator: Allocator<Real, Count, Count>,
{
let y = self.inner_pdf(xs);
vec_add(&y,self.fac.ln())
}
}
fn vec_exp<Real,Count>(v: &nalgebra::VectorN<Real,Count>) -> nalgebra::VectorN<Real,Count>
where
Real: RealField,
Count: Dim,
DefaultAllocator: Allocator<Real, Count>,
{
let nrows = Count::from_usize(v.nrows());
VectorN::from_iterator_generic(nrows, nalgebra::U1,
v.iter().map(|vi| vi.exp()))
}
fn vec_add<Real,Count>(v: &nalgebra::VectorN<Real,Count>, rhs: Real) -> nalgebra::VectorN<Real,Count>
where
Real: RealField,
Count: Dim,
DefaultAllocator: Allocator<Real, Count>,
{
let nrows = Count::from_usize(v.nrows());
VectorN::from_iterator_generic(nrows, nalgebra::U1,
v.iter().map(|vi| *vi + rhs))
}
fn broadcast_add<Real, R, C>(arr: &MatrixMN<Real,R,C>, vec: &VectorN<Real,C>) -> MatrixMN<Real,R,C>
where
Real: RealField,
R: Dim,
C: Dim,
DefaultAllocator: Allocator<Real, R, C>,
DefaultAllocator: Allocator<Real, C>,
{
let ndim = arr.nrows();
let nrows = R::from_usize(arr.nrows());
let ncols = C::from_usize(arr.ncols());
MatrixMN::from_iterator_generic( nrows, ncols,
arr.iter().enumerate().map(|(i,el)| {
let vi = i/ndim;
*el+vec[vi]
} )
)
}
#[cfg(test)]
mod tests {
use nalgebra as na;
use approx::relative_eq;
use crate::*;
fn sample_covariance<Real: RealField, N: Dim, K: Dim>(arr: &MatrixMN<Real,N,K>) -> nalgebra::MatrixN<Real,K>
where DefaultAllocator: Allocator<Real, N, K>,
DefaultAllocator: Allocator<Real, K, N>,
DefaultAllocator: Allocator<Real, K, K>,
DefaultAllocator: Allocator<Real, K>,
{
let mu: VectorN<Real,K> = mean_axis0(arr);
let y = broadcast_add(arr,&-mu);
let n: Real = Real::from_usize(arr.nrows()).unwrap();
let sigma = (y.transpose() * y) / (n - Real::one());
sigma
}
fn mean_axis0<Real, R, C>(arr: &MatrixMN<Real,R,C>) -> VectorN<Real,C>
where
Real: RealField,
R: Dim,
C: Dim,
DefaultAllocator: Allocator<Real, R, C>,
DefaultAllocator: Allocator<Real, C>,
{
let vec_dim: C = C::from_usize(arr.ncols());
let mut mu = VectorN::<Real,C>::zeros_generic(vec_dim, nalgebra::U1);
let scale: Real = Real::one()/na::convert(arr.nrows() as f64);
for j in 0..arr.ncols() {
let col_sum = arr
.column(j)
.iter()
.fold(Real::zero(), |acc, &x| acc + x);
mu[j] = col_sum * scale;
}
mu
}
#[test]
fn test_covar() {
use nalgebra::core::dimension::{U2, U3};
let arr = MatrixMN::<f64,U3,U2>::new(
1.0, 0.1,
2.0, 0.2,
3.0, 0.3,
);
let c = sample_covariance(&arr);
let expected = nalgebra::MatrixN::<f64,U2>::new(
0.2, 0.1,
0.1, 0.01,
);
relative_eq!(c, expected);
}
#[test]
fn test_mean_and_precision() {
let mu = na::Vector2::<f64>::new(0.0,0.0);
let precision = na::Matrix2::<f64>::new(1.0, 0.0, 0.0, 1.0);
let mvn = MultivariateNormal::from_mean_and_precision(&mu, &precision);
assert!(mu==mvn.mean());
assert!(&precision==mvn.precision());
}
#[test]
fn test_mean_axis0() {
use nalgebra::core::dimension::{U2, U4};
let a1 = MatrixMN::<f64,U2,U4>::new(1.0, 2.0, 3.0, 4.0,
5.0, 6.0, 7.0, 8.0);
let actual1: VectorN<f64,U4> = mean_axis0(&a1);
let expected1 = &[ 3.0, 4.0, 5.0, 6.0 ];
assert!(actual1.as_slice() == expected1);
let a2 = MatrixMN::<f64,U4,U2>::new(1.0, 2.0,
3.0, 4.0,
5.0, 6.0,
7.0, 8.0);
let actual2: VectorN<f64,U2> = mean_axis0(&a2);
let expected2 = &[ 4.0, 5.0 ];
assert!(actual2.as_slice() == expected2);
}
#[test]
fn test_broadcast_add() {
use nalgebra::core::dimension::{U3, U4};
let x = MatrixMN::<f64,U3,U4>::new(
1.0, 2.0, 3.0, 4.0,
5.0, 6.0, 7.0, 8.0,
100.0, 200.0, 300.0, 400.0,
);
let v = VectorN::<f64,U4>::new(-3.0, -4.0, -5.0, -3.0);
let actual = broadcast_add(&x, &v);
let expected = MatrixMN::<f64,U3,U4>::new(
-2.0, -2.0, -2.0, 1.0,
2.0, 2.0, 2.0, 5.0,
97.0, 196.0, 295.0, 397.0,
);
assert!(actual == expected);
}
#[test]
fn test_density() {
let mu = na::Vector2::<f64>::new(0.0,0.0);
let precision = na::Matrix2::<f64>::new(1.0, 0.0, 0.0, 1.0);
let mvn = MultivariateNormal::from_mean_and_precision(&mu, &precision);
let xs = na::MatrixMN::<f64,na::U2,na::U3>::new(
0.0, 1.0, 0.0,
0.0, 0.0, 1.0,
).transpose();
let results = mvn.pdf(&xs);
for i in 0..xs.nrows() {
let x = xs.row(i).clone_owned();
let di = mvn.pdf(&x)[0];
relative_eq!( di, results[i], epsilon = 1e-10 );
}
relative_eq!( results[0], 1.0/(2.0*std::f64::consts::PI).sqrt(), epsilon = 1e-10 );
relative_eq!( results[1], 1.0/(2.0*std::f64::consts::PI).sqrt() * (-0.5f64*1.0f64).exp(), epsilon = 1e-10 );
relative_eq!( results[2], 1.0/(2.0*std::f64::consts::PI).sqrt() * (-0.5f64*1.0f64).exp(), epsilon = 1e-10 );
}
}