Port deviation functions from StatsBase.jl

Author: munckymagik

src/deviation.rs

@@ -0,0 +1,386 @@
+use ndarray::{ArrayBase, Data, Dimension, Zip};
+use num_traits::{Float, Signed, ToPrimitive};
+use std::convert::Into;
+use std::ops::AddAssign;
+
+/// Extension trait for `ArrayBase` providing functions
+/// to compute different deviation measures.
+pub trait DeviationExt<A, S, D>
+where
+    S: Data<Elem = A>,
+    D: Dimension,
+{
+    fn count_eq(&self, other: &ArrayBase<S, D>) -> usize
+    where
+        A: PartialEq;
+
+    fn count_neq(&self, other: &ArrayBase<S, D>) -> usize
+    where
+        A: PartialEq;
+
+    fn sq_l2_dist(&self, other: &ArrayBase<S, D>) -> A
+    where
+        A: AddAssign + Clone + Signed;
+
+    fn l2_dist(&self, other: &ArrayBase<S, D>) -> f64
+    where
+        A: AddAssign + Clone + Signed + ToPrimitive;
+
+    fn l1_dist(&self, other: &ArrayBase<S, D>) -> A
+    where
+        A: AddAssign + Clone + Signed;
+
+    fn linf_dist(&self, other: &ArrayBase<S, D>) -> A
+    where
+        A: Clone + PartialOrd + Signed;
+
+    fn gkl_div(&self, other: &ArrayBase<S, D>) -> A
+    where
+        A: AddAssign + Clone + Float;
+
+    fn mean_abs_dev(&self, other: &ArrayBase<S, D>) -> f64
+    where
+        A: AddAssign + Clone + Signed + Into<f64>;
+
+    fn max_abs_dev(&self, other: &ArrayBase<S, D>) -> A
+    where
+        A: Clone + PartialOrd + Signed;
+
+    fn mean_sq_dev(&self, other: &ArrayBase<S, D>) -> f64
+    where
+        A: AddAssign + Clone + Signed + Into<f64>;
+
+    fn root_mean_sq_dev(&self, other: &ArrayBase<S, D>) -> f64
+    where
+        A: AddAssign + Clone + Signed + Into<f64>;
+
+    fn peak_signal_to_noise_ratio(&self, other: &ArrayBase<S, D>, maxv: A) -> f64
+    where
+        A: AddAssign + Clone + Signed + Into<f64>;
+}
+
+impl<A, S, D> DeviationExt<A, S, D> for ArrayBase<S, D>
+where
+    S: Data<Elem = A>,
+    D: Dimension,
+{
+    fn count_eq(&self, other: &ArrayBase<S, D>) -> usize
+    where
+        A: PartialEq,
+    {
+        let mut c = 0;
+
+        Zip::from(self).and(other).apply(|a, b| {
+            if a == b {
+                c += 1;
+            }
+        });
+
+        c
+    }
+
+    fn count_neq(&self, other: &ArrayBase<S, D>) -> usize
+    where
+        A: PartialEq,
+    {
+        self.len() - self.count_eq(other)
+    }
+
+    fn sq_l2_dist(&self, other: &ArrayBase<S, D>) -> A
+    where
+        A: AddAssign + Clone + Signed,
+    {
+        let mut r = A::zero();
+
+        Zip::from(self).and(other).apply(|self_i, other_i| {
+            let (a, b) = (self_i.clone(), other_i.clone());
+            let abs_diff = (a - b).abs();
+            r += abs_diff.clone() * abs_diff;
+        });
+
+        r
+    }
+
+    fn l2_dist(&self, other: &ArrayBase<S, D>) -> f64
+    where
+        A: AddAssign + Clone + Signed + ToPrimitive,
+    {
+        self.sq_l2_dist(other).to_f64().unwrap().sqrt()
+    }
+
+    fn l1_dist(&self, other: &ArrayBase<S, D>) -> A
+    where
+        A: AddAssign + Clone + Signed,
+    {
+        let mut r = A::zero();
+
+        Zip::from(self).and(other).apply(|self_i, other_i| {
+            let (a, b) = (self_i.clone(), other_i.clone());
+            r += (a - b).abs();
+        });
+
+        r
+    }
+
+    fn linf_dist(&self, other: &ArrayBase<S, D>) -> A
+    where
+        A: Clone + PartialOrd + Signed,
+    {
+        let mut max = A::zero();
+
+        Zip::from(self).and(other).apply(|self_i, other_i| {
+            let (a, b) = (self_i.clone(), other_i.clone());
+            let diff = (a - b).abs();
+            if diff > max {
+                max = diff;
+            }
+        });
+
+        max
+    }
+
+    fn gkl_div(&self, other: &ArrayBase<S, D>) -> A
+    where
+        A: AddAssign + Clone + Float,
+    {
+        let mut r = A::zero();
+
+        Zip::from(self).and(other).apply(|self_i, other_i| {
+            let (a, b) = (self_i.clone(), other_i.clone());
+            r += a * (a / b).ln() - a + b;
+        });
+
+        r
+    }
+
+    fn mean_abs_dev(&self, other: &ArrayBase<S, D>) -> f64
+    where
+        A: AddAssign + Clone + Signed + Into<f64>,
+    {
+        let a: f64 = self.l1_dist(other).into();
+        let b = self.len().to_f64().unwrap();
+        a / b
+    }
+
+    #[inline]
+    fn max_abs_dev(&self, other: &ArrayBase<S, D>) -> A
+    where
+        A: Clone + PartialOrd + Signed,
+    {
+        self.linf_dist(other)
+    }
+
+    fn mean_sq_dev(&self, other: &ArrayBase<S, D>) -> f64
+    where
+        A: AddAssign + Clone + Signed + Into<f64>,
+    {
+        let a: f64 = self.sq_l2_dist(other).into();
+        let b = self.len().to_f64().unwrap();
+        a / b
+    }
+
+    fn root_mean_sq_dev(&self, other: &ArrayBase<S, D>) -> f64
+    where
+        A: AddAssign + Clone + Signed + Into<f64>,
+    {
+        self.mean_sq_dev(other).sqrt()
+    }
+
+    fn peak_signal_to_noise_ratio(&self, other: &ArrayBase<S, D>, maxv: A) -> f64
+    where
+        A: AddAssign + Clone + Signed + Into<f64>,
+    {
+        let maxv_f: f64 = maxv.into();
+        10. * f64::log10(maxv_f * maxv_f / self.mean_sq_dev(&other))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use approx::assert_abs_diff_eq;
+    use ndarray::*;
+    use ndarray_rand::RandomExt;
+    use rand::distributions::Uniform;
+    use std::f64;
+
+    #[test]
+    fn test_count_eq() {
+        let a = array![0., 0.];
+        let b = array![1., 0.];
+        let c = array![0., 1.];
+        let d = array![1., 1.];
+
+        assert_eq!(a.count_eq(&a), 2);
+        assert_eq!(a.count_eq(&b), 1);
+        assert_eq!(a.count_eq(&c), 1);
+        assert_eq!(a.count_eq(&d), 0);
+    }
+
+    #[test]
+    fn test_count_neq() {
+        let a = array![0., 0.];
+        let b = array![1., 0.];
+        let c = array![0., 1.];
+        let d = array![1., 1.];
+
+        assert_eq!(a.count_neq(&a), 0);
+        assert_eq!(a.count_neq(&b), 1);
+        assert_eq!(a.count_neq(&c), 1);
+        assert_eq!(a.count_neq(&d), 2);
+    }
+
+    #[test]
+    fn test_sq_l2_dist() {
+        let a = array![0., 1., 4., 2.];
+        let b = array![1., 1., 2., 4.];
+
+        assert_eq!(a.sq_l2_dist(&b), 9.);
+    }
+
+    #[test]
+    fn test_l2_dist() {
+        let a = array![0., 1., 4., 2.];
+        let b = array![1., 1., 2., 4.];
+
+        assert_eq!(a.l2_dist(&b), 3.);
+    }
+
+    #[test]
+    fn test_l1_dist() {
+        let a = array![0., 1., 4., 2.];
+        let b = array![1., 1., 2., 4.];
+
+        assert_eq!(a.l1_dist(&b), 5.);
+    }
+
+    #[test]
+    fn test_linf_dist() {
+        let a = array![0., 0.];
+        let b = array![1., 0.];
+        let c = array![1., 2.];
+
+        assert_eq!(a.linf_dist(&a), 0.);
+
+        assert_eq!(a.linf_dist(&b), 1.);
+        assert_eq!(b.linf_dist(&a), 1.);
+
+        assert_eq!(a.linf_dist(&c), 2.);
+        assert_eq!(c.linf_dist(&a), 2.);
+    }
+
+    #[test]
+    fn test_gkl_div() {
+        let a = Array::random((5,), Uniform::new(0., 1.));
+        let b = Array::random((5,), Uniform::new(1., 2.));
+        let c = Array::random((5,), Uniform::new(-1., 0.));
+
+        assert_eq!(a.gkl_div(&a), 0.);
+        assert!(a.gkl_div(&b) > 0.);
+        assert!(b.gkl_div(&a) > 0.);
+        assert_ne!(a.gkl_div(&b), b.gkl_div(&a));
+
+        // TODO: what is the sign check logic doing in StatsBase.jl impl?
+        assert!(f64::is_nan(a.gkl_div(&c)));
+    }
+
+    #[test]
+    fn test_mean_abs_dev() {
+        let a = array![1., 1.];
+        let b = array![3., 5.];
+
+        assert_eq!(a.mean_abs_dev(&a), 0.);
+        assert_eq!(a.mean_abs_dev(&b), 3.);
+        assert_eq!(b.mean_abs_dev(&a), 3.);
+    }
+
+    #[test]
+    fn test_max_abs_dev() {
+        // This is effectively an alias for linf_dist, so not retesting deeply
+        let a = array![0., 0.];
+        let b = array![2., 4.];
+
+        assert_eq!(a.max_abs_dev(&a), 0.);
+        assert_eq!(a.max_abs_dev(&b), 4.);
+        assert_eq!(b.max_abs_dev(&a), 4.);
+    }
+
+    #[test]
+    fn test_mean_sq_dev() {
+        let a = array![1., 1.];
+        let b = array![3., 5.];
+
+        assert_eq!(a.mean_sq_dev(&a), 0.);
+        assert_eq!(a.mean_sq_dev(&b), 10.);
+        assert_eq!(b.mean_sq_dev(&a), 10.);
+    }
+
+    #[test]
+    fn test_root_mean_sq_dev() {
+        let a = array![1., 1.];
+        let b = array![3., 5.];
+
+        assert_eq!(a.root_mean_sq_dev(&a), 0.);
+        assert_abs_diff_eq!(a.root_mean_sq_dev(&b), 10.0.sqrt());
+        assert_abs_diff_eq!(b.root_mean_sq_dev(&a), 10.0.sqrt());
+    }
+
+    #[test]
+    fn test_peak_signal_to_noise_ratio() {
+        let a = array![1., 1.];
+        assert!(a.peak_signal_to_noise_ratio(&a, 1.).is_infinite());
+
+        let a = array![1., 2., 3., 4., 5., 6., 7.];
+        let b = array![1., 3., 3., 4., 6., 7., 8.];
+        let maxv = 8.;
+        let expected = 20. * Float::log10(maxv) - 10. * Float::log10(a.mean_sq_dev(&b));
+        let actual = a.peak_signal_to_noise_ratio(&b, maxv);
+
+        assert_abs_diff_eq!(actual, expected);
+    }
+
+    #[test]
+    fn test_deviations_with_n_by_m_ints() {
+        let a = array![[0, 1], [4, 2]];
+        let b = array![[1, 1], [2, 4]];
+
+        assert_eq!(a.count_eq(&a), 4);
+        assert_eq!(a.count_neq(&a), 0);
+        assert_eq!(a.sq_l2_dist(&b), 9);
+        assert_eq!(a.l2_dist(&b), 3.);
+        assert_eq!(a.l1_dist(&b), 5);
+        assert_eq!(a.linf_dist(&b), 2);
+
+        assert_abs_diff_eq!(a.mean_abs_dev(&b), 1.25);
+        assert_eq!(a.max_abs_dev(&b), 2);
+        assert_abs_diff_eq!(a.mean_sq_dev(&b), 2.25);
+        assert_abs_diff_eq!(a.root_mean_sq_dev(&b), 1.5);
+        assert_abs_diff_eq!(
+            a.peak_signal_to_noise_ratio(&b, 4),
+            8.519374645445623,
+            epsilon = f64::EPSILON
+        );
+
+        // TODO: gkl_div
+    }
+
+    #[test]
+    fn test_deviations_with_empty_inputs() {
+        let a: Array1<f64> = array![];
+
+        assert_eq!(a.count_eq(&a), 0);
+        assert_eq!(a.count_neq(&a), 0);
+        assert_eq!(a.sq_l2_dist(&a), 0.);
+        assert_eq!(a.l2_dist(&a), 0.);
+        assert_eq!(a.l1_dist(&a), 0.);
+        assert_eq!(a.linf_dist(&a), 0.);
+
+        assert!(a.mean_abs_dev(&a).is_nan());
+        assert_eq!(a.max_abs_dev(&a), 0.);
+        assert!(a.mean_sq_dev(&a).is_nan());
+        assert!(a.root_mean_sq_dev(&a).is_nan());
+        assert!(a.peak_signal_to_noise_ratio(&a, 0.).is_nan());
+
+        // TODO: gkl_div
+    }
+}