feature: vectorize metrics computation

etna/metrics/__init__.py

@@ -6,6 +6,8 @@
 
 from etna.metrics.base import Metric
 from etna.metrics.base import MetricAggregationMode
+from etna.metrics.base import MetricFunctionSignature
+from etna.metrics.functional_metrics import FunctionalMetricMode
 from etna.metrics.functional_metrics import mape
 from etna.metrics.functional_metrics import max_deviation
 from etna.metrics.functional_metrics import rmse

etna/metrics/base.py

@@ -8,6 +8,7 @@
 
 import numpy as np
 import pandas as pd
+from typing_extensions import assert_never
 
 from etna.core import BaseMixin
 from etna.datasets.tsdataset import TSDataset
@@ -27,6 +28,23 @@ def _missing_(cls, value):
         )
 
 
+class MetricFunctionSignature(str, Enum):
+    """Enum for different metric function signatures."""
+
+    #: function should expect arrays or y_pred and y_true with length ``n_timestamps`` and return scalar
+    array_to_scalar = "array_to_scalar"
+
+    #: function should expect matrices of y_pred and y_true with shape ``(n_timestamps, n_segments)``
+    #: and return vector of length ``n_segments``
+    matrix_to_array = "matrix_to_array"
+
+    @classmethod
+    def _missing_(cls, value):
+        raise NotImplementedError(
+            f"{value} is not a valid {cls.__name__}. Only {', '.join([repr(m.value) for m in cls])} signatures allowed"
+        )
+
+
 class AbstractMetric(ABC):
     """Abstract class for metric."""
 
@@ -66,6 +84,7 @@ def greater_is_better(self) -> Optional[bool]:
         pass
 
 
+# TODO: add tests on metric_fn_signature
 class Metric(AbstractMetric, BaseMixin):
     """
     Base class for all the multi-segment metrics.
@@ -74,7 +93,13 @@ class Metric(AbstractMetric, BaseMixin):
     dataset and aggregates it according to mode.
     """
 
-    def __init__(self, metric_fn: Callable[..., float], mode: str = MetricAggregationMode.per_segment, **kwargs):
+    def __init__(
+        self,
+        metric_fn: Callable,
+        mode: str = MetricAggregationMode.per_segment,
+        metric_fn_signature: str = "array_to_scalar",
+        **kwargs,
+    ):
         """
         Init Metric.
 
@@ -89,21 +114,29 @@ def __init__(self, metric_fn: Callable[..., float], mode: str = MetricAggregatio
 
             * if "per-segment" -- does not aggregate metrics
 
+        metric_fn_signature:
+            type of signature of ``metric_fn`` (see :py:class:`~etna.metrics.base.MetricFunctionSignature`)
         kwargs:
             functional metric's params
 
         Raises
         ------
         NotImplementedError:
-            it non existent mode is used
+            If non-existent ``mode`` is used.
+        NotImplementedError:
+            If non-existent ``metric_fn_signature`` is used.
         """
-        self.metric_fn = metric_fn
-        self.kwargs = kwargs
-        if MetricAggregationMode(mode) == MetricAggregationMode.macro:
+        if MetricAggregationMode(mode) is MetricAggregationMode.macro:
             self._aggregate_metrics = self._macro_average
-        elif MetricAggregationMode(mode) == MetricAggregationMode.per_segment:
+        elif MetricAggregationMode(mode) is MetricAggregationMode.per_segment:
             self._aggregate_metrics = self._per_segment_average
+
+        self._metric_fn_signature = MetricFunctionSignature(metric_fn_signature)
+
+        self.metric_fn = metric_fn
+        self.kwargs = kwargs
         self.mode = mode
+        self.metric_fn_signature = metric_fn_signature
 
     @property
     def name(self) -> str:
@@ -276,13 +309,20 @@ def __call__(self, y_true: TSDataset, y_pred: TSDataset) -> Union[float, Dict[st
         df_true = y_true[:, :, "target"].sort_index(axis=1)
         df_pred = y_pred[:, :, "target"].sort_index(axis=1)
 
-        metrics_per_segment = {}
         segments = df_true.columns.get_level_values("segment").unique()
 
-        for i, segment in enumerate(segments):
-            cur_y_true = df_true.iloc[:, i]
-            cur_y_pred = df_pred.iloc[:, i]
-            metrics_per_segment[segment] = self.metric_fn(y_true=cur_y_true, y_pred=cur_y_pred, **self.kwargs)
+        if self._metric_fn_signature is MetricFunctionSignature.array_to_scalar:
+            metrics_per_segment = {}
+            for i, segment in enumerate(segments):
+                cur_y_true = df_true.iloc[:, i].values
+                cur_y_pred = df_pred.iloc[:, i].values
+                metrics_per_segment[segment] = self.metric_fn(y_true=cur_y_true, y_pred=cur_y_pred, **self.kwargs)
+        elif self._metric_fn_signature is MetricFunctionSignature.matrix_to_array:
+            values = self.metric_fn(y_true=df_true.values, y_pred=df_pred.values, **self.kwargs)
+            metrics_per_segment = dict(zip(segments, values))
+        else:
+            assert_never(self._metric_fn_signature)
+
         metrics = self._aggregate_metrics(metrics_per_segment)
         return metrics
 

etna/metrics/functional_metrics.py

@@ -1,14 +1,33 @@
+from enum import Enum
 from functools import partial
-from typing import List
+from typing import Sequence
 from typing import Union
 
 import numpy as np
 from sklearn.metrics import mean_squared_error as mse
+from typing_extensions import assert_never
 
-ArrayLike = List[Union[float, List[float]]]
+ArrayLike = Union[float, Sequence[float], Sequence[Sequence[float]]]
 
 
-def mape(y_true: ArrayLike, y_pred: ArrayLike, eps: float = 1e-15) -> float:
+# TODO: протестировать новый режим
+class FunctionalMetricMode(str, Enum):
+    """Enum for different functional metric multioutput modes."""
+
+    #: Compute one scalar value taking into account all outputs.
+    joint = "joint"
+
+    #: Compute one value per each output.
+    per_output = "per_output"
+
+    @classmethod
+    def _missing_(cls, value):
+        raise NotImplementedError(
+            f"{value} is not a valid {cls.__name__}. Only {', '.join([repr(m.value) for m in cls])} options allowed"
+        )
+
+
+def mape(y_true: ArrayLike, y_pred: ArrayLike, eps: float = 1e-15, mode: str = "joint") -> ArrayLike:
     """Mean absolute percentage error.
 
     `Wikipedia entry on the Mean absolute percentage error
@@ -26,14 +45,19 @@ def mape(y_true: ArrayLike, y_pred: ArrayLike, eps: float = 1e-15) -> float:
 
         Estimated target values.
 
-    eps: float=1e-15
+    eps:
         MAPE is undefined for ``y_true[i]==0`` for any ``i``, so all zeros ``y_true[i]`` are
         clipped to ``max(eps, abs(y_true))``.
 
+    mode:
+        Defines aggregating of multiple output values
+        (see :py:class:`~etna.metrics.functional_metrics.FunctionalMetricMode`).
+
     Returns
     -------
-    float
-        A non-negative floating point value (the best value is 0.0).
+    :
+        A non-negative floating point value (the best value is 0.0), or an array of floating point values,
+        one for each individual target.
     """
     y_true_array, y_pred_array = np.asarray(y_true), np.asarray(y_pred)
 
@@ -42,10 +66,16 @@ def mape(y_true: ArrayLike, y_pred: ArrayLike, eps: float = 1e-15) -> float:
 
     y_true_array = y_true_array.clip(eps)
 
-    return np.mean(np.abs((y_true_array - y_pred_array) / y_true_array)) * 100
+    mode_enum = FunctionalMetricMode(mode)
+    if mode_enum is FunctionalMetricMode.joint:
+        return np.mean(np.abs((y_true_array - y_pred_array) / y_true_array)) * 100
+    elif mode_enum is FunctionalMetricMode.per_output:
+        return np.mean(np.abs((y_true_array - y_pred_array) / y_true_array), axis=0) * 100
+    else:
+        assert_never(mode_enum)
 
 
-def smape(y_true: ArrayLike, y_pred: ArrayLike, eps: float = 1e-15) -> float:
+def smape(y_true: ArrayLike, y_pred: ArrayLike, eps: float = 1e-15, mode: str = "joint") -> ArrayLike:
     """Symmetric mean absolute percentage error.
 
     `Wikipedia entry on the Symmetric mean absolute percentage error
@@ -70,22 +100,35 @@ def smape(y_true: ArrayLike, y_pred: ArrayLike, eps: float = 1e-15) -> float:
         SMAPE is undefined for ``y_true[i] + y_pred[i] == 0`` for any ``i``, so all zeros ``y_true[i] + y_pred[i]`` are
         clipped to ``max(eps, abs(y_true) + abs(y_pred))``.
 
+    mode:
+        Defines aggregating of multiple output values
+        (see :py:class:`~etna.metrics.functional_metrics.FunctionalMetricMode`).
+
     Returns
     -------
-    float
-        A non-negative floating point value (the best value is 0.0).
+    :
+        A non-negative floating point value (the best value is 0.0), or an array of floating point values,
+        one for each individual target.
     """
     y_true_array, y_pred_array = np.asarray(y_true), np.asarray(y_pred)
 
     if len(y_true_array.shape) != len(y_pred_array.shape):
         raise ValueError("Shapes of the labels must be the same")
 
-    return 100 * np.mean(
-        2 * np.abs(y_pred_array - y_true_array) / (np.abs(y_true_array) + np.abs(y_pred_array)).clip(eps)
-    )
+    mode_enum = FunctionalMetricMode(mode)
+    if mode_enum is FunctionalMetricMode.joint:
+        return 100 * np.mean(
+            2 * np.abs(y_pred_array - y_true_array) / (np.abs(y_true_array) + np.abs(y_pred_array)).clip(eps)
+        )
+    elif mode_enum is FunctionalMetricMode.per_output:
+        return 100 * np.mean(
+            2 * np.abs(y_pred_array - y_true_array) / (np.abs(y_true_array) + np.abs(y_pred_array)).clip(eps), axis=0
+        )
+    else:
+        assert_never(mode_enum)
 
 
-def sign(y_true: ArrayLike, y_pred: ArrayLike) -> float:
+def sign(y_true: ArrayLike, y_pred: ArrayLike, mode: str = "joint") -> ArrayLike:
     """Sign error metric.
 
     .. math::
@@ -103,20 +146,31 @@ def sign(y_true: ArrayLike, y_pred: ArrayLike) -> float:
 
         Estimated target values.
 
+    mode:
+        Defines aggregating of multiple output values
+        (see :py:class:`~etna.metrics.functional_metrics.FunctionalMetricMode`).
+
     Returns
     -------
-    float
-        A floating point value (the best value is 0.0).
+    :
+        A non-negative floating point value (the best value is 0.0), or an array of floating point values,
+        one for each individual target.
     """
     y_true_array, y_pred_array = np.asarray(y_true), np.asarray(y_pred)
 
     if len(y_true_array.shape) != len(y_pred_array.shape):
         raise ValueError("Shapes of the labels must be the same")
 
-    return np.mean(np.sign(y_true_array - y_pred_array))
+    mode_enum = FunctionalMetricMode(mode)
+    if mode_enum is FunctionalMetricMode.joint:
+        return np.mean(np.sign(y_true_array - y_pred_array))
+    elif mode_enum is FunctionalMetricMode.per_output:
+        return np.mean(np.sign(y_true_array - y_pred_array), axis=0)
+    else:
+        assert_never(mode_enum)
 
 
-def max_deviation(y_true: ArrayLike, y_pred: ArrayLike) -> float:
+def max_deviation(y_true: ArrayLike, y_pred: ArrayLike, mode: str = "joint") -> ArrayLike:
     """Max Deviation metric.
 
     Parameters
@@ -131,25 +185,36 @@ def max_deviation(y_true: ArrayLike, y_pred: ArrayLike) -> float:
 
         Estimated target values.
 
+    mode:
+        Defines aggregating of multiple output values
+        (see :py:class:`~etna.metrics.functional_metrics.FunctionalMetricMode`).
+
     Returns
     -------
-    float
-        A floating point value (the best value is 0.0).
+    :
+        A non-negative floating point value (the best value is 0.0), or an array of floating point values,
+        one for each individual target.
     """
     y_true_array, y_pred_array = np.asarray(y_true), np.asarray(y_pred)
 
     if len(y_true_array.shape) != len(y_pred_array.shape):
         raise ValueError("Shapes of the labels must be the same")
 
-    prefix_error_sum = np.cumsum(y_pred_array - y_true_array)
-
-    return max(np.abs(prefix_error_sum))
+    mode_enum = FunctionalMetricMode(mode)
+    if mode_enum is FunctionalMetricMode.joint:
+        prefix_error_sum = np.cumsum(y_pred_array - y_true_array)
+        return np.max(np.abs(prefix_error_sum))
+    elif mode_enum is FunctionalMetricMode.per_output:
+        prefix_error_sum = np.cumsum(y_pred_array - y_true_array, axis=0)
+        return np.max(np.abs(prefix_error_sum), axis=0)
+    else:
+        assert_never(mode_enum)
 
 
 rmse = partial(mse, squared=False)
 
 
-def wape(y_true: ArrayLike, y_pred: ArrayLike) -> float:
+def wape(y_true: ArrayLike, y_pred: ArrayLike, mode: str = "joint") -> ArrayLike:
     """Weighted average percentage Error metric.
 
     .. math::
@@ -167,14 +232,25 @@ def wape(y_true: ArrayLike, y_pred: ArrayLike) -> float:
 
         Estimated target values.
 
+    mode:
+        Defines aggregating of multiple output values
+        (see :py:class:`~etna.metrics.functional_metrics.FunctionalMetricMode`).
+
     Returns
     -------
-    float
-        A floating point value (the best value is 0.0).
+    :
+        A non-negative floating point value (the best value is 0.0), or an array of floating point values,
+        one for each individual target.
     """
     y_true_array, y_pred_array = np.asarray(y_true), np.asarray(y_pred)
 
     if len(y_true_array.shape) != len(y_pred_array.shape):
         raise ValueError("Shapes of the labels must be the same")
 
-    return np.sum(np.abs(y_true_array - y_pred_array)) / np.sum(np.abs(y_true_array))
+    mode_enum = FunctionalMetricMode(mode)
+    if mode_enum is FunctionalMetricMode.joint:
+        return np.sum(np.abs(y_true_array - y_pred_array)) / np.sum(np.abs(y_true_array))
+    elif mode_enum is FunctionalMetricMode.per_output:
+        return np.sum(np.abs(y_true_array - y_pred_array), axis=0) / np.sum(np.abs(y_true_array), axis=0)
+    else:
+        assert_never(mode_enum)