Handle multiclass precision/recall (#152)

* handle multiclass precision/recall

src/metrics/precision.rs

@@ -18,6 +18,8 @@
 //!
 //! <script src="https://polyfill.io/v3/polyfill.min.js?features=es6"></script>
 //! <script id="MathJax-script" async src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js"></script>
+use std::collections::HashSet;
+
 #[cfg(feature = "serde")]
 use serde::{Deserialize, Serialize};
 
@@ -42,34 +44,33 @@ impl Precision {
             );
         }
 
+        let mut classes = HashSet::new();
+        for i in 0..y_true.len() {
+            classes.insert(y_true.get(i).to_f64_bits());
+        }
+        let classes = classes.len();
+
         let mut tp = 0;
-        let mut p = 0;
-        let n = y_true.len();
-        for i in 0..n {
-            if y_true.get(i) != T::zero() && y_true.get(i) != T::one() {
-                panic!(
-                    "Precision can only be applied to binary classification: {}",
-                    y_true.get(i)
-                );
-            }
-
-            if y_pred.get(i) != T::zero() && y_pred.get(i) != T::one() {
-                panic!(
-                    "Precision can only be applied to binary classification: {}",
-                    y_pred.get(i)
-                );
-            }
-
-            if y_pred.get(i) == T::one() {
-                p += 1;
-
-                if y_true.get(i) == T::one() {
+        let mut fp = 0;
+        for i in 0..y_true.len() {
+            if y_pred.get(i) == y_true.get(i) {
+                if classes == 2 {
+                    if y_true.get(i) == T::one() {
+                        tp += 1;
+                    }
+                } else {
                     tp += 1;
                 }
+            } else if classes == 2 {
+                if y_true.get(i) == T::one() {
+                    fp += 1;
+                }
+            } else {
+                fp += 1;
             }
         }
 
-        T::from_i64(tp).unwrap() / T::from_i64(p).unwrap()
+        T::from_i64(tp).unwrap() / (T::from_i64(tp).unwrap() + T::from_i64(fp).unwrap())
     }
 }
 
@@ -88,5 +89,24 @@ mod tests {
 
         assert!((score1 - 0.5).abs() < 1e-8);
         assert!((score2 - 1.0).abs() < 1e-8);
+
+        let y_pred: Vec<f64> = vec![0., 0., 1., 1., 1., 1.];
+        let y_true: Vec<f64> = vec![0., 1., 1., 0., 1., 0.];
+
+        let score3: f64 = Precision {}.get_score(&y_pred, &y_true);
+        assert!((score3 - 0.5).abs() < 1e-8);
+    }
+
+    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
+    #[test]
+    fn precision_multiclass() {
+        let y_true: Vec<f64> = vec![0., 0., 0., 1., 1., 1., 2., 2., 2.];
+        let y_pred: Vec<f64> = vec![0., 1., 2., 0., 1., 2., 0., 1., 2.];
+
+        let score1: f64 = Precision {}.get_score(&y_pred, &y_true);
+        let score2: f64 = Precision {}.get_score(&y_pred, &y_pred);
+
+        assert!((score1 - 0.333333333).abs() < 1e-8);
+        assert!((score2 - 1.0).abs() < 1e-8);
     }
 }