Speed up metrics computation by optimizing segment validation

CHANGELOG.md

@@ -17,6 +17,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Changed
 - 
 - Add sorting by timestamp before the fit in `CatBoostPerSegmentModel` and `CatBoostMultiSegmentModel` ([#1337](https://github.com/tinkoff-ai/etna/pull/1337))
+- Speed up metrics computation by optimizing segment validation, forbid NaNs during metrics computation ([#1338](https://github.com/tinkoff-ai/etna/pull/1338))
 - Unify errors, warnings and checks in models ([#1312](https://github.com/tinkoff-ai/etna/pull/1312))
 - Remove upper limitation on version of numba ([#1321](https://github.com/tinkoff-ai/etna/pull/1321))
 

etna/metrics/base.py

@@ -111,9 +111,8 @@ def name(self) -> str:
  return self.__class__.__name__
 
  @staticmethod
- def _validate_segment_columns(y_true: TSDataset, y_pred: TSDataset):
- """
- Check if all the segments from ``y_true`` are in ``y_pred`` and vice versa.
+ def _validate_segments(y_true: TSDataset, y_pred: TSDataset):
+ """Check that segments in ``y_true`` and ``y_pred`` are the same.
 
  Parameters
  ----------
@@ -125,9 +124,7 @@ def _validate_segment_columns(y_true: TSDataset, y_pred: TSDataset):
  Raises
  ------
  ValueError:
- if there are mismatches in y_true and y_pred segments,
- ValueError:
- if one of segments in y_true or y_pred doesn't contain 'target' column.
+ if there are mismatches in y_true and y_pred segments
  """
  segments_true = set(y_true.df.columns.get_level_values("segment"))
  segments_pred = set(y_pred.df.columns.get_level_values("segment"))
@@ -144,33 +141,78 @@ def _validate_segment_columns(y_true: TSDataset, y_pred: TSDataset):
  f"There are segments in y_true that are not in y_pred, for example: "
  f"{', '.join(list(true_diff_pred)[:5])}"
  )
- for segment in segments_true:
+
+ @staticmethod
+ def _validate_target_columns(y_true: TSDataset, y_pred: TSDataset):
+ """Check that all the segments from ``y_true`` and ``y_pred`` has 'target' column.
+
+ Parameters
+ ----------
+ y_true:
+ y_true dataset
+ y_pred:
+ y_pred dataset
+
+ Raises
+ ------
+ ValueError:
+ if one of segments in y_true or y_pred doesn't contain 'target' column.
+ """
+ segments = set(y_true.df.columns.get_level_values("segment"))
+
+ for segment in segments:
  for name, dataset in zip(("y_true", "y_pred"), (y_true, y_pred)):
- if "target" not in dataset.loc[:, segment].columns:
+ if (segment, "target") not in dataset.columns:
  raise ValueError(
  f"All the segments in {name} should contain 'target' column. Segment {segment} doesn't."
  )
 
  @staticmethod
- def _validate_timestamp_columns(timestamp_true: pd.Series, timestamp_pred: pd.Series):
- """
- Check that ``y_true`` and ``y_pred`` have the same timestamp.
+ def _validate_index(y_true: TSDataset, y_pred: TSDataset):
+ """Check that ``y_true`` and ``y_pred`` have the same timestamps.
 
  Parameters
  ----------
- timestamp_true:
- y_true's timestamp column
- timestamp_pred:
- y_pred's timestamp column
+ y_true:
+ y_true dataset
+ y_pred:
+ y_pred dataset
 
  Raises
  ------
  ValueError:
  If there are mismatches in ``y_true`` and ``y_pred`` timestamps
  """
- if set(timestamp_pred) != set(timestamp_true):
+ if not y_true.index.equals(y_pred.index):
  raise ValueError("y_true and y_pred have different timestamps")
 
+ @staticmethod
+ def _validate_nans(y_true: TSDataset, y_pred: TSDataset):
+ """Check that ``y_true`` and ``y_pred`` doesn't have NaNs.
+
+ Parameters
+ ----------
+ y_true:
+ y_true dataset
+ y_pred:
+ y_pred dataset
+
+ Raises
+ ------
+ ValueError:
+ If there are NaNs in ``y_true`` or ``y_pred``
+ """
+ df_true = y_true.df.loc[:, pd.IndexSlice[:, "target"]]
+ df_pred = y_pred.df.loc[:, pd.IndexSlice[:, "target"]]
+
+ df_true_isna = df_true.isna().any().any()
+ if df_true_isna > 0:
+ raise ValueError("There are NaNs in y_true")
+
+ df_pred_isna = df_pred.isna().any().any()
+ if df_pred_isna > 0:
+ raise ValueError("There are NaNs in y_pred")
+
  @staticmethod
  def _macro_average(metrics_per_segments: Dict[str, float]) -> Union[float, Dict[str, float]]:
  """
@@ -226,18 +268,21 @@ def __call__(self, y_true: TSDataset, y_pred: TSDataset) -> Union[float, Dict[st
  metric's value aggregated over segments or not (depends on mode)
  """
  self._log_start()
- self._validate_segment_columns(y_true=y_true, y_pred=y_pred)
+ self._validate_segments(y_true=y_true, y_pred=y_pred)
+ self._validate_target_columns(y_true=y_true, y_pred=y_pred)
+ self._validate_index(y_true=y_true, y_pred=y_pred)
+ self._validate_nans(y_true=y_true, y_pred=y_pred)
+
+ df_true = y_true[:, :, "target"].sort_index(axis=1)
+ df_pred = y_pred[:, :, "target"].sort_index(axis=1)
 
- segments = set(y_true.df.columns.get_level_values("segment"))
  metrics_per_segment = {}
- for segment in segments:
- self._validate_timestamp_columns(
- timestamp_true=y_true[:, segment, "target"].dropna().index,
- timestamp_pred=y_pred[:, segment, "target"].dropna().index,
- )
- metrics_per_segment[segment] = self.metric_fn(
- y_true=y_true[:, segment, "target"].values, y_pred=y_pred[:, segment, "target"].values, **self.kwargs
- )
+ 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)
  metrics = self._aggregate_metrics(metrics_per_segment)
  return metrics
 

etna/metrics/intervals_metrics.py

@@ -67,16 +67,15 @@ def __call__(self, y_true: TSDataset, y_pred: TSDataset) -> Union[float, Dict[st
  -------
  metric's value aggregated over segments or not (depends on mode)
  """
- self._validate_segment_columns(y_true=y_true, y_pred=y_pred)
+ self._validate_segments(y_true=y_true, y_pred=y_pred)
+ self._validate_target_columns(y_true=y_true, y_pred=y_pred)
+ self._validate_index(y_true=y_true, y_pred=y_pred)
+ self._validate_nans(y_true=y_true, y_pred=y_pred)
  self._validate_tsdataset_quantiles(ts=y_pred, quantiles=self.quantiles)
 
  segments = set(y_true.df.columns.get_level_values("segment"))
  metrics_per_segment = {}
  for segment in segments:
- self._validate_timestamp_columns(
- timestamp_true=y_true[:, segment, "target"].dropna().index,
- timestamp_pred=y_pred[:, segment, "target"].dropna().index,
- )
  upper_quantile_flag = y_true[:, segment, "target"] <= y_pred[:, segment, f"target_{self.quantiles[1]:.4g}"]
  lower_quantile_flag = y_true[:, segment, "target"] >= y_pred[:, segment, f"target_{self.quantiles[0]:.4g}"]
 
@@ -135,16 +134,15 @@ def __call__(self, y_true: TSDataset, y_pred: TSDataset) -> Union[float, Dict[st
  -------
  metric's value aggregated over segments or not (depends on mode)
  """
- self._validate_segment_columns(y_true=y_true, y_pred=y_pred)
+ self._validate_segments(y_true=y_true, y_pred=y_pred)
+ self._validate_target_columns(y_true=y_true, y_pred=y_pred)
+ self._validate_index(y_true=y_true, y_pred=y_pred)
+ self._validate_nans(y_true=y_true, y_pred=y_pred)
  self._validate_tsdataset_quantiles(ts=y_pred, quantiles=self.quantiles)
 
  segments = set(y_true.df.columns.get_level_values("segment"))
  metrics_per_segment = {}
  for segment in segments:
- self._validate_timestamp_columns(
- timestamp_true=y_true[:, segment, "target"].dropna().index,
- timestamp_pred=y_pred[:, segment, "target"].dropna().index,
- )
  upper_quantile = y_pred[:, segment, f"target_{self.quantiles[1]:.4g}"]
  lower_quantile = y_pred[:, segment, f"target_{self.quantiles[0]:.4g}"]
 

tests/test_metrics/test_metrics.py

@@ -1,3 +1,6 @@
+from copy import deepcopy
+
+import numpy as np
 import pandas as pd
 import pytest
 
@@ -117,34 +120,60 @@ def test_metrics_invalid_aggregation(metric_class):
 @pytest.mark.parametrize(
  "metric_class", (MAE, MSE, RMSE, MedAE, MSLE, MAPE, SMAPE, R2, Sign, MaxDeviation, DummyMetric, WAPE)
 )
-def test_invalid_timestamps(metric_class, two_dfs_with_different_timestamps):
- """Check metrics behavior in case of invalid timeranges"""
+def test_invalid_segments(metric_class, two_dfs_with_different_segments_sets):
+ """Check metrics behavior in case of invalid segments sets"""
+ forecast_df, true_df = two_dfs_with_different_segments_sets
+ metric = metric_class()
+ with pytest.raises(ValueError, match="There are segments in .* that are not in .*"):
+ _ = metric(y_true=true_df, y_pred=forecast_df)
+
+
+@pytest.mark.parametrize(
+ "metric_class", (MAE, MSE, RMSE, MedAE, MSLE, MAPE, SMAPE, R2, Sign, MaxDeviation, DummyMetric, WAPE)
+)
+def test_invalid_target_columns(metric_class, train_test_dfs):
+ """Check metrics behavior in case of no target column in segment"""
+ forecast_df, true_df = train_test_dfs
+ columns = forecast_df.df.columns.to_list()
+ columns[0] = ("segment_1", "not_target")
+ forecast_df.df.columns = pd.MultiIndex.from_tuples(columns, names=["segment", "feature"])
+ metric = metric_class()
+ with pytest.raises(ValueError, match="All the segments in .* should contain 'target' column"):
+ _ = metric(y_true=true_df, y_pred=forecast_df)
+
+
+@pytest.mark.parametrize(
+ "metric_class", (MAE, MSE, RMSE, MedAE, MSLE, MAPE, SMAPE, R2, Sign, MaxDeviation, DummyMetric, WAPE)
+)
+def test_invalid_index(metric_class, two_dfs_with_different_timestamps):
+ """Check metrics behavior in case of invalid index"""
  forecast_df, true_df = two_dfs_with_different_timestamps
  metric = metric_class()
- with pytest.raises(ValueError):
+ with pytest.raises(ValueError, match="y_true and y_pred have different timestamps"):
  _ = metric(y_true=true_df, y_pred=forecast_df)
 
 
 @pytest.mark.parametrize(
  "metric_class", (MAE, MSE, RMSE, MedAE, MSLE, MAPE, SMAPE, R2, Sign, MaxDeviation, DummyMetric, WAPE)
 )
-def test_invalid_segments(metric_class, two_dfs_with_different_segments_sets):
- """Check metrics behavior in case of invalid segments sets"""
- forecast_df, true_df = two_dfs_with_different_segments_sets
+def test_invalid_nans_pred(metric_class, train_test_dfs):
+ """Check metrics behavior in case of nans in prediction."""
+ forecast_df, true_df = train_test_dfs
+ forecast_df.df.iloc[0, 0] = np.NaN
  metric = metric_class()
- with pytest.raises(ValueError):
+ with pytest.raises(ValueError, match="There are NaNs in y_pred"):
  _ = metric(y_true=true_df, y_pred=forecast_df)
 
 
 @pytest.mark.parametrize(
  "metric_class", (MAE, MSE, RMSE, MedAE, MSLE, MAPE, SMAPE, R2, Sign, MaxDeviation, DummyMetric, WAPE)
 )
-def test_invalid_segments_target(metric_class, train_test_dfs):
- """Check metrics behavior in case of no target column in segment"""
+def test_invalid_nans_true(metric_class, train_test_dfs):
+ """Check metrics behavior in case of nans in true values."""
  forecast_df, true_df = train_test_dfs
- forecast_df.df.drop(columns=[("segment_1", "target")], inplace=True)
+ true_df.df.iloc[0, 0] = np.NaN
  metric = metric_class()
- with pytest.raises(ValueError):
+ with pytest.raises(ValueError, match="There are NaNs in y_true"):
  _ = metric(y_true=true_df, y_pred=forecast_df)
 
 
@@ -181,6 +210,26 @@ def test_metrics_values(metric_class, metric_fn, train_test_dfs):
  assert value == true_metric_value
 
 
+@pytest.mark.parametrize(
+ "metric_class", (MAE, MSE, RMSE, MedAE, MSLE, MAPE, SMAPE, R2, Sign, MaxDeviation, DummyMetric, WAPE)
+)
+def test_metric_values_with_changed_segment_order(metric_class, train_test_dfs):
+ forecast_df, true_df = train_test_dfs
+ forecast_df_new, true_df_new = deepcopy(train_test_dfs)
+ segments = np.array(forecast_df.segments)
+
+ forecast_segment_order = segments[[3, 2, 0, 1, 4]]
+ forecast_df_new.df = forecast_df_new.df.loc[:, pd.IndexSlice[forecast_segment_order, :]]
+ true_segment_order = segments[[4, 1, 3, 2, 0]]
+ true_df_new.df = true_df_new.df.loc[:, pd.IndexSlice[true_segment_order, :]]
+
+ metric = metric_class(mode="per-segment")
+ metrics_initial = metric(y_pred=forecast_df, y_true=true_df)
+ metrics_changed_order = metric(y_pred=forecast_df_new, y_true=true_df_new)
+
+ assert metrics_initial == metrics_changed_order
+
+
 @pytest.mark.parametrize(
  "metric, greater_is_better",
  (