diff --git a/CHANGELOG.md b/CHANGELOG.md index 32dfcc8f9..71f1149a3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,7 +19,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed - - Fix `BaseReconciliator` to work on `pandas==1.1.5` ([#1229](https://github.com/tinkoff-ai/etna/pull/1229)) -- +- Fix `TSDataset.make_future` to handle hierarchy, quantiles, target components ([#1248](https://github.com/tinkoff-ai/etna/pull/1248)) - Fix warning during creation of `ResampleWithDistributionTransform` ([#1230](https://github.com/tinkoff-ai/etna/pull/1230)) - Add deep copy for copying attributes of `TSDataset` ([#1241](https://github.com/tinkoff-ai/etna/pull/1241)) - diff --git a/etna/datasets/tsdataset.py b/etna/datasets/tsdataset.py index 3ee0030b2..ea8e5f043 100644 --- a/etna/datasets/tsdataset.py +++ b/etna/datasets/tsdataset.py @@ -230,21 +230,23 @@ def __getitem__(self, item): def make_future( self, future_steps: int, transforms: Sequence["Transform"] = (), tail_steps: int = 0 ) -> "TSDataset": - """Return new TSDataset with future steps. + """Return new TSDataset with features extended into the future. + + The result dataset doesn't contain quantiles and target components. Parameters ---------- future_steps: - number of timestamp in the future to build features for. + number of steps to extend dataset into the future. transforms: sequence of transforms to be applied. tail_steps: - number of timestamp for context to build features for. + number of steps to keep from the tail of the original dataset. Returns ------- : - dataset with features in the future. + dataset with features extended into the. Examples -------- @@ -295,6 +297,14 @@ def make_future( f"NaN-s will be used for missing values" ) + # remove components and quantiles + # it should be done if we have quantiles and components in raw_df + # TODO: fix this after making quantiles to work like components, with special methods + if len(self.target_components_names) > 0: + df = df.drop(columns=list(self.target_components_names), level="feature") + if len(self.target_quantiles_names) > 0: + df = df.drop(columns=list(self.target_quantiles_names), level="feature") + # Here only df is required, other metadata is not necessary to build the dataset ts = TSDataset(df=df, freq=self.freq) for transform in transforms: @@ -305,7 +315,7 @@ def make_future( future_dataset = df.tail(future_steps + tail_steps).copy(deep=True) future_dataset = future_dataset.sort_index(axis=1, level=(0, 1)) - future_ts = TSDataset(df=future_dataset, freq=self.freq) + future_ts = TSDataset(df=future_dataset, freq=self.freq, hierarchical_structure=self.hierarchical_structure) # can't put known_future into constructor, _check_known_future fails with df_exog=None future_ts.known_future = deepcopy(self.known_future) diff --git a/tests/test_datasets/test_dataset.py b/tests/test_datasets/test_dataset.py index 3c796daf2..5862615b3 100644 --- a/tests/test_datasets/test_dataset.py +++ b/tests/test_datasets/test_dataset.py @@ -620,6 +620,24 @@ def test_make_future_inherits_regressors(df_and_regressors): assert ts_future.regressors == ts.regressors +def test_make_future_inherits_hierarchy(product_level_constant_forecast_with_quantiles): + ts = product_level_constant_forecast_with_quantiles + future = ts.make_future(future_steps=2) + assert future.hierarchical_structure is ts.hierarchical_structure + + +def test_make_future_removes_quantiles(product_level_constant_forecast_with_quantiles): + ts = product_level_constant_forecast_with_quantiles + future = ts.make_future(future_steps=2) + assert len(future.target_quantiles_names) == 0 + + +def test_make_future_removes_target_components(ts_with_target_components): + ts = ts_with_target_components + future = ts.make_future(future_steps=2) + assert len(future.target_components_names) == 0 + + def test_make_future_warn_not_enough_regressors(df_and_regressors): """Check that warning is thrown if regressors don't have enough values for the future.""" df, df_exog, known_future = df_and_regressors