Let eval metrics work w/ UnivariateTimeSeries. (#74)

salesforce · Feb 16, 2022 · 5a098ba · 5a098ba
1 parent 6aa1524
commit 5a098ba
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 16 deletions.
diff --git a/merlion/evaluate/anomaly.py b/merlion/evaluate/anomaly.py
@@ -1,5 +1,5 @@
 #
-# Copyright (c) 2021 salesforce.com, inc.
+# Copyright (c) 2022 salesforce.com, inc.
 # All rights reserved.
 # SPDX-License-Identifier: BSD-3-Clause
 # For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
@@ -16,7 +16,7 @@
 import pandas as pd
 
 from merlion.evaluate.base import EvaluatorBase, EvaluatorConfig
-from merlion.utils import TimeSeries
+from merlion.utils import TimeSeries, UnivariateTimeSeries
 
 
 def scaled_sigmoid(x, scale=2.5):
@@ -168,7 +168,11 @@ def nab_score(self, tp_weight=1.0, fp_weight=0.11, fn_weight=1.0, tn_weight=0.0)
 
 
 def accumulate_tsad_score(
-    ground_truth: TimeSeries, predict: TimeSeries, max_early_sec=None, max_delay_sec=None, metric=None
+    ground_truth: Union[TimeSeries, UnivariateTimeSeries],
+    predict: Union[TimeSeries, UnivariateTimeSeries],
+    max_early_sec=None,
+    max_delay_sec=None,
+    metric=None,
 ) -> Union[TSADScoreAccumulator, float]:
     """
     Computes the components required to compute multiple different types of
@@ -191,6 +195,8 @@ def accumulate_tsad_score(
         returns a ``float``. The `TSADScoreAccumulator` object is returned if
         ``metric`` is ``None``.
     """
+    ground_truth = ground_truth.to_ts() if isinstance(ground_truth, UnivariateTimeSeries) else ground_truth
+    predict = predict.to_ts() if isinstance(predict, UnivariateTimeSeries) else predict
     assert (
         ground_truth.dim == 1 and predict.dim == 1
     ), "Can only evaluate anomaly scores when ground truth and prediction are single-variable time series."

diff --git a/merlion/evaluate/forecast.py b/merlion/evaluate/forecast.py
@@ -1,5 +1,5 @@
 #
-# Copyright (c) 2021 salesforce.com, inc.
+# Copyright (c) 2022 salesforce.com, inc.
 # All rights reserved.
 # SPDX-License-Identifier: BSD-3-Clause
 # For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
@@ -16,7 +16,7 @@
 
 from merlion.evaluate.base import EvaluatorBase, EvaluatorConfig
 from merlion.models.forecast.base import ForecasterBase
-from merlion.utils import TimeSeries
+from merlion.utils import TimeSeries, UnivariateTimeSeries
 from merlion.utils.resample import granularity_str_to_seconds
 
 
@@ -29,26 +29,25 @@ class ForecastScoreAccumulator:
 
     def __init__(
         self,
-        ground_truth: TimeSeries,
-        predict: TimeSeries,
-        insample: TimeSeries = None,
+        ground_truth: Union[UnivariateTimeSeries, TimeSeries],
+        predict: Union[UnivariateTimeSeries, TimeSeries],
+        insample: Union[UnivariateTimeSeries, TimeSeries] = None,
         periodicity: int = 1,
         ub: TimeSeries = None,
         lb: TimeSeries = None,
     ):
         """
         :param ground_truth: ground truth time series
         :param predict: predicted truth time series
-        :param insample (optional): time series used for training model.
-            This value is used for computing MSES, MSIS
+        :param insample (optional): time series used for training model. This value is used for computing MSES, MSIS
         :param periodicity (optional): periodicity. m=1 indicates the non-seasonal time series,
-            whereas m>1 indicates seasonal time series.
-            This value is used for computing MSES, MSIS.
-        :param ub (optional): upper bound of 95% prediction interval. This value is used for
-            computing MSIS
-        :param lb (optional): lower bound of 95% prediction interval. This value is used for
-            computing MSIS
+            whereas m>1 indicates seasonal time series. This value is used for computing MSES, MSIS.
+        :param ub (optional): upper bound of 95% prediction interval. This value is used for computing MSIS
+        :param lb (optional): lower bound of 95% prediction interval. This value is used for computing MSIS
         """
+        ground_truth = ground_truth.to_ts() if isinstance(ground_truth, UnivariateTimeSeries) else ground_truth
+        predict = predict.to_ts() if isinstance(predict, UnivariateTimeSeries) else predict
+        insample = insample.to_ts() if isinstance(insample, UnivariateTimeSeries) else insample
         t0, tf = predict.t0, predict.tf
         ground_truth = ground_truth.window(t0, tf, include_tf=True).align()
         self.ground_truth = ground_truth