From a81f90f6958b2dc8bf9772a9e372c17acc2a9e5e Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Sat, 19 Sep 2026 11:29:38 +0200 Subject: [PATCH 1/2] Migrate SklearnWrapper to narwhals, add polars support SklearnWrapper (and its deprecated alias SklearnTransformerWrapper) now accepts pandas, polars and other narwhals-supported dataframes, and returns the same library it receives. - pandas keeps native inputs to scikit-learn; the wrapped transformer is set to return pandas output, which makes creators faster. - other backends pass numpy arrays to scikit-learn (None in text columns becomes NaN, as with pandas) and rebuild the output with narwhals, turning NaN back into null. - inverse_transform no longer modifies the input dataframe. - tests rewritten to the make_df conventions; user guide gains a polars example and fixes outdated examples. Co-Authored-By: Claude Opus 5 --- docs/user_guide/wrappers/Wrapper.rst | 73 +- feature_engine/wrappers/wrappers.py | 237 ++-- .../test_check_estimator_wrappers.py | 43 +- tests/test_wrappers/test_sklearn_wrapper.py | 1130 ++++++++--------- 4 files changed, 761 insertions(+), 722 deletions(-) diff --git a/docs/user_guide/wrappers/Wrapper.rst b/docs/user_guide/wrappers/Wrapper.rst index ff0722eb0..b3f2c01df 100644 --- a/docs/user_guide/wrappers/Wrapper.rst +++ b/docs/user_guide/wrappers/Wrapper.rst @@ -15,8 +15,8 @@ The :class:`SklearnWrapper()` offers similar functionality to the class available in scikit-learn. They differ in the implementation to select the variables and the output. -The :class:`SklearnWrapper()` returns a pandas dataframe with the variables -in the order of the original data. The +The :class:`SklearnWrapper()` returns a dataframe of the same library as the input, +for example pandas or polars, with the variables in the order of the original data. The `ColumnTransformer `_ returns a Numpy array, and the order of the variables may not coincide with that of the original dataset. @@ -170,7 +170,9 @@ list: .. code:: python - cols = [var for var in X_train.columns if X_train[var].dtypes !='O'] + from feature_engine.variable_handling import find_numerical_variables + + cols = find_numerical_variables(X_train) Now, we set up the wrapper with the SelectKBest, and fit it to the train set: @@ -211,7 +213,7 @@ imports and a function to load and clean the Titanic dataset: data = data.replace('?', np.nan) data['cabin'] = data['cabin'].astype(str).str[0] data['pclass'] = data['pclass'].astype('O') - data['embarked'].fillna('C', inplace=True) + data['embarked'] = data['embarked'].fillna('C') data.drop(["name", "home.dest", "ticket", "boat", "body"], axis=1, inplace=True) return data @@ -234,7 +236,7 @@ to transform the train and test sets: .. code:: python ohe = SklearnWrapper( - OneHotEncoder(sparse=False, drop='first'), + OneHotEncoder(sparse_output=False, drop='first'), variables = ['pclass','sex']) ohe.fit(X_train) @@ -253,11 +255,11 @@ The resulting dataframe is: .. code:: python age sibsp parch fare cabin embarked pclass_2 pclass_3 sex_male - 772 17 0 0 7.8958 n S 0.0 1.0 1.0 - 543 36 0 0 10.5 n S 1.0 0.0 1.0 + 772 17 0 0 7.8958 NaN S 0.0 1.0 1.0 + 543 36 0 0 10.5 NaN S 1.0 0.0 1.0 289 18 0 2 79.65 E S 0.0 0.0 0.0 10 47 1 0 227.525 C C 0.0 0.0 1.0 - 147 NaN 0 0 42.4 n S 0.0 0.0 1.0 + 147 NaN 0 0 42.4 NaN S 0.0 0.0 1.0 Let's say you want to use :class:`SklearnWrapper()` in a more complex @@ -323,7 +325,7 @@ We see the resulting dataframe, with the new interaction features at the end: 772 17.000000 0 0 7.8958 0 0 3.0 0.0 543 36.000000 0 0 10.5000 0 0 2.0 0.0 289 18.000000 0 2 79.6500 1 0 1.0 1.0 - 10 47.000000 1 0 227.5250 2 1 1.0 0.0 + 10 47.000000 1 0 227.5250 0 1 1.0 0.0 147 29.532738 0 0 42.4000 0 0 1.0 0.0 pclass sex @@ -334,6 +336,59 @@ We see the resulting dataframe, with the new interaction features at the end: 147 0.0 +With polars +^^^^^^^^^^^ + +:class:`SklearnWrapper()` works in the same way with a polars dataframe, and returns a +polars dataframe. In polars, missing values are `null`. The wrapper passes them to the +scikit-learn transformer as missing values, so imputers replace them, and it returns the +missing values that remain as `null`: + +.. code:: python + + import polars as pl + from sklearn.impute import SimpleImputer + from sklearn.preprocessing import StandardScaler + from feature_engine.wrappers import SklearnWrapper + + X = pl.DataFrame({ + "city": ["London", None, "Paris", "London"], + "age": [20, 30, None, 40], + "income": [1000.0, 2000.0, 1500.0, None], + }) + + imputer = SklearnWrapper( + transformer=SimpleImputer(strategy="most_frequent"), + variables=["city"], + ) + X_t = imputer.fit_transform(X) + + scaler = SklearnWrapper(transformer=StandardScaler()) + X_t = scaler.fit_transform(X_t) + + print(X_t) + +The imputer replaced the missing city with the most frequent one, and the scaler +standardised the numerical variables, leaving their missing values as `null`: + +.. code:: text + + shape: (4, 3) + ┌────────┬───────────┬───────────┐ + │ city ┆ age ┆ income │ + │ --- ┆ --- ┆ --- │ + │ str ┆ f64 ┆ f64 │ + ╞════════╪═══════════╪═══════════╡ + │ London ┆ -1.224745 ┆ -1.224745 │ + │ London ┆ 0.0 ┆ 1.224745 │ + │ Paris ┆ null ┆ 0.0 │ + │ London ┆ 1.224745 ┆ null │ + └────────┴───────────┴───────────┘ + +When you use a `FunctionTransformer`, its function receives the polars dataframe, so +write it with polars methods. + + More details ^^^^^^^^^^^^ diff --git a/feature_engine/wrappers/wrappers.py b/feature_engine/wrappers/wrappers.py index 211c00fa0..87c906c12 100644 --- a/feature_engine/wrappers/wrappers.py +++ b/feature_engine/wrappers/wrappers.py @@ -1,7 +1,10 @@ import warnings from typing import List, Optional, Union -import pandas as pd +import narwhals as nw +import narwhals.dependencies as nwd +import numpy as np +from narwhals.typing import IntoDataFrame, IntoSeries from sklearn.base import BaseEstimator, TransformerMixin, clone from sklearn.utils.validation import check_is_fitted @@ -150,7 +153,8 @@ class SklearnWrapper(TransformerMixin, BaseEstimator): ----- This transformer offers similar functionality to the ColumnTransformer from scikit-learn, but it allows entering the transformations directly into a - Pipeline and returns pandas dataframes. + Pipeline and returns a dataframe of the same library as the input, for + example, pandas or polars. See Also -------- @@ -195,6 +199,27 @@ class SklearnWrapper(TransformerMixin, BaseEstimator): 0 a 1.0 4.0 1.0 4.0 16.0 1 b 2.0 5.0 4.0 10.0 25.0 2 c 3.0 6.0 9.0 18.0 36.0 + + With polars: + + >>> import polars as pl + >>> from feature_engine.wrappers import SklearnWrapper + >>> from sklearn.preprocessing import OneHotEncoder + >>> X = pl.DataFrame(dict(x1 = ["a","b","c"], x2 = [1,2,3], x3 = [4,5,6])) + >>> skw = SklearnWrapper( + >>> OneHotEncoder(sparse_output = False), variables = "x1") + >>> skw.fit(X) + >>> skw.transform(X) + shape: (3, 5) + ┌─────┬─────┬──────┬──────┬──────┐ + │ x2 ┆ x3 ┆ x1_a ┆ x1_b ┆ x1_c │ + │ --- ┆ --- ┆ --- ┆ --- ┆ --- │ + │ i64 ┆ i64 ┆ f64 ┆ f64 ┆ f64 │ + ╞═════╪═════╪══════╪══════╪══════╡ + │ 1 ┆ 4 ┆ 1.0 ┆ 0.0 ┆ 0.0 │ + │ 2 ┆ 5 ┆ 0.0 ┆ 1.0 ┆ 0.0 │ + │ 3 ┆ 6 ┆ 0.0 ┆ 0.0 ┆ 1.0 │ + └─────┴─────┴──────┴──────┴──────┘ """ def __init__( @@ -204,10 +229,10 @@ def __init__( return_empty: bool = False, ) -> None: - if not issubclass(transformer.__class__, TransformerMixin): + if not isinstance(transformer, TransformerMixin): raise TypeError( "transformer expected a Scikit-learn transformer. " - f"got {transformer} instead. " + f"Got {transformer} instead." ) if transformer.__class__.__name__ not in _ALL_TRANSFORMERS: @@ -251,21 +276,22 @@ def __init__( self.variables = _check_variables_input_value(variables) self.return_empty = return_empty - def fit(self, X: pd.DataFrame, y: Optional[str] = None): + def fit(self, X: IntoDataFrame, y: Optional[IntoSeries] = None): """ Fits the scikit-learn transformer to the selected variables. Parameters ---------- - X: pandas DataFrame - The dataset to fit the transformer. + X: dataframe of shape = [n_samples, n_features] + The dataset to fit the transformer. Can be a pandas, polars, or any other + dataframe supported by narwhals. - y: pandas Series, default=None - The target variable. + y: Series, default=None + The target variable. Only needed by the transformers that use it, like + the feature selectors. """ - # check input dataframe - X = check_X(X) + nw_X = check_X(X) self.transformer_ = clone(self.transformer) @@ -290,27 +316,37 @@ def fit(self, X: pd.DataFrame, y: Optional[str] = None): else: self.variables_ = check_numerical_variables(X, self.variables) + if nwd.is_pandas_dataframe(X) is True: + self.feature_names_in_ = list(X.columns) + else: + self.feature_names_in_ = nw_X.columns + self.n_features_in_ = nw_X.shape[1] + if len(self.variables_) == 0: - # save input features - self.feature_names_in_ = X.columns.tolist() - self.n_features_in_ = X.shape[1] return self - self.transformer_.fit(X[self.variables_], y) + # set explicitly, so a global scikit-learn output config can't change the + # container that transform() expects. FunctionTransformer warns with + # "pandas" when its function returns an array. + if ( + nwd.is_pandas_dataframe(X) is True + and self.transformer_.__class__.__name__ != "FunctionTransformer" + ): + self.transformer_.set_output(transform="pandas") + else: + self.transformer_.set_output(transform="default") + + self.transformer_.fit(self._to_sklearn_input(X, nw_X), y) if self.transformer_.__class__.__name__ in _SELECTORS: - # Find features to drop. - selected = X[self.variables_].columns[self.transformer_.get_support()] + selected = [ + self.variables_[i] for i in self.transformer_.get_support(indices=True) + ] self.features_to_drop_ = [f for f in self.variables_ if f not in selected] - # save input features - self.feature_names_in_ = X.columns.tolist() - - self.n_features_in_ = X.shape[1] - return self - def transform(self, X: pd.DataFrame) -> pd.DataFrame: + def transform(self, X: IntoDataFrame) -> IntoDataFrame: """ Apply the transformation to the dataframe. Only the selected variables will be modified. @@ -326,56 +362,67 @@ def transform(self, X: pd.DataFrame) -> pd.DataFrame: Parameters ---------- - X: pandas DataFrame + X: dataframe of shape = [n_samples, n_features] The data to transform. Returns ------- - X_new: pandas DataFrame + X_new: dataframe The transformed dataset. """ check_is_fitted(self) - - # check that input is a dataframe - X = check_X(X) - - # Check that input data contains same number of columns than - # the dataframe used to fit the imputer. - + nw_X = check_X(X) _check_X_matches_training_df(X, self.n_features_in_) - # reorder df to match train set - X = X[self.feature_names_in_] - - # nothing to transform, e.g. when return_empty selected no variables if len(self.variables_) == 0: - return X - - # Transformers that add features: creators - if self.transformer_.__class__.__name__ in [ - "OneHotEncoder", - "PolynomialFeatures", - ]: - new_features_df = pd.DataFrame( - data=self.transformer_.transform(X[self.variables_]), - columns=self.transformer_.get_feature_names_out(self.variables_), - index=X.index, - ) - X = pd.concat([X.drop(columns=self.variables_), new_features_df], axis=1) + return self._select_columns(X, nw_X, self.feature_names_in_) # Feature selection: transformers that remove features - elif self.transformer_.__class__.__name__ in _SELECTORS: + if self.transformer_.__class__.__name__ in _SELECTORS: + return self._select_columns( + X, + nw_X, + [f for f in self.feature_names_in_ if f not in self.features_to_drop_], + ) - # return the dataframe with the selected features - X.drop(columns=self.features_to_drop_, inplace=True) + # Transformers that add features: creators + if self.transformer_.__class__.__name__ in _CREATORS: + X_remaining = self._select_columns( + X, + nw_X, + [f for f in self.feature_names_in_ if f not in self.variables_], + ) + X_new = self.transformer_.transform(self._to_sklearn_input(X, nw_X)) + # pandas input: set_output already returned a dataframe with X's index. + if nwd.is_pandas_dataframe(X) is True: + nw_new = nw.from_native(X_new, eager_only=True) + else: + nw_new = self._to_frame( + X_new, + list(self.transformer_.get_feature_names_out(self.variables_)), + nw_X, + ) + return nw.concat( + [nw.from_native(X_remaining, eager_only=True), nw_new], + how="horizontal", + ).to_native() # Transformers that modify existing features - else: + # pandas is faster than narwhals. + if nwd.is_pandas_dataframe(X) is True: + X = X[self.feature_names_in_] X[self.variables_] = self.transformer_.transform(X[self.variables_]) + return X + else: + X_new = self.transformer_.transform(self._to_sklearn_input(X, nw_X)) + nw_new = self._to_frame(X_new, self.variables_, nw_X) + return ( + nw_X.select(self.feature_names_in_) + .with_columns(*nw_new.iter_columns()) + .to_native() + ) - return X - - def inverse_transform(self, X: pd.DataFrame) -> pd.DataFrame: + def inverse_transform(self, X: IntoDataFrame) -> IntoDataFrame: """Convert the transformed variables back to the original values. Only implemented for the following scikit-learn transformers: @@ -388,19 +435,16 @@ def inverse_transform(self, X: pd.DataFrame) -> pd.DataFrame: Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features]. + X: dataframe of shape = [n_samples, n_features]. The transformed dataframe. Returns ------- - X_tr: pandas dataframe of shape = [n_samples, n_features]. + X_tr: dataframe of shape = [n_samples, n_features]. The dataframe with the original values. """ - # Check method fit has been called check_is_fitted(self) - - # check that input is a dataframe - X = check_X(X) + nw_X = check_X(X) if self.transformer_.__class__.__name__ not in _INVERSE_TRANSFORM: raise NotImplementedError( @@ -409,17 +453,66 @@ def inverse_transform(self, X: pd.DataFrame) -> pd.DataFrame: ", ".join(_INVERSE_TRANSFORM) ) ) - # For safety, we check that the transformer has the method implemented. - if hasattr(self.transformer_, "inverse_transform") and callable( - self.transformer_.inverse_transform - ): - X[self.variables_] = self.transformer_.inverse_transform(X[self.variables_]) + + X_inv = self.transformer_.inverse_transform(self._to_sklearn_input(X, nw_X)) + + # pandas is faster than narwhals. + if nwd.is_pandas_dataframe(X) is True: + # replacing whole columns leaves the user's dataframe untouched, so a + # shallow copy is enough. + X = X.copy(deep=False) + X[self.variables_] = X_inv + return X else: - raise NotImplementedError( - "This Scikit-learn transformer does not have the method " - "`inverse_transform` implemented." - ) - return X + nw_inv = self._to_frame(X_inv, self.variables_, nw_X) + return nw_X.with_columns(*nw_inv.iter_columns()).to_native() + + def _to_sklearn_input(self, X: IntoDataFrame, nw_X: nw.DataFrame): + """Return the variables to transform in the format passed to the + scikit-learn transformer.""" + if nwd.is_pandas_dataframe(X) is True: + return X[self.variables_] + + # the function in a FunctionTransformer expects the user's dataframe. + if self.transformer_.__class__.__name__ == "FunctionTransformer": + return nw_X.select(self.variables_).to_native() + + nw_vars = nw_X.select(self.variables_) + X_np = nw_vars.to_numpy() + # scikit-learn treats NaN, not None, as missing in text columns, which is + # what it gets from pandas. + if X_np.dtype == object and nw_vars.null_count().to_numpy().sum() > 0: + X_np[np.equal(X_np, None)] = np.nan + return X_np + + def _to_frame( + self, X_new, columns: List[Union[str, int]], nw_X: nw.DataFrame + ) -> nw.DataFrame: + """Return the output of the scikit-learn transformer as a narwhals + dataframe with the given column names, in the backend of nw_X.""" + # the function in a FunctionTransformer may return a dataframe. + if nwd.is_into_dataframe(X_new) is True: + nw_new = nw.from_native(X_new, eager_only=True) + return nw_new.rename(dict(zip(nw_new.columns, columns))) + + # scikit-learn marks missing values with NaN, polars and others with null. + if X_new.dtype == object: + X_new[np.not_equal(X_new, X_new)] = None + elif X_new.dtype.kind == "f" and bool(np.isnan(X_new).any()) is True: + return nw.from_numpy( + X_new, schema=columns, backend=nw_X.implementation + ).with_columns(nw.all().fill_nan(None)) + + return nw.from_numpy(X_new, schema=columns, backend=nw_X.implementation) + + def _select_columns( + self, X: IntoDataFrame, nw_X: nw.DataFrame, columns: List[Union[str, int]] + ) -> IntoDataFrame: + # pandas is faster than narwhals. + if nwd.is_pandas_dataframe(X) is True: + return X[columns] + else: + return nw_X.select(columns).to_native() def get_feature_names_out( self, input_features: Optional[List[Union[str, int]]] = None diff --git a/tests/test_wrappers/test_check_estimator_wrappers.py b/tests/test_wrappers/test_check_estimator_wrappers.py index 75ab4b9c0..2f0b44e81 100644 --- a/tests/test_wrappers/test_check_estimator_wrappers.py +++ b/tests/test_wrappers/test_check_estimator_wrappers.py @@ -1,4 +1,5 @@ -import pandas as pd +import re + import pytest from sklearn.impute import SimpleImputer from sklearn.preprocessing import OrdinalEncoder, StandardScaler @@ -40,39 +41,9 @@ def test_check_variables_assignment(): def test_raises_error_when_no_transformer_passed(): # this transformer needs an estimator as an input param. - with pytest.raises(TypeError): - SklearnWrapper() - - -def test_return_empty(): - X = pd.DataFrame({"var_cat": ["A", "B", "A"]}) - - transformer = SklearnWrapper( - transformer=StandardScaler(), variables=None, return_empty=False + msg = ( + "SklearnWrapper.__init__() missing 1 required positional argument: " + "'transformer'" ) - with pytest.raises( - TypeError, match="No numerical variables found in this dataframe" - ): - transformer.fit(X) - - transformer = SklearnWrapper( - transformer=StandardScaler(), variables=None, return_empty=True - ) - with pytest.warns( - UserWarning, - match="No numerical variables found in this dataframe. " - "Returning an empty list.", - ): - transformer.fit(X) - assert transformer.variables_ == [] - - # if return_empty=True, transformer should return same df - # after transformation - dft = transformer.transform(X) - pd.testing.assert_frame_equal(dft, X) - assert transformer.get_feature_names_out() == list(X.columns) - - # when wrapping a transformer that selects all variable types (e.g. - # OrdinalEncoder), find_all_variables always finds at least the 1 column - # present in a non-empty dataframe, so return_empty can't be exercised - # this way; there is no dataframe that reaches the "no variables" branch. + with pytest.raises(TypeError, match=re.escape(msg)): + SklearnWrapper() diff --git a/tests/test_wrappers/test_sklearn_wrapper.py b/tests/test_wrappers/test_sklearn_wrapper.py index 76e816c7e..9989a2960 100644 --- a/tests/test_wrappers/test_sklearn_wrapper.py +++ b/tests/test_wrappers/test_sklearn_wrapper.py @@ -1,19 +1,17 @@ +import re + +import narwhals as nw import numpy as np import pandas as pd +import polars as pl import pytest -from sklearn import __version__ as skl_version -from sklearn.base import clone -from sklearn.datasets import fetch_california_housing +from sklearn import config_context from sklearn.decomposition import PCA from sklearn.ensemble import RandomForestClassifier, VotingClassifier -from sklearn.feature_selection import ( - RFE, - SelectFromModel, - SelectKBest, - VarianceThreshold, - f_regression, -) -from sklearn.impute import MissingIndicator, SimpleImputer +from sklearn.exceptions import NotFittedError +from sklearn.experimental import enable_iterative_imputer # noqa: F401 +from sklearn.feature_selection import SelectKBest, VarianceThreshold, f_regression +from sklearn.impute import IterativeImputer, KNNImputer, MissingIndicator, SimpleImputer from sklearn.linear_model import Lasso from sklearn.model_selection import cross_val_score from sklearn.pipeline import Pipeline @@ -31,745 +29,667 @@ ) from feature_engine.wrappers import SklearnTransformerWrapper, SklearnWrapper - - -if skl_version < "1.7.0": - kbd = KBinsDiscretizer(n_bins=3, encode="ordinal") -else: - kbd = KBinsDiscretizer( - n_bins=3, encode="ordinal", quantile_method="averaged_inverted_cdf" - ) - -_transformers = [ - Binarizer(threshold=2), - kbd, - StandardScaler(), - MinMaxScaler(), - Normalizer(), - PowerTransformer(), - FunctionTransformer(np.cbrt, validate=True), - OrdinalEncoder(), +from tests.backend_helpers import frame_to_dict, make_series, null_count + +DATA = { + "num_1": [1, 2, 3, 4], + "num_2": [2.0, 4.0, 6.0, 8.0], + "cat": ["a", "b", "a", "c"], +} + +DATA_NA = { + "num": [1.0, None, 3.0, 5.0], + "cat": ["a", None, "a", "b"], + "other": [1, 2, 3, 4], +} + +# x1 follows the target, x2 does not +DATA_SELECTION = { + "x1": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0], + "x2": [6.0, 1.0, 5.0, 2.0, 4.0, 3.0], + "cat": ["a", "b", "a", "b", "a", "b"], +} +TARGET = [1.1, 2.0, 2.9, 4.2, 5.0, 6.1] + +# StandardScaler of [1, 2, 3, 4], or of any linear transformation of it +SCALED = [ + -1.3416407864998738, + -0.4472135954999579, + 0.4472135954999579, + 1.3416407864998738, ] -_selectors = [ - SelectFromModel(Lasso(random_state=1)), - SelectKBest(f_regression, k=2), - VarianceThreshold(), - RFE(Lasso(random_state=1)), -] +NOT_FITTED_MSG = ( + "This SklearnWrapper instance is not fitted yet. Call 'fit' with " + "appropriate arguments before using this estimator." +) -def _OneHotEncoder(sparse, drop=None, dtype=np.float64) -> OneHotEncoder: - return OneHotEncoder(sparse_output=sparse, drop=drop, dtype=dtype) +# init parameters +@pytest.mark.parametrize( + "transformer", [Lasso(), RandomForestClassifier(), "StandardScaler", 1, None] +) +def test_error_if_transformer_is_not_a_transformer(transformer): + msg = ( + "transformer expected a Scikit-learn transformer. " + f"Got {transformer} instead." + ) + with pytest.raises(TypeError, match=re.escape(msg)): + SklearnWrapper(transformer=transformer) @pytest.mark.parametrize( "transformer", - [ - SimpleImputer(), - _OneHotEncoder(sparse=False), - StandardScaler(), - SelectKBest(), - ], + [PCA(), VotingClassifier([("rf", RandomForestClassifier())]), MissingIndicator()], ) -def test_permitted_param_transformer(transformer, df_na): - tr = SklearnWrapper(transformer=transformer) - assert tr.transformer == transformer - - -@pytest.mark.parametrize("transformer", [Lasso(), RandomForestClassifier()]) -def test_error_when_transformer_is_estimator(transformer, df_na): - with pytest.raises(TypeError): +def test_error_if_transformer_is_not_supported(transformer): + msg = ( + "This transformer is not compatible with the wrapper. Supported " + "transformers are GenericUnivariateSelect, RFE, RFECV, SelectFdr, " + "SelectFpr, SelectFromModel, SelectFwe, SelectKBest, SelectPercentile, " + "SequentialFeatureSelector, VarianceThreshold, OneHotEncoder, " + "PolynomialFeatures, Binarizer, FunctionTransformer, KBinsDiscretizer, " + "PowerTransformer, QuantileTransformer, SimpleImputer, IterativeImputer, " + "KNNImputer, OrdinalEncoder, MaxAbsScaler, MinMaxScaler, StandardScaler, " + "RobustScaler, Normalizer." + ) + with pytest.raises(NotImplementedError, match=re.escape(msg)): SklearnWrapper(transformer=transformer) -if skl_version < "1.7.0": - kbd = KBinsDiscretizer(encode="one_hot") -else: - kbd = KBinsDiscretizer(encode="one_hot", quantile_method="averaged_inverted_cdf") - - @pytest.mark.parametrize( "transformer", [ - PCA(), - VotingClassifier(RandomForestClassifier()), - MissingIndicator(), - kbd, SimpleImputer(add_indicator=True), - _OneHotEncoder(sparse=True), + KNNImputer(add_indicator=True), + IterativeImputer(add_indicator=True), ], ) -def test_error_not_implemented_transformer(transformer, df_na): - with pytest.raises(NotImplementedError): +def test_error_if_imputer_adds_indicator(transformer): + msg = ( + "The imputer is only compatible with the wrapper when the " + "parameter `add_indicator` is False. " + ) + with pytest.raises(NotImplementedError, match=re.escape(msg)): SklearnWrapper(transformer=transformer) -@pytest.mark.parametrize("transformer", _selectors) -def test_wrap_selectors(transformer): - # load data - X = fetch_california_housing(as_frame=True).frame - y = X["MedHouseVal"] - X = X.drop(["MedHouseVal"], axis=1) - - # prepare selectors - sel = clone(transformer) - sel_wrap = SklearnWrapper(transformer=transformer) - - # Test: - # When passing variable list - varlist = ["MedInc", "HouseAge", "AveRooms", "AveBedrms"] - sel_wrap.set_params(variables=varlist) - - Xt = pd.DataFrame( - sel.fit_transform(X[varlist], y), - columns=X[varlist].columns[(sel.get_support())], +@pytest.mark.parametrize("encode", ["onehot", "onehot-dense"]) +def test_error_if_kbins_discretizer_encoding_is_not_ordinal(encode): + msg = ( + "The KBinsDiscretizer is only compatible with the wrapper when the " + "parameter `encode` is `ordinal`. " ) - Xw = sel_wrap.fit_transform(X, y) - - selected = X[varlist].columns[(sel.get_support())] - remaining = [f for f in X.columns if f not in varlist] + with pytest.raises(NotImplementedError, match=re.escape(msg)): + SklearnWrapper(transformer=KBinsDiscretizer(encode=encode)) - pd.testing.assert_frame_equal(Xt, Xw[selected]) - pd.testing.assert_frame_equal(X[remaining], Xw[remaining]) - assert Xw.shape[1] == len(remaining) + len(selected) - # when variable list is None - sel_wrap.set_params(variables=None) +def test_error_if_one_hot_encoder_output_is_sparse(): + msg = "SklearnWrapper can only wrap OneHotEncoder if the sparse is set to False." + with pytest.raises(NotImplementedError, match=re.escape(msg)): + SklearnWrapper(transformer=OneHotEncoder(sparse_output=True)) - Xt = pd.DataFrame(sel.fit_transform(X, y), columns=X.columns[(sel.get_support())]) - Xw = sel_wrap.fit_transform(X, y) - pd.testing.assert_frame_equal(Xt, Xw) - - -@pytest.mark.parametrize("transformer", _transformers) -def test_wrap_transformers(transformer): - # load data - X = fetch_california_housing(as_frame=True).frame - - # prepare selectors - tr = clone(transformer) - tr_wrap = SklearnWrapper(transformer=transformer) - - # Test: - # When passing variable list - varlist = ["MedInc", "HouseAge", "AveRooms", "AveBedrms"] - tr_wrap.set_params(variables=varlist) +@pytest.mark.parametrize( + "transformer", + [ + SimpleImputer(), + OneHotEncoder(sparse_output=False), + StandardScaler(), + SelectKBest(), + KBinsDiscretizer(encode="ordinal"), + ], +) +def test_init_param_assignment(transformer): + wrapper = SklearnWrapper(transformer=transformer) + assert wrapper.transformer is transformer - Xt = pd.DataFrame(tr.fit_transform(X[varlist]), columns=X[varlist].columns) - Xw = tr_wrap.fit_transform(X) - remaining = [f for f in X.columns if f not in varlist] +def test_sklearn_transformer_wrapper_is_deprecated(make_df): + msg = ( + "SklearnTransformerWrapper was deprecated in favour of SklearnWrapper in " + "version 2.0.0 and will be removed in version 2.1.0. To silence this " + "warning, use SklearnWrapper instead." + ) + with pytest.warns(FutureWarning, match=re.escape(msg)): + wrapper = SklearnTransformerWrapper( + transformer=StandardScaler(), variables=["num_1"] + ) + assert isinstance(wrapper, SklearnWrapper) - assert Xt.shape[1] == 4 - assert Xw.shape[1] == 9 - pd.testing.assert_frame_equal(Xt, Xw[varlist]) - pd.testing.assert_frame_equal(X[remaining], Xw[remaining]) + Xt = wrapper.fit_transform(make_df(DATA)) + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == {**DATA, "num_1": pytest.approx(SCALED)} - # when variable list is None - tr_wrap.set_params(variables=None) - Xt = pd.DataFrame(tr.fit_transform(X), columns=X.columns) - Xw = tr_wrap.fit_transform(X) +# fit and transform +@pytest.mark.parametrize( + "transformer, expected", + [ + (StandardScaler(), {"num_1": SCALED, "num_2": SCALED}), + ( + MinMaxScaler(), + {"num_1": [0.0, 1 / 3, 2 / 3, 1.0], "num_2": [0.0, 1 / 3, 2 / 3, 1.0]}, + ), + (Binarizer(threshold=2), {"num_1": [0, 0, 1, 1], "num_2": [0, 1, 1, 1]}), + ( + FunctionTransformer(np.cbrt, validate=True), + { + "num_1": [1.0, 1.2599210498948732, 1.4422495703074083, 1.5874010519681], + "num_2": [ + 1.2599210498948732, + 1.5874010519681, + 1.8171205928321397, + 2.0, + ], + }, + ), + ], +) +def test_transformers_replace_the_variables(make_df, transformer, expected): + wrapper = SklearnWrapper(transformer=transformer, variables=["num_1", "num_2"]) + Xt = wrapper.fit_transform(make_df(DATA)) - pd.testing.assert_frame_equal(Xt, Xw) + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == { + "num_1": pytest.approx(expected["num_1"]), + "num_2": pytest.approx(expected["num_2"]), + "cat": ["a", "b", "a", "c"], + } -def test_wrap_polynomial_features(): - # load data - X = fetch_california_housing(as_frame=True).frame +def test_variables_none_selects_numerical_variables(make_df): + wrapper = SklearnWrapper(transformer=StandardScaler()) + Xt = wrapper.fit_transform(make_df(DATA)) - # prepare selectors - tr = PolynomialFeatures() - tr_wrap = SklearnWrapper(transformer=PolynomialFeatures()) + assert wrapper.variables_ == ["num_1", "num_2"] + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == { + "num_1": pytest.approx(SCALED), + "num_2": pytest.approx(SCALED), + "cat": ["a", "b", "a", "c"], + } - # Test: - # When passing variable list - varlist = ["MedInc", "HouseAge", "AveRooms", "AveBedrms"] - tr_wrap.set_params(variables=varlist) - Xt = pd.DataFrame( - tr.fit_transform(X[varlist]), columns=tr.get_feature_names_out(varlist) - ) - Xw = tr_wrap.fit_transform(X) +def test_variables_none_selects_all_variables_with_ordinal_encoder(make_df): + wrapper = SklearnWrapper(transformer=OrdinalEncoder()) + Xt = wrapper.fit_transform(make_df(DATA)) - pd.testing.assert_frame_equal(Xw, pd.concat([X.drop(columns=varlist), Xt], axis=1)) - assert Xw.shape[1] == len(X.drop(columns=varlist).columns) + len( - tr.get_feature_names_out(varlist) - ) + assert wrapper.variables_ == ["num_1", "num_2", "cat"] + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == { + "num_1": [0.0, 1.0, 2.0, 3.0], + "num_2": [0.0, 1.0, 2.0, 3.0], + "cat": [0.0, 1.0, 0.0, 2.0], + } - # when variable list is None - tr_wrap.set_params(variables=None) - Xt = pd.DataFrame(tr.fit_transform(X), columns=tr.get_feature_names_out()) - Xw = tr_wrap.fit_transform(X) +def test_transform_returns_variables_in_the_order_seen_in_fit(make_df): + wrapper = SklearnWrapper(transformer=StandardScaler(), variables=["num_1"]) + wrapper.fit(make_df(DATA)) + Xt = wrapper.transform(make_df({k: DATA[k] for k in ["cat", "num_2", "num_1"]})) - pd.testing.assert_frame_equal(Xw, Xt) - assert Xw.shape[1] == len(tr.get_feature_names_out(X.columns)) + assert isinstance(Xt, make_df) + assert list(Xt.columns) == ["num_1", "num_2", "cat"] + assert frame_to_dict(Xt) == {**DATA, "num_1": pytest.approx(SCALED)} -def test_wrap_polynomial_features_get_features_name_out(): - X = fetch_california_housing(as_frame=True).frame +def test_simple_imputer_with_numerical_variables(make_df): + wrapper = SklearnWrapper(transformer=SimpleImputer(), variables=["num"]) + Xt = wrapper.fit_transform(make_df(DATA_NA)) - varlist = ["MedInc", "HouseAge", "AveRooms", "AveBedrms"] - tr_wrap = SklearnWrapper( - transformer=PolynomialFeatures(), variables=varlist - ) + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == {**DATA_NA, "num": [1.0, 3.0, 3.0, 5.0]} - tr_wrap.fit(X) - expected_features_all = [ - "Population", - "AveOccup", - "Latitude", - "Longitude", - "MedHouseVal", - "1", - "MedInc", - "HouseAge", - "AveRooms", - "AveBedrms", - "MedInc^2", - "MedInc HouseAge", - "MedInc AveRooms", - "MedInc AveBedrms", - "HouseAge^2", - "HouseAge AveRooms", - "HouseAge AveBedrms", - "AveRooms^2", - "AveRooms AveBedrms", - "AveBedrms^2", - ] - expected_features_varlist = [ - "1", - "MedInc", - "HouseAge", - "AveRooms", - "AveBedrms", - "MedInc^2", - "MedInc HouseAge", - "MedInc AveRooms", - "MedInc AveBedrms", - "HouseAge^2", - "HouseAge AveRooms", - "HouseAge AveBedrms", - "AveRooms^2", - "AveRooms AveBedrms", - "AveBedrms^2", - ] - assert tr_wrap.get_feature_names_out() == expected_features_all - assert tr_wrap.get_feature_names_out(varlist) == expected_features_varlist +@pytest.mark.parametrize( + "transformer, expected", + [ + (SimpleImputer(strategy="most_frequent"), ["a", "a", "a", "b"]), + ( + SimpleImputer(strategy="constant", fill_value="missing"), + ["a", "missing", "a", "b"], + ), + ], +) +def test_simple_imputer_with_categorical_variables(make_df, transformer, expected): + wrapper = SklearnWrapper(transformer=transformer, variables=["cat"]) + Xt = wrapper.fit_transform(make_df(DATA_NA)) + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == {**DATA_NA, "cat": expected} -# SimpleImputer -def test_wrap_simple_imputer(df_na): - variables_to_impute = ["Age", "Marks"] - na_variables_left_after_imputation = [ - col - for col in df_na.loc[:, df_na.isna().any()].columns - if col not in variables_to_impute - ] - transformer = SklearnWrapper( - transformer=SimpleImputer(fill_value=-999, strategy="constant"), - variables=variables_to_impute, - ) +def test_missing_values_stay_missing(make_df): + wrapper = SklearnWrapper(transformer=StandardScaler(), variables=["num"]) + Xt = wrapper.fit_transform(make_df(DATA_NA)) - # transformed dataframe - ref = df_na.copy() - ref[variables_to_impute] = ref[variables_to_impute].fillna(-999) + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == { + **DATA_NA, + "num": [pytest.approx(-1.224744871391589), None, 0.0, 1.224744871391589], + } + assert null_count(Xt, "num") == 1 - dataframe_na_transformed = transformer.fit_transform(df_na) - # transformed output - assert all( - dataframe_na_transformed[na_variables_left_after_imputation].isna().sum() != 0 - ) - assert all(dataframe_na_transformed[variables_to_impute].isna().sum() == 0) - pd.testing.assert_frame_equal(ref, dataframe_na_transformed) +def test_ordinal_encoder_leaves_missing_values_missing(make_df): + wrapper = SklearnWrapper(transformer=OrdinalEncoder(), variables=["cat"]) + Xt = wrapper.fit_transform(make_df(DATA_NA)) + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == {**DATA_NA, "cat": [0.0, None, 0.0, 1.0]} + assert null_count(Xt, "cat") == 1 -def test_sklearn_imputer_object_with_constant(df_na): - variables_to_impute = ["Name", "City"] - na_variables_left_after_imputation = [ - col - for col in df_na.loc[:, df_na.isna().any()].columns - if col not in variables_to_impute - ] - transformer = SklearnWrapper( - transformer=SimpleImputer(fill_value="missing", strategy="constant"), - variables=variables_to_impute, - ) +@pytest.mark.parametrize( + "transformer, expected", + [ + ( + OneHotEncoder(sparse_output=False), + { + "cat_a": [1.0, 0.0, 1.0, 0.0], + "cat_b": [0.0, 1.0, 0.0, 0.0], + "cat_c": [0.0, 0.0, 0.0, 1.0], + }, + ), + ( + OneHotEncoder(sparse_output=False, drop="first", dtype=np.int64), + {"cat_b": [0, 1, 0, 0], "cat_c": [0, 0, 0, 1]}, + ), + ], +) +def test_one_hot_encoder_adds_variables_at_the_end(make_df, transformer, expected): + wrapper = SklearnWrapper(transformer=transformer, variables="cat") + Xt = wrapper.fit_transform(make_df(DATA)) - # transformed dataframe - ref = df_na.copy() - ref[variables_to_impute] = ref[variables_to_impute].fillna("missing") + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == { + "num_1": [1, 2, 3, 4], + "num_2": [2.0, 4.0, 6.0, 8.0], + **expected, + } - dataframe_na_transformed = transformer.fit_transform(df_na) - # transformed output - assert all( - dataframe_na_transformed[na_variables_left_after_imputation].isna().sum() != 0 +def test_one_hot_encoder_encodes_missing_values_as_a_category(make_df): + wrapper = SklearnWrapper( + transformer=OneHotEncoder(sparse_output=False), variables=["cat"] ) - assert all(dataframe_na_transformed[variables_to_impute].isna().sum() == 0) - pd.testing.assert_frame_equal(ref, dataframe_na_transformed) + Xt = wrapper.fit_transform(make_df(DATA_NA)) + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == { + "num": [1.0, None, 3.0, 5.0], + "other": [1, 2, 3, 4], + "cat_a": [1.0, 0.0, 1.0, 0.0], + "cat_b": [0.0, 0.0, 0.0, 1.0], + "cat_nan": [0.0, 1.0, 0.0, 0.0], + } -def test_sklearn_imputer_allfeatures_with_constant(df_na): - transformer = SklearnWrapper( - transformer=SimpleImputer(fill_value="missing", strategy="constant") - ) - - # transformed dataframe - ref = df_na.copy() - ref = ref.fillna("missing") - dataframe_na_transformed = transformer.fit_transform(df_na) +@pytest.mark.parametrize( + "include_bias, expected", + [ + ( + True, + { + "cat": ["a", "b", "a", "c"], + "1": [1.0, 1.0, 1.0, 1.0], + "num_1": [1.0, 2.0, 3.0, 4.0], + "num_2": [2.0, 4.0, 6.0, 8.0], + "num_1^2": [1.0, 4.0, 9.0, 16.0], + "num_1 num_2": [2.0, 8.0, 18.0, 32.0], + "num_2^2": [4.0, 16.0, 36.0, 64.0], + }, + ), + ( + False, + { + "cat": ["a", "b", "a", "c"], + "num_1": [1.0, 2.0, 3.0, 4.0], + "num_2": [2.0, 4.0, 6.0, 8.0], + "num_1^2": [1.0, 4.0, 9.0, 16.0], + "num_1 num_2": [2.0, 8.0, 18.0, 32.0], + "num_2^2": [4.0, 16.0, 36.0, 64.0], + }, + ), + ], +) +def test_polynomial_features_adds_variables_at_the_end(make_df, include_bias, expected): + wrapper = SklearnWrapper(transformer=PolynomialFeatures(include_bias=include_bias)) + Xt = wrapper.fit_transform(make_df(DATA)) - # transformed output - assert all(dataframe_na_transformed.isna().sum() == 0) - pd.testing.assert_frame_equal(ref, dataframe_na_transformed) + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == expected + assert wrapper.get_feature_names_out() == list(expected.keys()) -# One Hot Encoder -def test_sklearn_ohe_object_one_feature(df_vartypes): - variables_to_encode = ["Name"] +def test_polynomial_features_when_all_variables_are_transformed(make_df): + wrapper = SklearnWrapper(transformer=PolynomialFeatures(include_bias=False)) + Xt = wrapper.fit_transform(make_df({"num_1": [1, 2, 3, 4]})) - transformer = SklearnWrapper( - transformer=_OneHotEncoder(sparse=False, dtype=np.int64), - variables=variables_to_encode, - ) + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == { + "num_1": [1.0, 2.0, 3.0, 4.0], + "num_1^2": [1.0, 4.0, 9.0, 16.0], + } - ref = pd.DataFrame( - { - "Name_jack": [0, 0, 0, 1], - "Name_krish": [0, 0, 1, 0], - "Name_nick": [0, 1, 0, 0], - "Name_tom": [1, 0, 0, 0], - } - ) - transformed_df = transformer.fit_transform(df_vartypes[variables_to_encode]) +def test_variance_threshold_drops_variables(make_df): + wrapper = SklearnWrapper(transformer=VarianceThreshold(threshold=2)) + Xt = wrapper.fit_transform(make_df(DATA)) - # TODO: Remove pandas < 3 support when dropping older pandas versions - if pd.__version__ >= "3": - # Pandas 3 uses microseconds format - transformed_df.columns = [ - c.replace(".000000000", "").replace(".000000", "") - for c in transformed_df.columns - ] - ref.columns = [ - c.replace(".000000000", "").replace(".000000", "") for c in ref.columns - ] - else: - # Pandas 2 uses nanoseconds format - transformed_df.columns = [ - c.replace(".000000000", "").replace(".000000", "") - for c in transformed_df.columns - ] - ref.columns = [ - c.replace(".000000000", "").replace(".000000", "") for c in ref.columns - ] - pd.testing.assert_frame_equal(ref, transformed_df) + assert wrapper.features_to_drop_ == ["num_1"] + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == {"num_2": DATA["num_2"], "cat": DATA["cat"]} -def test_sklearn_ohe_object_many_features(df_vartypes): - variables_to_encode = ["Name", "City"] +def test_select_k_best_uses_the_target(make_df): + wrapper = SklearnWrapper(transformer=SelectKBest(f_regression, k=1)) + Xt = wrapper.fit_transform(make_df(DATA_SELECTION), make_series(make_df, TARGET)) - transformer = SklearnWrapper( - transformer=_OneHotEncoder(sparse=False, dtype=np.int64), - variables=variables_to_encode, - ) + assert wrapper.features_to_drop_ == ["x2"] + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == { + "x1": DATA_SELECTION["x1"], + "cat": DATA_SELECTION["cat"], + } - ref = pd.DataFrame( - { - "Name_jack": [0, 0, 0, 1], - "Name_krish": [0, 0, 1, 0], - "Name_nick": [0, 1, 0, 0], - "Name_tom": [1, 0, 0, 0], - "City_Bristol": [0, 0, 0, 1], - "City_Liverpool": [0, 0, 1, 0], - "City_London": [1, 0, 0, 0], - "City_Manchester": [0, 1, 0, 0], - } - ) - transformed_df = transformer.fit_transform(df_vartypes[variables_to_encode]) +@pytest.mark.parametrize("y", [TARGET, np.array(TARGET)]) +def test_target_as_list_or_array(make_df, y): + wrapper = SklearnWrapper(transformer=SelectKBest(f_regression, k=1)) + Xt = wrapper.fit_transform(make_df(DATA_SELECTION), y) - # TODO: Remove pandas < 3 support when dropping older pandas versions - if pd.__version__ >= "3": - # Pandas 3 uses microseconds format - transformed_df.columns = [ - c.replace(".000000000", "").replace(".000000", "") - for c in transformed_df.columns - ] - ref.columns = [ - c.replace(".000000000", "").replace(".000000", "") for c in ref.columns - ] - else: - # Pandas 2 uses nanoseconds format - transformed_df.columns = [ - c.replace(".000000000", "").replace(".000000", "") - for c in transformed_df.columns - ] - ref.columns = [ - c.replace(".000000000", "").replace(".000000", "") for c in ref.columns - ] - pd.testing.assert_frame_equal(ref, transformed_df) + assert wrapper.features_to_drop_ == ["x2"] + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == { + "x1": DATA_SELECTION["x1"], + "cat": DATA_SELECTION["cat"], + } -def test_sklearn_ohe_numeric(df_vartypes): - variables_to_encode = ["Age"] +def test_function_transformer_receives_the_input_dataframe(make_df): + def to_float(X): + # the function gets the user's dataframe, pandas or polars + return nw.from_native(X).with_columns(nw.all().cast(nw.Float64)).to_native() - transformer = SklearnWrapper( - transformer=_OneHotEncoder(sparse=False, dtype=np.int64), - variables=variables_to_encode, + wrapper = SklearnWrapper( + transformer=FunctionTransformer(to_float), variables=["col1"] ) + Xt = wrapper.fit_transform(make_df({"col1": ["1", "2", "3"], "col2": [1, 2, 3]})) - ref = pd.DataFrame( - { - "Age_18": [0, 0, 0, 1], - "Age_19": [0, 0, 1, 0], - "Age_20": [1, 0, 0, 0], - "Age_21": [0, 1, 0, 0], - } - ) + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == {"col1": [1.0, 2.0, 3.0], "col2": [1, 2, 3]} - transformed_df = transformer.fit_transform(df_vartypes[variables_to_encode]) - # TODO: Remove pandas < 3 support when dropping older pandas versions - if pd.__version__ >= "3": - # Pandas 3 uses microseconds format - transformed_df.columns = [ - c.replace(".000000000", "").replace(".000000", "") - for c in transformed_df.columns - ] - ref.columns = [ - c.replace(".000000000", "").replace(".000000", "") for c in ref.columns - ] - else: - # Pandas 2 uses nanoseconds format - transformed_df.columns = [ - c.replace(".000000000", "").replace(".000000", "") - for c in transformed_df.columns - ] - ref.columns = [ - c.replace(".000000000", "").replace(".000000", "") for c in ref.columns - ] - pd.testing.assert_frame_equal(ref, transformed_df) +def test_function_transformer_with_numerical_variables(make_df): + wrapper = SklearnWrapper( + transformer=FunctionTransformer(lambda x: x + 1), variables=["col1"] + ) + Xt = wrapper.fit_transform(make_df({"col1": [1, 2, 3], "col2": ["a", "b", "c"]})) + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == {"col1": [2, 3, 4], "col2": ["a", "b", "c"]} -def test_sklearn_ohe_all_features(df_vartypes): - transformer = SklearnWrapper( - transformer=_OneHotEncoder(sparse=False, dtype=np.int64) - ) - ref = pd.DataFrame( +@pytest.mark.parametrize( + "transform_output, output_type", + [("pandas", pd.DataFrame), ("polars", pl.DataFrame)], +) +def test_sklearn_transform_output_config(make_df, transform_output, output_type): + with config_context(transform_output=transform_output): + Xt = SklearnWrapper( + transformer=OneHotEncoder(sparse_output=False), variables=["cat"] + ).fit_transform(make_df(DATA)) + + assert isinstance(Xt, output_type) + assert frame_to_dict(Xt) == { + "num_1": [1, 2, 3, 4], + "num_2": [2.0, 4.0, 6.0, 8.0], + "cat_a": [1.0, 0.0, 1.0, 0.0], + "cat_b": [0.0, 1.0, 0.0, 0.0], + "cat_c": [0.0, 0.0, 0.0, 1.0], + } + + +def test_wrapper_in_pipeline_with_cross_validation(make_df): + # regression test for issue #368 + rng = np.random.RandomState(0) + X = make_df( { - "Name_jack": [0, 0, 0, 1], - "Name_krish": [0, 0, 1, 0], - "Name_nick": [0, 1, 0, 0], - "Name_tom": [1, 0, 0, 0], - "City_Bristol": [0, 0, 0, 1], - "City_Liverpool": [0, 0, 1, 0], - "City_London": [1, 0, 0, 0], - "City_Manchester": [0, 1, 0, 0], - "Age_18": [0, 0, 0, 1], - "Age_19": [0, 0, 1, 0], - "Age_20": [1, 0, 0, 0], - "Age_21": [0, 1, 0, 0], - "Marks_0.6": [0, 0, 0, 1], - "Marks_0.7": [0, 0, 1, 0], - "Marks_0.8": [0, 1, 0, 0], - "Marks_0.9": [1, 0, 0, 0], - "dob_2020-02-24T00:00:00.000000000": [1, 0, 0, 0], - "dob_2020-02-24T00:01:00.000000000": [0, 1, 0, 0], - "dob_2020-02-24T00:02:00.000000000": [0, 0, 1, 0], - "dob_2020-02-24T00:03:00.000000000": [0, 0, 0, 1], + "num": rng.normal(size=30).tolist(), + "cat": rng.choice(["a", "b", "c"], size=30).tolist(), } ) - - transformed_df = transformer.fit_transform(df_vartypes) - - # TODO: Remove pandas < 3 support when dropping older pandas versions - if pd.__version__ >= "3": - # Pandas 3 uses microseconds format - transformed_df.columns = [ - c.replace(".000000000", "").replace(".000000", "") - for c in transformed_df.columns - ] - ref.columns = [ - c.replace(".000000000", "").replace(".000000", "") for c in ref.columns - ] - else: - # Pandas 2 uses nanoseconds format - transformed_df.columns = [ - c.replace(".000000000", "").replace(".000000", "") - for c in transformed_df.columns - ] - ref.columns = [ - c.replace(".000000000", "").replace(".000000", "") for c in ref.columns - ] - pd.testing.assert_frame_equal(ref, transformed_df) - - -def test_sklearn_ohe_with_crossvalidation(): - """ - Created 2022-02-14 to test fix to issue # 368 - """ - - # Set up test pipeline with wrapped OneHotEncoder, with simple regression model - # to be able to run cross-validation; use sklearn CA housing data - df = fetch_california_housing(as_frame=True).frame - y = df["MedHouseVal"] - X = ( - df[["HouseAge", "AveBedrms"]] - .assign( - AveBedrms_cat=lambda x: pd.cut(x.AveBedrms, [0, 1, 2, 3, 4, np.inf]).astype( - str - ) - ) - .drop(columns="AveBedrms") - ) - pipeline: Pipeline = Pipeline( + y = make_series(make_df, rng.normal(size=30).tolist()) + pipeline = Pipeline( steps=[ ( - "encode_cat", + "encoder", SklearnWrapper( - transformer=_OneHotEncoder(drop="first", sparse=False), - variables=["AveBedrms_cat"], + transformer=OneHotEncoder(sparse_output=False, drop="first"), + variables=["cat"], ), ), ("model", Lasso()), ] ) - # Run cross-validation - results: np.ndarray = cross_val_score( - pipeline, X, y, scoring="neg_mean_squared_error", cv=3 - ) - assert not any(np.isnan(i) for i in results) - - -def test_wrap_one_hot_encoder_get_features_name_out(df_vartypes): - ohe_wrap = SklearnWrapper(transformer=_OneHotEncoder(sparse=False)) - ohe_wrap.fit(df_vartypes) - - expected_features_all = [ - "Name_jack", - "Name_krish", - "Name_nick", - "Name_tom", - "City_Bristol", - "City_Liverpool", - "City_London", - "City_Manchester", - "Age_18", - "Age_19", - "Age_20", - "Age_21", - "Marks_0.6", - "Marks_0.7", - "Marks_0.8", - "Marks_0.9", - "dob_2020-02-24T00:00:00.000000000", - "dob_2020-02-24T00:01:00.000000000", - "dob_2020-02-24T00:02:00.000000000", - "dob_2020-02-24T00:03:00.000000000", - ] - - # TODO: Remove pandas < 3 support when dropping older pandas versions - if pd.__version__ >= "3": - # Pandas 3 uses microseconds format - actual_features = [ - f.replace(".000000000", "").replace(".000000", "") - for f in ohe_wrap.get_feature_names_out() - ] - expected_features = [ - f.replace(".000000000", "").replace(".000000", "") - for f in expected_features_all - ] - else: - # Pandas 2 uses nanoseconds format - actual_features = [ - f.replace(".000000000", "").replace(".000000", "") - for f in ohe_wrap.get_feature_names_out() - ] - expected_features = [ - f.replace(".000000000", "").replace(".000000", "") - for f in expected_features_all - ] - assert actual_features == expected_features + results = cross_val_score(pipeline, X, y, scoring="neg_mean_squared_error", cv=3) + assert np.isfinite(results).all() @pytest.mark.parametrize( "transformer", - [PowerTransformer(), OrdinalEncoder(), MinMaxScaler(), StandardScaler()], + [StandardScaler(), MinMaxScaler(), PowerTransformer(), OrdinalEncoder()], ) -def test_inverse_transform(transformer): - X = fetch_california_housing(as_frame=True).frame - X = X.drop(["Longitude"], axis=1) - - tr_wrap = SklearnWrapper(transformer=transformer) +def test_inverse_transform(make_df, transformer): + wrapper = SklearnWrapper(transformer=transformer, variables=["num_1", "num_2"]) + Xt = wrapper.fit_transform(make_df(DATA)) + Xt_before = frame_to_dict(Xt) + X_inv = wrapper.inverse_transform(Xt) - # When passing variable list - varlist = ["MedInc", "HouseAge", "AveRooms", "AveBedrms"] - tr_wrap.set_params(variables=varlist) - X_tr = tr_wrap.fit_transform(X) - X_inv = tr_wrap.inverse_transform(X_tr) + assert isinstance(X_inv, make_df) + assert frame_to_dict(X_inv) == { + "num_1": pytest.approx(DATA["num_1"]), + "num_2": pytest.approx(DATA["num_2"]), + "cat": DATA["cat"], + } + # the input dataframe is not modified + assert frame_to_dict(Xt) == Xt_before - pd.testing.assert_frame_equal(X_inv, X) - # when variable list is None - tr_wrap.set_params(variables=None) +def test_inverse_transform_with_ordinal_encoder_and_categorical_variables(make_df): + wrapper = SklearnWrapper(transformer=OrdinalEncoder(), variables=["cat"]) + X_inv = wrapper.inverse_transform(wrapper.fit_transform(make_df(DATA))) - X_tr = tr_wrap.fit_transform(X) - X_inv = tr_wrap.inverse_transform(X_tr) - - pd.testing.assert_frame_equal(X_inv, X) + assert isinstance(X_inv, make_df) + assert frame_to_dict(X_inv) == DATA @pytest.mark.parametrize( "transformer", [ - SelectKBest(f_regression, k=2), + SelectKBest(f_regression, k=1), PolynomialFeatures(), - SimpleImputer(), + SimpleImputer(strategy="most_frequent"), ], ) -def test_error_when_inverse_transform_not_implemented(transformer): - X = fetch_california_housing(as_frame=True).frame - y = X["MedHouseVal"] - X = X.drop(["MedHouseVal"], axis=1) - - tr_wrap = SklearnWrapper(transformer=transformer) - tr_wrap.fit(X, y) - X_tr = tr_wrap.transform(X) +def test_error_if_inverse_transform_is_not_supported(make_df, transformer): + msg = ( + "The method `inverse_transform` is not implemented for this transformer. " + "Supported transformers are PowerTransformer, QuantileTransformer, " + "OrdinalEncoder, MaxAbsScaler, MinMaxScaler, StandardScaler, RobustScaler." + ) + X = make_df(DATA_SELECTION) + wrapper = SklearnWrapper(transformer=transformer).fit( + X, make_series(make_df, TARGET) + ) - with pytest.raises(NotImplementedError): - tr_wrap.inverse_transform(X_tr) + with pytest.raises(NotImplementedError, match=re.escape(msg)): + wrapper.inverse_transform(wrapper.transform(X)) +@pytest.mark.parametrize("variables", [["num_1", "num_2"], None]) @pytest.mark.parametrize( - "varlist", [["MedInc", "HouseAge", "AveRooms", "AveBedrms"], None] + "transformer", + [ + StandardScaler(), + Binarizer(threshold=2), + Normalizer(), + MinMaxScaler(), + ], ) -@pytest.mark.parametrize("transformer", _transformers) -def test_get_feature_names_out_transformers(varlist, transformer): - X = fetch_california_housing(as_frame=True).frame - tr_wrap = SklearnWrapper(transformer=transformer, variables=varlist) - Xw = tr_wrap.fit_transform(X) +def test_get_feature_names_out_transformers(make_df, transformer, variables): + wrapper = SklearnWrapper(transformer=transformer, variables=variables) + Xt = wrapper.fit_transform(make_df(DATA)) - assert Xw.columns.to_list() == tr_wrap.get_feature_names_out() - assert Xw.columns.to_list() == tr_wrap.get_feature_names_out(["MedInc", "HouseAge"]) + assert wrapper.get_feature_names_out() == ["num_1", "num_2", "cat"] + assert wrapper.get_feature_names_out() == list(Xt.columns) + # input_features is ignored + assert wrapper.get_feature_names_out(["num_1"]) == ["num_1", "num_2", "cat"] -@pytest.mark.parametrize( - "varlist", [["MedInc", "HouseAge", "AveRooms", "AveBedrms"], None] -) -@pytest.mark.parametrize("transformer", _selectors) -def test_get_feature_names_out_selectors(varlist, transformer): - X = fetch_california_housing(as_frame=True).frame - y = X["MedHouseVal"] - X = X.drop(["MedHouseVal"], axis=1) - tr_wrap = SklearnWrapper(transformer=transformer, variables=varlist) - Xw = tr_wrap.fit_transform(X, y) +@pytest.mark.parametrize("variables", [["x1", "x2"], None]) +def test_get_feature_names_out_selectors(make_df, variables): + wrapper = SklearnWrapper( + transformer=SelectKBest(f_regression, k=1), variables=variables + ) + Xt = wrapper.fit_transform(make_df(DATA_SELECTION), make_series(make_df, TARGET)) - assert Xw.columns.to_list() == tr_wrap.get_feature_names_out() - assert Xw.columns.to_list() == tr_wrap.get_feature_names_out(["MedInc", "HouseAge"]) + assert wrapper.get_feature_names_out() == ["x1", "cat"] + assert wrapper.get_feature_names_out() == list(Xt.columns) + # input_features is ignored + assert wrapper.get_feature_names_out(["x1"]) == ["x1", "cat"] -@pytest.mark.parametrize( - "varlist", [["MedInc", "HouseAge", "AveRooms", "AveBedrms"], None] -) -def test_get_feature_names_out_polynomialfeatures(varlist): - X = fetch_california_housing(as_frame=True).frame - tr_wrap = SklearnWrapper( - transformer=PolynomialFeatures(), variables=varlist +def test_get_feature_names_out_one_hot_encoder(make_df): + wrapper = SklearnWrapper( + transformer=OneHotEncoder(sparse_output=False), variables=["cat", "num_1"] ) - Xw = tr_wrap.fit_transform(X) - assert Xw.columns.tolist() == tr_wrap.get_feature_names_out() - - if varlist is not None: - output_feat = [ - "1", - "MedInc", - "HouseAge", - "AveRooms", - "AveBedrms", - "MedInc^2", - "MedInc HouseAge", - "MedInc AveRooms", - "MedInc AveBedrms", - "HouseAge^2", - "HouseAge AveRooms", - "HouseAge AveBedrms", - "AveRooms^2", - "AveRooms AveBedrms", - "AveBedrms^2", - ] + Xt = wrapper.fit_transform(make_df(DATA)) + new_features = [ + "cat_a", + "cat_b", + "cat_c", + "num_1_1", + "num_1_2", + "num_1_3", + "num_1_4", + ] - assert output_feat == tr_wrap.get_feature_names_out(varlist) + assert wrapper.get_feature_names_out() == ["num_2"] + new_features + assert wrapper.get_feature_names_out() == list(Xt.columns) + assert wrapper.get_feature_names_out(["cat", "num_1"]) == new_features -@pytest.mark.parametrize("varlist", [["Name", "City"], None]) -def test_get_feature_names_out_ohe(varlist, df_vartypes): - transformer = SklearnWrapper( - transformer=_OneHotEncoder(sparse=False, dtype=np.int64), - variables=varlist, +def test_get_feature_names_out_polynomial_features(make_df): + wrapper = SklearnWrapper( + transformer=PolynomialFeatures(), variables=["num_1", "num_2"] ) + Xt = wrapper.fit_transform(make_df(DATA)) + new_features = ["1", "num_1", "num_2", "num_1^2", "num_1 num_2", "num_2^2"] - df_tr = transformer.fit_transform(df_vartypes) + assert wrapper.get_feature_names_out() == ["cat"] + new_features + assert wrapper.get_feature_names_out() == list(Xt.columns) + assert wrapper.get_feature_names_out(["num_1", "num_2"]) == new_features - assert df_tr.columns.to_list() == transformer.get_feature_names_out() - - if varlist is not None: - output_feat = [ - "Name_jack", - "Name_krish", - "Name_nick", - "Name_tom", - "City_Bristol", - "City_Liverpool", - "City_London", - "City_Manchester", - ] - assert output_feat == transformer.get_feature_names_out(varlist) +def test_return_empty_when_no_numerical_variables(make_df): + msg = "No numerical variables found in this dataframe. Returning an empty list." + X = make_df({"cat": ["a", "b", "a"]}) + wrapper = SklearnWrapper(transformer=StandardScaler(), return_empty=True) + with pytest.warns(UserWarning, match=re.escape(msg)): + wrapper.fit(X) + Xt = wrapper.transform(X) -def test_function_transformer_works_with_categoricals(): - X = pd.DataFrame({"col1": ["1", "2", "3"], "col2": ["a", "b", "c"]}) + assert wrapper.variables_ == [] + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == {"cat": ["a", "b", "a"]} + assert wrapper.get_feature_names_out() == ["cat"] - X_expected = pd.DataFrame({"col1": [1.0, 2.0, 3.0], "col2": ["a", "b", "c"]}) - transformer = SklearnWrapper( - FunctionTransformer(lambda x: x.astype(np.float64)), variables=["col1"] +def test_error_if_no_numerical_variables_and_return_empty_is_false(make_df): + msg = ( + "No numerical variables found in this dataframe. Check variable dtypes " + "or set return_empty to True to return an empty list instead." ) + wrapper = SklearnWrapper(transformer=StandardScaler()) + with pytest.raises(TypeError, match=re.escape(msg)): + wrapper.fit(make_df({"cat": ["a", "b", "a"]})) - X_tf = transformer.fit_transform(X) - pd.testing.assert_frame_equal(X_expected, X_tf) +def test_error_if_transform_df_has_different_number_of_columns(make_df): + msg = ( + "The number of columns in this dataset is different from the one used to " + "fit this transformer (when using the fit() method)." + ) + wrapper = SklearnWrapper(transformer=StandardScaler()).fit(make_df(DATA)) + with pytest.raises(ValueError, match=re.escape(msg)): + wrapper.transform(make_df({"num_1": DATA["num_1"]})) -def test_function_transformer_works_with_numericals(): - X = pd.DataFrame({"col1": [1, 2, 3], "col2": ["a", "b", "c"]}) +@pytest.mark.parametrize( + "method", ["transform", "inverse_transform", "get_feature_names_out"] +) +def test_error_if_not_fitted(make_df, method): + wrapper = SklearnWrapper(transformer=StandardScaler()) + with pytest.raises(NotFittedError, match=re.escape(NOT_FITTED_MSG)): + if method == "get_feature_names_out": + wrapper.get_feature_names_out() + else: + getattr(wrapper, method)(make_df(DATA)) + + +# pandas-specific behaviour +def test_pandas_integer_column_names(): + X = pd.DataFrame( + {0: [1, 2, 3, 4], 1: [2.0, 4.0, 6.0, 8.0], 2: ["a", "b", "a", "c"]} + ) - X_expected = pd.DataFrame({"col1": [2, 3, 4], "col2": ["a", "b", "c"]}) + wrapper = SklearnWrapper(transformer=StandardScaler()) + Xt = wrapper.fit_transform(X) + expected = pd.DataFrame({0: SCALED, 1: SCALED, 2: ["a", "b", "a", "c"]}) + pd.testing.assert_frame_equal(Xt, expected) + pd.testing.assert_frame_equal(wrapper.inverse_transform(Xt), X, check_dtype=False) + + wrapper = SklearnWrapper(transformer=VarianceThreshold(threshold=2)) + pd.testing.assert_frame_equal(wrapper.fit_transform(X), X[[1, 2]]) - transformer = SklearnWrapper( - FunctionTransformer(lambda x: x + 1), variables=["col1"] - ) - X_tf = transformer.fit_transform(X) +@pytest.mark.parametrize( + "transformer, expected", + [ + (StandardScaler(), {"num_1": SCALED, "num_2": SCALED, "cat": DATA["cat"]}), + ( + OneHotEncoder(sparse_output=False), + { + "num_1": DATA["num_1"], + "num_2": DATA["num_2"], + "cat_a": [1.0, 0.0, 1.0, 0.0], + "cat_b": [0.0, 1.0, 0.0, 0.0], + "cat_c": [0.0, 0.0, 0.0, 1.0], + }, + ), + (VarianceThreshold(threshold=2), {"num_2": DATA["num_2"], "cat": DATA["cat"]}), + ], +) +def test_pandas_index_is_kept(transformer, expected): + index = [10, 10, 5, 7] + variables = ["cat"] if isinstance(transformer, OneHotEncoder) else None + wrapper = SklearnWrapper(transformer=transformer, variables=variables) + Xt = wrapper.fit_transform(pd.DataFrame(DATA, index=index)) - pd.testing.assert_frame_equal(X_expected, X_tf) + pd.testing.assert_frame_equal(Xt, pd.DataFrame(expected, index=index)) -def test_sklearn_transformer_wrapper_is_deprecated(): - """SklearnTransformerWrapper should emit a FutureWarning and still work.""" - with pytest.warns(FutureWarning, match="SklearnTransformerWrapper was deprecated"): - transformer = SklearnTransformerWrapper(transformer=StandardScaler()) - assert isinstance(transformer, SklearnWrapper) - assert isinstance(transformer.transformer, StandardScaler) +def test_pandas_one_hot_encoder_with_category_dtype(): + X = pd.DataFrame({"cat": pd.Categorical(["a", "b", "a"]), "num": [1, 2, 3]}) + wrapper = SklearnWrapper( + transformer=OneHotEncoder(sparse_output=False), variables=["cat"] + ) + expected = pd.DataFrame( + {"num": [1, 2, 3], "cat_a": [1.0, 0.0, 1.0], "cat_b": [0.0, 1.0, 0.0]} + ) + pd.testing.assert_frame_equal(wrapper.fit_transform(X), expected) From 05e9a730329a01c2586f4c1a6ea14b30c2c05555 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Sat, 19 Sep 2026 11:33:29 +0200 Subject: [PATCH 2/2] Shorten a comment in SklearnWrapper Co-Authored-By: Claude Opus 5 --- feature_engine/wrappers/wrappers.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/feature_engine/wrappers/wrappers.py b/feature_engine/wrappers/wrappers.py index 87c906c12..ee966ea69 100644 --- a/feature_engine/wrappers/wrappers.py +++ b/feature_engine/wrappers/wrappers.py @@ -325,9 +325,8 @@ def fit(self, X: IntoDataFrame, y: Optional[IntoSeries] = None): if len(self.variables_) == 0: return self - # set explicitly, so a global scikit-learn output config can't change the - # container that transform() expects. FunctionTransformer warns with - # "pandas" when its function returns an array. + # set explicitly, so a global sklearn output config can't change the container + # transform() expects. FunctionTransformer warns if its function returns arrays. if ( nwd.is_pandas_dataframe(X) is True and self.transformer_.__class__.__name__ != "FunctionTransformer"