Fix bug with indexing in _forecast_segment

tinkoff-ai · Jul 8, 2022 · 1f0eee4 · 1f0eee4
1 parent 116709e
commit 1f0eee4
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/etna/models/base.py b/etna/models/base.py
@@ -78,7 +78,7 @@ def forecast(
 
     @staticmethod
     def _forecast_segment(model, segment: Union[str, List[str]], ts: TSDataset) -> pd.DataFrame:
-        segment_features = ts[:, segment, :]
+        segment_features = ts.df.loc[:, pd.IndexSlice[segment, :]]
         segment_features = segment_features.droplevel("segment", axis=1)
         segment_features = segment_features.reset_index()
         dates = segment_features["timestamp"]
@@ -251,7 +251,7 @@ def get_model(self) -> Dict[str, Any]:
     @staticmethod
     def _forecast_segment(model: Any, segment: str, ts: TSDataset, *args, **kwargs) -> pd.DataFrame:
         """Make predictions for one segment."""
-        segment_features = ts[:, segment, :]
+        segment_features = ts.df.loc[:, pd.IndexSlice[segment, :]]
         segment_features = segment_features.droplevel("segment", axis=1)
         segment_features = segment_features.reset_index()
         dates = segment_features["timestamp"]

diff --git a/tests/test_models/test_common.py b/tests/test_models/test_common.py
@@ -30,7 +30,7 @@
 @pytest.mark.parametrize(
     "model, transforms, num_skip_timestamps",
     [
-        (CatBoostModelPerSegment(), [LagTransform(in_column="target", lags=[2, 3])], 2),
+        (CatBoostModelPerSegment(), [LagTransform(in_column="target", lags=[2, 3])], 0),
         (CatBoostModelMultiSegment(), [LagTransform(in_column="target", lags=[2, 3])], 0),
         (LinearPerSegmentModel(), [LagTransform(in_column="target", lags=[2, 3])], 3),
         (LinearMultiSegmentModel(), [LagTransform(in_column="target", lags=[2, 3])], 3),