Merge pull request #71 from CompML/feature/#70/fix_prepare_data_assert

fix _prepare_data
CompML · Feb 17, 2021 · 0ec80c1 · 0ec80c1
2 parents dcd5764 + 4047c96
commit 0ec80c1
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 2 deletions.
diff --git a/prts/base/time_series_metrics.py b/prts/base/time_series_metrics.py
@@ -127,8 +127,12 @@ def _shift(self, arr, num, fill_value=np.nan):
     def _prepare_data(self, values_real, values_pred):
 
         assert len(values_real) == len(values_pred)
-        assert np.allclose(np.unique(values_real), np.array([0, 1]))
-        assert np.allclose(np.unique(values_pred), np.array([0, 1]))
+        assert np.allclose(np.unique(values_real), np.array([0, 1])) or np.allclose(
+            np.unique(values_real), np.array([1])
+        )
+        assert np.allclose(np.unique(values_pred), np.array([0, 1])) or np.allclose(
+            np.unique(values_pred), np.array([1])
+        )
 
         predicted_anomalies_ = np.argwhere(values_pred == 1).ravel()
         predicted_anomalies_shift_forward = self._shift(predicted_anomalies_, 1, fill_value=predicted_anomalies_[0])

diff --git a/tests/test_precision.py b/tests/test_precision.py
@@ -135,3 +135,23 @@ def test_precision_function_with_invalid_bias(self):
 
         with self.assertRaises(Exception):
             ts_precision(real, pred, bias="Invalid")
+
+    def test_precision_function_with_all_zeros(self):
+        """Test of ts_precision function with all zero values
+        """
+
+        real = np.array([0, 0, 0, 0, 0])
+        pred = np.array([0, 0, 0, 0, 0])
+
+        with self.assertRaises(Exception):
+            ts_precision(real, pred)
+
+    def test_precision_function_with_all_ones(self):
+        """Test of ts_precision function with all zero values
+        """
+
+        real = np.array([1, 1, 1, 1, 1])
+        pred = np.array([1, 1, 1, 1, 1])
+
+        self.assertEqual(ts_precision(real, pred), 1.0)
+