Merge pull request #66 from CompML/features/#64_allow_list_type_fscore

modify assertion in fscore
CompML · Jan 4, 2021 · e9b5f20 · e9b5f20
2 parents a63f0e9 + 82c58b8
commit e9b5f20
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/prts/time_series_metrics/fscore.py b/prts/time_series_metrics/fscore.py
@@ -42,8 +42,13 @@ def score(self, real: np.ndarray, pred: np.ndarray) -> float:
             float: fbeta
         """
 
-        assert isinstance(real, np.ndarray)
-        assert isinstance(pred, np.ndarray)
+        assert isinstance(real, np.ndarray) or isinstance(real, list)
+        assert isinstance(pred, np.ndarray) or isinstance(pred, list)
+
+        if not isinstance(real, np.ndarray):
+            real = np.array(real)
+        if not isinstance(pred, np.ndarray):
+            pred = np.array(pred)
 
         precision = TimeSeriesPrecision(self.p_alpha, self.cardinality, self.p_bias).score(real, pred)
         recall = TimeSeriesRecall(self.r_alpha, self.cardinality, self.r_bias).score(real, pred)

diff --git a/tests/test_fscore.py b/tests/test_fscore.py
@@ -41,3 +41,13 @@ def test_fscore_function(self):
 
         score = ts_fscore(real, pred)
         self.assertEqual(score, 0.5)
+
+    def test_fscore_function_with_list(self):
+        """Teest of ts_fscore function with list type arguments.
+        """
+
+        real = [1, 1, 1, 0, 0]
+        pred = [0, 1, 0, 0, 0]
+
+        score = ts_fscore(real, pred)
+        self.assertEqual(score, 0.5)