deviation: Implement count_eq

Author: munckymagik

src/deviation.rs

@@ -0,0 +1,53 @@
+use ndarray::{ArrayBase, Data, Dimension, Zip};
+
+/// 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;
+}
+
+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
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use ndarray::array;
+
+    #[test]
+    fn test_count_eq() {
+        let a = array![1., 2., 3., 4., 5., 6., 7.];
+        let b = array![1., 3., 3., 4., 6., 7., 8.];
+        let c = array![2., 4., 4., 5., 7., 8., 9.];
+        let d = array![[1, 2], [3, 4], [5, 6]];
+        let e = array![[1, 2], [4, 3], [5, 6]];
+
+        assert_eq!(a.count_eq(&b), 3);
+        assert_eq!(b.count_eq(&c), 0);
+        assert_eq!(d.count_eq(&e), 4);
+    }
+}

src/lib.rs

@@ -28,6 +28,7 @@
 //! [`StatsBase.jl`]: https://juliastats.github.io/StatsBase.jl/latest/
 
 pub use crate::correlation::CorrelationExt;
+pub use crate::deviation::DeviationExt;
 pub use crate::entropy::EntropyExt;
 pub use crate::histogram::HistogramExt;
 pub use crate::maybe_nan::{MaybeNan, MaybeNanExt};
@@ -69,6 +70,7 @@ mod private {
 }
 
 mod correlation;
+mod deviation;
 mod entropy;
 pub mod errors;
 pub mod histogram;