Repository: tinkoff-ai/etna (language: python)

Issue:
Shift of exogenous data
### 🚀 Feature Request

Add class `ExogShiftTransform` for automatic shift of exogenous data based on lag transform.

### Proposal

```
class ExogShiftTransform(IrreversibleTransform, FutureMixin):
    def __init__(self, lags: Union[int, Literal["auto"]], horizon: Optional[int] = None):
        pass

   def _fit(self, df: pd.DataFrame) -> "ExogShiftTransform":
       pass
   
   def _transform(self, df: pd.DataFrame) -> pd.DataFrame:
       pass
```

* `lags: int` -- use passed value to shift all exog columns
* `lags="auto"` -- compute the size of the shift automatically. In this case, estimate for each series the minimal number of steps required to cover all `horizon`. It may be possible that each exog series will have its own shift size. Computation may be performed using the last available date of the exog variable and the prediction horizon.

When `lags="auto"` parameter `horizon` is mandatory.

Name shifted columns in format `{exog_name}_shift_{exog_shift_size}`.

### Test cases

* Test all necessary parameters set when initialising the transform.
* Test that names for transformed data are correct.
* Test all exog data correctly shifted for both modes.
* Test that transform works within `Pipeline` for both modes.
* Add transform to inference tests.


### Additional context

_No response_


Discussion / hints:
It seems like the task #1072 is relevant here.
It seems like the task #946 is relevant here.

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2c3d952ca..d3ae89966 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,7 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## Unreleased
 ### Added
 - Notebook `forecast_interpretation.ipynb` with forecast decomposition ([#1220](https://github.com/tinkoff-ai/etna/pull/1220))
--
+- Exogenous variables shift transform `ExogShiftTransform`([#1254](https://github.com/tinkoff-ai/etna/pull/1254))
 -
 - 
 ### Changed
diff --git a/etna/transforms/__init__.py b/etna/transforms/__init__.py
index 00c67b6ce..4c8332cae 100644
--- a/etna/transforms/__init__.py
+++ b/etna/transforms/__init__.py
@@ -26,6 +26,7 @@
 from etna.transforms.math import AddConstTransform
 from etna.transforms.math import BoxCoxTransform
 from etna.transforms.math import DifferencingTransform
+from etna.transforms.math import ExogShiftTransform
 from etna.transforms.math import LagTransform
 from etna.transforms.math import LambdaTransform
 from etna.transforms.math import LogTransform
diff --git a/etna/transforms/math/__init__.py b/etna/transforms/math/__init__.py
index 71877775d..dc3df9b09 100644
--- a/etna/transforms/math/__init__.py
+++ b/etna/transforms/math/__init__.py
@@ -1,6 +1,7 @@
 from etna.transforms.math.add_constant import AddConstTransform
 from etna.transforms.math.apply_lambda import LambdaTransform
 from etna.transforms.math.differencing import DifferencingTransform
+from etna.transforms.math.lags import ExogShiftTransform
 from etna.transforms.math.lags import LagTransform
 from etna.transforms.math.log import LogTransform
 from etna.transforms.math.power import BoxCoxTransform
diff --git a/etna/transforms/math/lags.py b/etna/transforms/math/lags.py
index aae433624..38d8f4f40 100644
--- a/etna/transforms/math/lags.py
+++ b/etna/transforms/math/lags.py
@@ -1,9 +1,13 @@
+from typing import Dict
 from typing import List
+from typing import Literal
 from typing import Optional
 from typing import Union
 
 import pandas as pd
 
+from etna.datasets import TSDataset
+from etna.models.utils import determine_num_steps
 from etna.transforms.base import FutureMixin
 from etna.transforms.base import IrreversibleTransform
 
@@ -96,3 +100,198 @@ def _transform(self, df: pd.DataFrame) -> pd.DataFrame:
     def get_regressors_info(self) -> List[str]:
         """Return the list with regressors created by the transform."""
         return [self._get_column_name(lag) for lag in self.lags]
+
+
+class ExogShiftTransform(IrreversibleTransform, FutureMixin):
+    """Shifts exogenous variables from a given dataframe."""
+
+    def __init__(self, lag: Union[int, Literal["auto"]], horizon: Optional[int] = None):
+        """Create instance of ExogShiftTransform.
+
+        Parameters
+        ----------
+        lag:
+            value for shift estimation
+
+            * if set to `int` all exogenous variables will be shifted `lag` steps forward;
+
+            * if set to `auto` minimal shift will be estimated for each variable based on
+              the prediction horizon and available timeline
+
+        horizon:
+            prediction horizon. Mandatory when set to `lag="auto"`, ignored otherwise
+        """
+        super().__init__(required_features="all")
+
+        self.lag: Optional[int] = None
+        self.horizon: Optional[int] = None
+        self._auto = False
+
+        self._freq: Optional[str] = None
+        self._created_regressors: Optional[List[str]] = None
+        self._exog_shifts: Optional[Dict[str, int]] = None
+        self._exog_last_date: Optional[Dict[str, pd.Timestamp]] = None
+        self._filter_out_columns = {"target"}
+
+        if isinstance(lag, int):
+            if lag <= 0:
+                raise ValueError(f"{self.__class__.__name__} works only with positive lags values, {lag} given")
+            self.lag = lag
+
+        else:
+            if horizon is None:
+                raise ValueError("Value of `horizon` should be specified when using `auto`!")
+
+            if horizon < 1:
+                raise ValueError(f"{self.__class__.__name__} works only with positive horizon values, {horizon} given")
+
+            self.horizon = horizon
+            self._auto = True
+
+    def _save_exog_last_date(self, df_exog: Optional[pd.DataFrame] = None):
+        """Save last available date of each exogenous variable."""
+        self._exog_last_date = {}
+        if df_exog is not None:
+            exog_names = set(df_exog.columns.get_level_values("feature"))
+
+            for name in exog_names:
+                feature = df_exog.loc[:, pd.IndexSlice[:, name]]
+
+                na_mask = pd.isna(feature).any(axis=1)
+                last_date = feature.index[~na_mask].max()
+
+                self._exog_last_date[name] = last_date
+
+    def fit(self, ts: TSDataset) -> "ExogShiftTransform":
+        """Fit the transform.
+
+        Parameters
+        ----------
+        ts:
+            Dataset to fit the transform on.
+
+        Returns
+        -------
+        :
+            The fitted transform instance.
+        """
+        self._freq = ts.freq
+        self._save_exog_last_date(df_exog=ts.df_exog)
+
+        super().fit(ts=ts)
+
+        return self
+
+    def _fit(self, df: pd.DataFrame) -> "ExogShiftTransform":
+        """Estimate shifts for exogenous variables.
+
+        Parameters
+        ----------
+        df:
+            dataframe with data.
+
+        Returns
+        -------
+        :
+            Fitted `ExogShiftTransform` instance.
+        """
+        feature_names = self._get_feature_names(df=df)
+
+        self._exog_shifts = dict()
+        self._created_regressors = []
+
+        for feature_name in feature_names:
+            shift = self._estimate_shift(df=df, feature_name=feature_name)
+            self._exog_shifts[feature_name] = shift
+
+            if shift > 0:
+                self._created_regressors.append(f"{feature_name}_shift_{shift}")
+
+        return self
+
+    def _get_feature_names(self, df: pd.DataFrame) -> List[str]:
+        """Return the names of exogenous variables."""
+        feature_names = []
+        if self._exog_last_date is not None:
+            feature_names = list(self._exog_last_date.keys())
+
+        df_columns = df.columns.get_level_values("feature")
+        for name in feature_names:
+            if name not in df_columns:
+                raise ValueError(f"Feature `{name}` is expected to be in the dataframe!")
+
+        return feature_names
+
+    def _estimate_shift(self, df: pd.DataFrame, feature_name: str) -> int:
+        """Estimate shift value for exogenous variable."""
+        if not self._auto:
+            return self.lag  # type: ignore
+
+        if self._exog_last_date is None or self._freq is None:
+            raise ValueError("Call `fit()` method before estimating exog shifts!")
+
+        last_date = df.index.max()
+        last_feature_date = self._exog_last_date[feature_name]
+
+        if last_feature_date > last_date:
+            delta = -determine_num_steps(start_timestamp=last_date, end_timestamp=last_feature_date, freq=self._freq)
+
+        elif last_feature_date < last_date:
+            delta = determine_num_steps(start_timestamp=last_feature_date, end_timestamp=last_date, freq=self._freq)
+
+        else:
+            delta = 0
+
+        shift = max(0, delta + self.horizon)  # type: ignore
+
+        return shift
+
+    def _transform(self, df: pd.DataFrame) -> pd.DataFrame:
+        """Shift exogenous variables.
+
+        Parameters
+        ----------
+        df:
+            dataframe with data to transform.
+
+        Returns
+        -------
+        :
+            Transformed dataframe.
+        """
+        if self._exog_shifts is None:
+            raise ValueError("Transform is not fitted!")
+
+        result = df
+        segments = sorted(set(df.columns.get_level_values("segment")))
+        feature_names = self._get_feature_names(df=df)
+
+        shifted_features = []
+        features_to_remove = []
+        for feature_name in feature_names:
+            shift = self._exog_shifts[feature_name]
+
+            feature = df.loc[:, pd.IndexSlice[:, feature_name]]
+
+            if shift > 0:
+                shifted_feature = feature.shift(shift, freq=self._freq)
+
+                column_name = f"{feature_name}_shift_{shift}"
+                shifted_feature.columns = pd.MultiIndex.from_product([segments, [column_name]])
+
+                shifted_features.append(shifted_feature)
+                features_to_remove.append(feature_name)
+
+        if len(features_to_remove) > 0:
+            result = result.drop(columns=pd.MultiIndex.from_product([segments, features_to_remove]))
+
+        result = pd.concat([result] + shifted_features, axis=1)
+        result.sort_index(axis=1, inplace=True)
+        return result
+
+    def get_regressors_info(self) -> List[str]:
+        """Return the list with regressors created by the transform."""
+        if self._created_regressors is None:
+            raise ValueError("Fit the transform to get the regressors info!")
+
+        return self._created_regressors


Test patch:
diff --git a/tests/test_transforms/test_math/test_exog_shift_transform.py b/tests/test_transforms/test_math/test_exog_shift_transform.py
new file mode 100644
index 000000000..3d13319fe
--- /dev/null
+++ b/tests/test_transforms/test_math/test_exog_shift_transform.py
@@ -0,0 +1,211 @@
+import numpy as np
+import pandas as pd
+import pytest
+
+from etna.datasets import TSDataset
+from etna.metrics import MAE
+from etna.models import LinearPerSegmentModel
+from etna.pipeline import Pipeline
+from etna.transforms import ExogShiftTransform
+
+
+def build_df_exog_with_nans(freq="D"):
+    df = pd.DataFrame(
+        {
+            "timestamp": list(pd.date_range("2023-01-01", periods=5, freq=freq)) * 2,
+            "segment": ["A"] * 5 + ["B"] * 5,
+            "feat1": [1, 2, 3, 4, None] + [1, 2, 3, 4, 5],
+            "feat2": [1, 2, 3, None, None] + [1, 2, 3, None, None],
+            "feat3": [1, 2, 3, 4, 5] + [1, 2, 3, 4, 5],
+        }
+    )
+
+    return TSDataset.to_dataset(df=df)
+
+
+def build_ts_with_exogs(freq="D"):
+    df = pd.DataFrame(
+        {
+            "timestamp": list(pd.date_range("2023-01-01", periods=4, freq=freq)) * 2,
+            "segment": ["A"] * 4 + ["B"] * 4,
+            "target": list(3 * np.arange(1, 5)) * 2,
+        }
+    )
+
+    df = TSDataset.to_dataset(df=df)
+    ts = TSDataset(df=df, df_exog=build_df_exog_with_nans(freq=freq), freq=freq)
+    return ts
+
+
+@pytest.fixture()
+def df_exog_with_nans():
+    return build_df_exog_with_nans(freq="D")
+
+
+@pytest.fixture()
+def ts_with_exogs():
+    return build_ts_with_exogs(freq="D")
+
+
+@pytest.fixture()
+def ts_with_exogs_ms_freq():
+    return build_ts_with_exogs(freq="MS")
+
+
+@pytest.mark.parametrize(
+    "expected",
+    (
+        {
+            "feat1": pd.Timestamp("2023-01-04"),
+            "feat2": pd.Timestamp("2023-01-03"),
+            "feat3": pd.Timestamp("2023-01-05 00:00:00"),
+        },
+    ),
+)
+def test_save_exog_last_date(df_exog_with_nans, expected):
+    t = ExogShiftTransform(lag=1)
+    t._save_exog_last_date(df_exog=df_exog_with_nans)
+    assert t._exog_last_date == expected
+
+
+def test_negative_lag():
+    with pytest.raises(ValueError, match=".* works only with positive lags"):
+        ExogShiftTransform(lag=-1)
+
+
+def test_horizon_not_set():
+    with pytest.raises(ValueError, match="`horizon` should be specified"):
+        ExogShiftTransform(lag="auto")
+
+
+def test_negative_horizon_set():
+    with pytest.raises(ValueError, match=".* works only with positive horizon"):
+        ExogShiftTransform(lag="auto", horizon=-1)
+
+
+def test_regressors_info_not_fit():
+    with pytest.raises(ValueError, match="Fit the transform"):
+        ExogShiftTransform(lag=1).get_regressors_info()
+
+
+def test_estimate_shift_not_fit():
+    with pytest.raises(ValueError, match=".* method before estimating exog shifts!"):
+        ExogShiftTransform(lag="auto", horizon=1)._estimate_shift(None, None)
+
+
+def test_transform_not_fit(ts_with_exogs):
+    t = ExogShiftTransform(lag=1)
+    with pytest.raises(ValueError, match="Transform is not fitted!"):
+        t.transform(ts=ts_with_exogs)
+
+
+def test_get_feature_names(example_reg_tsds, expected={"regressor_exog_weekend"}):
+    t = ExogShiftTransform(lag=1)
+    t.fit(ts=example_reg_tsds)
+    feature_names = t._get_feature_names(example_reg_tsds.df)
+    assert set(feature_names) == expected
+
+
+def test_get_feature_names_no_exog_error(ts_with_exogs):
+    t = ExogShiftTransform(lag=1)
+    t.fit(ts=ts_with_exogs)
+    df = ts_with_exogs.df.drop(columns=["feat3"], level=1)
+
+    with pytest.raises(ValueError, match="Feature `feat3` is expected to be in the dataframe!"):
+        t._get_feature_names(df)
+
+
+@pytest.mark.parametrize(
+    "ts_name", ("toy_dataset_with_mean_shift_in_target", "product_level_constant_forecast_with_target_components")
+)
+def test_ts_with_quantiles_and_components(ts_name, request):
+    ts = request.getfixturevalue(ts_name)
+    t = ExogShiftTransform(lag=1)
+    t.fit(ts=ts)
+    assert t.get_regressors_info() == []
+
+
+@pytest.mark.parametrize(
+    "horizon,expected",
+    ((1, {"feat1_shift_1", "feat2_shift_2"}), (2, {"feat1_shift_2", "feat2_shift_3", "feat3_shift_1"})),
+)
+def test_regressors_info(ts_with_exogs, horizon, expected):
+    t = ExogShiftTransform(lag="auto", horizon=horizon)
+    t.fit(ts=ts_with_exogs)
+    assert set(t.get_regressors_info()) == expected
+
+
+@pytest.mark.parametrize(
+    "lag,horizon,expected",
+    (
+        (1, None, {"feat1": 1, "feat2": 1, "feat3": 1}),
+        ("auto", 1, {"feat1": 1, "feat2": 2, "feat3": 0}),
+        ("auto", 2, {"feat1": 2, "feat2": 3, "feat3": 1}),
+    ),
+)
+def test_estimate_shift(ts_with_exogs, lag, horizon, expected):
+    t = ExogShiftTransform(lag=lag, horizon=horizon)
+    t.fit(ts=ts_with_exogs)
+    assert t._exog_shifts == expected
+
+
+@pytest.mark.parametrize("lag", (1, "auto"))
+def test_shift_no_exog(simple_df, lag, expected={"target"}):
+    t = ExogShiftTransform(lag=lag, horizon=1)
+    transformed = t.fit_transform(simple_df)
+    assert set(transformed.df.columns.get_level_values("feature")) == expected
+
+
+@pytest.mark.parametrize(
+    "lag,horizon,expected",
+    (
+        (1, None, {"feat1_shift_1", "feat2_shift_1", "feat3_shift_1", "target"}),
+        ("auto", 1, {"feat1_shift_1", "feat2_shift_2", "feat3", "target"}),
+        ("auto", 2, {"feat1_shift_2", "feat2_shift_3", "feat3_shift_1", "target"}),
+    ),
+)
+@pytest.mark.parametrize(
+    "ts_name",
+    (
+        "ts_with_exogs",
+        "ts_with_exogs_ms_freq",
+    ),
+)
+def test_transformed_names(ts_name, lag, horizon, expected, request):
+    ts = request.getfixturevalue(ts_name)
+
+    t = ExogShiftTransform(lag=lag, horizon=horizon)
+    transformed = t.fit_transform(ts=ts)
+    column_names = transformed.df.columns.get_level_values("feature")
+    assert set(column_names) == expected
+
+
+@pytest.mark.parametrize("lag", (3, "auto"))
+@pytest.mark.parametrize("horizon", range(1, 3))
+@pytest.mark.parametrize(
+    "ts_name",
+    (
+        "ts_with_exogs",
+        "ts_with_exogs_ms_freq",
+    ),
+)
+def test_pipeline_forecast(ts_name, lag, horizon, request):
+    ts = request.getfixturevalue(ts_name)
+
+    pipeline = Pipeline(
+        transforms=[ExogShiftTransform(lag=lag, horizon=horizon)], model=LinearPerSegmentModel(), horizon=horizon
+    )
+    pipeline.fit(ts)
+    pipeline.forecast()
+
+
+@pytest.mark.parametrize("lag", (3, "auto"))
+@pytest.mark.parametrize("horizon", range(7, 10))
+def test_pipeline_backtest(example_reg_tsds, lag, horizon):
+    ts = example_reg_tsds
+
+    pipeline = Pipeline(
+        transforms=[ExogShiftTransform(lag=lag, horizon=horizon)], model=LinearPerSegmentModel(), horizon=horizon
+    )
+
+    pipeline.backtest(ts=ts, metrics=[MAE()])