New feature: Lag or windows features grouped by

feature_engine/_docstrings/init_parameters/all_trasnformers.py

@@ -22,3 +22,10 @@
         contain missing values. If `'ignore'`, missing data will be ignored when
         learning parameters or performing the transformation.
         """.rstrip()
+
+_group_by_docstring = """group_by: str, int, or list of strings or integers,default=None
+        A group_by operation involves some combination of splitting the object,
+        applying a function, and combining the results.
+        This can be used to group large amounts of data and
+        compute operations on these groups.
+        """.rstrip()

feature_engine/selection/drop_psi_features.py

@@ -1,5 +1,5 @@
 import datetime
-from typing import List, Union
+from typing import Dict, List, Union
 
 import numpy as np
 import pandas as pd
@@ -475,7 +475,7 @@ def fit(self, X: pd.DataFrame, y: pd.Series = None):
                 threshold_cat = self.threshold
 
         # Compute the PSI by looping over the features
-        self.psi_values_ = {}
+        self.psi_values_: Dict = {}
         self.features_to_drop_ = []
 
         # Compute PSI for numerical features

feature_engine/timeseries/forecasting/base_forecast_transformers.py

@@ -14,6 +14,7 @@
 )
 from feature_engine._docstrings.init_parameters.all_trasnformers import (
     _drop_original_docstring,
+    _group_by_docstring,
     _missing_values_docstring,
 )
 from feature_engine._docstrings.methods import _fit_not_learn_docstring
@@ -37,6 +38,7 @@
     feature_names_in_=_feature_names_in_docstring,
     fit=_fit_not_learn_docstring,
     n_features_in_=_n_features_in_docstring,
+    group_by=_group_by_docstring,
 )
 class BaseForecastTransformer(BaseEstimator, TransformerMixin, GetFeatureNamesOutMixin):
     """
@@ -51,6 +53,8 @@ class BaseForecastTransformer(BaseEstimator, TransformerMixin, GetFeatureNamesOu
 
     {drop_original}
 
+    {group_by}
+
     Attributes
     ----------
     {feature_names_in_}
@@ -64,6 +68,7 @@ def __init__(
         variables: Union[None, int, str, List[Union[str, int]]] = None,
         missing_values: str = "raise",
         drop_original: bool = False,
+        group_by: Union[None, int, str, List[Union[str, int]]] = None,
     ) -> None:
 
         if missing_values not in ["raise", "ignore"]:
@@ -81,6 +86,7 @@ def __init__(
         self.variables = _check_variables_input_value(variables)
         self.missing_values = missing_values
         self.drop_original = drop_original
+        self.group_by = group_by
 
     def _check_index(self, X: pd.DataFrame):
         """

feature_engine/timeseries/forecasting/expanding_window_features.py

@@ -3,7 +3,7 @@
 
 from __future__ import annotations
 
-from typing import List
+from typing import List, Union
 
 import pandas as pd
 
@@ -13,6 +13,7 @@
 )
 from feature_engine._docstrings.init_parameters.all_trasnformers import (
     _drop_original_docstring,
+    _group_by_docstring,
     _missing_values_docstring,
     _variables_numerical_docstring,
 )
@@ -34,6 +35,7 @@
     n_features_in_=_n_features_in_docstring,
     fit=_fit_not_learn_docstring,
     fit_transform=_fit_transform_docstring,
+    group_by=_group_by_docstring,
 )
 class ExpandingWindowFeatures(BaseForecastTransformer):
     """
@@ -93,6 +95,8 @@ class ExpandingWindowFeatures(BaseForecastTransformer):
 
     {drop_original}
 
+    {group_by}
+
     Attributes
     ----------
     variables_:
@@ -151,6 +155,7 @@ def __init__(
         sort_index: bool = True,
         missing_values: str = "raise",
         drop_original: bool = False,
+        group_by: Union[None, int, str, List[Union[str, int]]] = None,
     ) -> None:
 
         if not isinstance(functions, (str, list)) or not all(
@@ -168,7 +173,7 @@ def __init__(
                 f"periods must be a non-negative integer. Got {periods} instead."
             )
 
-        super().__init__(variables, missing_values, drop_original)
+        super().__init__(variables, missing_values, drop_original, group_by)
 
         self.min_periods = min_periods
         self.functions = functions
@@ -193,12 +198,21 @@ def transform(self, X: pd.DataFrame) -> pd.DataFrame:
         # Common dataframe checks and setting up.
         X = self._check_transform_input_and_state(X)
 
-        tmp = (
-            X[self.variables_]
-            .expanding(min_periods=self.min_periods)
-            .agg(self.functions)
-            .shift(periods=self.periods, freq=self.freq)
-        )
+        if self.group_by:
+            original_index = X.index
+            tmp = X.groupby(self.group_by, as_index=False).apply(
+                self._agg_expanding_window_features,
+                include_groups=False,
+            )
+            tmp = tmp.set_index(original_index)
+            tmp = tmp.reindex(original_index)
+        else:
+            tmp = (
+                X[self.variables_]
+                .expanding(min_periods=self.min_periods)
+                .agg(self.functions)
+                .shift(periods=self.periods, freq=self.freq)
+            )
 
         tmp.columns = self._get_new_features_name()
 
@@ -224,3 +238,25 @@ def _get_new_features_name(self) -> List:
         ]
 
         return feature_names
+
+    def _agg_expanding_window_features(
+        self,
+        grouped_df: pd.core.groupby.generic.DataFrameGroupBy,
+    ) -> Union[pd.Series, pd.DataFrame]:
+        """generate expanding window features based on groups
+        Parameters
+        ----------
+        grouped_df : pd.core.groupby.generic.DataFrameGroupBy
+            dataframe of groups
+
+        Returns
+        -------
+        Union[pd.Series, pd.DataFrame]
+            returned expanding window features
+        """
+        return (
+            grouped_df[self.variables_]
+            .expanding(min_periods=self.min_periods)
+            .agg(self.functions)
+            .shift(periods=self.periods, freq=self.freq)
+        )

feature_engine/timeseries/forecasting/lag_features.py

@@ -11,6 +11,7 @@
 )
 from feature_engine._docstrings.init_parameters.all_trasnformers import (
     _drop_original_docstring,
+    _group_by_docstring,
     _missing_values_docstring,
     _variables_numerical_docstring,
 )
@@ -32,6 +33,7 @@
     n_features_in_=_n_features_in_docstring,
     fit=_fit_not_learn_docstring,
     fit_transform=_fit_transform_docstring,
+    group_by=_group_by_docstring,
 )
 class LagFeatures(BaseForecastTransformer):
     """
@@ -74,6 +76,8 @@ class LagFeatures(BaseForecastTransformer):
 
     {drop_original}
 
+    {group_by}
+
     Attributes
     ----------
     variables_:
@@ -127,6 +131,7 @@ def __init__(
         sort_index: bool = True,
         missing_values: str = "raise",
         drop_original: bool = False,
+        group_by: Union[None, int, str, List[Union[str, int]]] = None,
     ) -> None:
 
         if not (
@@ -151,7 +156,7 @@ def __init__(
                 "sort_index takes values True and False." f"Got {sort_index} instead."
             )
 
-        super().__init__(variables, missing_values, drop_original)
+        super().__init__(variables, missing_values, drop_original, group_by)
 
         self.periods = periods
         self.freq = freq
@@ -180,35 +185,57 @@ def transform(self, X: pd.DataFrame) -> pd.DataFrame:
             if isinstance(self.freq, list):
                 df_ls = []
                 for fr in self.freq:
-                    tmp = X[self.variables_].shift(
-                        freq=fr,
-                        axis=0,
-                    )
+                    if self.group_by:
+                        tmp = self._agg_freq_lags(
+                            grouped_df=X.groupby(self.group_by),
+                            freq=fr,
+                        )
+                    else:
+                        tmp = X[self.variables_].shift(
+                            freq=fr,
+                            axis=0,
+                        )
                     df_ls.append(tmp)
                 tmp = pd.concat(df_ls, axis=1)
 
             else:
-                tmp = X[self.variables_].shift(
-                    freq=self.freq,
-                    axis=0,
-                )
+                if self.group_by:
+                    tmp = self._agg_freq_lags(
+                        grouped_df=X.groupby(self.group_by),
+                        freq=self.freq,
+                    )
+                else:
+                    tmp = X[self.variables_].shift(
+                        freq=self.freq,
+                        axis=0,
+                    )
 
         else:
             if isinstance(self.periods, list):
                 df_ls = []
                 for pr in self.periods:
-                    tmp = X[self.variables_].shift(
-                        periods=pr,
-                        axis=0,
-                    )
+                    if self.group_by:
+                        tmp = X.groupby(self.group_by)[self.variables_].shift(
+                            periods=pr,
+                        )
+                    else:
+                        tmp = X[self.variables_].shift(
+                            periods=pr,
+                            axis=0,
+                        )
                     df_ls.append(tmp)
                 tmp = pd.concat(df_ls, axis=1)
 
             else:
-                tmp = X[self.variables_].shift(
-                    periods=self.periods,
-                    axis=0,
-                )
+                if self.group_by:
+                    tmp = X.groupby(self.group_by)[self.variables_].shift(
+                        periods=self.periods,
+                    )
+                else:
+                    tmp = X[self.variables_].shift(
+                        periods=self.periods,
+                        axis=0,
+                    )
 
         tmp.columns = self._get_new_features_name()
 
@@ -243,3 +270,30 @@ def _get_new_features_name(self) -> List:
             ]
 
         return feature_names
+
+    def _agg_freq_lags(
+        self,
+        grouped_df: pd.core.groupby.generic.DataFrameGroupBy,
+        freq: Union[str, List[str]],
+    ) -> Union[pd.Series, pd.DataFrame]:
+        """_summary_
+
+        Parameters
+        ----------
+        grouped_df : pd.core.groupby.generic.DataFrameGroupBy
+            dataframe of groups
+        freq : Union[str, List[str]]
+            Offset to use from the tseries module or time rule.
+
+        Returns
+        -------
+        Union[pd.Series, pd.DataFrame]
+            lag feature or dataframe of lag features
+        """
+        tmp_data = []
+        for _, group in grouped_df:
+            original_idx = group.index
+            tmp = group[self.variables_].shift(freq=freq).reindex(original_idx)
+            tmp_data.append(tmp)
+        tmp = pd.concat(tmp_data).sort_index()
+        return tmp