Feat/multivariate dimension generic: electric boogaloo

src/distribution/dirichlet.rs

@@ -345,10 +345,15 @@ fn is_valid_alpha(a: &[f64]) -> bool {
 #[rustfmt::skip]
 #[cfg(test)]
 mod tests {
-    use nalgebra::{dvector, vector, DimMin, OVector};
+    use std::fmt::{Debug, Display};
 
-    use super::*;
-    use crate::distribution::Continuous;
+    use nalgebra::{dmatrix, dvector, vector, DimMin, OVector};
+
+    use super::is_valid_alpha;
+    use crate::{
+        distribution::{Continuous, Dirichlet},
+        statistics::{MeanN, VarianceN},
+    };
 
     fn try_create<D>(alpha: OVector<f64, D>) -> Dirichlet<D>
     where
@@ -369,15 +374,16 @@ mod tests {
         assert!(dd.is_err());
     }
 
-    fn test_almost<F, D>(alpha: OVector<f64, D>, expected: f64, acc: f64, eval: F)
+    fn test_almost<F, T, D>(alpha: OVector<f64, D>, expected: T, acc: f64, eval: F)
     where
-        F: FnOnce(Dirichlet<D>) -> f64,
+        T: Debug + Display + approx::RelativeEq<Epsilon = f64>,
+        F: FnOnce(Dirichlet<D>) -> T,
         D: DimMin<D, Output = D>,
         nalgebra::DefaultAllocator: nalgebra::allocator::Allocator<f64, D>,
     {
         let dd = try_create(alpha);
         let x = eval(dd);
-        assert_almost_eq!(expected, x, acc);
+        assert_relative_eq!(expected, x, epsilon = acc);
     }
 
     #[test]
@@ -406,26 +412,51 @@ mod tests {
         bad_create_case(vector![0.001, f64::INFINITY, 3756.0]); // moved to bad case as this is degenerate
     }
 
-    // #[test]
-    // fn test_mean() {
-    //     let n = Dirichlet::new_with_param(0.3, 5).unwrap();
-    //     let res = n.mean();
-    //     for x in res {
-    //         assert_eq!(x, 0.3 / 1.5);
-    //     }
-    // }
+    #[test]
+    fn test_mean() {
+        let mean = |dd: Dirichlet<_>| dd.mean().unwrap();
 
-    // #[test]
-    // fn test_variance() {
-    //     let alpha = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0];
-    //     let sum = alpha.iter().fold(0.0, |acc, x| acc + x);
-    //     let n = Dirichlet::new(&alpha).unwrap();
-    //     let res = n.variance();
-    //     for i in 1..11 {
-    //         let f = i as f64;
-    //         assert_almost_eq!(res[i-1], f * (sum - f) / (sum * sum * (sum + 1.0)), 1e-15);
-    //     }
-    // }
+        test_almost(vec![0.5; 5].into(), vec![1.0 / 5.0; 5].into(), 1e-15, mean);
+
+        test_almost(
+            dvector![0.1, 0.2, 0.3, 0.4],
+            dvector![0.1, 0.2, 0.3, 0.4],
+            1e-15,
+            mean,
+        );
+
+        test_almost(
+            dvector![1.0, 2.0, 3.0, 4.0],
+            dvector![0.1, 0.2, 0.3, 0.4],
+            1e-15,
+            mean,
+        );
+    }
+
+    #[test]
+    fn test_variance() {
+        let variance = |dd: Dirichlet<_>| dd.variance().unwrap();
+
+        test_almost(
+            dvector![1.0, 2.0],
+            dmatrix![0.055555555555555, -0.055555555555555;
+                    -0.055555555555555,  0.055555555555555;
+            ],
+            1e-15,
+            variance,
+        );
+
+        test_almost(
+            dvector![0.1, 0.2, 0.3, 0.4],
+            dmatrix![0.045, -0.010, -0.015, -0.020;
+                    -0.010,  0.080, -0.030, -0.040;
+                    -0.015, -0.030,  0.105, -0.060;
+                    -0.020, -0.040, -0.060,  0.120;
+            ],
+            1e-15,
+            variance,
+        );
+    }
 
     // #[test]
     // fn test_std_dev() {

src/distribution/internal.rs

@@ -1,5 +1,8 @@
+use nalgebra::{Dim, OVector};
 use num_traits::Num;
 
+use crate::StatsError;
+
 /// Returns true if there are no elements in `x` in `arr`
 /// such that `x <= 0.0` or `x` is `f64::NAN` and `sum(arr) > 0.0`.
 /// IF `incl_zero` is true, it tests for `x < 0.0` instead of `x <= 0.0`
@@ -14,6 +17,37 @@ pub fn is_valid_multinomial(arr: &[f64], incl_zero: bool) -> bool {
     sum != 0.0
 }
 
+pub fn check_multinomial<D>(arr: &OVector<f64, D>, accept_zeroes: bool) -> crate::Result<()>
+where
+    D: Dim,
+    nalgebra::DefaultAllocator: nalgebra::allocator::Allocator<f64, D>,
+{
+    if arr.len() < 2 {
+        return Err(StatsError::BadParams);
+    }
+    let mut sum = 0.0;
+    for &x in arr.iter() {
+        #[allow(clippy::if_same_then_else)]
+        if x.is_nan() {
+            return Err(StatsError::BadParams);
+        } else if x.is_infinite() {
+            return Err(StatsError::BadParams);
+        } else if x < 0.0 {
+            return Err(StatsError::BadParams);
+        } else if x == 0.0 && !accept_zeroes {
+            return Err(StatsError::BadParams);
+        } else {
+            sum += x;
+        }
+    }
+
+    if sum != 0.0 {
+        Ok(())
+    } else {
+        Err(StatsError::BadParams)
+    }
+}
+
 /// Implements univariate function bisection searching for criteria
 /// ```text
 /// smallest k such that f(k) >= z
@@ -225,12 +259,16 @@ pub mod test {
 
         let invalid = [1.0, f64::NAN, 3.0];
         assert!(!is_valid_multinomial(&invalid, true));
+        assert!(check_multinomial(&invalid.to_vec().into(), true).is_err());
         let invalid2 = [-2.0, 5.0, 1.0, 6.2];
         assert!(!is_valid_multinomial(&invalid2, true));
+        assert!(check_multinomial(&invalid2.to_vec().into(), true).is_err());
         let invalid3 = [0.0, 0.0, 0.0];
         assert!(!is_valid_multinomial(&invalid3, true));
+        assert!(check_multinomial(&invalid3.to_vec().into(), true).is_err());
         let valid = [5.2, 0.0, 1e-15, 1000000.12];
         assert!(is_valid_multinomial(&valid, true));
+        assert!(check_multinomial(&valid.to_vec().into(), true).is_ok());
     }
 
     #[test]