From f5dac1d8473978d5cee331f3b88928127199a3d0 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Sat, 19 Sep 2026 11:22:09 +0200 Subject: [PATCH 1/3] Migrate BaseForecastTransformer to narwhals, add polars support The base of the forecasting transformers takes pandas and polars dataframes. With pandas the index keeps ordering the rows in time; with polars the rows are used in the order given, sort_index does not apply and freq raises an error. The base now provides transform() and transform_x_y() on top of an _add_features() hook for the subclasses. Also sets AGENTS.md's priorities to speed, readability and simplicity. Co-Authored-By: Claude Opus 5 --- AGENTS.md | 6 +- .../timeseries/forecasting/index.rst | 18 + .../forecasting/base_forecast_transformers.py | 229 ++++++++--- .../forecasting/expanding_window_features.py | 2 +- .../timeseries/forecasting/lag_features.py | 2 +- .../timeseries/forecasting/window_features.py | 2 +- .../test_base_forecast_transformer.py | 367 ++++++++++++++++++ 7 files changed, 564 insertions(+), 62 deletions(-) create mode 100644 tests/test_time_series/test_forecasting/test_base_forecast_transformer.py diff --git a/AGENTS.md b/AGENTS.md index 96c9a3ebc..3e8bc1309 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,8 +1,8 @@ # AGENTS.md -Conventions for working in this repo. Optimize for readability and speed, -in that order of how you decide, but don't ship a slow default when a -fast one is free. +Conventions for working in this repo. Optimize for speed, readability and +simplicity, in that order of how you decide: never ship a slow default when a +fast one is available. ## Inputs diff --git a/docs/user_guide/timeseries/forecasting/index.rst b/docs/user_guide/timeseries/forecasting/index.rst index c05d47fee..431f882db 100644 --- a/docs/user_guide/timeseries/forecasting/index.rst +++ b/docs/user_guide/timeseries/forecasting/index.rst @@ -96,6 +96,24 @@ methods like mean, sum, standard deviation, min, max, among others. Unlike rolli don't need to specify a window size as the window is expanded automatically at each time step. +Time order of the rows +~~~~~~~~~~~~~~~~~~~~~~ + +Lag and window features take values from previous rows, so the rows must be in time +order. How the transformers find that order depends on the dataframe you pass. + +With **pandas** dataframes, the index gives the time order. The index must have unique +values and no missing data. With `sort_index=True`, the default, the transformers sort +the rows by the index before creating the features. The parameter `freq`, and windows +given as a time span, like `"3D"`, need a DatetimeIndex. + +**polars** dataframes have no index, so the transformers use the rows in the order in +which they are in the dataframe. Sort the dataframe by time before passing it to the +transformers, for example with `df.sort("date")`. The parameter `sort_index` has no +effect with polars. The parameter `freq`, and windows given as a time span, are only +supported with pandas dataframes; with polars, use a number of rows instead. + + Datetime Features ----------------- diff --git a/feature_engine/timeseries/forecasting/base_forecast_transformers.py b/feature_engine/timeseries/forecasting/base_forecast_transformers.py index f7ec468bf..502c66fd0 100644 --- a/feature_engine/timeseries/forecasting/base_forecast_transformers.py +++ b/feature_engine/timeseries/forecasting/base_forecast_transformers.py @@ -1,13 +1,12 @@ from typing import List, Optional, Union -import pandas as pd +import narwhals as nw +import narwhals.dependencies as nwd +from narwhals.typing import IntoDataFrame, IntoSeries from sklearn.base import BaseEstimator, TransformerMixin from sklearn.utils.validation import check_is_fitted -from feature_engine._base_transformers.mixins import ( - GetFeatureNamesOutMixin, - TransformXyMixin, -) +from feature_engine._base_transformers.mixins import GetFeatureNamesOutMixin from feature_engine._check_init_parameters.check_init_input_params import ( _check_return_empty_is_bool, ) @@ -29,6 +28,7 @@ _check_contains_na, _check_X_matches_training_df, check_X, + check_X_y, ) from feature_engine.tags import _return_tags from feature_engine.variable_handling import ( @@ -44,12 +44,16 @@ fit=_fit_not_learn_docstring, n_features_in_=_n_features_in_docstring, ) -class BaseForecastTransformer( - TransformerMixin, BaseEstimator, GetFeatureNamesOutMixin, TransformXyMixin -): +class BaseForecastTransformer(TransformerMixin, BaseEstimator, GetFeatureNamesOutMixin): """ Shared methods across time-series forecasting transformers. + With pandas, the rows are ordered in time by the dataframe's index. Other + dataframes, like polars, have no index: their rows are taken in the order given. + + Subclasses define the parameters `freq` and `sort_index`, and the methods + `_add_features()` and `_get_new_features_name()`. + Parameters ---------- variables: str, int, or list of strings or integers, default=None. @@ -79,7 +83,10 @@ def __init__( drop_na: bool = False, ) -> None: - if missing_values not in ["raise", "ignore"]: + if not isinstance(missing_values, str) or missing_values not in [ + "raise", + "ignore", + ]: raise ValueError( "missing_values takes only values 'raise' or 'ignore'. " f"Got {missing_values} instead." @@ -105,38 +112,48 @@ def __init__( self.drop_original = drop_original self.drop_na = drop_na - def _check_index(self, X: pd.DataFrame): + def _check_index(self, X: IntoDataFrame): """ - Check that the index does not have missing data and its values are unique. + Checks that the rows of the dataframe can be ordered in time. With pandas, + the index must be unique and not contain missing data. Other dataframes have + no index, so they can't use `freq`. Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The dataset. """ - if X.index.isnull().any(): - raise NotImplementedError( - "The dataframe's index contains NaN values or missing data. " - "Only dataframes with complete indexes are compatible with " - "this transformer." - ) - - if X.index.is_unique is False: + if nwd.is_pandas_dataframe(X) is True: + if X.index.isnull().any(): + raise NotImplementedError( + "The dataframe's index contains NaN values or missing data. " + "Only dataframes with complete indexes are compatible with " + "this transformer." + ) + + if X.index.is_unique is False: + raise NotImplementedError( + "The dataframe's index does not contain unique values. " + "Only dataframes with unique values in the index are " + "compatible with this transformer." + ) + + elif self.freq is not None: raise NotImplementedError( - "The dataframe's index does not contain unique values. " - "Only dataframes with unique values in the index are " - "compatible with this transformer." + "freq is only supported with pandas dataframes, because it uses the " + "dataframe's DatetimeIndex. With other dataframes, leave freq=None " + f"to shift the rows by periods. Got {self.freq} instead." ) return self - def _check_na_and_inf(self, X: pd.DataFrame): + def _check_na_and_inf(self, X: IntoDataFrame): """ Checks that the dataframe does not contain NaN or Infinite values. Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The dataset for training or transformation. """ _check_contains_na(X, self.variables_) @@ -144,41 +161,41 @@ def _check_na_and_inf(self, X: pd.DataFrame): return self - def _get_feature_names_in(self, X: pd.DataFrame): + def _get_feature_names_in(self, X: IntoDataFrame): """ Finds the number and name of the features in the training set. Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The dataset for training or transformation. """ - - self.feature_names_in_ = X.columns.tolist() + if nwd.is_pandas_dataframe(X) is True: + self.feature_names_in_ = list(X.columns) + else: + self.feature_names_in_ = nw.from_native(X, eager_only=True).columns self.n_features_in_ = X.shape[1] return self - def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): + def fit(self, X: IntoDataFrame, y: Optional[IntoSeries] = None): """ This transformer does not learn parameters. Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The training dataset. - y: pandas Series, default=None + y: Series, default=None y is not needed in this transformer. You can pass None or y. """ - # check input dataframe - X = check_X(X) + check_X(X) - # We need the dataframes to have unique values in the index and no missing data. - # Otherwise, when we merge the new features we will duplicate rows. + # With pandas, the new features are aligned to the rows through the index, + # so an index with duplicates or missing data would duplicate rows. self._check_index(X) - # find or check for numerical variables if self.variables is None: self.variables_ = find_numerical_variables( X, return_empty=self.return_empty @@ -186,7 +203,6 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): else: self.variables_ = check_numerical_variables(X, self.variables) - # check if dataset contains na if self.missing_values == "raise": self._check_na_and_inf(X) @@ -194,45 +210,146 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): return self - def _check_transform_input_and_state(self, X: pd.DataFrame) -> pd.DataFrame: + def _check_transform_input_and_state(self, X: IntoDataFrame) -> nw.DataFrame: """ Common checks performed before the feature transformation. Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The data to transform. Returns ------- - X: pandas dataframe of shape = [n_samples, n_features] - The data to transform. + nw_X: narwhals dataframe of shape = [n_samples, n_features] + The data to transform, with the variables in the same order as in the + train set. With pandas, the rows are sorted by the index if + `sort_index=True`. """ - # check method fit has been called check_is_fitted(self) - # check if 'X' is a dataframe - X = check_X(X) + nw_X = check_X(X) - # check if input data contains the same number of columns as the fitted - # dataframe. - _check_X_matches_training_df(X, self.n_features_in_) + _check_X_matches_training_df(nw_X, self.n_features_in_) - # Dataframes must have unique values in the index and no missing data. - # Otherwise, when we merge the created features we will duplicate rows. self._check_index(X) - # check if dataset contains na if self.missing_values == "raise": self._check_na_and_inf(X) - # reorder variables to match train set - X = X[self.feature_names_in_] + if nwd.is_pandas_dataframe(X) is True: + # pandas is faster than narwhals, and only pandas has an index to sort. + X = X[self.feature_names_in_] + if self.sort_index is True: + X.sort_index(inplace=True) + return nw.from_native(X, eager_only=True) + + # without an index, the rows are used in the order given: sort_index is ignored. + return nw_X.select(self.feature_names_in_) + + def _add_features(self, nw_X: nw.DataFrame) -> nw.DataFrame: + """ + Adds the new features, named as `_get_new_features_name()` returns, after + the columns of `nw_X`. Rows must keep their order. + + Parameters + ---------- + nw_X: narwhals dataframe of shape = [n_samples, n_features] + The data returned by `_check_transform_input_and_state()`. + """ + raise NotImplementedError + + def _transform(self, X: IntoDataFrame) -> nw.DataFrame: + """ + Checks the input, adds the new features and drops the original variables if + requested. Rows with missing data are not dropped. + + Parameters + ---------- + X: dataframe of shape = [n_samples, n_features] + The data to transform. + """ + nw_X = self._add_features(self._check_transform_input_and_state(X)) + + if self.drop_original is True: + nw_X = nw_X.drop(self.variables_) - if self.sort_index is True: - X.sort_index(inplace=True) + return nw_X + + def transform(self, X: IntoDataFrame) -> IntoDataFrame: + """ + Adds the new features. + + Parameters + ---------- + X: dataframe of shape = [n_samples, n_features] + The data to transform. + + Returns + ------- + X_new: dataframe of shape = [n_samples, n_features + new_features] + The dataframe with the original plus the new variables. If + `drop_na=True`, rows with missing data in the new features are removed. + """ + nw_X = self._transform(X) + + if self.drop_na is True: + new_features = self._get_new_features_name() + X = nw_X.to_native() + if nwd.is_pandas_dataframe(X) is True: + # a numpy mask is faster than pandas' dropna() and narwhals. + return X[~X[new_features].isna().to_numpy().any(axis=1)] + nw_X = nw_X.drop_nulls(subset=new_features) + + return nw_X.to_native() + + def transform_x_y(self, X: IntoDataFrame, y: IntoSeries): + """ + Adds the new features to X and, if `drop_na=True`, removes the rows with + missing data in the new features from both X and y. + + Parameters + ---------- + X: dataframe of shape = [n_samples, n_features] + The dataframe to transform. + + y: Series or Dataframe of length = n_samples + The target variable to transform. Can be multi-output. + + Returns + ------- + X_new: dataframe + The transformed dataframe. It may contain less rows than the original + dataset. + + y_new: Series or DataFrame + The target variable, with as many rows as those left in X_new. + """ + _, y = check_X_y(X, y) + + if nwd.is_pandas_dataframe(X) is True: + # the rows may be sorted by the index, so y follows the index of X. + X = self.transform(X) + return X, y.loc[X.index] + + nw_X = self._transform(X) + new_features = self._get_new_features_name() + + if self.drop_na is False or len(new_features) == 0: + return nw_X.to_native(), y + + # without an index, a mask of the rows to keep subsets both X and y. + keep = nw_X.select( + ( + ~nw.any_horizontal(nw.col(new_features).is_null(), ignore_nulls=True) + ).alias("__keep__") + ).get_column("__keep__") + if nwd.is_into_series(y): + y = nw.from_native(y, series_only=True).filter(keep).to_native() + else: + y = nw.from_native(y, eager_only=True).filter(keep).to_native() - return X + return nw_X.filter(keep).to_native(), y def _more_tags(self): tags_dict = _return_tags() diff --git a/feature_engine/timeseries/forecasting/expanding_window_features.py b/feature_engine/timeseries/forecasting/expanding_window_features.py index c121734e1..2b56bf943 100644 --- a/feature_engine/timeseries/forecasting/expanding_window_features.py +++ b/feature_engine/timeseries/forecasting/expanding_window_features.py @@ -206,7 +206,7 @@ def transform(self, X: pd.DataFrame) -> pd.DataFrame: The dataframe with the original plus the new variables. """ # Common dataframe checks and setting up. - X = self._check_transform_input_and_state(X) + X = self._check_transform_input_and_state(X).to_native() tmp = ( X[self.variables_] diff --git a/feature_engine/timeseries/forecasting/lag_features.py b/feature_engine/timeseries/forecasting/lag_features.py index 578802d43..d73426718 100644 --- a/feature_engine/timeseries/forecasting/lag_features.py +++ b/feature_engine/timeseries/forecasting/lag_features.py @@ -195,7 +195,7 @@ def transform(self, X: pd.DataFrame) -> pd.DataFrame: The dataframe with the original plus the new variables. """ # Common dataframe checks and setting up. - X = self._check_transform_input_and_state(X) + X = self._check_transform_input_and_state(X).to_native() # if freq is not None, it overrides periods. if self.freq is not None: diff --git a/feature_engine/timeseries/forecasting/window_features.py b/feature_engine/timeseries/forecasting/window_features.py index f5f73c7f7..7a86aaadf 100644 --- a/feature_engine/timeseries/forecasting/window_features.py +++ b/feature_engine/timeseries/forecasting/window_features.py @@ -214,7 +214,7 @@ def transform(self, X: pd.DataFrame) -> pd.DataFrame: The dataframe with the original plus the new variables. """ # Common dataframe checks and setting up. - X = self._check_transform_input_and_state(X) + X = self._check_transform_input_and_state(X).to_native() if isinstance(self.window, list): df_ls = [] diff --git a/tests/test_time_series/test_forecasting/test_base_forecast_transformer.py b/tests/test_time_series/test_forecasting/test_base_forecast_transformer.py new file mode 100644 index 000000000..b21899e3a --- /dev/null +++ b/tests/test_time_series/test_forecasting/test_base_forecast_transformer.py @@ -0,0 +1,367 @@ +import re + +import numpy as np +import pandas as pd +import polars as pl +import pytest +from sklearn.exceptions import NotFittedError + +from feature_engine.timeseries.forecasting.base_forecast_transformers import ( + BaseForecastTransformer, +) +from tests.backend_helpers import frame_to_dict, make_series + +MSG_NA = ( + "Some of the variables in the dataset contain NaN. Check and " + "remove those before using this transformer." +) +MSG_INF = ( + "Some of the variables to transform contain inf values. Check and " + "remove those before using this transformer." +) +MSG_INDEX_NA = ( + "The dataframe's index contains NaN values or missing data. Only dataframes " + "with complete indexes are compatible with this transformer." +) +MSG_INDEX_DUPLICATED = ( + "The dataframe's index does not contain unique values. Only dataframes with " + "unique values in the index are compatible with this transformer." +) + +DATA = { + "a": [1.0, 2.0, 3.0, 4.0], + "b": [10, 20, 30, 40], + "c": ["w", "x", "y", "z"], +} + + +class MockLag(BaseForecastTransformer): + # adds the value of the previous row, to test the shared logic + def __init__( + self, + variables=None, + return_empty=False, + missing_values="raise", + drop_original=False, + drop_na=False, + freq=None, + sort_index=True, + ): + super().__init__( + variables, return_empty, missing_values, drop_original, drop_na + ) + self.freq = freq + self.sort_index = sort_index + + def _add_features(self, nw_X): + # get_column also works with the integer column names of pandas + return nw_X.with_columns( + nw_X.get_column(var).shift(1).alias(name) + for var, name in zip(self.variables_, self._get_new_features_name()) + ) + + def _get_new_features_name(self): + return [f"{var}_lag_1" for var in self.variables_] + + +# init parameters +@pytest.mark.parametrize("missing_values", ["other", "Raise", 1, True, None, ["raise"]]) +def test_error_if_missing_values_not_permitted(missing_values): + msg = ( + "missing_values takes only values 'raise' or 'ignore'. " + f"Got {missing_values} instead." + ) + with pytest.raises(ValueError, match=re.escape(msg)): + BaseForecastTransformer(missing_values=missing_values) + + +@pytest.mark.parametrize("drop_original", ["True", 1, 0, None, [True]]) +def test_error_if_drop_original_not_bool(drop_original): + msg = ( + "drop_original takes only boolean values True and False. " + f"Got {drop_original} instead." + ) + with pytest.raises(ValueError, match=re.escape(msg)): + BaseForecastTransformer(drop_original=drop_original) + + +@pytest.mark.parametrize("drop_na", ["True", 1, 0, None, [True]]) +def test_error_if_drop_na_not_bool(drop_na): + msg = f"drop_na takes only boolean values True and False. Got {drop_na} instead." + with pytest.raises(ValueError, match=re.escape(msg)): + BaseForecastTransformer(drop_na=drop_na) + + +@pytest.mark.parametrize("return_empty", ["True", 1, 0, None, [True]]) +def test_error_if_return_empty_not_bool(return_empty): + msg = ( + "return_empty takes only boolean values True and False. " + f"Got {return_empty} instead." + ) + with pytest.raises(ValueError, match=re.escape(msg)): + BaseForecastTransformer(return_empty=return_empty) + + +@pytest.mark.parametrize( + "missing_values, drop_original, drop_na", + [("raise", False, False), ("ignore", True, False), ("raise", False, True)], +) +def test_init_param_assignment(missing_values, drop_original, drop_na): + transformer = BaseForecastTransformer( + missing_values=missing_values, + drop_original=drop_original, + drop_na=drop_na, + ) + assert transformer.missing_values == missing_values + assert transformer.drop_original == drop_original + assert transformer.drop_na == drop_na + + +# fit and transform +def test_fit_finds_numerical_variables(make_df): + transformer = MockLag().fit(make_df(DATA)) + + assert transformer.variables_ == ["a", "b"] + assert transformer.feature_names_in_ == ["a", "b", "c"] + assert transformer.n_features_in_ == 3 + + +def test_fit_checks_variables_entered_by_user(make_df): + transformer = MockLag(variables="b").fit(make_df(DATA)) + assert transformer.variables_ == ["b"] + + +@pytest.mark.parametrize("y", [[1, 2, 3, 4], np.array([1, 2, 3, 4])]) +def test_fit_accepts_y_as_list_and_array(make_df, y): + transformer = MockLag().fit(make_df(DATA), y) + assert transformer.variables_ == ["a", "b"] + + +def test_fit_and_transform_raise_error_if_na(make_df): + X_na = make_df({**DATA, "a": [1.0, None, 3.0, 4.0]}) + with pytest.raises(ValueError, match=re.escape(MSG_NA)): + MockLag().fit(X_na) + + transformer = MockLag().fit(make_df(DATA)) + with pytest.raises(ValueError, match=re.escape(MSG_NA)): + transformer.transform(X_na) + + +def test_fit_and_transform_raise_error_if_inf(make_df): + X_inf = make_df({**DATA, "a": [1.0, np.inf, 3.0, 4.0]}) + with pytest.raises(ValueError, match=re.escape(MSG_INF)): + MockLag().fit(X_inf) + + transformer = MockLag().fit(make_df(DATA)) + with pytest.raises(ValueError, match=re.escape(MSG_INF)): + transformer.transform(X_inf) + + +def test_transform_adds_features(make_df): + X = make_df(DATA) + Xt = MockLag().fit(X).transform(X) + + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == { + **DATA, + "a_lag_1": [None, 1.0, 2.0, 3.0], + "b_lag_1": [None, 10, 20, 30], + } + + +def test_transform_keeps_na_when_ignored(make_df): + X = make_df({**DATA, "a": [1.0, None, 3.0, 4.0]}) + Xt = MockLag(variables="a", missing_values="ignore").fit(X).transform(X) + + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt)["a_lag_1"] == [None, 1.0, None, 3.0] + + +def test_transform_reorders_columns_to_match_fit(make_df): + transformer = MockLag().fit(make_df(DATA)) + Xt = transformer.transform(make_df({k: DATA[k] for k in ["c", "b", "a"]})) + + assert isinstance(Xt, make_df) + assert list(Xt.columns) == ["a", "b", "c", "a_lag_1", "b_lag_1"] + + +def test_drop_original(make_df): + X = make_df(DATA) + Xt = MockLag(drop_original=True).fit(X).transform(X) + + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == { + "c": ["w", "x", "y", "z"], + "a_lag_1": [None, 1.0, 2.0, 3.0], + "b_lag_1": [None, 10, 20, 30], + } + + +def test_drop_na(make_df): + X = make_df({**DATA, "a": [1.0, 2.0, None, 4.0]}) + transformer = MockLag(missing_values="ignore", drop_na=True).fit(X) + Xt = transformer.transform(X) + + assert isinstance(Xt, make_df) + # only missing data in the new features removes a row + assert frame_to_dict(Xt) == { + "a": [2.0, None], + "b": [20, 30], + "c": ["x", "y"], + "a_lag_1": [1.0, 2.0], + "b_lag_1": [10, 20], + } + + +def test_drop_na_without_variables_keeps_all_rows(make_df): + X = make_df({"c": DATA["c"]}) + transformer = MockLag(return_empty=True, drop_na=True).fit(X) + y = make_series(make_df, [1, 2, 3, 4]) + + Xt, yt = transformer.transform_x_y(X, y) + + assert isinstance(Xt, make_df) + assert frame_to_dict(transformer.transform(X)) == {"c": DATA["c"]} + assert frame_to_dict(Xt) == {"c": DATA["c"]} + assert list(yt) == [1, 2, 3, 4] + + +def test_transform_raises_error_if_different_number_of_columns(make_df): + transformer = MockLag().fit(make_df(DATA)) + msg = ( + "The number of columns in this dataset is different from the one used to " + "fit this transformer (when using the fit() method)." + ) + with pytest.raises(ValueError, match=re.escape(msg)): + transformer.transform(make_df({"a": DATA["a"], "b": DATA["b"]})) + + +@pytest.mark.parametrize("method", ["transform", "transform_x_y"]) +def test_raises_non_fitted_error(make_df, method): + X = make_df(DATA) + msg = ( + "This MockLag instance is not fitted yet. Call 'fit' with " + "appropriate arguments before using this estimator." + ) + with pytest.raises(NotFittedError, match=re.escape(msg)): + if method == "transform": + MockLag().transform(X) + else: + MockLag().transform_x_y(X, make_series(make_df, [1, 2, 3, 4])) + + +@pytest.mark.parametrize( + "drop_na, expected", [(True, [2, 3, 4]), (False, [1, 2, 3, 4])] +) +def test_transform_x_y(make_df, drop_na, expected): + X = make_df(DATA) + y = make_series(make_df, [1, 2, 3, 4], name="y") + transformer = MockLag(drop_na=drop_na).fit(X) + + Xt, yt = transformer.transform_x_y(X, y) + + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == frame_to_dict(transformer.transform(X)) + assert list(yt) == expected + + +def test_transform_x_y_with_multioutput_target(make_df): + X = make_df(DATA) + y = make_df({"y1": [1, 2, 3, 4], "y2": [5, 6, 7, 8]}) + + Xt, yt = MockLag(drop_na=True).fit(X).transform_x_y(X, y) + + assert isinstance(yt, make_df) + assert frame_to_dict(yt) == {"y1": [2, 3, 4], "y2": [6, 7, 8]} + assert Xt.shape == (3, 5) + + +# pandas: the index orders the rows in time +DATES = pd.to_datetime(["2020-01-01", "2020-01-02", "2020-01-03", "2020-01-04"]) + + +@pytest.fixture +def df_shuffled(): + # the rows are not in time order + return pd.DataFrame(DATA, index=DATES).iloc[[2, 0, 3, 1]] + + +def test_sort_index_orders_rows_by_the_index_with_pandas(df_shuffled): + Xt = MockLag(variables="a").fit(df_shuffled).transform(df_shuffled) + + expected = pd.DataFrame({**DATA, "a_lag_1": [np.nan, 1.0, 2.0, 3.0]}, index=DATES) + pd.testing.assert_frame_equal(Xt, expected) + # the input dataframe is not sorted in place + assert list(df_shuffled.index) == list(DATES[[2, 0, 3, 1]]) + + +def test_rows_keep_their_order_without_sort_index_with_pandas(df_shuffled): + transformer = MockLag(variables="a", sort_index=False).fit(df_shuffled) + Xt = transformer.transform(df_shuffled) + + expected = df_shuffled.assign(a_lag_1=[np.nan, 3.0, 1.0, 4.0]) + pd.testing.assert_frame_equal(Xt, expected) + + +def test_transform_x_y_aligns_y_on_the_index_with_pandas(df_shuffled): + y = pd.Series([3, 1, 4, 2], index=df_shuffled.index, name="y") + transformer = MockLag(variables="a", drop_na=True).fit(df_shuffled) + + Xt, yt = transformer.transform_x_y(df_shuffled, y) + + pd.testing.assert_frame_equal(Xt, transformer.transform(df_shuffled)) + pd.testing.assert_series_equal(yt, pd.Series([2, 3, 4], index=DATES[1:], name="y")) + + +@pytest.mark.parametrize( + "index, msg", + [ + (list(DATES[:3]) + [pd.NaT], MSG_INDEX_NA), + (list(DATES[:3]) + [DATES[0]], MSG_INDEX_DUPLICATED), + ], +) +def test_error_if_index_incomplete_or_duplicated_with_pandas(index, msg): + X_wrong = pd.DataFrame(DATA, index=index) + with pytest.raises(NotImplementedError, match=re.escape(msg)): + MockLag().fit(X_wrong) + + transformer = MockLag().fit(pd.DataFrame(DATA, index=DATES)) + with pytest.raises(NotImplementedError, match=re.escape(msg)): + transformer.transform(X_wrong) + + +def test_integer_column_names_with_pandas(): + X = pd.DataFrame({0: [1.0, 2.0, 3.0], 1: [4, 5, 6], "c": ["x", "y", "z"]}) + transformer = MockLag(drop_original=True, drop_na=True).fit(X) + Xt = transformer.transform(X[["c", 1, 0]]) + + expected = pd.DataFrame( + {"c": ["y", "z"], "0_lag_1": [1.0, 2.0], "1_lag_1": [4.0, 5.0]}, index=[1, 2] + ) + # the columns started as a mix of integers and strings + expected.columns = expected.columns.astype(object) + pd.testing.assert_frame_equal(Xt, expected) + + +# polars: rows are used in the order given +def test_sort_index_is_ignored_with_polars(): + X = pl.DataFrame({"a": [3.0, 1.0, 4.0, 2.0], "b": [30, 10, 40, 20]}) + Xt = MockLag(variables="a", sort_index=True).fit(X).transform(X) + + assert frame_to_dict(Xt)["a_lag_1"] == [None, 3.0, 1.0, 4.0] + + +def test_error_if_freq_with_polars(): + X = pl.DataFrame(DATA) + msg = ( + "freq is only supported with pandas dataframes, because it uses the " + "dataframe's DatetimeIndex. With other dataframes, leave freq=None " + "to shift the rows by periods. Got 1D instead." + ) + with pytest.raises(NotImplementedError, match=re.escape(msg)): + MockLag(freq="1D").fit(X) + + # a transformer fitted on pandas can't transform polars with freq either + transformer = MockLag(freq="1D").fit(pd.DataFrame(DATA, index=DATES)) + with pytest.raises(NotImplementedError, match=re.escape(msg)): + transformer.transform(X) From 4c77509cd5a4e9595f3290878e494bcb2f717d7c Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Sat, 19 Sep 2026 12:06:06 +0200 Subject: [PATCH 2/3] Migrate WindowFeatures to narwhals, add polars support WindowFeatures takes pandas and polars dataframes. With pandas it keeps rolling on the index, including time-span windows, freq and all rolling functions; it now calls the rolling methods instead of agg(), which is faster with many variables. With polars the rows are used in the order given, windows are numbers of rows, and mean, sum, std, var and count use narwhals, while min, max, median, skew and kurt use polars expressions. Time-span windows, other functions and min_periods=0 raise a clear error with polars. Also stops return_empty=True from failing with lists of windows or functions when there are no numerical variables. Co-Authored-By: Claude Opus 5 --- .../timeseries/forecasting/WindowFeatures.rst | 101 +- .../timeseries/forecasting/window_features.py | 233 ++++- .../test_forecasting/test_window_features.py | 882 +++++++----------- 3 files changed, 621 insertions(+), 595 deletions(-) diff --git a/docs/user_guide/timeseries/forecasting/WindowFeatures.rst b/docs/user_guide/timeseries/forecasting/WindowFeatures.rst index f601afcb6..ddec26a5d 100644 --- a/docs/user_guide/timeseries/forecasting/WindowFeatures.rst +++ b/docs/user_guide/timeseries/forecasting/WindowFeatures.rst @@ -33,7 +33,7 @@ by executing: .. code:: python - X[["var_1", "var_2"].rolling(window=3).agg(["max", "mean"]) + X[["var_1", "var_2"]].rolling(window=3).agg(["max", "mean"]) With the previous command, we create 2 window features for each variable, `var_1` and `var_2`, by taking the maximum and average value of the current and 2 previous rows of data. @@ -43,7 +43,7 @@ algorithms, we also need to shift the window forward with pandas method `shift`: .. code:: python - X[["var_1", "var_2"].rolling(window=3).agg(["max", "mean"]).shift(period=1) + X[["var_1", "var_2"]].rolling(window=3).agg(["max", "mean"]).shift(periods=1) Shifting is important to ensure that we are using values strictly in the past, with respect to the point that we want to forecast. @@ -72,15 +72,15 @@ weeks data, leaving a gap of a window in between the window feature and the fore WindowFeatures: under the hood ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -:class:`WindowFeatures` works on top of `pandas.rolling`, `pandas.aggregate` and +With pandas, :class:`WindowFeatures` works on top of `pandas.rolling` and `pandas.shift`. With `pandas.rolling`, :class:`WindowFeatures` determines the size of the windows for the operations. With `pandas.rolling` we can specify the window size with an integer, a string or a function. With :class:`WindowFeatures`, in addition, we can pass a list of integers, strings or functions, to perform computations over multiple window sizes. -:class:`WindowFeatures` uses `pandas.aggregate` to perform the mathematical operations -over the windows. Therefore, you can use any operation supported +:class:`WindowFeatures` performs the mathematical operations over the windows with the +methods of `pandas.rolling`. Therefore, you can use any operation supported by pandas. For supported aggregation functions, see Rolling Window `Functions `_. @@ -94,8 +94,10 @@ forward the value 2 weeks forward. original dataframe. It also has the methods `fit()` and `transform()` that make it compatible with the scikit-learn's `Pipeline` and cross-validation functions. -Note that, in the current implementation, :class:`WindowFeatures` only works with dataframes whose index, -containing the time series timestamp, contains unique values and no NaN. +:class:`WindowFeatures` works with pandas and polars dataframes. With pandas, the rows are +ordered in time by the dataframe's index, which should contain the time series timestamp, +with unique values and no NaN. Polars dataframes have no index, so the rows are used in the +order given. See :ref:`Window features with polars ` below. .. attention:: @@ -504,7 +506,7 @@ This is the resulting output: 2020-05-15 13:15:00 32.50 2020-05-15 13:30:00 32.52 2020-05-15 13:45:00 32.68 - Freq: 15T, Name: ambient_temp, dtype: float64 + Freq: 15min, Name: ambient_temp, dtype: float64 We can use :class:`WindowFeatures` to create, for example, 2 new window features by finding the mean and maximum value within a 45 minute windows of a pandas Series if we convert it @@ -604,6 +606,89 @@ We obtain the names of the original variables plus the new window features: 'module_temp_window_3_mean', 'irradiation_window_3_mean'] +.. _window_features_polars: + +Window features with polars +~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +:class:`WindowFeatures` also creates window features from polars dataframes. Polars +dataframes have no index, so the rows are used in the order given: make sure they are +sorted in time before creating the features. For the same reason, with polars: + +- `window` takes the number of rows in each window, as an integer or a list of integers. + Time spans, like `"30min"`, are not supported. +- `freq` is not supported: the window features are shifted by `periods` rows. +- `functions` takes `"count"`, `"kurt"`, `"max"`, `"mean"`, `"median"`, `"min"`, + `"skew"`, `"std"`, `"sum"` and `"var"`. +- `min_periods` must be greater than 0. +- `sort_index` does not apply. + +Let's create a polars dataframe with the ambient temperature and irradiation of our +previous example, which were recorded every 15 minutes: + +.. code:: python + + import polars as pl + from feature_engine.timeseries.forecasting import WindowFeatures + + X = pl.DataFrame({ + "ambient_temp": [31.31, 31.51, 32.15, 32.39, 32.62, 32.5, 32.52, 32.68], + "irradiation": [0.51, 0.79, 0.65, 0.76, 0.42, 0.49, 0.57, 0.56], + }) + +Windows of 2 and 4 rows cover 30 and 60 minutes of data: + +.. code:: python + + win_f = WindowFeatures(window=[2, 4], functions="mean") + + X_tr = win_f.fit_transform(X) + + print(X_tr) + +We see the window features on the right of the dataframe. The first rows, which don't +have enough past data to fill the window, contain null values: + +.. code:: python + + shape: (8, 6) + ┌──────────────┬─────────────┬─────────────────┬─────────────────┬────────────────┬────────────────┐ + │ ambient_temp ┆ irradiation ┆ ambient_temp_wi ┆ irradiation_win ┆ ambient_temp_w ┆ irradiation_wi │ + │ --- ┆ --- ┆ ndow_2_mean ┆ dow_2_mean ┆ indow_4_mean ┆ ndow_4_mean │ + │ f64 ┆ f64 ┆ --- ┆ --- ┆ --- ┆ --- │ + │ ┆ ┆ f64 ┆ f64 ┆ f64 ┆ f64 │ + ╞══════════════╪═════════════╪═════════════════╪═════════════════╪════════════════╪════════════════╡ + │ 31.31 ┆ 0.51 ┆ null ┆ null ┆ null ┆ null │ + │ 31.51 ┆ 0.79 ┆ null ┆ null ┆ null ┆ null │ + │ 32.15 ┆ 0.65 ┆ 31.41 ┆ 0.65 ┆ null ┆ null │ + │ 32.39 ┆ 0.76 ┆ 31.83 ┆ 0.72 ┆ null ┆ null │ + │ 32.62 ┆ 0.42 ┆ 32.27 ┆ 0.705 ┆ 31.84 ┆ 0.6775 │ + │ 32.5 ┆ 0.49 ┆ 32.505 ┆ 0.59 ┆ 32.1675 ┆ 0.655 │ + │ 32.52 ┆ 0.57 ┆ 32.56 ┆ 0.455 ┆ 32.415 ┆ 0.58 │ + │ 32.68 ┆ 0.56 ┆ 32.51 ┆ 0.53 ┆ 32.5075 ┆ 0.56 │ + └──────────────┴─────────────┴─────────────────┴─────────────────┴────────────────┴────────────────┘ + +To take windows with fewer rows at the start of the series, set `min_periods`. + +We can also drop the rows with null values in the window features from the dataframe and +the target: + +.. code:: python + + y = pl.Series("y", [1, 2, 3, 4, 5, 6, 7, 8]) + + win_f = WindowFeatures(window=[2, 4], functions=["mean", "max"], drop_na=True) + + X_tr, y_tr = win_f.fit(X).transform_x_y(X, y) + + print(X.shape, y.shape, X_tr.shape, y_tr.shape) + +The first 4 rows were removed from both: + +.. code:: python + + (8, 2) (8,) (4, 10) (4,) + Windows from the target vs windows from predictor variables ----------------------------------------------------------- diff --git a/feature_engine/timeseries/forecasting/window_features.py b/feature_engine/timeseries/forecasting/window_features.py index 7a86aaadf..3fb9733b6 100644 --- a/feature_engine/timeseries/forecasting/window_features.py +++ b/feature_engine/timeseries/forecasting/window_features.py @@ -1,6 +1,9 @@ +from itertools import product from typing import Callable, List, Union -import pandas as pd +import narwhals as nw +import narwhals.dependencies as nwd +from narwhals.typing import IntoDataFrame from feature_engine._docstrings.fit_attributes import ( _feature_names_in_docstring, @@ -21,6 +24,34 @@ BaseForecastTransformer, ) +# pandas' rolling functions that other dataframes support: with narwhals, or with +# polars' own expressions when narwhals has none (they are faster than numpy). +_NW_ROLLING = { + "count": lambda col, win, min_samples: (~col.is_null()) + .cast(nw.Float64) + .rolling_sum(win, min_samples=min_samples), + "mean": lambda col, win, min_samples: col.rolling_mean( + win, min_samples=min_samples + ), + "std": lambda col, win, min_samples: col.rolling_std(win, min_samples=min_samples), + "sum": lambda col, win, min_samples: col.rolling_sum(win, min_samples=min_samples), + "var": lambda col, win, min_samples: col.rolling_var(win, min_samples=min_samples), +} +_POLARS_ROLLING = { + # bias=False gives the sample skewness and kurtosis, as pandas. + "kurt": lambda col, win, min_samples: col.rolling_kurtosis( + win, bias=False, min_samples=min_samples + ), + "max": lambda col, win, min_samples: col.rolling_max(win, min_samples=min_samples), + "median": lambda col, win, min_samples: col.rolling_median( + win, min_samples=min_samples + ), + "min": lambda col, win, min_samples: col.rolling_min(win, min_samples=min_samples), + "skew": lambda col, win, min_samples: col.rolling_skew( + win, bias=False, min_samples=min_samples + ), +} + @Substitution( variables=_variables_numerical_docstring, @@ -41,11 +72,12 @@ class WindowFeatures(BaseForecastTransformer): example, the mean value of the previous 3 months of data is a window feature. The maximum value of the previous three rows of data is another window feature. - WindowFeatures uses pandas functions `rolling()`, `agg()` and `shift()`. With - `rolling()`, it creates rolling windows. With `agg()` it applies multiple functions - within those windows. With `shift()` it allocates the values to the correct rows. + WindowFeatures works with pandas and polars dataframes. - For supported aggregation functions, see Rolling Window + With pandas, WindowFeatures uses pandas functions `rolling()` and `shift()`. With + `rolling()`, it creates rolling windows and applies the functions within those + windows. With `shift()` it allocates the values to the correct rows. For supported + aggregation functions, see Rolling Window `Functions `_. With pandas `rolling()` we can perform rolling operations over 1 window size at a @@ -53,8 +85,10 @@ class WindowFeatures(BaseForecastTransformer): be derived from multiple window sizes, and the created features will be automatically concatenated to the original dataframe. - To be compatible with WindowFeatures, the dataframe's index must have unique values - and no missing data. + With pandas, the rows are ordered in time by the dataframe's index, which must have + unique values and no missing data. Polars dataframes have no index: their rows are + taken in the order given, so they must be sorted in time. With polars, the windows + are a number of rows, and `freq` is not supported. WindowFeatures works only with numerical variables. You can pass a list of variables to use as input for the windows. Alternatively, WindowFeatures will automatically @@ -78,25 +112,31 @@ class WindowFeatures(BaseForecastTransformer): the above specified values, in which case, features will be created for each one of the windows specified in the list. + With polars, `window` takes only integers or lists of integers. + min_periods: int, default None. Minimum number of observations in the window required to have a value; otherwise, the result is np.nan. See parameter `min_periods` in pandas - `rolling()` documentation for more details. + `rolling()` documentation for more details. With polars, `min_periods` must be + greater than 0. functions: string or list of strings, default = 'mean' The functions to apply within the window. Valid functions can be found - `here `_. + `here `_. With polars, + the valid functions are 'count', 'kurt', 'max', 'mean', 'median', 'min', + 'skew', 'std', 'sum' and 'var'. - periods: int, list of ints, default=1 + periods: int, default=1 Number of periods to shift. Can be a positive integer. See param `periods` in pandas `shift()`. - freq: str, list of str, default=None + freq: str, default=None Offset to use from the tseries module or time rule. See parameter `freq` in - pandas `shift()`. + pandas `shift()`. Only supported with pandas. sort_index: bool, default=True - Whether to order the index of the dataframe before creating the features. + Whether to order the index of the dataframe before creating the features. It + does not apply to polars dataframes, which have no index. {missing_values} @@ -153,6 +193,25 @@ class WindowFeatures(BaseForecastTransformer): 2 2022-09-20 3 8 1.5 6.5 3 2022-09-21 4 9 2.5 7.5 4 2022-09-22 5 10 3.5 8.5 + + With polars, the rows are taken in the order given: + + >>> import polars as pl + >>> X = pl.DataFrame(dict(x1 = [1,2,3,4,5], x2 = [6,7,8,9,10])) + >>> wf = WindowFeatures(window = 2) + >>> wf.fit_transform(X) + shape: (5, 4) + ┌─────┬─────┬──────────────────┬──────────────────┐ + │ x1 ┆ x2 ┆ x1_window_2_mean ┆ x2_window_2_mean │ + │ --- ┆ --- ┆ --- ┆ --- │ + │ i64 ┆ i64 ┆ f64 ┆ f64 │ + ╞═════╪═════╪══════════════════╪══════════════════╡ + │ 1 ┆ 6 ┆ null ┆ null │ + │ 2 ┆ 7 ┆ null ┆ null │ + │ 3 ┆ 8 ┆ 1.5 ┆ 6.5 │ + │ 4 ┆ 9 ┆ 2.5 ┆ 7.5 │ + │ 5 ┆ 10 ┆ 3.5 ┆ 8.5 │ + └─────┴─────┴──────────────────┴──────────────────┘ """ def __init__( @@ -199,54 +258,134 @@ def __init__( self.freq = freq self.sort_index = sort_index - def transform(self, X: pd.DataFrame) -> pd.DataFrame: + def _check_index(self, X: IntoDataFrame): + """ + Checks that the rows of the dataframe can be ordered in time and, with + dataframes other than pandas, that `window`, `functions` and `min_periods` + are supported. + + Parameters + ---------- + X: dataframe of shape = [n_samples, n_features] + The dataset. + """ + super()._check_index(X) + + if nwd.is_pandas_dataframe(X) is False: + windows = self.window if isinstance(self.window, list) else [self.window] + if not all(isinstance(win, int) for win in windows): + raise NotImplementedError( + "Time spans in window, like '3D', are only supported with pandas " + "dataframes, because they use the dataframe's DatetimeIndex. With " + "other dataframes, window takes integers, the number of rows in " + f"each window. Got {self.window} instead." + ) + + functions = ( + self.functions if isinstance(self.functions, list) else [self.functions] + ) + if not all( + function in _NW_ROLLING or function in _POLARS_ROLLING + for function in functions + ): + raise NotImplementedError( + "With dataframes other than pandas, functions takes only " + f"{sorted([*_NW_ROLLING, *_POLARS_ROLLING])}. " + f"Got {self.functions} instead." + ) + + if self.min_periods == 0: + raise NotImplementedError( + "min_periods=0 is only supported with pandas dataframes. With " + "other dataframes, min_periods takes integers greater than 0 or " + f"None. Got {self.min_periods} instead." + ) + + return self + + def transform(self, X: IntoDataFrame) -> IntoDataFrame: """ Adds window features. Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The data to transform. Returns ------- - X_new: pandas dataframe, shape = [n_samples, n_features + window_features] - The dataframe with the original plus the new variables. + X_new: dataframe of shape = [n_samples, n_features + window_features] + The dataframe with the original plus the new variables. If + `drop_na=True`, rows with missing data in the new features are removed. """ - # Common dataframe checks and setting up. - X = self._check_transform_input_and_state(X).to_native() - - if isinstance(self.window, list): - df_ls = [] - for win in self.window: - tmp = ( - X[self.variables_] - .rolling(window=win, min_periods=self.min_periods) - .agg(self.functions) - .shift(periods=self.periods, freq=self.freq) - ) - df_ls.append(tmp) - tmp = pd.concat(df_ls, axis=1, sort=False) - - else: - tmp = ( - X[self.variables_] - .rolling(window=self.window, min_periods=self.min_periods) - .agg(self.functions) - .shift(periods=self.periods, freq=self.freq) - ) - - tmp.columns = self._get_new_features_name() + return super().transform(X) - X = X.merge(tmp, left_index=True, right_index=True, how="left") + def _add_features(self, nw_X: nw.DataFrame) -> nw.DataFrame: + """ + Adds the window features after the columns of `nw_X`. - if self.drop_original: - X = X.drop(self.variables_, axis=1) + Parameters + ---------- + nw_X: narwhals dataframe of shape = [n_samples, n_features] + The data returned by `_check_transform_input_and_state()`. + """ + if len(self.variables_) == 0: + return nw_X - if self.drop_na: - X = X.dropna(subset=tmp.columns, axis=0) + windows = self.window if isinstance(self.window, list) else [self.window] + functions = ( + self.functions if isinstance(self.functions, list) else [self.functions] + ) + new_features = self._get_new_features_name() + X = nw_X.to_native() + + if nwd.is_pandas_dataframe(X) is True: + # pandas is faster than narwhals, and time spans and freq need its index. + tmp = [] + for win in windows: + rolling = X[self.variables_].rolling( + window=win, min_periods=self.min_periods + ) + for function in functions: + # the rolling methods are faster than agg() with many variables. + rolled = getattr(rolling, function)().shift( + periods=self.periods, freq=self.freq + ) + rolled.columns = [ + f"{var}_window_{win}_{function}" for var in self.variables_ + ] + tmp.append(rolled) + # with freq, the index of the features moves: merge aligns them to X. + X = X.merge( + tmp[0].join(tmp[1:])[new_features], + left_index=True, + right_index=True, + how="left", + ) + return nw.from_native(X, eager_only=True) - return X + plx = nw.get_native_namespace(nw_X) + nw_exprs, native_exprs = [], [] + for name, (win, var, function) in zip( + new_features, product(windows, self.variables_, functions) + ): + # like pandas, windows of integers need all their rows by default. + min_samples = win if self.min_periods is None else self.min_periods + if function in _NW_ROLLING: + rolled = _NW_ROLLING[function](nw.col(var), win, min_samples) + nw_exprs.append(rolled.shift(self.periods).alias(name)) + else: + rolled = _POLARS_ROLLING[function](plx.col(var), win, min_samples) + native_exprs.append(rolled.shift(self.periods).alias(name)) + + columns = nw_X.columns + nw_X = nw_X.with_columns(nw_exprs) + if len(native_exprs) > 0: + # the polars features were added last: restore the order of the names. + X = nw_X.to_native().with_columns(native_exprs) + nw_X = nw.from_native(X, eager_only=True).select(columns + new_features) + + return nw_X def _get_new_features_name(self) -> List: """Get names of the lag features.""" diff --git a/tests/test_time_series/test_forecasting/test_window_features.py b/tests/test_time_series/test_forecasting/test_window_features.py index 1ec83ed8e..6c475e305 100644 --- a/tests/test_time_series/test_forecasting/test_window_features.py +++ b/tests/test_time_series/test_forecasting/test_window_features.py @@ -1,634 +1,436 @@ +import re + import numpy as np import pandas as pd +import polars as pl import pytest -from pandas.testing import assert_frame_equal from feature_engine.timeseries.forecasting import WindowFeatures +from tests.backend_helpers import frame_to_dict, make_series -_date_time = [ - pd.Timestamp("2020-05-15 12:00:00"), - pd.Timestamp("2020-05-15 12:15:00"), - pd.Timestamp("2020-05-15 12:30:00"), - pd.Timestamp("2020-05-15 12:45:00"), - pd.Timestamp("2020-05-15 13:00:00"), - pd.Timestamp("2020-05-15 13:15:00"), - pd.Timestamp("2020-05-15 13:30:00"), - pd.Timestamp("2020-05-15 13:45:00"), - pd.Timestamp("2020-05-15 14:00:00"), -] +DATA = { + "x1": [1.0, 2.0, 4.0, 7.0, 11.0, 16.0, 22.0], + "x2": [10, 20, 30, 40, 50, 60, 70], + "cat": ["a", "b", "c", "d", "e", "f", "g"], +} +MSG_POLARS_FUNCTIONS = ( + "With dataframes other than pandas, functions takes only ['count', 'kurt', " + "'max', 'mean', 'median', 'min', 'skew', 'std', 'sum', 'var']. " +) -@pytest.mark.parametrize("_window", [[2, 2, 3], ["45min", "45min", "30min"]]) -def test_error_when_duplicated_windows(_window): - with pytest.raises(ValueError): - WindowFeatures(window=_window) +# init parameters +@pytest.mark.parametrize("window", [[2, 2, 3], ["45min", "45min", "30min"]]) +def test_error_if_window_duplicated(window): + msg = f"There are duplicated windows in the list: {window}" + with pytest.raises(ValueError, match=re.escape(msg)): + WindowFeatures(window=window) -@pytest.mark.parametrize("_periods", [1, 2, 3]) -def test_permitted_param_periods(_periods): - transformer = WindowFeatures(periods=_periods) - assert transformer.periods == _periods +@pytest.mark.parametrize("functions", [3.33, [1, "mean"], None, ("mean", "sum")]) +def test_error_if_functions_not_string_or_list(functions): + msg = f"functions must be a string or a list of strings. Got {functions} instead." + with pytest.raises(ValueError, match=re.escape(msg)): + WindowFeatures(functions=functions) -@pytest.mark.parametrize("_periods", ["pizza", 3.33, ["mate", "cumbia"], None]) -def test_error_when_non_permitted_param_periods(_periods): - with pytest.raises(ValueError): - WindowFeatures(periods=_periods) +def test_error_if_functions_duplicated(): + msg = "There are duplicated functions in the list: ['sum', 'sum', 'mean']" + with pytest.raises(ValueError, match=re.escape(msg)): + WindowFeatures(functions=["sum", "sum", "mean"]) -@pytest.mark.parametrize( - "_functions", ["mean", "std", ["sum", "mean"], ["sum", "mean", "count"]] -) -def test_permitted_param_functions(_functions): - transformer = WindowFeatures(functions=_functions) - assert transformer.functions == _functions + +@pytest.mark.parametrize("periods", ["pizza", 3.33, ["mate", "cumbia"], None, 0, -1]) +def test_error_if_periods_not_positive_integer(periods): + msg = f"periods must be a positive integer. Got {periods} instead." + with pytest.raises(ValueError, match=re.escape(msg)): + WindowFeatures(periods=periods) @pytest.mark.parametrize( - "_functions", [3.33, [1, "mean"], None, ["sum", "sum", "mean"]] + "window, min_periods, functions, periods, freq, sort_index", + [ + (3, None, "mean", 1, None, True), + ([2, 3], 1, ["sum", "mean", "count"], 2, "15min", False), + (["45min", "30min"], 2, ["std"], 3, "1h", True), + ], ) -def test_error_when_non_permitted_param_functions(_functions): - with pytest.raises(ValueError): - WindowFeatures(functions=_functions) - - -@pytest.mark.parametrize("_drop_or", [-1, [0], None, 7, "hola"]) -def test_error_when_non_permitted_param_drop_original(_drop_or): - with pytest.raises(ValueError): - WindowFeatures(drop_original=_drop_or) +def test_init_param_assignment( + window, min_periods, functions, periods, freq, sort_index +): + transformer = WindowFeatures( + window=window, + min_periods=min_periods, + functions=functions, + periods=periods, + freq=freq, + sort_index=sort_index, + missing_values="ignore", + drop_original=True, + drop_na=True, + ) + assert transformer.window == window + assert transformer.min_periods == min_periods + assert transformer.functions == functions + assert transformer.periods == periods + assert transformer.freq == freq + assert transformer.sort_index == sort_index + assert transformer.missing_values == "ignore" + assert transformer.drop_original is True + assert transformer.drop_na is True + + +# fit and transform +def test_get_feature_names_out(make_df): + X = make_df(DATA) + original = ["x1", "x2", "cat"] + + transformer = WindowFeatures(window=3, functions=["mean", "sum"]).fit(X) + expected = original + [ + "x1_window_3_mean", + "x1_window_3_sum", + "x2_window_3_mean", + "x2_window_3_sum", + ] + assert transformer.get_feature_names_out() == expected + assert transformer.get_feature_names_out(input_features=original) == expected + + transformer = WindowFeatures(window=[2, 3], functions="mean").fit(X) + assert transformer.get_feature_names_out() == original + [ + "x1_window_2_mean", + "x2_window_2_mean", + "x1_window_3_mean", + "x2_window_3_mean", + ] + transformer = WindowFeatures(variables="x2", drop_original=True).fit(X) + assert transformer.get_feature_names_out() == ["x1", "cat", "x2_window_3_mean"] + + +@pytest.mark.parametrize("input_features", [["x1"], ["cat"]]) +def test_error_if_input_features_differ_from_fit(make_df, input_features): + transformer = WindowFeatures().fit(make_df(DATA)) + msg = "input_features is not equal to feature_names_in_" + with pytest.raises(ValueError, match=re.escape(msg)): + transformer.get_feature_names_out(input_features=input_features) + + +def test_error_if_input_features_not_list(make_df): + transformer = WindowFeatures().fit(make_df(DATA)) + msg = "input_features must be a list or an array. Got x1 instead." + with pytest.raises(ValueError, match=re.escape(msg)): + transformer.get_feature_names_out(input_features="x1") + + +def test_multiple_windows_and_functions(make_df): + X = make_df(DATA) + Xt = WindowFeatures(window=[2, 3], functions=["sum", "mean"]).fit_transform(X) + + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == { + **DATA, + "x1_window_2_sum": [None, None, 3.0, 6.0, 11.0, 18.0, 27.0], + "x1_window_2_mean": [None, None, 1.5, 3.0, 5.5, 9.0, 13.5], + "x2_window_2_sum": [None, None, 30, 50, 70, 90, 110], + "x2_window_2_mean": [None, None, 15.0, 25.0, 35.0, 45.0, 55.0], + "x1_window_3_sum": [None, None, None, 7.0, 13.0, 22.0, 34.0], + "x1_window_3_mean": pytest.approx( + [None, None, None, 7 / 3, 13 / 3, 22 / 3, 34 / 3] + ), + "x2_window_3_sum": [None, None, None, 60, 90, 120, 150], + "x2_window_3_mean": [None, None, None, 20.0, 30.0, 40.0, 50.0], + } -@pytest.mark.parametrize("_drop_na", [-1, [0], None, 7, "hola"]) -def test_error_when_non_permitted_param_drop_na(_drop_na): - with pytest.raises(ValueError): - WindowFeatures(drop_na=_drop_na) +@pytest.mark.parametrize( + "function, expected", + [ + ("count", [4.0, 4.0, 4.0]), + ("sum", [14.0, 24.0, 38.0]), + ("mean", [3.5, 6.0, 9.5]), + ("median", [3.0, 5.5, 9.0]), + ("var", [7.0, 15.333333, 27.0]), + ("std", [2.645751, 3.915780, 5.196152]), + ("min", [1.0, 2.0, 4.0]), + ("max", [7.0, 11.0, 16.0]), + ("skew", [0.863919, 0.599581, 0.456178]), + ("kurt", [-0.285714, -0.768431, -0.951989]), + ], +) +def test_functions(make_df, function, expected): + X = make_df(DATA) + Xt = WindowFeatures(variables="x1", window=4, functions=function).fit_transform(X) + + assert isinstance(Xt, make_df) + # the first window with 4 rows is shifted to the fifth row + assert frame_to_dict(Xt)[f"x1_window_4_{function}"] == pytest.approx( + [None, None, None, None] + expected, abs=1e-6 + ) -def test_get_feature_names_out(df_time): - # input features - original_features = ["ambient_temp", "module_temp", "irradiation", "color"] - # case 1: 1 window, multiple variables, multiple functions - tr = WindowFeatures(window=3, functions=["mean", "sum"]) - tr.fit(df_time) +def test_periods(make_df): + X = make_df(DATA) + Xt = WindowFeatures(window=2, functions="max", periods=3).fit_transform(X) - # expected - output = [ - "ambient_temp_window_3_mean", - "ambient_temp_window_3_sum", - "module_temp_window_3_mean", - "module_temp_window_3_sum", - "irradiation_window_3_mean", - "irradiation_window_3_sum", - ] - output = original_features + output - assert tr.get_feature_names_out(input_features=None) == output - assert tr.get_feature_names_out(input_features=original_features) == output + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == { + **DATA, + "x1_window_2_max": [None, None, None, None, 2.0, 4.0, 7.0], + "x2_window_2_max": [None, None, None, None, 20, 30, 40], + } - with pytest.raises(ValueError): - # get error when a user passes a string instead of list - tr.get_feature_names_out(input_features=["ambient_temp"]) - with pytest.raises(ValueError): - # assert error when uses passes features that were not transformed - tr.get_feature_names_out(input_features=["color"]) +def test_min_periods(make_df): + X = make_df(DATA) + Xt = WindowFeatures(window=3, min_periods=1).fit_transform(X) - # case 2: 1 window, 1 variable, multiple functions - tr = WindowFeatures( - variables="ambient_temp", window=2, functions=["sum", "mean", "count"] - ) - tr.fit(df_time) + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == { + **DATA, + "x1_window_3_mean": pytest.approx( + [None, 1.0, 1.5, 7 / 3, 13 / 3, 22 / 3, 34 / 3] + ), + "x2_window_3_mean": [None, 10.0, 15.0, 20.0, 30.0, 40.0, 50.0], + } - # expected - output = [ - "ambient_temp_window_2_sum", - "ambient_temp_window_2_mean", - "ambient_temp_window_2_count", - ] - output = original_features + output - assert tr.get_feature_names_out(input_features=None) == output - assert tr.get_feature_names_out(input_features=df_time.columns) == output - - # case 3: multiple windows, multiple variables, multiple functions - tr = WindowFeatures(window=[2, 3], functions=["mean", "sum"]) - tr.fit(df_time) - - # expected - output = [ - "ambient_temp_window_2_mean", - "ambient_temp_window_2_sum", - "module_temp_window_2_mean", - "module_temp_window_2_sum", - "irradiation_window_2_mean", - "irradiation_window_2_sum", - "ambient_temp_window_3_mean", - "ambient_temp_window_3_sum", - "module_temp_window_3_mean", - "module_temp_window_3_sum", - "irradiation_window_3_mean", - "irradiation_window_3_sum", - ] - output = original_features + output - assert tr.get_feature_names_out(input_features=None) == output - assert tr.get_feature_names_out(input_features=original_features) == output - - # case 4: 1 window, multiple variables, 1 function - tr = WindowFeatures(window=2, functions=["mean"]) - tr.fit(df_time) - - # expected - output = [ - "ambient_temp_window_2_mean", - "module_temp_window_2_mean", - "irradiation_window_2_mean", - ] - output = original_features + output - assert tr.get_feature_names_out(input_features=None) == output - - # case 5: When function is a string - tr = WindowFeatures(window=2, functions="mean") - tr.fit(df_time) - - # expected - output = [ - "ambient_temp_window_2_mean", - "module_temp_window_2_mean", - "irradiation_window_2_mean", - ] - output = original_features + output - assert tr.get_feature_names_out(input_features=None) == output - assert tr.get_feature_names_out(input_features=original_features) == output - - # case 6: drop_original is True - tr = WindowFeatures(window=2, functions="mean", drop_original=True) - tr.fit(df_time) - - # expected - output = [ - "ambient_temp_window_2_mean", - "module_temp_window_2_mean", - "irradiation_window_2_mean", - ] - assert tr.get_feature_names_out(input_features=None) == ["color"] + output - assert ( - tr.get_feature_names_out(input_features=original_features) == ["color"] + output - ) +def test_min_periods_with_multiple_windows(make_df): + X = make_df(DATA) + Xt = WindowFeatures(variables="x2", window=[2, 3], min_periods=2).fit_transform(X) -def test_single_window_when_using_periods(df_time): - - expected_results = { - "ambient_temp": [31.31, 31.51, 32.15, 32.39, 32.62, 32.5, 32.52, 32.68, 33.76], - "module_temp": [49.18, 49.84, 52.35, 50.63, 49.61, 47.01, 46.67, 47.52, 49.8], - "irradiation": [0.51, 0.79, 0.65, 0.76, 0.42, 0.49, 0.57, 0.56, 0.74], - "color": [ - "blue", - "blue", - "blue", - "blue", - "blue", - "blue", - "blue", - "blue", - "blue", - ], - "ambient_temp_window_3_median": [ - np.nan, - np.nan, - np.nan, - np.nan, - 31.51, - 32.15, - 32.39, - 32.5, - 32.52, - ], - "module_temp_window_3_median": [ - np.nan, - np.nan, - np.nan, - np.nan, - 49.84, - 50.63, - 50.63, - 49.61, - 47.01, - ], - "irradiation_window_3_median": [ - np.nan, - np.nan, - np.nan, - np.nan, - 0.65, - 0.76, - 0.65, - 0.49, - 0.49, - ], + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == { + **DATA, + "x2_window_2_mean": [None, None, 15.0, 25.0, 35.0, 45.0, 55.0], + "x2_window_3_mean": [None, None, 15.0, 20.0, 30.0, 40.0, 50.0], } - expected_results_df = pd.DataFrame(data=expected_results, index=_date_time) - - # Case 1: automatically select variables - transformer = WindowFeatures(window=3, functions=["median"], periods=2) - df_tr = transformer.fit_transform(df_time) - assert df_tr.head(9).equals(expected_results_df) - # Case 2: when drop_original is true +@pytest.mark.parametrize( + "min_periods, mean, count", + [ + ( + None, + [None, None, None, None, None, None, 34 / 3], + [None, None, None, 2.0, 2.0, 2.0, 3.0], + ), + ( + 2, + [None, None, 1.5, 1.5, 4.5, 9.0, 34 / 3], + [None, None, 2.0, 2.0, 2.0, 2.0, 3.0], + ), + ], +) +def test_missing_values_ignored(make_df, min_periods, mean, count): + X = make_df({**DATA, "x1": [1.0, 2.0, None, 7.0, 11.0, 16.0, 22.0]}) transformer = WindowFeatures( - window=3, functions=["median"], periods=2, drop_original=True + variables="x1", + functions=["mean", "count"], + min_periods=min_periods, + missing_values="ignore", ) - df_tr = transformer.fit_transform(df_time) + Xt = transformer.fit_transform(X) - assert df_tr.head(9).equals( - expected_results_df.drop(["ambient_temp", "module_temp", "irradiation"], axis=1) - ) + assert isinstance(Xt, make_df) + # windows with missing data need min_periods rows with data + assert frame_to_dict(Xt)["x1_window_3_mean"] == pytest.approx(mean) + assert frame_to_dict(Xt)["x1_window_3_count"] == count - # Case 3: user indicates multiple variables - transformer = WindowFeatures( - variables=["module_temp", "irradiation"], - window=3, - functions="median", - periods=2, - ) - transformer.fit(df_time) - df_tr = transformer.transform(df_time) - assert df_tr.head(9).equals( - expected_results_df.drop(["ambient_temp_window_3_median"], axis=1) - ) +def test_drop_original(make_df): + X = make_df(DATA) + Xt = WindowFeatures(window=2, drop_original=True).fit_transform(X) - # Case 4: user indicates 1 variable - transformer = WindowFeatures( - variables=["module_temp"], window=3, functions=["median"], periods=2 - ) - transformer.fit(df_time) - df_tr = transformer.transform(df_time) + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == { + "cat": DATA["cat"], + "x1_window_2_mean": [None, None, 1.5, 3.0, 5.5, 9.0, 13.5], + "x2_window_2_mean": [None, None, 15.0, 25.0, 35.0, 45.0, 55.0], + } - assert df_tr.head(9).equals( - expected_results_df.drop( - ["ambient_temp_window_3_median", "irradiation_window_3_median"], axis=1 - ) - ) +def test_drop_na(make_df): + X = make_df(DATA) + Xt = WindowFeatures(window=[2, 4], functions="min", drop_na=True).fit_transform(X) -def test_single_window_when_using_freq(df_time): - - expected_results = { - "ambient_temp": [31.31, 31.51, 32.15, 32.39, 32.62, 32.5, 32.52, 32.68, 33.76], - "module_temp": [49.18, 49.84, 52.35, 50.63, 49.61, 47.01, 46.67, 47.52, 49.8], - "irradiation": [0.51, 0.79, 0.65, 0.76, 0.42, 0.49, 0.57, 0.56, 0.74], - "color": [ - "blue", - "blue", - "blue", - "blue", - "blue", - "blue", - "blue", - "blue", - "blue", - ], - "ambient_temp_window_2_sum": [ - np.nan, - np.nan, - np.nan, - np.nan, - 62.82, - 63.66, - 64.54, - 65.01, - 65.12, - ], - "module_temp_window_2_sum": [ - np.nan, - np.nan, - np.nan, - np.nan, - 99.02, - 102.19, - 102.98, - 100.24, - 96.62, - ], - "irradiation_window_2_sum": [ - np.nan, - np.nan, - np.nan, - np.nan, - 1.3, - 1.44, - 1.41, - 1.18, - 0.91, - ], + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == { + "x1": [11.0, 16.0, 22.0], + "x2": [50, 60, 70], + "cat": ["e", "f", "g"], + "x1_window_2_min": [4.0, 7.0, 11.0], + "x2_window_2_min": [30, 40, 50], + "x1_window_4_min": [1.0, 2.0, 4.0], + "x2_window_4_min": [10, 20, 30], } - expected_results_df = pd.DataFrame(data=expected_results, index=_date_time) - # Case 1: automatically select variables - transformer = WindowFeatures(window=2, functions=["sum"], freq="45min") - df_tr = transformer.fit_transform(df_time) +def test_transform_x_y(make_df): + X = make_df(DATA) + y = make_series(make_df, [1, 2, 3, 4, 5, 6, 7], name="y") + transformer = WindowFeatures(window=4, drop_na=True).fit(X) + + Xt, yt = transformer.transform_x_y(X, y) + + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == frame_to_dict(transformer.transform(X)) + assert list(yt) == [5, 6, 7] - assert df_tr.head(9).round(3).equals(expected_results_df) - # Case 2: when drop_original is true +def test_return_empty_without_numerical_variables(make_df): + X = make_df({"cat": DATA["cat"]}) transformer = WindowFeatures( - window=2, functions=["sum"], freq="45min", drop_original=True + window=[2, 3], functions=["mean", "max"], return_empty=True ) - df_tr = transformer.fit_transform(df_time) - - assert ( - df_tr.head(9) - .round(3) - .equals( - expected_results_df.drop( - ["ambient_temp", "module_temp", "irradiation"], axis=1 - ) - ) + Xt = transformer.fit_transform(X) + + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == {"cat": DATA["cat"]} + + +# pandas: the index orders the rows in time and allows time spans and freq +@pytest.fixture +def df_time_expected(): + return pd.DataFrame( + { + "ambient_temp": [31.31, 31.51, 32.15, 32.39, 32.62, 32.5, 32.52, 32.68], + "module_temp": [49.18, 49.84, 52.35, 50.63, 49.61, 47.01, 46.67, 47.52], + "irradiation": [0.51, 0.79, 0.65, 0.76, 0.42, 0.49, 0.57, 0.56], + "color": ["blue"] * 8, + }, + index=pd.date_range("2020-05-15 12:00:00", periods=8, freq="15min"), ) - # Case 3: user indicates multiple variables + +def test_freq_with_pandas(df_time, df_time_expected): transformer = WindowFeatures( variables=["ambient_temp", "irradiation"], window=2, functions="sum", freq="45min", ) - df_tr = transformer.fit_transform(df_time) + Xt = transformer.fit_transform(df_time).head(8) - assert ( - df_tr.head(9) - .round(3) - .equals(expected_results_df.drop(["module_temp_window_2_sum"], axis=1)) + expected = df_time_expected.assign( + ambient_temp_window_2_sum=[np.nan] * 4 + [62.82, 63.66, 64.54, 65.01], + irradiation_window_2_sum=[np.nan] * 4 + [1.3, 1.44, 1.41, 1.18], ) + pd.testing.assert_frame_equal(Xt, expected, check_freq=False) - # Case 4: user indicates 1 variable - transformer = WindowFeatures( - variables=["irradiation"], window=2, functions=["sum"], freq="45min" - ) - df_tr = transformer.fit_transform(df_time) - - assert ( - df_tr.head(9) - .round(3) - .equals( - expected_results_df.drop( - ["ambient_temp_window_2_sum", "module_temp_window_2_sum"], axis=1 - ) - ) - ) - - -def test_multiple_windows(df_time): - # Case 1: automatically select variables +def test_time_span_windows_with_pandas(df_time, df_time_expected): transformer = WindowFeatures( - window=[2, 3], functions=["sum", "mean"], periods=15, freq="min" + variables="irradiation", window=["30min", "45min"], functions="max" ) - df_time_tr = transformer.fit_transform(df_time) - - # Expected - X = df_time.copy() - num_vars = ["ambient_temp", "module_temp", "irradiation"] - tmp = X[num_vars].rolling(2).agg(["sum", "mean"]).shift(periods=15, freq="min") - tmp.columns = tmp.columns.droplevel() - X_tr = X.merge(tmp, left_index=True, right_index=True, how="left") - tmp = X[num_vars].rolling(3).agg(["sum", "mean"]).shift(periods=15, freq="min") - tmp.columns = tmp.columns.droplevel() - X_tr = X_tr.merge(tmp, left_index=True, right_index=True, how="left") - X_tr.columns = transformer.get_feature_names_out() - - assert df_time_tr.equals(X_tr) - - # Case 2: user indicates multiple variables - transformer = WindowFeatures( - variables=["ambient_temp", "irradiation"], - window=[2, 3], - functions=["sum", "mean"], - freq="30min", - ) - df_time_tr = transformer.fit_transform(df_time) - - # Expected - X = df_time.copy() - tmp = ( - X[["ambient_temp", "irradiation"]] - .rolling(2) - .agg(["sum", "mean"]) - .shift(freq="30min") - ) - tmp.columns = tmp.columns.droplevel() - X_tr = X.merge(tmp, left_index=True, right_index=True, how="left") - tmp = ( - X[["ambient_temp", "irradiation"]] - .rolling(3) - .agg(["sum", "mean"]) - .shift(freq="30min") - ) - tmp.columns = tmp.columns.droplevel() - X_tr = X_tr.merge(tmp, left_index=True, right_index=True, how="left") - X_tr.columns = transformer.get_feature_names_out() + Xt = transformer.fit_transform(df_time).head(8) - assert df_time_tr.equals(X_tr) - - # Case 4: user indicates 1 variable - transformer = WindowFeatures( - variables=["irradiation"], - window=[2, 3], - functions="std", - periods=1, + # windows of time spans take the rows available + expected = df_time_expected.assign( + irradiation_window_30min_max=[np.nan, 0.51, 0.79, 0.79, 0.76, 0.76, 0.49, 0.57], + irradiation_window_45min_max=[np.nan, 0.51, 0.79, 0.79, 0.79, 0.76, 0.76, 0.57], ) - df_time_tr = transformer.fit_transform(df_time) - - # Expected - X = df_time.copy() - tmp = X["irradiation"].rolling(2).agg("std").shift(periods=1) - X_tr = X.merge(tmp, left_index=True, right_index=True, how="left") - tmp = X["irradiation"].rolling(3).agg("std").shift(periods=1) - X_tr = X_tr.merge(tmp, left_index=True, right_index=True, how="left") - X_tr.columns = transformer.get_feature_names_out() - - assert df_time_tr.equals(X_tr) + pd.testing.assert_frame_equal(Xt, expected, check_freq=False) -def test_min_periods_is_used(df_time): - # min_periods was accepted and documented but never reached pandas - # rolling(), so the leading rows stayed NaN whatever the user asked for. - variables = ["ambient_temp", "module_temp", "irradiation"] +def test_functions_only_in_pandas(df_time): + transformer = WindowFeatures(variables="irradiation", functions=["sem", "last"]) + Xt = transformer.fit_transform(df_time) - transformer = WindowFeatures(window=3, min_periods=1) - df_tr = transformer.fit_transform(df_time) + expected = df_time["irradiation"].rolling(3).agg(["sem", "last"]).shift(1) + expected.columns = ["irradiation_window_3_sem", "irradiation_window_3_last"] + pd.testing.assert_frame_equal(Xt[expected.columns], expected) - expected = ( - df_time[variables].rolling(window=3, min_periods=1).agg("mean").shift(periods=1) - ) - expected.columns = [f"{var}_window_3_mean" for var in variables] - assert_frame_equal(df_tr[expected.columns], expected) +def test_sort_index_with_pandas(df_time): + X = df_time.sample(frac=1, random_state=1) - # with min_periods=1 only the very first row (shifted out) remains NaN - assert df_tr["ambient_temp_window_3_mean"].isna().sum() == 1 + Xt = WindowFeatures(sort_index=True).fit_transform(X) + pd.testing.assert_frame_equal(Xt, WindowFeatures().fit_transform(df_time)) - # the default is unchanged: pandas requires a full window - df_default = WindowFeatures(window=3).fit_transform(df_time) - assert df_default["ambient_temp_window_3_mean"].isna().sum() == 3 + transformer = WindowFeatures(sort_index=False) + Xt = transformer.fit_transform(X) + pd.testing.assert_frame_equal(Xt[df_time.columns], X) -def test_min_periods_is_used_with_multiple_windows(df_time): - transformer = WindowFeatures(window=[2, 3], min_periods=1) - df_tr = transformer.fit_transform(df_time) +def test_transform_x_y_aligns_y_on_the_index_with_pandas(df_time): + X = df_time.sample(frac=1, random_state=1) + y = pd.Series(np.arange(len(X)), index=X.index) + transformer = WindowFeatures(window="45min", drop_na=True).fit(X) - for win in (2, 3): - expected = ( - df_time["ambient_temp"].rolling(window=win, min_periods=1).mean().shift(1) - ) - assert_frame_equal( - df_tr[[f"ambient_temp_window_{win}_mean"]], - expected.to_frame(f"ambient_temp_window_{win}_mean"), - ) + Xt, yt = transformer.transform_x_y(X, y) + pd.testing.assert_frame_equal(Xt, transformer.transform(X)) + pd.testing.assert_series_equal(yt, y.loc[Xt.index]) -def test_sort_index(df_time): - # Shuffle dataframe - Xs = df_time.copy() - Xs = Xs.sample(frac=1) - transformer = WindowFeatures(sort_index=False) - df_tr = transformer.fit_transform(Xs) - assert_frame_equal(df_tr[transformer.variables_], Xs[transformer.variables_]) +def test_integer_column_names_with_pandas(): + X = pd.DataFrame({0: DATA["x1"], 1: DATA["x2"], "cat": DATA["cat"]}) + transformer = WindowFeatures(window=[2, 3], drop_original=True, drop_na=True) + Xt = transformer.fit_transform(X) - transformer = WindowFeatures(sort_index=True) - df_tr = transformer.fit_transform(Xs) - assert_frame_equal( - df_tr[transformer.variables_], Xs[transformer.variables_].sort_index() + expected = pd.DataFrame( + { + "cat": ["d", "e", "f", "g"], + "0_window_2_mean": [3.0, 5.5, 9.0, 13.5], + "1_window_2_mean": [25.0, 35.0, 45.0, 55.0], + "0_window_3_mean": [7 / 3, 13 / 3, 22 / 3, 34 / 3], + "1_window_3_mean": [20.0, 30.0, 40.0, 50.0], + }, + index=[3, 4, 5, 6], ) + # the columns started as a mix of integers and strings + expected.columns = expected.columns.astype(object) + pd.testing.assert_frame_equal(Xt, expected) -def test_drop_na(df_time): - df = df_time.head(9).copy() - - expected_results = { - "ambient_temp": [32.62, 32.5, 32.52, 32.68, 33.76], - "module_temp": [49.61, 47.01, 46.67, 47.52, 49.8], - "irradiation": [0.42, 0.49, 0.57, 0.56, 0.74], - "color": [ - "blue", - "blue", - "blue", - "blue", - "blue", - ], - "ambient_temp_window_3_median": [ - 31.51, - 32.15, - 32.39, - 32.5, - 32.52, - ], - "module_temp_window_3_median": [ - 49.84, - 50.63, - 50.63, - 49.61, - 47.01, - ], - "irradiation_window_3_median": [ - 0.65, - 0.76, - 0.65, - 0.49, - 0.49, - ], - } - expected_results_df = pd.DataFrame(data=expected_results, index=_date_time[4:]) +# polars: the rows are used in the order given and the windows are numbers of rows +def test_rows_keep_their_order_with_polars(): + X = pl.DataFrame({"x1": [4.0, 1.0, 7.0, 2.0]}) + Xt = WindowFeatures(window=2, functions="max").fit_transform(X) - # Case 1: automatically select variables - transformer = WindowFeatures( - window=3, functions=["median"], periods=2, drop_na=True - ) - df_tr = transformer.fit_transform(df) - - assert df_tr.equals(expected_results_df) + assert frame_to_dict(Xt)["x1_window_2_max"] == [None, None, 4.0, 7.0] - # Case 2: when drop_original is true - transformer = WindowFeatures( - window=3, - functions=["median"], - periods=2, - drop_original=True, - drop_na=True, - ) - df_tr = transformer.fit_transform(df) - assert df_tr.equals( - expected_results_df.drop(["ambient_temp", "module_temp", "irradiation"], axis=1) +@pytest.mark.parametrize("window", ["3D", ["45min", "30min"], [2, "30min"]]) +def test_error_if_time_span_window_with_polars(window): + msg = ( + "Time spans in window, like '3D', are only supported with pandas " + "dataframes, because they use the dataframe's DatetimeIndex. With other " + "dataframes, window takes integers, the number of rows in each window. " + f"Got {window} instead." ) + with pytest.raises(NotImplementedError, match=re.escape(msg)): + WindowFeatures(window=window).fit(pl.DataFrame(DATA)) - # Case 3: user indicates multiple variables - transformer = WindowFeatures( - variables=["module_temp", "irradiation"], - window=3, - functions="median", - periods=2, - drop_na=True, + # a transformer fitted on pandas can't transform polars with time spans either + transformer = WindowFeatures(window=window).fit( + pd.DataFrame(DATA, index=pd.date_range("2020-01-01", periods=7)) ) - transformer.fit(df) - df_tr = transformer.transform(df) + with pytest.raises(NotImplementedError, match=re.escape(msg)): + transformer.transform(pl.DataFrame(DATA)) - assert df_tr.equals( - expected_results_df.drop(["ambient_temp_window_3_median"], axis=1) - ) +@pytest.mark.parametrize("functions", ["sem", ["mean", "last"]]) +def test_error_if_function_not_supported_with_polars(functions): + msg = MSG_POLARS_FUNCTIONS + f"Got {functions} instead." + with pytest.raises(NotImplementedError, match=re.escape(msg)): + WindowFeatures(functions=functions).fit(pl.DataFrame(DATA)) -def test_transform_x_y(df_time): - df = df_time.head(9).copy() - y = pd.Series(np.zeros(len(df)), index=df.index) - - expected_results = { - "ambient_temp": [32.62, 32.5, 32.52, 32.68, 33.76], - "module_temp": [49.61, 47.01, 46.67, 47.52, 49.8], - "irradiation": [0.42, 0.49, 0.57, 0.56, 0.74], - "color": [ - "blue", - "blue", - "blue", - "blue", - "blue", - ], - "ambient_temp_window_3_median": [ - 31.51, - 32.15, - 32.39, - 32.5, - 32.52, - ], - "module_temp_window_3_median": [ - 49.84, - 50.63, - 50.63, - 49.61, - 47.01, - ], - "irradiation_window_3_median": [ - 0.65, - 0.76, - 0.65, - 0.49, - 0.49, - ], - } - expected_results_df = pd.DataFrame(data=expected_results, index=_date_time[4:]) - transformer = WindowFeatures( - window=3, functions=["median"], periods=2, drop_na=True - ) + transformer = WindowFeatures(functions=functions).fit(pd.DataFrame(DATA)) + with pytest.raises(NotImplementedError, match=re.escape(msg)): + transformer.transform(pl.DataFrame(DATA)) - df_tr = transformer.fit_transform(df) - assert df_tr.equals(expected_results_df) - assert len(df_tr) != len(y) +def test_error_if_min_periods_is_zero_with_polars(): + msg = ( + "min_periods=0 is only supported with pandas dataframes. With other " + "dataframes, min_periods takes integers greater than 0 or None. " + "Got 0 instead." + ) + with pytest.raises(NotImplementedError, match=re.escape(msg)): + WindowFeatures(min_periods=0).fit(pl.DataFrame(DATA)) - Xt, yt = transformer.transform_x_y(df, y) - assert len(Xt) == len(yt) - assert len(y) != len(yt) - assert (Xt.index == yt.index).all() + transformer = WindowFeatures(min_periods=0).fit(pd.DataFrame(DATA)) + with pytest.raises(NotImplementedError, match=re.escape(msg)): + transformer.transform(pl.DataFrame(DATA)) From f6b794bcb1697f2f184d70c469360a8cbadfc7a5 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Sat, 19 Sep 2026 12:50:15 +0200 Subject: [PATCH 3/3] Mark polars output blocks in the user guide as text Co-Authored-By: Claude Opus 5 --- docs/user_guide/timeseries/forecasting/WindowFeatures.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/user_guide/timeseries/forecasting/WindowFeatures.rst b/docs/user_guide/timeseries/forecasting/WindowFeatures.rst index ddec26a5d..5ca3624a5 100644 --- a/docs/user_guide/timeseries/forecasting/WindowFeatures.rst +++ b/docs/user_guide/timeseries/forecasting/WindowFeatures.rst @@ -649,7 +649,7 @@ Windows of 2 and 4 rows cover 30 and 60 minutes of data: We see the window features on the right of the dataframe. The first rows, which don't have enough past data to fill the window, contain null values: -.. code:: python +.. code:: text shape: (8, 6) ┌──────────────┬─────────────┬─────────────────┬─────────────────┬────────────────┬────────────────┐