From 85508ba8db7b14bfbb3e9d938f804e7ef18a62af Mon Sep 17 00:00:00 2001 From: alex-hse-repository Date: Tue, 7 Mar 2023 13:28:27 +0100 Subject: [PATCH 01/15] Add return_components flag to the signature of the models --- etna/models/base.py | 38 +++++++++++++++++++++++++++++------ etna/models/mixins.py | 46 +++++++++++++++++++++++++++++++++++++------ 2 files changed, 72 insertions(+), 12 deletions(-) diff --git a/etna/models/base.py b/etna/models/base.py index 607b921dd..478656a52 100644 --- a/etna/models/base.py +++ b/etna/models/base.py @@ -93,13 +93,15 @@ def context_size(self) -> int: return 0 @abstractmethod - def forecast(self, ts: TSDataset) -> TSDataset: + def forecast(self, ts: TSDataset, return_components: bool = False) -> TSDataset: """Make predictions. Parameters ---------- ts: Dataset with features + return_components: + If True additionally returns forecast components Returns ------- @@ -109,13 +111,15 @@ def forecast(self, ts: TSDataset) -> TSDataset: pass @abstractmethod - def predict(self, ts: TSDataset) -> TSDataset: + def predict(self, ts: TSDataset, return_components: bool = False) -> TSDataset: """Make predictions with using true values as autoregression context if possible (teacher forcing). Parameters ---------- ts: Dataset with features + return_components: + If True additionally returns prediction components Returns ------- @@ -129,7 +133,7 @@ class NonPredictionIntervalContextRequiredAbstractModel(AbstractModel): """Interface for models that don't support prediction intervals and need context for prediction.""" @abstractmethod - def forecast(self, ts: TSDataset, prediction_size: int) -> TSDataset: + def forecast(self, ts: TSDataset, prediction_size: int, return_components: bool = False) -> TSDataset: """Make predictions. Parameters @@ -139,6 +143,8 @@ def forecast(self, ts: TSDataset, prediction_size: int) -> TSDataset: prediction_size: Number of last timestamps to leave after making prediction. Previous timestamps will be used as a context for models that require it. + return_components: + If True additionally returns forecast components Returns ------- @@ -148,7 +154,7 @@ def forecast(self, ts: TSDataset, prediction_size: int) -> TSDataset: pass @abstractmethod - def predict(self, ts: TSDataset, prediction_size: int) -> TSDataset: + def predict(self, ts: TSDataset, prediction_size: int, return_components: bool = False) -> TSDataset: """Make predictions with using true values as autoregression context if possible (teacher forcing). Parameters @@ -158,6 +164,8 @@ def predict(self, ts: TSDataset, prediction_size: int) -> TSDataset: prediction_size: Number of last timestamps to leave after making prediction. Previous timestamps will be used as a context for models that require it. + return_components: + If True additionally returns prediction components Returns ------- @@ -180,7 +188,11 @@ def context_size(self) -> int: @abstractmethod def forecast( - self, ts: TSDataset, prediction_interval: bool = False, quantiles: Sequence[float] = (0.025, 0.975) + self, + ts: TSDataset, + prediction_interval: bool = False, + quantiles: Sequence[float] = (0.025, 0.975), + return_components: bool = False, ) -> TSDataset: """Make predictions. @@ -192,6 +204,8 @@ def forecast( If True returns prediction interval for forecast quantiles: Levels of prediction distribution. By default 2.5% and 97.5% are taken to form a 95% prediction interval + return_components: + If True additionally returns forecast components Returns ------- @@ -202,7 +216,11 @@ def forecast( @abstractmethod def predict( - self, ts: TSDataset, prediction_interval: bool = False, quantiles: Sequence[float] = (0.025, 0.975) + self, + ts: TSDataset, + prediction_interval: bool = False, + quantiles: Sequence[float] = (0.025, 0.975), + return_components: bool = False, ) -> TSDataset: """Make predictions with using true values as autoregression context if possible (teacher forcing). @@ -214,6 +232,8 @@ def predict( If True returns prediction interval for forecast quantiles: Levels of prediction distribution. By default 2.5% and 97.5% are taken to form a 95% prediction interval + return_components: + If True additionally returns prediction components Returns ------- @@ -233,6 +253,7 @@ def forecast( prediction_size: int, prediction_interval: bool = False, quantiles: Sequence[float] = (0.025, 0.975), + return_components: bool = False, ) -> TSDataset: """Make predictions. @@ -247,6 +268,8 @@ def forecast( If True returns prediction interval for forecast quantiles: Levels of prediction distribution. By default 2.5% and 97.5% are taken to form a 95% prediction interval + return_components: + If True additionally returns forecast components Returns ------- @@ -262,6 +285,7 @@ def predict( prediction_size: int, prediction_interval: bool = False, quantiles: Sequence[float] = (0.025, 0.975), + return_components: bool = False, ) -> TSDataset: """Make predictions with using true values as autoregression context if possible (teacher forcing). @@ -276,6 +300,8 @@ def predict( If True returns prediction interval for forecast quantiles: Levels of prediction distribution. By default 2.5% and 97.5% are taken to form a 95% prediction interval + return_components: + If True additionally returns prediction components Returns ------- diff --git a/etna/models/mixins.py b/etna/models/mixins.py index 22cb1fe3b..80a7d3c30 100644 --- a/etna/models/mixins.py +++ b/etna/models/mixins.py @@ -28,17 +28,27 @@ def _forecast(self, **kwargs) -> TSDataset: def _predict(self, **kwargs) -> TSDataset: pass + @abstractmethod + def _forecast_components(self, **kwargs) -> pd.DataFrame: + pass + + @abstractmethod + def _predict_components(self, **kwargs) -> pd.DataFrame: + pass + class NonPredictionIntervalContextIgnorantModelMixin(ModelForecastingMixin): """Mixin for models that don't support prediction intervals and don't need context for prediction.""" - def forecast(self, ts: TSDataset) -> TSDataset: + def forecast(self, ts: TSDataset, return_components: bool = False) -> TSDataset: """Make predictions. Parameters ---------- ts: Dataset with features + return_components: + If True additionally returns forecast components Returns ------- @@ -47,13 +57,15 @@ def forecast(self, ts: TSDataset) -> TSDataset: """ return self._forecast(ts=ts) - def predict(self, ts: TSDataset) -> TSDataset: + def predict(self, ts: TSDataset, return_components: bool = False) -> TSDataset: """Make predictions with using true values as autoregression context if possible (teacher forcing). Parameters ---------- ts: Dataset with features + return_components: + If True additionally returns prediction components Returns ------- @@ -66,7 +78,7 @@ def predict(self, ts: TSDataset) -> TSDataset: class NonPredictionIntervalContextRequiredModelMixin(ModelForecastingMixin): """Mixin for models that don't support prediction intervals and need context for prediction.""" - def forecast(self, ts: TSDataset, prediction_size: int) -> TSDataset: + def forecast(self, ts: TSDataset, prediction_size: int, return_components: bool = False) -> TSDataset: """Make predictions. Parameters @@ -76,6 +88,8 @@ def forecast(self, ts: TSDataset, prediction_size: int) -> TSDataset: prediction_size: Number of last timestamps to leave after making prediction. Previous timestamps will be used as a context for models that require it. + return_components: + If True additionally returns forecast components Returns ------- @@ -84,7 +98,7 @@ def forecast(self, ts: TSDataset, prediction_size: int) -> TSDataset: """ return self._forecast(ts=ts, prediction_size=prediction_size) - def predict(self, ts: TSDataset, prediction_size: int) -> TSDataset: + def predict(self, ts: TSDataset, prediction_size: int, return_components: bool = False) -> TSDataset: """Make predictions with using true values as autoregression context if possible (teacher forcing). Parameters @@ -94,6 +108,8 @@ def predict(self, ts: TSDataset, prediction_size: int) -> TSDataset: prediction_size: Number of last timestamps to leave after making prediction. Previous timestamps will be used as a context for models that require it. + return_components: + If True additionally returns prediction components Returns ------- @@ -107,7 +123,11 @@ class PredictionIntervalContextIgnorantModelMixin(ModelForecastingMixin): """Mixin for models that support prediction intervals and don't need context for prediction.""" def forecast( - self, ts: TSDataset, prediction_interval: bool = False, quantiles: Sequence[float] = (0.025, 0.975) + self, + ts: TSDataset, + prediction_interval: bool = False, + quantiles: Sequence[float] = (0.025, 0.975), + return_components: bool = False, ) -> TSDataset: """Make predictions. @@ -119,6 +139,8 @@ def forecast( If True returns prediction interval for forecast quantiles: Levels of prediction distribution. By default 2.5% and 97.5% are taken to form a 95% prediction interval + return_components: + If True additionally returns forecast components Returns ------- @@ -128,7 +150,11 @@ def forecast( return self._forecast(ts=ts, prediction_interval=prediction_interval, quantiles=quantiles) def predict( - self, ts: TSDataset, prediction_interval: bool = False, quantiles: Sequence[float] = (0.025, 0.975) + self, + ts: TSDataset, + prediction_interval: bool = False, + quantiles: Sequence[float] = (0.025, 0.975), + return_components: bool = False, ) -> TSDataset: """Make predictions with using true values as autoregression context if possible (teacher forcing). @@ -140,6 +166,8 @@ def predict( If True returns prediction interval for forecast quantiles: Levels of prediction distribution. By default 2.5% and 97.5% are taken to form a 95% prediction interval + return_components: + If True additionally returns prediction components Returns ------- @@ -158,6 +186,7 @@ def forecast( prediction_size: int, prediction_interval: bool = False, quantiles: Sequence[float] = (0.025, 0.975), + return_components: bool = False, ) -> TSDataset: """Make predictions. @@ -172,6 +201,8 @@ def forecast( If True returns prediction interval for forecast quantiles: Levels of prediction distribution. By default 2.5% and 97.5% are taken to form a 95% prediction interval + return_components: + If True additionally returns forecast components Returns ------- @@ -188,6 +219,7 @@ def predict( prediction_size: int, prediction_interval: bool = False, quantiles: Sequence[float] = (0.025, 0.975), + return_components: bool = False, ) -> TSDataset: """Make predictions with using true values as autoregression context if possible (teacher forcing). @@ -202,6 +234,8 @@ def predict( If True returns prediction interval for forecast quantiles: Levels of prediction distribution. By default 2.5% and 97.5% are taken to form a 95% prediction interval + return_components: + If True additionally returns prediction components Returns ------- From 8c48a86d84ee29f9adc3549e9d9193341a35d973 Mon Sep 17 00:00:00 2001 From: alex-hse-repository Date: Tue, 7 Mar 2023 14:36:23 +0100 Subject: [PATCH 02/15] Add components merging logic --- etna/models/mixins.py | 48 +++++++++++++++++++++++++++++++++++-------- 1 file changed, 40 insertions(+), 8 deletions(-) diff --git a/etna/models/mixins.py b/etna/models/mixins.py index 80a7d3c30..0f08e0cba 100644 --- a/etna/models/mixins.py +++ b/etna/models/mixins.py @@ -55,7 +55,11 @@ def forecast(self, ts: TSDataset, return_components: bool = False) -> TSDataset: : Dataset with predictions """ - return self._forecast(ts=ts) + forecast = self._forecast(ts=ts) + if return_components: + forecast_components_df = self._forecast_components(ts=ts) + forecast.add_target_components(target_components_df=forecast_components_df) + return forecast def predict(self, ts: TSDataset, return_components: bool = False) -> TSDataset: """Make predictions with using true values as autoregression context if possible (teacher forcing). @@ -72,7 +76,11 @@ def predict(self, ts: TSDataset, return_components: bool = False) -> TSDataset: : Dataset with predictions """ - return self._predict(ts=ts) + prediction = self._predict(ts=ts) + if return_components: + prediction_components_df = self._predict_components(ts=ts) + prediction.add_target_components(target_components_df=prediction_components_df) + return prediction class NonPredictionIntervalContextRequiredModelMixin(ModelForecastingMixin): @@ -96,7 +104,11 @@ def forecast(self, ts: TSDataset, prediction_size: int, return_components: bool : Dataset with predictions """ - return self._forecast(ts=ts, prediction_size=prediction_size) + forecast = self._forecast(ts=ts, prediction_size=prediction_size) + if return_components: + forecast_components_df = self._forecast_components(ts=ts) + forecast.add_target_components(target_components_df=forecast_components_df) + return forecast def predict(self, ts: TSDataset, prediction_size: int, return_components: bool = False) -> TSDataset: """Make predictions with using true values as autoregression context if possible (teacher forcing). @@ -116,7 +128,11 @@ def predict(self, ts: TSDataset, prediction_size: int, return_components: bool = : Dataset with predictions """ - return self._predict(ts=ts, prediction_size=prediction_size) + prediction = self._predict(ts=ts, prediction_size=prediction_size) + if return_components: + prediction_components_df = self._predict_components(ts=ts) + prediction.add_target_components(target_components_df=prediction_components_df) + return prediction class PredictionIntervalContextIgnorantModelMixin(ModelForecastingMixin): @@ -147,7 +163,11 @@ def forecast( : Dataset with predictions """ - return self._forecast(ts=ts, prediction_interval=prediction_interval, quantiles=quantiles) + forecast = self._forecast(ts=ts, prediction_interval=prediction_interval, quantiles=quantiles) + if return_components: + forecast_components_df = self._forecast_components(ts=ts) + forecast.add_target_components(target_components_df=forecast_components_df) + return forecast def predict( self, @@ -174,7 +194,11 @@ def predict( : Dataset with predictions """ - return self._predict(ts=ts, prediction_interval=prediction_interval, quantiles=quantiles) + prediction = self._predict(ts=ts, prediction_interval=prediction_interval, quantiles=quantiles) + if return_components: + prediction_components_df = self._predict_components(ts=ts) + prediction.add_target_components(target_components_df=prediction_components_df) + return prediction class PredictionIntervalContextRequiredModelMixin(ModelForecastingMixin): @@ -209,9 +233,13 @@ def forecast( : Dataset with predictions """ - return self._forecast( + forecast = self._forecast( ts=ts, prediction_size=prediction_size, prediction_interval=prediction_interval, quantiles=quantiles ) + if return_components: + forecast_components_df = self._forecast_components(ts=ts) + forecast.add_target_components(target_components_df=forecast_components_df) + return forecast def predict( self, @@ -242,9 +270,13 @@ def predict( : Dataset with predictions """ - return self._predict( + prediction = self._predict( ts=ts, prediction_size=prediction_size, prediction_interval=prediction_interval, quantiles=quantiles ) + if return_components: + prediction_components_df = self._predict_components(ts=ts) + prediction.add_target_components(target_components_df=prediction_components_df) + return prediction class PerSegmentModelMixin(ModelForecastingMixin): From d66b3efa39bccb0b285338ac0b5d48956ba6a518 Mon Sep 17 00:00:00 2001 From: alex-hse-repository Date: Tue, 7 Mar 2023 14:36:49 +0100 Subject: [PATCH 03/15] Add tests for components merging logic --- tests/test_models/test_mixins.py | 100 +++++++++++++++++++++++++++++++ 1 file changed, 100 insertions(+) diff --git a/tests/test_models/test_mixins.py b/tests/test_models/test_mixins.py index 0b7f8c9d3..d617c6938 100644 --- a/tests/test_models/test_mixins.py +++ b/tests/test_models/test_mixins.py @@ -1,6 +1,7 @@ import json import pathlib from unittest.mock import MagicMock +from unittest.mock import Mock from unittest.mock import patch from zipfile import ZipFile @@ -8,12 +9,17 @@ import pytest from etna import SETTINGS +from etna.datasets import TSDataset if SETTINGS.torch_required: import torch from etna.models.mixins import MultiSegmentModelMixin +from etna.models.mixins import NonPredictionIntervalContextIgnorantModelMixin +from etna.models.mixins import NonPredictionIntervalContextRequiredModelMixin from etna.models.mixins import PerSegmentModelMixin +from etna.models.mixins import PredictionIntervalContextIgnorantModelMixin +from etna.models.mixins import PredictionIntervalContextRequiredModelMixin from etna.models.mixins import SaveNNMixin @@ -124,3 +130,97 @@ def test_save_mixin_load_warning(get_version_mock, save_version, load_version, t ): get_version_mock.return_value = load_version _ = DummyNN.load(path) + + +@pytest.mark.parametrize( + "mixin_constructor, call_params", + [ + (PredictionIntervalContextIgnorantModelMixin, {}), + (NonPredictionIntervalContextIgnorantModelMixin, {}), + (PredictionIntervalContextRequiredModelMixin, {"prediction_size": 10}), + (NonPredictionIntervalContextRequiredModelMixin, {"prediction_size": 10}), + ], +) +def test_model_mixins_calls_forecast_components_in_forecast(mixin_constructor, call_params): + with patch.multiple(mixin_constructor, __abstractmethods__=set()): + ts = Mock() + forecast_ts = Mock(spec=TSDataset) + target_components_df = Mock() + mixin = mixin_constructor() + mixin._forecast = Mock(return_value=forecast_ts) + mixin._forecast_components = Mock(return_value=target_components_df) + + _ = mixin.forecast(ts=ts, return_components=True, **call_params) + + mixin._forecast_components.assert_called_with(ts=ts) + forecast_ts.add_target_components.assert_called_with(target_components_df=target_components_df) + + +@pytest.mark.parametrize( + "mixin_constructor, call_params", + [ + (PredictionIntervalContextIgnorantModelMixin, {}), + (NonPredictionIntervalContextIgnorantModelMixin, {}), + (PredictionIntervalContextRequiredModelMixin, {"prediction_size": 10}), + (NonPredictionIntervalContextRequiredModelMixin, {"prediction_size": 10}), + ], +) +def test_model_mixins_not_calls_forecast_components_in_forecast(mixin_constructor, call_params): + with patch.multiple(mixin_constructor, __abstractmethods__=set()): + ts = Mock() + forecast_ts = Mock(spec=TSDataset) + target_components_df = Mock() + mixin = mixin_constructor() + mixin._forecast = Mock(return_value=forecast_ts) + mixin._forecast_components = Mock(return_value=target_components_df) + + _ = mixin.forecast(ts=ts, return_components=False, **call_params) + + mixin._forecast_components.assert_not_called() + forecast_ts.add_target_components.assert_not_called() + + +@pytest.mark.parametrize( + "mixin_constructor, call_params", + [ + (PredictionIntervalContextIgnorantModelMixin, {}), + (NonPredictionIntervalContextIgnorantModelMixin, {}), + (PredictionIntervalContextRequiredModelMixin, {"prediction_size": 10}), + (NonPredictionIntervalContextRequiredModelMixin, {"prediction_size": 10}), + ], +) +def test_model_mixins_calls_predict_components_in_predict(mixin_constructor, call_params): + with patch.multiple(mixin_constructor, __abstractmethods__=set()): + ts = Mock() + predict_ts = Mock(spec=TSDataset) + target_components_df = Mock() + mixin = mixin_constructor() + mixin._predict = Mock(return_value=predict_ts) + mixin._predict_components = Mock(return_value=target_components_df) + _ = mixin.predict(ts=ts, return_components=True, **call_params) + + mixin._predict_components.assert_called_with(ts=ts) + predict_ts.add_target_components.assert_called_with(target_components_df=target_components_df) + + +@pytest.mark.parametrize( + "mixin_constructor, call_params", + [ + (PredictionIntervalContextIgnorantModelMixin, {}), + (NonPredictionIntervalContextIgnorantModelMixin, {}), + (PredictionIntervalContextRequiredModelMixin, {"prediction_size": 10}), + (NonPredictionIntervalContextRequiredModelMixin, {"prediction_size": 10}), + ], +) +def test_model_mixins_not_calls_predict_components_in_predict(mixin_constructor, call_params): + with patch.multiple(mixin_constructor, __abstractmethods__=set()): + ts = Mock() + predict_ts = Mock(spec=TSDataset) + target_components_df = Mock() + mixin = mixin_constructor() + mixin._predict = Mock(return_value=predict_ts) + mixin._predict_components = Mock(return_value=target_components_df) + _ = mixin.predict(ts=ts, return_components=False, **call_params) + + mixin._predict_components.assert_not_called() + predict_ts.add_target_components.assert_not_called() From 2cf3fa3e0f71ab684aa92df316730c951ee96acb Mon Sep 17 00:00:00 2001 From: alex-hse-repository Date: Tue, 7 Mar 2023 15:04:03 +0100 Subject: [PATCH 04/15] Add return_components to nns --- etna/models/base.py | 16 ++++++++++++++-- etna/models/nn/deepar.py | 9 +++++++++ etna/models/nn/tft.py | 9 +++++++++ tests/test_models/nn/test_deepar.py | 7 +++++++ tests/test_models/nn/test_tft.py | 7 +++++++ tests/test_models/test_base.py | 7 ++++++- 6 files changed, 52 insertions(+), 3 deletions(-) diff --git a/etna/models/base.py b/etna/models/base.py index 478656a52..8e8f0675c 100644 --- a/etna/models/base.py +++ b/etna/models/base.py @@ -630,7 +630,7 @@ def raw_predict(self, torch_dataset: "Dataset") -> Dict[Tuple[str, str], np.ndar return predictions_dict @log_decorator - def forecast(self, ts: "TSDataset", prediction_size: int) -> "TSDataset": + def forecast(self, ts: "TSDataset", prediction_size: int, return_components: bool = False) -> "TSDataset": """Make predictions. This method will make autoregressive predictions. @@ -642,12 +642,17 @@ def forecast(self, ts: "TSDataset", prediction_size: int) -> "TSDataset": prediction_size: Number of last timestamps to leave after making prediction. Previous timestamps will be used as a context. + return_components: + If True additionally returns forecast components Returns ------- : Dataset with predictions """ + if return_components: + raise NotImplementedError("This mode isn't currently implemented!") + test_dataset = ts.to_torch_dataset( make_samples=functools.partial( self.net.make_samples, encoder_length=self.encoder_length, decoder_length=prediction_size @@ -662,7 +667,12 @@ def forecast(self, ts: "TSDataset", prediction_size: int) -> "TSDataset": return future_ts @log_decorator - def predict(self, ts: "TSDataset", prediction_size: int) -> "TSDataset": + def predict( + self, + ts: "TSDataset", + prediction_size: int, + return_components: bool = False, + ) -> "TSDataset": """Make predictions. This method will make predictions using true values instead of predicted on a previous step. @@ -675,6 +685,8 @@ def predict(self, ts: "TSDataset", prediction_size: int) -> "TSDataset": prediction_size: Number of last timestamps to leave after making prediction. Previous timestamps will be used as a context. + return_components: + If True additionally returns prediction components Returns ------- diff --git a/etna/models/nn/deepar.py b/etna/models/nn/deepar.py index c0c45ed08..15edc9fae 100644 --- a/etna/models/nn/deepar.py +++ b/etna/models/nn/deepar.py @@ -149,6 +149,7 @@ def forecast( prediction_size: int, prediction_interval: bool = False, quantiles: Sequence[float] = (0.025, 0.975), + return_components: bool = False, ) -> TSDataset: """Make predictions. @@ -165,12 +166,17 @@ def forecast( If True returns prediction interval for forecast quantiles: Levels of prediction distribution. By default 2.5% and 97.5% are taken to form a 95% prediction interval + return_components: + If True additionally returns forecast components Returns ------- TSDataset TSDataset with predictions. """ + if return_components: + raise NotImplementedError("This mode isn't currently implemented!") + ts, prediction_dataloader = self._make_target_prediction(ts, prediction_size) if prediction_interval: @@ -203,6 +209,7 @@ def predict( prediction_size: int, prediction_interval: bool = False, quantiles: Sequence[float] = (0.025, 0.975), + return_components: bool = False, ) -> TSDataset: """Make predictions. @@ -220,6 +227,8 @@ def predict( If True returns prediction interval for forecast quantiles: Levels of prediction distribution. By default 2.5% and 97.5% are taken to form a 95% prediction interval + return_components: + If True additionally returns prediction components Returns ------- diff --git a/etna/models/nn/tft.py b/etna/models/nn/tft.py index 10d771bc2..ce27d448e 100644 --- a/etna/models/nn/tft.py +++ b/etna/models/nn/tft.py @@ -153,6 +153,7 @@ def forecast( prediction_size: int, prediction_interval: bool = False, quantiles: Sequence[float] = (0.025, 0.975), + return_components: bool = False, ) -> TSDataset: """Make predictions. @@ -169,12 +170,17 @@ def forecast( If True returns prediction interval for forecast quantiles: Levels of prediction distribution. By default 2.5% and 97.5% are taken to form a 95% prediction interval + return_components: + If True additionally returns forecast components Returns ------- TSDataset TSDataset with predictions. """ + if return_components: + raise NotImplementedError("This mode isn't currently implemented!") + ts, prediction_dataloader = self._make_target_prediction(ts, prediction_size) if prediction_interval: @@ -232,6 +238,7 @@ def predict( prediction_size: int, prediction_interval: bool = False, quantiles: Sequence[float] = (0.025, 0.975), + return_components: bool = False, ) -> TSDataset: """Make predictions. @@ -249,6 +256,8 @@ def predict( If True returns prediction interval for forecast quantiles: Levels of prediction distribution. By default 2.5% and 97.5% are taken to form a 95% prediction interval + return_components: + If True additionally returns prediction components Returns ------- diff --git a/tests/test_models/nn/test_deepar.py b/tests/test_models/nn/test_deepar.py index 9eb31b75a..7ca8750c7 100644 --- a/tests/test_models/nn/test_deepar.py +++ b/tests/test_models/nn/test_deepar.py @@ -1,3 +1,5 @@ +from unittest.mock import Mock + import pandas as pd import pytest from pytorch_forecasting.data import GroupNormalizer @@ -184,3 +186,8 @@ def test_repr(): "train_batch_size = 64, test_batch_size = 64, lr = 0.1, cell_type = 'LSTM', hidden_size = 10, rnn_layers = 2, " "dropout = 0.1, loss = NormalDistributionLoss(), trainer_params = {'max_epochs': 2, 'gpus': 0}, quantiles_kwargs = {}, )" ) + + +def test_deepar_forecast_throw_error_on_return_components(): + with pytest.raises(NotImplementedError, match="This mode isn't currently implemented!"): + DeepARModel.forecast(self=Mock(), ts=Mock(), prediction_size=Mock(), return_components=True) diff --git a/tests/test_models/nn/test_tft.py b/tests/test_models/nn/test_tft.py index 636b58ad0..7f0d7b737 100644 --- a/tests/test_models/nn/test_tft.py +++ b/tests/test_models/nn/test_tft.py @@ -1,3 +1,5 @@ +from unittest.mock import Mock + import pandas as pd import pytest @@ -189,3 +191,8 @@ def test_repr(): "attention_head_size = 4, dropout = 0.1, hidden_continuous_size = 8, " "loss = QuantileLoss(), trainer_params = {'max_epochs': 2, 'gpus': 0}, quantiles_kwargs = {}, )" ) + + +def test_tft_forecast_throw_error_on_return_components(): + with pytest.raises(NotImplementedError, match="This mode isn't currently implemented!"): + TFTModel.forecast(self=Mock(), ts=Mock(), prediction_size=Mock(), return_components=True) diff --git a/tests/test_models/test_base.py b/tests/test_models/test_base.py index 400b6a3a1..faf8f65c8 100644 --- a/tests/test_models/test_base.py +++ b/tests/test_models/test_base.py @@ -13,7 +13,7 @@ @pytest.fixture() def deep_base_model_mock(): - model = MagicMock() + model = MagicMock(spec=DeepBaseModel) model.train_batch_size = 32 model.train_dataloader_params = {} model.val_dataloader_params = {} @@ -156,3 +156,8 @@ def test_deep_base_model_forecast_loop(simple_df, deep_base_model_mock): np.testing.assert_allclose( future.df.loc[:, pd.IndexSlice["B", "target"]], raw_predict[("B", "target")][:horizon, 0] ) + + +def test_deep_base_model_forecast_throw_error_on_return_components(): + with pytest.raises(NotImplementedError, match="This mode isn't currently implemented!"): + DeepBaseModel.forecast(self=Mock(), ts=Mock(), prediction_size=Mock(), return_components=True) From be2fc4b5bab659b57de4b14ba5d0bde057b75288 Mon Sep 17 00:00:00 2001 From: alex-hse-repository Date: Thu, 9 Mar 2023 10:27:46 +0100 Subject: [PATCH 05/15] Add components logic to multi-segment mixin --- etna/models/mixins.py | 50 ++++++++++++++++++++++++++++++++ tests/test_models/test_mixins.py | 35 ++++++++++++++++++++++ 2 files changed, 85 insertions(+) diff --git a/etna/models/mixins.py b/etna/models/mixins.py index 0f08e0cba..dfb566526 100644 --- a/etna/models/mixins.py +++ b/etna/models/mixins.py @@ -424,6 +424,16 @@ def _forecast(self, ts: TSDataset, **kwargs) -> TSDataset: def _predict(self, ts: TSDataset, **kwargs) -> TSDataset: return self._make_predictions(ts=ts, prediction_method=self._base_model.__class__.predict, **kwargs) + @log_decorator + def _forecast_components(self, **kwargs) -> pd.DataFrame: + if hasattr(self._base_model, "forecast_components"): + return self._base_model.forecast_components() + return self._base_model.predict_components() + + @log_decorator + def _predict_components(self, **kwargs) -> pd.DataFrame: + return self._base_model.predict_components() + class MultiSegmentModelMixin(ModelForecastingMixin): """Mixin for holding methods for multi-segment prediction. @@ -484,6 +494,30 @@ def _make_predictions(self, ts: TSDataset, prediction_method: Callable, **kwargs ts.loc[:, pd.IndexSlice[:, "target"]] = y return ts + def _make_component_predictions(self, ts: TSDataset, prediction_method: Callable, **kwargs) -> pd.DataFrame: + """Make target component predictions. + + Parameters + ---------- + ts: + Dataset with features + prediction_method: + Method for making components predictions + + Returns + ------- + : + Dataset with predicted components + """ + features_df = ts.to_pandas(flatten=True) + segment_column = features_df["segment"].values + features_df = features_df.drop(["segment"], axis=1) + + target_components_df = prediction_method(self=self._base_model, df=features_df, **kwargs) + target_components_df["segment"] = segment_column + target_components_df = TSDataset.to_dataset(target_components_df) + return target_components_df + @log_decorator def _forecast(self, ts: TSDataset, **kwargs) -> TSDataset: if hasattr(self._base_model, "forecast"): @@ -494,6 +528,22 @@ def _forecast(self, ts: TSDataset, **kwargs) -> TSDataset: def _predict(self, ts: TSDataset, **kwargs) -> TSDataset: return self._make_predictions(ts=ts, prediction_method=self._base_model.__class__.predict, **kwargs) + @log_decorator + def _forecast_components(self, ts: TSDataset, **kwargs) -> pd.DataFrame: + if hasattr(self._base_model, "forecast_components"): + return self._make_component_predictions( + ts=ts, prediction_method=self._base_model.__class__.forecast_components, **kwargs + ) + return self._make_component_predictions( + ts=ts, prediction_method=self._base_model.__class__.predict_components, **kwargs + ) + + @log_decorator + def _predict_components(self, ts: TSDataset, **kwargs) -> pd.DataFrame: + return self._make_component_predictions( + ts=ts, prediction_method=self._base_model.__class__.predict_components, **kwargs + ) + def get_model(self) -> Any: """Get internal model that is used inside etna class. diff --git a/tests/test_models/test_mixins.py b/tests/test_models/test_mixins.py index d617c6938..ae4f2bbf4 100644 --- a/tests/test_models/test_mixins.py +++ b/tests/test_models/test_mixins.py @@ -14,6 +14,8 @@ if SETTINGS.torch_required: import torch +import pandas as pd + from etna.models.mixins import MultiSegmentModelMixin from etna.models.mixins import NonPredictionIntervalContextIgnorantModelMixin from etna.models.mixins import NonPredictionIntervalContextRequiredModelMixin @@ -43,6 +45,27 @@ def autoregression_base_model_mock(): return model +@pytest.fixture +def target_components_df(): + timestamp = pd.date_range("2021-01-01", "2021-01-15") + df_1 = pd.DataFrame({"timestamp": timestamp, "target_component_a": 1, "target_component_b": 2, "segment": 1}) + df_2 = pd.DataFrame({"timestamp": timestamp, "target_component_a": 3, "target_component_b": 4, "segment": 2}) + df = pd.concat([df_1, df_2]) + df = TSDataset.to_dataset(df) + return df + + +@pytest.fixture +def ts_without_target_components(): + timestamp = pd.date_range("2021-01-01", "2021-01-15") + df_1 = pd.DataFrame({"timestamp": timestamp, "target": 3, "segment": 1}) + df_2 = pd.DataFrame({"timestamp": timestamp, "target": 7, "segment": 2}) + df = pd.concat([df_1, df_2]) + df = TSDataset.to_dataset(df) + ts = TSDataset(df=df, freq="D") + return ts + + @pytest.mark.parametrize("mixin_constructor", [PerSegmentModelMixin, MultiSegmentModelMixin]) @pytest.mark.parametrize( "base_model_name, called_method_name, expected_method_name", @@ -224,3 +247,15 @@ def test_model_mixins_not_calls_predict_components_in_predict(mixin_constructor, mixin._predict_components.assert_not_called() predict_ts.add_target_components.assert_not_called() + + +@pytest.mark.parametrize("mixin_constructor", [PerSegmentModelMixin, MultiSegmentModelMixin]) +def test_make_components_prediction(mixin_constructor, target_components_df, ts_without_target_components): + mixin = mixin_constructor(base_model=Mock()) + target_components_df_model_format = TSDataset.to_flatten(target_components_df).drop(columns=["segment"]) + prediction_method = Mock(return_value=target_components_df_model_format) + + target_components_pred = mixin._make_component_predictions( + ts=ts_without_target_components, prediction_method=prediction_method + ) + pd.testing.assert_frame_equal(target_components_pred, target_components_df) From 22a64869111f05288a198dc0e20c06ba228499fe Mon Sep 17 00:00:00 2001 From: alex-hse-repository Date: Thu, 9 Mar 2023 11:07:53 +0100 Subject: [PATCH 06/15] Add components prediction logic to per-segment mixin --- etna/models/mixins.py | 45 +++++++++++++++++++++++++++----- tests/test_models/test_mixins.py | 28 +++++++++++++++----- 2 files changed, 61 insertions(+), 12 deletions(-) diff --git a/etna/models/mixins.py b/etna/models/mixins.py index dfb566526..0b7dc82ac 100644 --- a/etna/models/mixins.py +++ b/etna/models/mixins.py @@ -414,6 +414,33 @@ def _make_predictions(self, ts: TSDataset, prediction_method: Callable, **kwargs ts.df = ts.df.iloc[-prediction_size:] return ts + def _make_component_predictions(self, ts: TSDataset, prediction_method: Callable, **kwargs) -> pd.DataFrame: + """Make target component predictions. + + Parameters + ---------- + ts: + Dataset with features + prediction_method: + Method for making components predictions + + Returns + ------- + : + Dataset with predicted components + """ + features_df = ts.to_pandas() + result_list = list() + for segment, model in self._get_model().items(): + segment_predict = self._make_predictions_segment( + model=model, segment=segment, df=features_df, prediction_method=prediction_method, **kwargs + ) + result_list.append(segment_predict) + + target_components_df = pd.concat(result_list, ignore_index=True) + target_components_df = TSDataset.to_dataset(target_components_df) + return target_components_df + @log_decorator def _forecast(self, ts: TSDataset, **kwargs) -> TSDataset: if hasattr(self._base_model, "forecast"): @@ -425,14 +452,20 @@ def _predict(self, ts: TSDataset, **kwargs) -> TSDataset: return self._make_predictions(ts=ts, prediction_method=self._base_model.__class__.predict, **kwargs) @log_decorator - def _forecast_components(self, **kwargs) -> pd.DataFrame: + def _forecast_components(self, ts: TSDataset, **kwargs) -> pd.DataFrame: if hasattr(self._base_model, "forecast_components"): - return self._base_model.forecast_components() - return self._base_model.predict_components() + return self._make_component_predictions( + ts=ts, prediction_method=self._base_model.__class__.forecast_components, **kwargs + ) + return self._make_component_predictions( + ts=ts, prediction_method=self._base_model.__class__.predict_components, **kwargs + ) @log_decorator - def _predict_components(self, **kwargs) -> pd.DataFrame: - return self._base_model.predict_components() + def _predict_components(self, ts: TSDataset, **kwargs) -> pd.DataFrame: + return self._make_component_predictions( + ts=ts, prediction_method=self._base_model.__class__.predict_components, **kwargs + ) class MultiSegmentModelMixin(ModelForecastingMixin): @@ -512,7 +545,7 @@ def _make_component_predictions(self, ts: TSDataset, prediction_method: Callable features_df = ts.to_pandas(flatten=True) segment_column = features_df["segment"].values features_df = features_df.drop(["segment"], axis=1) - + # TODO: make it work with prediction intervals and context target_components_df = prediction_method(self=self._base_model, df=features_df, **kwargs) target_components_df["segment"] = segment_column target_components_df = TSDataset.to_dataset(target_components_df) diff --git a/tests/test_models/test_mixins.py b/tests/test_models/test_mixins.py index ae4f2bbf4..9e4d12f8c 100644 --- a/tests/test_models/test_mixins.py +++ b/tests/test_models/test_mixins.py @@ -48,9 +48,7 @@ def autoregression_base_model_mock(): @pytest.fixture def target_components_df(): timestamp = pd.date_range("2021-01-01", "2021-01-15") - df_1 = pd.DataFrame({"timestamp": timestamp, "target_component_a": 1, "target_component_b": 2, "segment": 1}) - df_2 = pd.DataFrame({"timestamp": timestamp, "target_component_a": 3, "target_component_b": 4, "segment": 2}) - df = pd.concat([df_1, df_2]) + df = pd.DataFrame({"timestamp": timestamp, "target_component_a": 1, "target_component_b": 2, "segment": 1}) df = TSDataset.to_dataset(df) return df @@ -58,9 +56,7 @@ def target_components_df(): @pytest.fixture def ts_without_target_components(): timestamp = pd.date_range("2021-01-01", "2021-01-15") - df_1 = pd.DataFrame({"timestamp": timestamp, "target": 3, "segment": 1}) - df_2 = pd.DataFrame({"timestamp": timestamp, "target": 7, "segment": 2}) - df = pd.concat([df_1, df_2]) + df = pd.DataFrame({"timestamp": timestamp, "target": 3, "segment": 1}) df = TSDataset.to_dataset(df) ts = TSDataset(df=df, freq="D") return ts @@ -249,9 +245,29 @@ def test_model_mixins_not_calls_predict_components_in_predict(mixin_constructor, predict_ts.add_target_components.assert_not_called() +@pytest.mark.parametrize("mixin_constructor", [PerSegmentModelMixin]) +def test_make_prediction_segment_with_components(mixin_constructor, target_components_df, ts_without_target_components): + mixin = mixin_constructor(base_model=Mock()) + target_components_df_model_format = TSDataset.to_flatten(target_components_df).drop(columns=["segment"]) + prediction_method = Mock(return_value=target_components_df_model_format) + + target_components_pred = mixin._make_predictions_segment( + model=mixin._base_model, + segment="1", + df=ts_without_target_components.to_pandas(), + prediction_method=prediction_method, + ) + + pd.testing.assert_frame_equal( + target_components_pred.set_index(["timestamp", "segment"]), + TSDataset.to_flatten(target_components_df).set_index(["timestamp", "segment"]), + ) + + @pytest.mark.parametrize("mixin_constructor", [PerSegmentModelMixin, MultiSegmentModelMixin]) def test_make_components_prediction(mixin_constructor, target_components_df, ts_without_target_components): mixin = mixin_constructor(base_model=Mock()) + mixin.fit(ts_without_target_components) target_components_df_model_format = TSDataset.to_flatten(target_components_df).drop(columns=["segment"]) prediction_method = Mock(return_value=target_components_df_model_format) From 02319b7fc01977b1da9cf6a60fe306205de6f113 Mon Sep 17 00:00:00 2001 From: alex-hse-repository Date: Thu, 9 Mar 2023 11:14:41 +0100 Subject: [PATCH 07/15] Update changelog --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index af890e306..c4ded829d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## Unreleased ### Added +- Target components logic base classes of models ([#1158](https://github.com/tinkoff-ai/etna/pull/1158)) - Target components logic to TSDataset ([#1153](https://github.com/tinkoff-ai/etna/pull/1153)) - Methods `save` and `load` to HierarchicalPipeline ([#1096](https://github.com/tinkoff-ai/etna/pull/1096)) - New data access methods in `TSDataset` : `update_columns_from_pandas`, `add_columns_from_pandas`, `drop_features` ([#809](https://github.com/tinkoff-ai/etna/pull/809)) From 5905c0ba3cccb003ea8825a3b90d82521ade0e9b Mon Sep 17 00:00:00 2001 From: alex-hse-repository Date: Mon, 13 Mar 2023 08:17:07 +0100 Subject: [PATCH 08/15] Fix changelog --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c4ded829d..591bae09b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,7 +7,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## Unreleased ### Added -- Target components logic base classes of models ([#1158](https://github.com/tinkoff-ai/etna/pull/1158)) +- Target components logic into base classes of models ([#1158](https://github.com/tinkoff-ai/etna/pull/1158)) - Target components logic to TSDataset ([#1153](https://github.com/tinkoff-ai/etna/pull/1153)) - Methods `save` and `load` to HierarchicalPipeline ([#1096](https://github.com/tinkoff-ai/etna/pull/1096)) - New data access methods in `TSDataset` : `update_columns_from_pandas`, `add_columns_from_pandas`, `drop_features` ([#809](https://github.com/tinkoff-ai/etna/pull/809)) From 84de9d7e50fe045a23fb21c9117db8f81b6331c6 Mon Sep 17 00:00:00 2001 From: alex-hse-repository Date: Mon, 13 Mar 2023 08:17:32 +0100 Subject: [PATCH 09/15] Review fixes --- etna/models/mixins.py | 83 ++++++++++++++++++++++---------- tests/test_models/test_base.py | 2 +- tests/test_models/test_mixins.py | 81 ++++++++----------------------- 3 files changed, 79 insertions(+), 87 deletions(-) diff --git a/etna/models/mixins.py b/etna/models/mixins.py index 0b7dc82ac..d335cfbfe 100644 --- a/etna/models/mixins.py +++ b/etna/models/mixins.py @@ -36,6 +36,13 @@ def _forecast_components(self, **kwargs) -> pd.DataFrame: def _predict_components(self, **kwargs) -> pd.DataFrame: pass + def _add_target_components( + self, ts: TSDataset, predictions: TSDataset, components_prediction_method: Callable, return_components: bool + ): + if return_components: + target_components_df = components_prediction_method(self=self, ts=ts) + predictions.add_target_components(target_components_df=target_components_df) + class NonPredictionIntervalContextIgnorantModelMixin(ModelForecastingMixin): """Mixin for models that don't support prediction intervals and don't need context for prediction.""" @@ -56,9 +63,12 @@ def forecast(self, ts: TSDataset, return_components: bool = False) -> TSDataset: Dataset with predictions """ forecast = self._forecast(ts=ts) - if return_components: - forecast_components_df = self._forecast_components(ts=ts) - forecast.add_target_components(target_components_df=forecast_components_df) + self._add_target_components( + ts=ts, + predictions=forecast, + components_prediction_method=self._forecast_components, + return_components=return_components, + ) return forecast def predict(self, ts: TSDataset, return_components: bool = False) -> TSDataset: @@ -77,9 +87,12 @@ def predict(self, ts: TSDataset, return_components: bool = False) -> TSDataset: Dataset with predictions """ prediction = self._predict(ts=ts) - if return_components: - prediction_components_df = self._predict_components(ts=ts) - prediction.add_target_components(target_components_df=prediction_components_df) + self._add_target_components( + ts=ts, + predictions=prediction, + components_prediction_method=self._predict_components, + return_components=return_components, + ) return prediction @@ -105,9 +118,12 @@ def forecast(self, ts: TSDataset, prediction_size: int, return_components: bool Dataset with predictions """ forecast = self._forecast(ts=ts, prediction_size=prediction_size) - if return_components: - forecast_components_df = self._forecast_components(ts=ts) - forecast.add_target_components(target_components_df=forecast_components_df) + self._add_target_components( + ts=ts, + predictions=forecast, + components_prediction_method=self._forecast_components, + return_components=return_components, + ) return forecast def predict(self, ts: TSDataset, prediction_size: int, return_components: bool = False) -> TSDataset: @@ -129,9 +145,12 @@ def predict(self, ts: TSDataset, prediction_size: int, return_components: bool = Dataset with predictions """ prediction = self._predict(ts=ts, prediction_size=prediction_size) - if return_components: - prediction_components_df = self._predict_components(ts=ts) - prediction.add_target_components(target_components_df=prediction_components_df) + self._add_target_components( + ts=ts, + predictions=prediction, + components_prediction_method=self._predict_components, + return_components=return_components, + ) return prediction @@ -164,9 +183,12 @@ def forecast( Dataset with predictions """ forecast = self._forecast(ts=ts, prediction_interval=prediction_interval, quantiles=quantiles) - if return_components: - forecast_components_df = self._forecast_components(ts=ts) - forecast.add_target_components(target_components_df=forecast_components_df) + self._add_target_components( + ts=ts, + predictions=forecast, + components_prediction_method=self._forecast_components, + return_components=return_components, + ) return forecast def predict( @@ -195,9 +217,12 @@ def predict( Dataset with predictions """ prediction = self._predict(ts=ts, prediction_interval=prediction_interval, quantiles=quantiles) - if return_components: - prediction_components_df = self._predict_components(ts=ts) - prediction.add_target_components(target_components_df=prediction_components_df) + self._add_target_components( + ts=ts, + predictions=prediction, + components_prediction_method=self._predict_components, + return_components=return_components, + ) return prediction @@ -236,9 +261,12 @@ def forecast( forecast = self._forecast( ts=ts, prediction_size=prediction_size, prediction_interval=prediction_interval, quantiles=quantiles ) - if return_components: - forecast_components_df = self._forecast_components(ts=ts) - forecast.add_target_components(target_components_df=forecast_components_df) + self._add_target_components( + ts=ts, + predictions=forecast, + components_prediction_method=self._forecast_components, + return_components=return_components, + ) return forecast def predict( @@ -273,9 +301,12 @@ def predict( prediction = self._predict( ts=ts, prediction_size=prediction_size, prediction_interval=prediction_interval, quantiles=quantiles ) - if return_components: - prediction_components_df = self._predict_components(ts=ts) - prediction.add_target_components(target_components_df=prediction_components_df) + self._add_target_components( + ts=ts, + predictions=prediction, + components_prediction_method=self._predict_components, + return_components=return_components, + ) return prediction @@ -427,7 +458,7 @@ def _make_component_predictions(self, ts: TSDataset, prediction_method: Callable Returns ------- : - Dataset with predicted components + DataFrame with predicted components """ features_df = ts.to_pandas() result_list = list() @@ -540,7 +571,7 @@ def _make_component_predictions(self, ts: TSDataset, prediction_method: Callable Returns ------- : - Dataset with predicted components + DataFrame with predicted components """ features_df = ts.to_pandas(flatten=True) segment_column = features_df["segment"].values diff --git a/tests/test_models/test_base.py b/tests/test_models/test_base.py index faf8f65c8..17612b7b3 100644 --- a/tests/test_models/test_base.py +++ b/tests/test_models/test_base.py @@ -13,7 +13,7 @@ @pytest.fixture() def deep_base_model_mock(): - model = MagicMock(spec=DeepBaseModel) + model = MagicMock() model.train_batch_size = 32 model.train_dataloader_params = {} model.val_dataloader_params = {} diff --git a/tests/test_models/test_mixins.py b/tests/test_models/test_mixins.py index 9e4d12f8c..ca306d815 100644 --- a/tests/test_models/test_mixins.py +++ b/tests/test_models/test_mixins.py @@ -160,19 +160,23 @@ def test_save_mixin_load_warning(get_version_mock, save_version, load_version, t (NonPredictionIntervalContextRequiredModelMixin, {"prediction_size": 10}), ], ) -def test_model_mixins_calls_forecast_components_in_forecast(mixin_constructor, call_params): +@pytest.mark.parametrize("return_components", (True, False)) +def test_model_mixins_calls_add_target_components_in_forecast(mixin_constructor, return_components, call_params): with patch.multiple(mixin_constructor, __abstractmethods__=set()): ts = Mock() forecast_ts = Mock(spec=TSDataset) - target_components_df = Mock() mixin = mixin_constructor() mixin._forecast = Mock(return_value=forecast_ts) - mixin._forecast_components = Mock(return_value=target_components_df) + mixin._add_target_components = Mock() - _ = mixin.forecast(ts=ts, return_components=True, **call_params) + _ = mixin.forecast(ts=ts, return_components=return_components, **call_params) - mixin._forecast_components.assert_called_with(ts=ts) - forecast_ts.add_target_components.assert_called_with(target_components_df=target_components_df) + mixin._add_target_components.assert_called_with( + ts=ts, + predictions=forecast_ts, + components_prediction_method=mixin._forecast_components, + return_components=return_components, + ) @pytest.mark.parametrize( @@ -184,65 +188,22 @@ def test_model_mixins_calls_forecast_components_in_forecast(mixin_constructor, c (NonPredictionIntervalContextRequiredModelMixin, {"prediction_size": 10}), ], ) -def test_model_mixins_not_calls_forecast_components_in_forecast(mixin_constructor, call_params): - with patch.multiple(mixin_constructor, __abstractmethods__=set()): - ts = Mock() - forecast_ts = Mock(spec=TSDataset) - target_components_df = Mock() - mixin = mixin_constructor() - mixin._forecast = Mock(return_value=forecast_ts) - mixin._forecast_components = Mock(return_value=target_components_df) - - _ = mixin.forecast(ts=ts, return_components=False, **call_params) - - mixin._forecast_components.assert_not_called() - forecast_ts.add_target_components.assert_not_called() - - -@pytest.mark.parametrize( - "mixin_constructor, call_params", - [ - (PredictionIntervalContextIgnorantModelMixin, {}), - (NonPredictionIntervalContextIgnorantModelMixin, {}), - (PredictionIntervalContextRequiredModelMixin, {"prediction_size": 10}), - (NonPredictionIntervalContextRequiredModelMixin, {"prediction_size": 10}), - ], -) -def test_model_mixins_calls_predict_components_in_predict(mixin_constructor, call_params): +@pytest.mark.parametrize("return_components", (True, False)) +def test_model_mixins_calls_add_target_components_in_predict(mixin_constructor, return_components, call_params): with patch.multiple(mixin_constructor, __abstractmethods__=set()): ts = Mock() predict_ts = Mock(spec=TSDataset) - target_components_df = Mock() mixin = mixin_constructor() mixin._predict = Mock(return_value=predict_ts) - mixin._predict_components = Mock(return_value=target_components_df) - _ = mixin.predict(ts=ts, return_components=True, **call_params) - - mixin._predict_components.assert_called_with(ts=ts) - predict_ts.add_target_components.assert_called_with(target_components_df=target_components_df) - - -@pytest.mark.parametrize( - "mixin_constructor, call_params", - [ - (PredictionIntervalContextIgnorantModelMixin, {}), - (NonPredictionIntervalContextIgnorantModelMixin, {}), - (PredictionIntervalContextRequiredModelMixin, {"prediction_size": 10}), - (NonPredictionIntervalContextRequiredModelMixin, {"prediction_size": 10}), - ], -) -def test_model_mixins_not_calls_predict_components_in_predict(mixin_constructor, call_params): - with patch.multiple(mixin_constructor, __abstractmethods__=set()): - ts = Mock() - predict_ts = Mock(spec=TSDataset) - target_components_df = Mock() - mixin = mixin_constructor() - mixin._predict = Mock(return_value=predict_ts) - mixin._predict_components = Mock(return_value=target_components_df) - _ = mixin.predict(ts=ts, return_components=False, **call_params) - - mixin._predict_components.assert_not_called() - predict_ts.add_target_components.assert_not_called() + mixin._add_target_components = Mock() + _ = mixin.predict(ts=ts, return_components=return_components, **call_params) + + mixin._add_target_components.assert_called_with( + ts=ts, + predictions=predict_ts, + components_prediction_method=mixin._predict_components, + return_components=return_components, + ) @pytest.mark.parametrize("mixin_constructor", [PerSegmentModelMixin]) From 317ce1beff3ba280ec2550541039733dc8733ba9 Mon Sep 17 00:00:00 2001 From: alex-hse-repository Date: Tue, 14 Mar 2023 11:03:16 +0100 Subject: [PATCH 10/15] Rework tests for forecast and predict with target components --- tests/test_models/test_mixins.py | 152 ++++++++++++++++++++----------- 1 file changed, 97 insertions(+), 55 deletions(-) diff --git a/tests/test_models/test_mixins.py b/tests/test_models/test_mixins.py index ca306d815..c0e66b739 100644 --- a/tests/test_models/test_mixins.py +++ b/tests/test_models/test_mixins.py @@ -24,6 +24,39 @@ from etna.models.mixins import PredictionIntervalContextRequiredModelMixin from etna.models.mixins import SaveNNMixin +class DummyModelBase: + def _forecast(self, ts: TSDataset, **kwargs) -> TSDataset: + ts.loc[pd.IndexSlice[:], pd.IndexSlice[:, "target"]] = 100 + return ts + + def _predict(self, ts: TSDataset, **kwargs) -> TSDataset: + ts.loc[pd.IndexSlice[:], pd.IndexSlice[:, "target"]] = 200 + return ts + + def _forecast_components(self, ts: TSDataset, **kwargs) -> pd.DataFrame: + df = ts.to_pandas(flatten=True, features=["target"]) + df["target_component_a"] = 10 + df["target_component_b"] = 90 + df = df.drop(columns=["target"]) + df = TSDataset.to_dataset(df) + return df + + def _predict_components(self, ts: TSDataset, **kwargs) -> pd.DataFrame: + df = ts.to_pandas(flatten=True, features=["target"]) + df["target_component_a"] = 20 + df["target_component_b"] = 180 + df = df.drop(columns=["target"]) + df = TSDataset.to_dataset(df) + return df + +class NonPredictionIntervalContextIgnorantDummyModel(DummyModelBase, NonPredictionIntervalContextIgnorantModelMixin): + pass +class NonPredictionIntervalContextRequiredDummyModel(DummyModelBase, NonPredictionIntervalContextRequiredModelMixin): + pass +class PredictionIntervalContextIgnorantDummyModel(DummyModelBase, PredictionIntervalContextIgnorantModelMixin): + pass +class PredictionIntervalContextRequiredDummyModel(DummyModelBase, PredictionIntervalContextRequiredModelMixin): + pass @pytest.fixture() def regression_base_model_mock(): @@ -151,61 +184,6 @@ def test_save_mixin_load_warning(get_version_mock, save_version, load_version, t _ = DummyNN.load(path) -@pytest.mark.parametrize( - "mixin_constructor, call_params", - [ - (PredictionIntervalContextIgnorantModelMixin, {}), - (NonPredictionIntervalContextIgnorantModelMixin, {}), - (PredictionIntervalContextRequiredModelMixin, {"prediction_size": 10}), - (NonPredictionIntervalContextRequiredModelMixin, {"prediction_size": 10}), - ], -) -@pytest.mark.parametrize("return_components", (True, False)) -def test_model_mixins_calls_add_target_components_in_forecast(mixin_constructor, return_components, call_params): - with patch.multiple(mixin_constructor, __abstractmethods__=set()): - ts = Mock() - forecast_ts = Mock(spec=TSDataset) - mixin = mixin_constructor() - mixin._forecast = Mock(return_value=forecast_ts) - mixin._add_target_components = Mock() - - _ = mixin.forecast(ts=ts, return_components=return_components, **call_params) - - mixin._add_target_components.assert_called_with( - ts=ts, - predictions=forecast_ts, - components_prediction_method=mixin._forecast_components, - return_components=return_components, - ) - - -@pytest.mark.parametrize( - "mixin_constructor, call_params", - [ - (PredictionIntervalContextIgnorantModelMixin, {}), - (NonPredictionIntervalContextIgnorantModelMixin, {}), - (PredictionIntervalContextRequiredModelMixin, {"prediction_size": 10}), - (NonPredictionIntervalContextRequiredModelMixin, {"prediction_size": 10}), - ], -) -@pytest.mark.parametrize("return_components", (True, False)) -def test_model_mixins_calls_add_target_components_in_predict(mixin_constructor, return_components, call_params): - with patch.multiple(mixin_constructor, __abstractmethods__=set()): - ts = Mock() - predict_ts = Mock(spec=TSDataset) - mixin = mixin_constructor() - mixin._predict = Mock(return_value=predict_ts) - mixin._add_target_components = Mock() - _ = mixin.predict(ts=ts, return_components=return_components, **call_params) - - mixin._add_target_components.assert_called_with( - ts=ts, - predictions=predict_ts, - components_prediction_method=mixin._predict_components, - return_components=return_components, - ) - - @pytest.mark.parametrize("mixin_constructor", [PerSegmentModelMixin]) def test_make_prediction_segment_with_components(mixin_constructor, target_components_df, ts_without_target_components): mixin = mixin_constructor(base_model=Mock()) @@ -236,3 +214,67 @@ def test_make_components_prediction(mixin_constructor, target_components_df, ts_ ts=ts_without_target_components, prediction_method=prediction_method ) pd.testing.assert_frame_equal(target_components_pred, target_components_df) + +@pytest.mark.parametrize( + "mixin_constructor, call_params", + [ + (NonPredictionIntervalContextIgnorantDummyModel, {}), + (NonPredictionIntervalContextRequiredDummyModel, {"prediction_size": 10}), + (PredictionIntervalContextIgnorantDummyModel, {}), + (PredictionIntervalContextRequiredDummyModel, {"prediction_size": 10}), + ], +) +def test_model_mixins_forecast_without_target_components(ts_without_target_components, mixin_constructor, call_params, expected_target=100, expected_columns=["timestamp", "segment", "target"]): + mixin = mixin_constructor() + forecast = mixin.forecast(ts=ts_without_target_components, return_components=False, **call_params).to_pandas(flatten=True) + assert sorted(forecast.columns) == sorted(expected_columns) + assert (forecast["target"] == expected_target).all() + +@pytest.mark.parametrize( + "mixin_constructor, call_params", + [ + (NonPredictionIntervalContextIgnorantDummyModel, {}), + (NonPredictionIntervalContextRequiredDummyModel, {"prediction_size": 10}), + (PredictionIntervalContextIgnorantDummyModel, {}), + (PredictionIntervalContextRequiredDummyModel, {"prediction_size": 10}), + ], +) +def test_model_mixins_predict_without_target_components(ts_without_target_components, mixin_constructor, call_params, expected_target=200, expected_columns=["timestamp", "segment", "target"]): + mixin = mixin_constructor() + forecast = mixin.predict(ts=ts_without_target_components, return_components=False, **call_params).to_pandas(flatten=True) + assert sorted(forecast.columns) == sorted(expected_columns) + assert (forecast["target"] == expected_target).all() + +@pytest.mark.parametrize( + "mixin_constructor, call_params", + [ + (NonPredictionIntervalContextIgnorantDummyModel, {}), + (NonPredictionIntervalContextRequiredDummyModel, {"prediction_size": 10}), + (PredictionIntervalContextIgnorantDummyModel, {}), + (PredictionIntervalContextRequiredDummyModel, {"prediction_size": 10}), + ], +) +def test_model_mixins_forecast_with_target_components(ts_without_target_components, mixin_constructor, call_params, expected_target=100, expected_component_a=10, expected_component_b=90, expected_columns=["timestamp", "segment", "target", "target_component_a", "target_component_b"]): + mixin = mixin_constructor() + forecast = mixin.forecast(ts=ts_without_target_components, return_components=True, **call_params).to_pandas(flatten=True) + assert sorted(forecast.columns) == sorted(expected_columns) + assert (forecast["target"] == expected_target).all() + assert (forecast["target_component_a"] == expected_component_a).all() + assert (forecast["target_component_b"] == expected_component_b).all() + +@pytest.mark.parametrize( + "mixin_constructor, call_params", + [ + (NonPredictionIntervalContextIgnorantDummyModel, {}), + (NonPredictionIntervalContextRequiredDummyModel, {"prediction_size": 10}), + (PredictionIntervalContextIgnorantDummyModel, {}), + (PredictionIntervalContextRequiredDummyModel, {"prediction_size": 10}), + ], +) +def test_model_mixins_predict_with_target_components(ts_without_target_components, mixin_constructor, call_params, expected_target=200, expected_component_a=20, expected_component_b=180, expected_columns=["timestamp", "segment", "target", "target_component_a", "target_component_b"]): + mixin = mixin_constructor() + forecast = mixin.predict(ts=ts_without_target_components, return_components=True, **call_params).to_pandas(flatten=True) + assert sorted(forecast.columns) == sorted(expected_columns) + assert (forecast["target"] == expected_target).all() + assert (forecast["target_component_a"] == expected_component_a).all() + assert (forecast["target_component_b"] == expected_component_b).all() From c3363d32a6480de4e7948e4e66a469d82ca2013c Mon Sep 17 00:00:00 2001 From: alex-hse-repository Date: Tue, 14 Mar 2023 11:30:54 +0100 Subject: [PATCH 11/15] Add tests for mixin implementations interface --- tests/test_models/test_mixins.py | 103 ++++++++++++++++--------------- 1 file changed, 52 insertions(+), 51 deletions(-) diff --git a/tests/test_models/test_mixins.py b/tests/test_models/test_mixins.py index c0e66b739..43fdfc8ac 100644 --- a/tests/test_models/test_mixins.py +++ b/tests/test_models/test_mixins.py @@ -1,7 +1,6 @@ import json import pathlib from unittest.mock import MagicMock -from unittest.mock import Mock from unittest.mock import patch from zipfile import ZipFile @@ -23,6 +22,34 @@ from etna.models.mixins import PredictionIntervalContextIgnorantModelMixin from etna.models.mixins import PredictionIntervalContextRequiredModelMixin from etna.models.mixins import SaveNNMixin +from etna.models.base import BaseAdapter + +class DummyAdapter(BaseAdapter): + def fit(self, df: pd.DataFrame, **kwargs) -> "DummyAdapter": + return self + + def forecast(self, df: pd.DataFrame, **kwargs) -> pd.DataFrame: + df["target"] = 100 + return df + + def predict(self, df: pd.DataFrame, **kwargs) -> pd.DataFrame: + df["target"] = 200 + return df + + def forecast_components(self, df: pd.DataFrame, **kwargs) -> pd.DataFrame: + df["target_component_a"] = 10 + df["target_component_b"] = 90 + df = df.drop(columns=["target"]) + return df + + def predict_components(self, df: pd.DataFrame, **kwargs) -> pd.DataFrame: + df["target_component_a"] = 20 + df["target_component_b"] = 180 + df = df.drop(columns=["target"]) + return df + + def get_model(self) -> "DummyAdapter": + return self class DummyModelBase: def _forecast(self, ts: TSDataset, **kwargs) -> TSDataset: @@ -78,23 +105,6 @@ def autoregression_base_model_mock(): return model -@pytest.fixture -def target_components_df(): - timestamp = pd.date_range("2021-01-01", "2021-01-15") - df = pd.DataFrame({"timestamp": timestamp, "target_component_a": 1, "target_component_b": 2, "segment": 1}) - df = TSDataset.to_dataset(df) - return df - - -@pytest.fixture -def ts_without_target_components(): - timestamp = pd.date_range("2021-01-01", "2021-01-15") - df = pd.DataFrame({"timestamp": timestamp, "target": 3, "segment": 1}) - df = TSDataset.to_dataset(df) - ts = TSDataset(df=df, freq="D") - return ts - - @pytest.mark.parametrize("mixin_constructor", [PerSegmentModelMixin, MultiSegmentModelMixin]) @pytest.mark.parametrize( "base_model_name, called_method_name, expected_method_name", @@ -184,37 +194,6 @@ def test_save_mixin_load_warning(get_version_mock, save_version, load_version, t _ = DummyNN.load(path) -@pytest.mark.parametrize("mixin_constructor", [PerSegmentModelMixin]) -def test_make_prediction_segment_with_components(mixin_constructor, target_components_df, ts_without_target_components): - mixin = mixin_constructor(base_model=Mock()) - target_components_df_model_format = TSDataset.to_flatten(target_components_df).drop(columns=["segment"]) - prediction_method = Mock(return_value=target_components_df_model_format) - - target_components_pred = mixin._make_predictions_segment( - model=mixin._base_model, - segment="1", - df=ts_without_target_components.to_pandas(), - prediction_method=prediction_method, - ) - - pd.testing.assert_frame_equal( - target_components_pred.set_index(["timestamp", "segment"]), - TSDataset.to_flatten(target_components_df).set_index(["timestamp", "segment"]), - ) - - -@pytest.mark.parametrize("mixin_constructor", [PerSegmentModelMixin, MultiSegmentModelMixin]) -def test_make_components_prediction(mixin_constructor, target_components_df, ts_without_target_components): - mixin = mixin_constructor(base_model=Mock()) - mixin.fit(ts_without_target_components) - target_components_df_model_format = TSDataset.to_flatten(target_components_df).drop(columns=["segment"]) - prediction_method = Mock(return_value=target_components_df_model_format) - - target_components_pred = mixin._make_component_predictions( - ts=ts_without_target_components, prediction_method=prediction_method - ) - pd.testing.assert_frame_equal(target_components_pred, target_components_df) - @pytest.mark.parametrize( "mixin_constructor, call_params", [ @@ -271,10 +250,32 @@ def test_model_mixins_forecast_with_target_components(ts_without_target_componen (PredictionIntervalContextRequiredDummyModel, {"prediction_size": 10}), ], ) -def test_model_mixins_predict_with_target_components(ts_without_target_components, mixin_constructor, call_params, expected_target=200, expected_component_a=20, expected_component_b=180, expected_columns=["timestamp", "segment", "target", "target_component_a", "target_component_b"]): +def test_model_mixins_predict_with_target_components(example_tsds, mixin_constructor, call_params, expected_target=200, expected_component_a=20, expected_component_b=180, expected_columns=["timestamp", "segment", "target", "target_component_a", "target_component_b"]): mixin = mixin_constructor() - forecast = mixin.predict(ts=ts_without_target_components, return_components=True, **call_params).to_pandas(flatten=True) + forecast = mixin.predict(ts=example_tsds, return_components=True, **call_params).to_pandas(flatten=True) assert sorted(forecast.columns) == sorted(expected_columns) assert (forecast["target"] == expected_target).all() assert (forecast["target_component_a"] == expected_component_a).all() assert (forecast["target_component_b"] == expected_component_b).all() + + +@pytest.mark.parametrize("mixin_constructor", [PerSegmentModelMixin, MultiSegmentModelMixin]) +@pytest.mark.parametrize("method_name, expected_target", [("_forecast", 100), ("_predict", 200)]) +def test_mixin_implementations_prediction_methods(example_tsds, mixin_constructor, method_name, expected_target, expected_columns=["timestamp", "segment", "target"]): + mixin = mixin_constructor(base_model=DummyAdapter()) + mixin = mixin.fit(ts=example_tsds) + to_call = getattr(mixin, method_name) + forecast = to_call(ts=example_tsds).to_pandas(flatten=True) + assert sorted(forecast.columns) == sorted(expected_columns) + assert (forecast["target"] == expected_target).all() + +@pytest.mark.parametrize("mixin_constructor", [PerSegmentModelMixin, MultiSegmentModelMixin]) +@pytest.mark.parametrize("method_name, expected_component_a, expected_component_b", [("_forecast_components", 10, 90), ("_predict_components", 20, 180)]) +def test_mixin_implementations_prediction_components_methods(example_tsds, mixin_constructor, method_name, expected_component_a, expected_component_b, expected_columns=["timestamp", "segment", "target_component_a", "target_component_b"]): + mixin = mixin_constructor(base_model=DummyAdapter()) + mixin = mixin.fit(ts=example_tsds) + to_call = getattr(mixin, method_name) + forecast = TSDataset.to_flatten(to_call(ts=example_tsds)) + assert sorted(forecast.columns) == sorted(expected_columns) + assert (forecast["target_component_a"] == expected_component_a).all() + assert (forecast["target_component_b"] == expected_component_b).all() From dac65e8e184d571635fae0aadb7a131e762785fc Mon Sep 17 00:00:00 2001 From: alex-hse-repository Date: Tue, 14 Mar 2023 11:36:56 +0100 Subject: [PATCH 12/15] Refactor tests --- tests/test_models/test_mixins.py | 89 +++++++++++++++++++------------- 1 file changed, 52 insertions(+), 37 deletions(-) diff --git a/tests/test_models/test_mixins.py b/tests/test_models/test_mixins.py index 43fdfc8ac..78d42b24b 100644 --- a/tests/test_models/test_mixins.py +++ b/tests/test_models/test_mixins.py @@ -15,6 +15,7 @@ import pandas as pd +from etna.models.base import BaseAdapter from etna.models.mixins import MultiSegmentModelMixin from etna.models.mixins import NonPredictionIntervalContextIgnorantModelMixin from etna.models.mixins import NonPredictionIntervalContextRequiredModelMixin @@ -22,7 +23,7 @@ from etna.models.mixins import PredictionIntervalContextIgnorantModelMixin from etna.models.mixins import PredictionIntervalContextRequiredModelMixin from etna.models.mixins import SaveNNMixin -from etna.models.base import BaseAdapter + class DummyAdapter(BaseAdapter): def fit(self, df: pd.DataFrame, **kwargs) -> "DummyAdapter": @@ -51,6 +52,7 @@ def predict_components(self, df: pd.DataFrame, **kwargs) -> pd.DataFrame: def get_model(self) -> "DummyAdapter": return self + class DummyModelBase: def _forecast(self, ts: TSDataset, **kwargs) -> TSDataset: ts.loc[pd.IndexSlice[:], pd.IndexSlice[:, "target"]] = 100 @@ -76,15 +78,23 @@ def _predict_components(self, ts: TSDataset, **kwargs) -> pd.DataFrame: df = TSDataset.to_dataset(df) return df + class NonPredictionIntervalContextIgnorantDummyModel(DummyModelBase, NonPredictionIntervalContextIgnorantModelMixin): pass + + class NonPredictionIntervalContextRequiredDummyModel(DummyModelBase, NonPredictionIntervalContextRequiredModelMixin): pass + + class PredictionIntervalContextIgnorantDummyModel(DummyModelBase, PredictionIntervalContextIgnorantModelMixin): pass + + class PredictionIntervalContextRequiredDummyModel(DummyModelBase, PredictionIntervalContextRequiredModelMixin): pass + @pytest.fixture() def regression_base_model_mock(): cls = MagicMock() @@ -203,26 +213,21 @@ def test_save_mixin_load_warning(get_version_mock, save_version, load_version, t (PredictionIntervalContextRequiredDummyModel, {"prediction_size": 10}), ], ) -def test_model_mixins_forecast_without_target_components(ts_without_target_components, mixin_constructor, call_params, expected_target=100, expected_columns=["timestamp", "segment", "target"]): +@pytest.mark.parametrize("method_name, expected_target", [("forecast", 100), ("predict", 200)]) +def test_model_mixins_predict_without_target_components( + example_tsds, + mixin_constructor, + call_params, + method_name, + expected_target, + expected_columns=["timestamp", "segment", "target"], +): mixin = mixin_constructor() - forecast = mixin.forecast(ts=ts_without_target_components, return_components=False, **call_params).to_pandas(flatten=True) + to_call = getattr(mixin, method_name) + forecast = to_call(ts=example_tsds, return_components=False, **call_params).to_pandas(flatten=True) assert sorted(forecast.columns) == sorted(expected_columns) assert (forecast["target"] == expected_target).all() -@pytest.mark.parametrize( - "mixin_constructor, call_params", - [ - (NonPredictionIntervalContextIgnorantDummyModel, {}), - (NonPredictionIntervalContextRequiredDummyModel, {"prediction_size": 10}), - (PredictionIntervalContextIgnorantDummyModel, {}), - (PredictionIntervalContextRequiredDummyModel, {"prediction_size": 10}), - ], -) -def test_model_mixins_predict_without_target_components(ts_without_target_components, mixin_constructor, call_params, expected_target=200, expected_columns=["timestamp", "segment", "target"]): - mixin = mixin_constructor() - forecast = mixin.predict(ts=ts_without_target_components, return_components=False, **call_params).to_pandas(flatten=True) - assert sorted(forecast.columns) == sorted(expected_columns) - assert (forecast["target"] == expected_target).all() @pytest.mark.parametrize( "mixin_constructor, call_params", @@ -233,26 +238,23 @@ def test_model_mixins_predict_without_target_components(ts_without_target_compon (PredictionIntervalContextRequiredDummyModel, {"prediction_size": 10}), ], ) -def test_model_mixins_forecast_with_target_components(ts_without_target_components, mixin_constructor, call_params, expected_target=100, expected_component_a=10, expected_component_b=90, expected_columns=["timestamp", "segment", "target", "target_component_a", "target_component_b"]): - mixin = mixin_constructor() - forecast = mixin.forecast(ts=ts_without_target_components, return_components=True, **call_params).to_pandas(flatten=True) - assert sorted(forecast.columns) == sorted(expected_columns) - assert (forecast["target"] == expected_target).all() - assert (forecast["target_component_a"] == expected_component_a).all() - assert (forecast["target_component_b"] == expected_component_b).all() - @pytest.mark.parametrize( - "mixin_constructor, call_params", - [ - (NonPredictionIntervalContextIgnorantDummyModel, {}), - (NonPredictionIntervalContextRequiredDummyModel, {"prediction_size": 10}), - (PredictionIntervalContextIgnorantDummyModel, {}), - (PredictionIntervalContextRequiredDummyModel, {"prediction_size": 10}), - ], + "method_name, expected_target, expected_component_a, expected_component_b", + [("forecast", 100, 10, 90), ("predict", 200, 20, 180)], ) -def test_model_mixins_predict_with_target_components(example_tsds, mixin_constructor, call_params, expected_target=200, expected_component_a=20, expected_component_b=180, expected_columns=["timestamp", "segment", "target", "target_component_a", "target_component_b"]): +def test_model_mixins_prediction_methods_with_target_components( + example_tsds, + mixin_constructor, + call_params, + method_name, + expected_target, + expected_component_a, + expected_component_b, + expected_columns=["timestamp", "segment", "target", "target_component_a", "target_component_b"], +): mixin = mixin_constructor() - forecast = mixin.predict(ts=example_tsds, return_components=True, **call_params).to_pandas(flatten=True) + to_call = getattr(mixin, method_name) + forecast = to_call(ts=example_tsds, return_components=True, **call_params).to_pandas(flatten=True) assert sorted(forecast.columns) == sorted(expected_columns) assert (forecast["target"] == expected_target).all() assert (forecast["target_component_a"] == expected_component_a).all() @@ -261,7 +263,9 @@ def test_model_mixins_predict_with_target_components(example_tsds, mixin_constru @pytest.mark.parametrize("mixin_constructor", [PerSegmentModelMixin, MultiSegmentModelMixin]) @pytest.mark.parametrize("method_name, expected_target", [("_forecast", 100), ("_predict", 200)]) -def test_mixin_implementations_prediction_methods(example_tsds, mixin_constructor, method_name, expected_target, expected_columns=["timestamp", "segment", "target"]): +def test_mixin_implementations_prediction_methods( + example_tsds, mixin_constructor, method_name, expected_target, expected_columns=["timestamp", "segment", "target"] +): mixin = mixin_constructor(base_model=DummyAdapter()) mixin = mixin.fit(ts=example_tsds) to_call = getattr(mixin, method_name) @@ -269,9 +273,20 @@ def test_mixin_implementations_prediction_methods(example_tsds, mixin_constructo assert sorted(forecast.columns) == sorted(expected_columns) assert (forecast["target"] == expected_target).all() + @pytest.mark.parametrize("mixin_constructor", [PerSegmentModelMixin, MultiSegmentModelMixin]) -@pytest.mark.parametrize("method_name, expected_component_a, expected_component_b", [("_forecast_components", 10, 90), ("_predict_components", 20, 180)]) -def test_mixin_implementations_prediction_components_methods(example_tsds, mixin_constructor, method_name, expected_component_a, expected_component_b, expected_columns=["timestamp", "segment", "target_component_a", "target_component_b"]): +@pytest.mark.parametrize( + "method_name, expected_component_a, expected_component_b", + [("_forecast_components", 10, 90), ("_predict_components", 20, 180)], +) +def test_mixin_implementations_prediction_components_methods( + example_tsds, + mixin_constructor, + method_name, + expected_component_a, + expected_component_b, + expected_columns=["timestamp", "segment", "target_component_a", "target_component_b"], +): mixin = mixin_constructor(base_model=DummyAdapter()) mixin = mixin.fit(ts=example_tsds) to_call = getattr(mixin, method_name) From e4e1a2f62f3d91ddd15c64cbba8a03013c955e13 Mon Sep 17 00:00:00 2001 From: alex-hse-repository Date: Tue, 14 Mar 2023 11:47:14 +0100 Subject: [PATCH 13/15] Small fixes --- etna/models/mixins.py | 2 +- tests/test_models/test_mixins.py | 9 +++++---- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/etna/models/mixins.py b/etna/models/mixins.py index d335cfbfe..6e462167d 100644 --- a/etna/models/mixins.py +++ b/etna/models/mixins.py @@ -40,7 +40,7 @@ def _add_target_components( self, ts: TSDataset, predictions: TSDataset, components_prediction_method: Callable, return_components: bool ): if return_components: - target_components_df = components_prediction_method(self=self, ts=ts) + target_components_df = components_prediction_method(ts=ts) predictions.add_target_components(target_components_df=target_components_df) diff --git a/tests/test_models/test_mixins.py b/tests/test_models/test_mixins.py index 78d42b24b..417b5709e 100644 --- a/tests/test_models/test_mixins.py +++ b/tests/test_models/test_mixins.py @@ -5,6 +5,7 @@ from zipfile import ZipFile import dill +import numpy as np import pytest from etna import SETTINGS @@ -29,13 +30,13 @@ class DummyAdapter(BaseAdapter): def fit(self, df: pd.DataFrame, **kwargs) -> "DummyAdapter": return self - def forecast(self, df: pd.DataFrame, **kwargs) -> pd.DataFrame: + def forecast(self, df: pd.DataFrame, **kwargs) -> np.ndarray: df["target"] = 100 - return df + return df["target"].values - def predict(self, df: pd.DataFrame, **kwargs) -> pd.DataFrame: + def predict(self, df: pd.DataFrame, **kwargs) -> np.ndarray: df["target"] = 200 - return df + return df["target"].values def forecast_components(self, df: pd.DataFrame, **kwargs) -> pd.DataFrame: df["target_component_a"] = 10 From 1ffa2e579631c18854d72ef953b3a023aea6809c Mon Sep 17 00:00:00 2001 From: alex-hse-repository Date: Tue, 14 Mar 2023 13:04:12 +0100 Subject: [PATCH 14/15] Add test for adapter withot frorecast --- tests/test_models/test_mixins.py | 58 ++++++++++++++++++++++---------- 1 file changed, 41 insertions(+), 17 deletions(-) diff --git a/tests/test_models/test_mixins.py b/tests/test_models/test_mixins.py index 417b5709e..4bfaaba74 100644 --- a/tests/test_models/test_mixins.py +++ b/tests/test_models/test_mixins.py @@ -26,24 +26,14 @@ from etna.models.mixins import SaveNNMixin -class DummyAdapter(BaseAdapter): +class DummyPredictAdapter(BaseAdapter): def fit(self, df: pd.DataFrame, **kwargs) -> "DummyAdapter": return self - def forecast(self, df: pd.DataFrame, **kwargs) -> np.ndarray: - df["target"] = 100 - return df["target"].values - def predict(self, df: pd.DataFrame, **kwargs) -> np.ndarray: df["target"] = 200 return df["target"].values - def forecast_components(self, df: pd.DataFrame, **kwargs) -> pd.DataFrame: - df["target_component_a"] = 10 - df["target_component_b"] = 90 - df = df.drop(columns=["target"]) - return df - def predict_components(self, df: pd.DataFrame, **kwargs) -> pd.DataFrame: df["target_component_a"] = 20 df["target_component_b"] = 180 @@ -54,6 +44,21 @@ def get_model(self) -> "DummyAdapter": return self +class DummyForecastPredictAdapter(DummyPredictAdapter): + def fit(self, df: pd.DataFrame, **kwargs) -> "DummyAdapter": + return self + + def forecast(self, df: pd.DataFrame, **kwargs) -> np.ndarray: + df["target"] = 100 + return df["target"].values + + def forecast_components(self, df: pd.DataFrame, **kwargs) -> pd.DataFrame: + df["target_component_a"] = 10 + df["target_component_b"] = 90 + df = df.drop(columns=["target"]) + return df + + class DummyModelBase: def _forecast(self, ts: TSDataset, **kwargs) -> TSDataset: ts.loc[pd.IndexSlice[:], pd.IndexSlice[:, "target"]] = 100 @@ -263,11 +268,24 @@ def test_model_mixins_prediction_methods_with_target_components( @pytest.mark.parametrize("mixin_constructor", [PerSegmentModelMixin, MultiSegmentModelMixin]) -@pytest.mark.parametrize("method_name, expected_target", [("_forecast", 100), ("_predict", 200)]) +@pytest.mark.parametrize( + "method_name, adapter_constructor, expected_target", + [ + ("_forecast", DummyForecastPredictAdapter, 100), + ("_predict", DummyForecastPredictAdapter, 200), + ("_forecast", DummyPredictAdapter, 200), + ("_predict", DummyPredictAdapter, 200), + ], +) def test_mixin_implementations_prediction_methods( - example_tsds, mixin_constructor, method_name, expected_target, expected_columns=["timestamp", "segment", "target"] + example_tsds, + mixin_constructor, + method_name, + adapter_constructor, + expected_target, + expected_columns=["timestamp", "segment", "target"], ): - mixin = mixin_constructor(base_model=DummyAdapter()) + mixin = mixin_constructor(base_model=adapter_constructor()) mixin = mixin.fit(ts=example_tsds) to_call = getattr(mixin, method_name) forecast = to_call(ts=example_tsds).to_pandas(flatten=True) @@ -277,18 +295,24 @@ def test_mixin_implementations_prediction_methods( @pytest.mark.parametrize("mixin_constructor", [PerSegmentModelMixin, MultiSegmentModelMixin]) @pytest.mark.parametrize( - "method_name, expected_component_a, expected_component_b", - [("_forecast_components", 10, 90), ("_predict_components", 20, 180)], + "method_name, adapter_constructor, expected_component_a, expected_component_b", + [ + ("_forecast_components", DummyForecastPredictAdapter, 10, 90), + ("_predict_components", DummyForecastPredictAdapter, 20, 180), + ("_forecast_components", DummyPredictAdapter, 20, 180), + ("_predict_components", DummyPredictAdapter, 20, 180), + ], ) def test_mixin_implementations_prediction_components_methods( example_tsds, mixin_constructor, method_name, + adapter_constructor, expected_component_a, expected_component_b, expected_columns=["timestamp", "segment", "target_component_a", "target_component_b"], ): - mixin = mixin_constructor(base_model=DummyAdapter()) + mixin = mixin_constructor(base_model=adapter_constructor()) mixin = mixin.fit(ts=example_tsds) to_call = getattr(mixin, method_name) forecast = TSDataset.to_flatten(to_call(ts=example_tsds)) From 4f7231bb3d58dd7483acbe9dd650a42bcbab7f2c Mon Sep 17 00:00:00 2001 From: alex-hse-repository Date: Tue, 14 Mar 2023 16:37:39 +0100 Subject: [PATCH 15/15] Fix typing --- tests/test_models/test_mixins.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_models/test_mixins.py b/tests/test_models/test_mixins.py index 4bfaaba74..099040720 100644 --- a/tests/test_models/test_mixins.py +++ b/tests/test_models/test_mixins.py @@ -27,7 +27,7 @@ class DummyPredictAdapter(BaseAdapter): - def fit(self, df: pd.DataFrame, **kwargs) -> "DummyAdapter": + def fit(self, df: pd.DataFrame, **kwargs) -> "DummyPredictAdapter": return self def predict(self, df: pd.DataFrame, **kwargs) -> np.ndarray: @@ -40,12 +40,12 @@ def predict_components(self, df: pd.DataFrame, **kwargs) -> pd.DataFrame: df = df.drop(columns=["target"]) return df - def get_model(self) -> "DummyAdapter": + def get_model(self) -> "DummyPredictAdapter": return self class DummyForecastPredictAdapter(DummyPredictAdapter): - def fit(self, df: pd.DataFrame, **kwargs) -> "DummyAdapter": + def fit(self, df: pd.DataFrame, **kwargs) -> "DummyForecastPredictAdapter": return self def forecast(self, df: pd.DataFrame, **kwargs) -> np.ndarray: