From 905557743062b79ec36dd992e37ae11bbc3c275a Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Sat, 19 Sep 2026 11:28:30 +0200 Subject: [PATCH 1/2] Migrate BaseTargetMeanEstimator to narwhals, add polars support Compute the target mean per bin and per category directly instead of through a Pipeline of discretiser and MeanEncoders, reusing the fitted discretiser's bin edges and labels. Fit and _predict accept pandas and polars dataframes, and pandas integer column names now work. Co-Authored-By: Claude Opus 5 --- feature_engine/_prediction/base_predictor.py | 297 ++++++++--------- tests/test_prediction/test_base_predictor.py | 313 ++++++++++++++++++ .../test_check_estimator_prediction.py | 15 +- 3 files changed, 451 insertions(+), 174 deletions(-) create mode 100644 tests/test_prediction/test_base_predictor.py diff --git a/feature_engine/_prediction/base_predictor.py b/feature_engine/_prediction/base_predictor.py index 819b6a5f0..92031a40a 100644 --- a/feature_engine/_prediction/base_predictor.py +++ b/feature_engine/_prediction/base_predictor.py @@ -1,9 +1,10 @@ from typing import List, Union +import narwhals as nw +import narwhals.dependencies as nwd import numpy as np -import pandas as pd +from narwhals.typing import IntoDataFrame, IntoSeries from sklearn.base import BaseEstimator -from sklearn.pipeline import Pipeline from sklearn.utils.validation import check_is_fitted from feature_engine._check_init_parameters.check_variables import ( @@ -20,7 +21,7 @@ EqualFrequencyDiscretiser, EqualWidthDiscretiser, ) -from feature_engine.encoding import MeanEncoder +from feature_engine.encoding._helper_functions import TARGET_NAME, add_target_to_X from feature_engine.tags import _return_tags from feature_engine.variable_handling import find_categorical_and_numerical_variables @@ -42,7 +43,7 @@ class BaseTargetMeanEstimator(BaseEstimator): the values will be sorted. strategy: str, default='equal_width' - Whether the bins should of equal width ('equal_width') or equal frequency + Whether the bins should be of equal width ('equal_width') or equal frequency ('equal_frequency'). Attributes @@ -87,10 +88,13 @@ def __init__( strategy: str = "equal_width", ): - if not isinstance(bins, int): - raise ValueError(f"bins must be an integer. Got {bins} instead.") + if not isinstance(bins, int) or bins < 1: + raise ValueError(f"bins must be a positive integer. Got {bins} instead.") - if strategy not in ["equal_width", "equal_frequency"]: + if not isinstance(strategy, str) or strategy not in [ + "equal_width", + "equal_frequency", + ]: raise ValueError( "strategy takes only values 'equal_width' or 'equal_frequency'. " f"Got {strategy} instead." @@ -100,201 +104,172 @@ def __init__( self.bins = bins self.strategy = strategy - def fit(self, X: pd.DataFrame, y: pd.Series): + def fit( + self, + X: IntoDataFrame, + y: Union[IntoSeries, np.ndarray, List], + ): """ Learn the mean target value per category or bin. Parameters ---------- - X : pandas dataframe of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The training input samples. - y : pandas series of shape = [n_samples,] + y: Series, numpy array or list of shape = [n_samples,] The target variable. """ - # check if 'X' is a dataframe - X, y = check_X_y(X, y) + nw_X, y = check_X_y(X, y) - # find categorical and numerical variables ( - self.variables_categorical_, - self.variables_numerical_, + variables_categorical_, + variables_numerical_, ) = find_categorical_and_numerical_variables(X, self.variables) - # check for missing values - _check_contains_na(X, self.variables_numerical_) - _check_contains_na(X, self.variables_categorical_) - - # check inf - _check_contains_inf(X, self.variables_numerical_) - - # Create pipelines - if self.variables_categorical_ and self.variables_numerical_: - self._pipeline = self._make_combined_pipeline() - - elif self.variables_categorical_: - self._pipeline = self._make_categorical_pipeline() - + _check_contains_na(X, variables_numerical_ + variables_categorical_) + _check_contains_inf(X, variables_numerical_) + + nw_Xy = add_target_to_X(nw_X, y) + if nwd.is_pandas_dataframe(X) is True: + y_pd = nw_Xy.get_column(TARGET_NAME).to_native() + + encoder_dict_ = {} + bin_means = {} + + if len(variables_numerical_) > 0: + discretiser = self._make_discretiser(variables_numerical_).fit(X) + binner_dict_ = discretiser.binner_dict_ + for var in variables_numerical_: + edges = np.asarray(binner_dict_[var], dtype=float) + codes, _ = discretiser._digitize(nw_X.get_column(var).to_numpy(), edges) + # pandas is faster than narwhals. + if nwd.is_pandas_dataframe(X) is True: + means_per_bin = y_pd.groupby(codes).mean() + bins_seen = means_per_bin.index.to_numpy() + means = means_per_bin.to_numpy() + else: + stats = ( + nw_Xy.select(TARGET_NAME) + .with_columns( + nw.new_series("__bin__", codes, backend=nw_X.implementation) + ) + .group_by("__bin__") + .agg(nw.col(TARGET_NAME).mean()) + .sort("__bin__") + ) + bins_seen = stats.get_column("__bin__").to_numpy() + means = stats.get_column(TARGET_NAME).to_numpy() + # NaN marks the bins without training observations, which _predict + # treats as unseen values. + bin_means[var] = np.full(len(edges) - 1, np.nan) + bin_means[var][bins_seen] = means + labels = discretiser._format_bin_labels(edges, discretiser.precision) + encoder_dict_[var] = { + labels[code]: mean + for code, mean in zip(bins_seen.tolist(), means.tolist()) + } + self._discretiser = discretiser else: - self._pipeline = self._make_numerical_pipeline() - - # Train pipeline - self._pipeline.fit(X, y) - - # Assign attributes (useful to interpret features) - # Use dict() to make a copy of the dictionary. Otherwise, like in pandas, - # it is just another view of the same data, mind-blowing. - if self.variables_categorical_ and self.variables_numerical_: - self.binner_dict_ = dict( - self._pipeline.named_steps["discretiser"].binner_dict_ - ) - self.encoder_dict_ = dict( - self._pipeline.named_steps["encoder_num"].encoder_dict_ - ) - tmp_dict = dict(self._pipeline.named_steps["encoder_cat"].encoder_dict_) - self.encoder_dict_.update(tmp_dict) - - elif self.variables_categorical_: - self.binner_dict_ = {} - self.encoder_dict_ = dict(self._pipeline.encoder_dict_) - - else: - self.binner_dict_ = dict( - self._pipeline.named_steps["discretiser"].binner_dict_ - ) - self.encoder_dict_ = dict( - self._pipeline.named_steps["encoder"].encoder_dict_ - ) - - # store input features - self.n_features_in_ = X.shape[1] - self.feature_names_in_ = list(X.columns) + binner_dict_ = {} + + for var in variables_categorical_: + # pandas is faster than narwhals. + if nwd.is_pandas_dataframe(X) is True: + encoder_dict_[var] = ( + y_pd.groupby(X[var], observed=True, dropna=False).mean().to_dict() + ) + else: + stats = nw_Xy.group_by(var).agg(nw.col(TARGET_NAME).mean()) + encoder_dict_[var] = dict( + zip( + stats.get_column(var).to_list(), + stats.get_column(TARGET_NAME).to_list(), + ) + ) + + self.variables_categorical_ = variables_categorical_ + self.variables_numerical_ = variables_numerical_ + self.binner_dict_ = binner_dict_ + self.encoder_dict_ = encoder_dict_ + self._bin_means = bin_means + self.feature_names_in_ = nw_X.columns + self.n_features_in_ = nw_X.shape[1] return self - def _make_numerical_pipeline(self): - """ - Create pipeline for a dataframe solely comprised of numerical variables - using a discretiser and an encoder. - """ - encoder = MeanEncoder(variables=self.variables_numerical_, unseen="raise") - - pipeline = Pipeline( - [ - ("discretiser", self._make_discretiser()), - ("encoder", encoder), - ] - ) - - return pipeline - - def _make_categorical_pipeline(self): - """ - Instantiate the target mean encoder. Used when all variables are categorical. - """ - - pipeline = MeanEncoder(variables=self.variables_categorical_, unseen="raise") - - return pipeline - - def _make_combined_pipeline(self): - - encoder_num = MeanEncoder(variables=self.variables_numerical_, unseen="raise") - encoder_cat = MeanEncoder(variables=self.variables_categorical_, unseen="raise") - - pipeline = Pipeline( - [ - ("discretiser", self._make_discretiser()), - ("encoder_num", encoder_num), - ("encoder_cat", encoder_cat), - ] - ) - - return pipeline - - def _make_discretiser(self): + def _make_discretiser(self, variables: List[Union[str, int]]): """ Instantiate the EqualWidthDiscretiser or EqualFrequencyDiscretiser. """ if self.strategy == "equal_width": - discretiser = EqualWidthDiscretiser( - bins=self.bins, - variables=self.variables_numerical_, - return_boundaries=True, - ) + discretiser = EqualWidthDiscretiser(bins=self.bins, variables=variables) else: - discretiser = EqualFrequencyDiscretiser( - q=self.bins, - variables=self.variables_numerical_, - return_boundaries=True, - ) + discretiser = EqualFrequencyDiscretiser(q=self.bins, variables=variables) return discretiser - def _transform(self, X: pd.DataFrame) -> pd.DataFrame: + def _predict(self, X: IntoDataFrame) -> np.ndarray: """ - Replace original values by the average of the target mean value per bin or - category in each one of the variables. + Predict using the average of the target mean value across variables. Parameters ---------- - X : pandas dataframe of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The input samples. Returns ------- - X_new: pandas dataframe of shape = [n_samples, n_features] - The transformed data with the discrete variables. + y_pred: numpy array of shape = (n_samples, ) + The mean target value per observation. """ - # check method fit has been called check_is_fitted(self) - - # check that input is a dataframe - X = check_X(X) - - # Check input data contains same number of columns as df used to fit + nw_X = check_X(X) _check_X_matches_training_df(X, self.n_features_in_) - - # check for missing values - _check_contains_na(X, self.variables_numerical_) - _check_contains_na(X, self.variables_categorical_) - - # check inf + _check_contains_na(X, self.variables_numerical_ + self.variables_categorical_) _check_contains_inf(X, self.variables_numerical_) - # reorder dataframe to match train set - X = X[self.feature_names_in_] + predictions = np.zeros(nw_X.shape[0]) - # transform dataframe - X_tr = self._pipeline.transform(X) - - return X_tr - - def _predict(self, X: pd.DataFrame) -> np.ndarray: - """ - Predict using the average of the target mean value across variables. - - Parameters - ---------- - X : pandas dataframe of shape = [n_samples, n_features] - The input samples. - - Returns - ------- - y_pred: numpy array of shape = (n_samples, ) - The mean target value per observation. - """ - # transform dataframe - X_tr = self._transform(X) - - # calculate the average for each observation - predictions = ( - X_tr[self.variables_numerical_ + self.variables_categorical_] - .mean(axis=1) - .to_numpy() + unseen = [] + for var in self.variables_numerical_: + codes, _ = self._discretiser._digitize( + nw_X.get_column(var).to_numpy(), + np.asarray(self.binner_dict_[var], dtype=float), + ) + encoded = self._bin_means[var][codes] + if np.isnan(encoded).any(): + unseen.append(var) + predictions += encoded + self._raise_if_unseen(unseen) + + for var in self.variables_categorical_: + mapping = self.encoder_dict_[var] + # pandas is faster than narwhals. + if nwd.is_pandas_dataframe(X) is True: + codes, categories = X[var].factorize(use_na_sentinel=False) + encoded = np.array([mapping.get(c, np.nan) for c in categories])[codes] + else: + encoded = ( + nw_X.get_column(var) + .replace_strict(mapping, default=None, return_dtype=nw.Float64) + .to_numpy() + ) + if np.isnan(encoded).any(): + unseen.append(var) + predictions += encoded + self._raise_if_unseen(unseen) + + return predictions / ( + len(self.variables_numerical_) + len(self.variables_categorical_) ) - return predictions + def _raise_if_unseen(self, variables: List[Union[str, int]]): + if len(variables) > 0: + raise ValueError( + "During the encoding, NaN values were introduced in the feature(s) " + f"{', '.join(str(var) for var in variables)}." + ) def _more_tags(self): return _return_tags() diff --git a/tests/test_prediction/test_base_predictor.py b/tests/test_prediction/test_base_predictor.py new file mode 100644 index 000000000..b57c02b90 --- /dev/null +++ b/tests/test_prediction/test_base_predictor.py @@ -0,0 +1,313 @@ +import re +from datetime import datetime + +import numpy as np +import pandas as pd +import pytest +from sklearn.exceptions import NotFittedError + +from feature_engine._prediction.base_predictor import BaseTargetMeanEstimator +from tests.backend_helpers import make_series + +INF = float("inf") + +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." +) + +DATA = { + "num": [1, 2, 3, 4, 5, 6, 7, 8, 9, 100], + "cat": ["a", "a", "b", "b", "b", "c", "c", "c", "c", "a"], +} +TARGET = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] +ENC_DICT_CAT = {"a": 13 / 3, "b": 4.0, "c": 7.5} +# num has an outlier, so 2 equal-width bins split it at 50.5 and 2 equal-frequency +# bins at the median, 5.5. +BINNER_DICT_NUM = { + "equal_width": [-INF, 50.5, INF], + "equal_frequency": [-INF, 5.5, INF], +} +ENC_DICT_NUM = { + "equal_width": {"(-inf, 50.5]": 5.0, "(50.5, inf]": 10.0}, + "equal_frequency": {"(-inf, 5.5]": 3.0, "(5.5, inf]": 8.0}, +} +PREDICTIONS = { + "equal_width": [14 / 3, 14 / 3, 4.5, 4.5, 4.5, 6.25, 6.25, 6.25, 6.25, 43 / 6], + "equal_frequency": [ + 11 / 3, + 11 / 3, + 3.5, + 3.5, + 3.5, + 7.75, + 7.75, + 7.75, + 7.75, + 37 / 6, + ], +} + + +# init parameters +@pytest.mark.parametrize("bins", [0, -1, 2.5, "3", None, [5]]) +def test_error_if_bins_not_positive_integer(bins): + msg = f"bins must be a positive integer. Got {bins} instead." + with pytest.raises(ValueError, match=re.escape(msg)): + BaseTargetMeanEstimator(bins=bins) + + +@pytest.mark.parametrize( + "strategy", ["arbitrary", "Equal_width", "", 1, None, ["equal_width"]] +) +def test_error_if_strategy_not_permitted(strategy): + msg = ( + "strategy takes only values 'equal_width' or 'equal_frequency'. " + f"Got {strategy} instead." + ) + with pytest.raises(ValueError, match=re.escape(msg)): + BaseTargetMeanEstimator(strategy=strategy) + + +@pytest.mark.parametrize( + "bins, strategy", [(1, "equal_width"), (5, "equal_frequency"), (10, "equal_width")] +) +def test_init_param_assignment(bins, strategy): + estimator = BaseTargetMeanEstimator(bins=bins, strategy=strategy) + assert estimator.bins == bins + assert estimator.strategy == strategy + + +# fit and transform +@pytest.mark.parametrize("strategy", ["equal_width", "equal_frequency"]) +def test_fit_attributes(make_df, strategy): + X = make_df(DATA) + y = make_series(make_df, TARGET) + + estimator = BaseTargetMeanEstimator(bins=2, strategy=strategy) + estimator.fit(X, y) + + assert estimator.variables_numerical_ == ["num"] + assert estimator.variables_categorical_ == ["cat"] + assert estimator.binner_dict_ == {"num": BINNER_DICT_NUM[strategy]} + assert estimator.encoder_dict_["num"] == pytest.approx(ENC_DICT_NUM[strategy]) + assert estimator.encoder_dict_["cat"] == pytest.approx(ENC_DICT_CAT) + assert estimator.feature_names_in_ == ["num", "cat"] + assert estimator.n_features_in_ == 2 + + +@pytest.mark.parametrize("strategy", ["equal_width", "equal_frequency"]) +def test_predict(make_df, strategy): + X = make_df(DATA) + y = make_series(make_df, TARGET) + + estimator = BaseTargetMeanEstimator(bins=2, strategy=strategy).fit(X, y) + y_pred = estimator._predict(X) + + assert isinstance(y_pred, np.ndarray) + assert y_pred.tolist() == pytest.approx(PREDICTIONS[strategy]) + + +@pytest.mark.parametrize("to_target", [list, np.array]) +def test_target_as_list_or_array(make_df, to_target): + X = make_df(DATA) + y = to_target(TARGET) + + estimator = BaseTargetMeanEstimator(bins=2).fit(X, y) + + assert estimator.encoder_dict_["num"] == pytest.approx(ENC_DICT_NUM["equal_width"]) + assert estimator.encoder_dict_["cat"] == pytest.approx(ENC_DICT_CAT) + assert estimator._predict(X).tolist() == pytest.approx(PREDICTIONS["equal_width"]) + + +def test_only_numerical_variables(make_df): + X = make_df({"num": DATA["num"]}) + y = make_series(make_df, TARGET) + + estimator = BaseTargetMeanEstimator(bins=2).fit(X, y) + + assert estimator.variables_categorical_ == [] + assert estimator.encoder_dict_["num"] == pytest.approx(ENC_DICT_NUM["equal_width"]) + assert estimator._predict(X).tolist() == pytest.approx([5.0] * 9 + [10.0]) + + +def test_only_categorical_variables(make_df): + X = make_df({"cat": DATA["cat"]}) + y = make_series(make_df, TARGET) + + estimator = BaseTargetMeanEstimator(bins=2).fit(X, y) + + assert estimator.variables_numerical_ == [] + assert estimator.binner_dict_ == {} + assert estimator.encoder_dict_["cat"] == pytest.approx(ENC_DICT_CAT) + assert estimator._predict(X).tolist() == pytest.approx( + [13 / 3, 13 / 3, 4.0, 4.0, 4.0, 7.5, 7.5, 7.5, 7.5, 13 / 3] + ) + + +@pytest.mark.parametrize( + "variables, numerical, categorical", + [ + ("num", ["num"], []), + (["cat"], [], ["cat"]), + (["cat", "num"], ["num"], ["cat"]), + (None, ["num"], ["cat"]), + ], +) +def test_variable_selection(make_df, variables, numerical, categorical): + # datetime variables are never used for prediction + data = {**DATA, "date": [datetime(2020, 1, day) for day in range(1, 11)]} + X = make_df(data) + y = make_series(make_df, TARGET) + + estimator = BaseTargetMeanEstimator(bins=2, variables=variables).fit(X, y) + + assert estimator.variables_numerical_ == numerical + assert estimator.variables_categorical_ == categorical + assert estimator.feature_names_in_ == ["num", "cat", "date"] + assert estimator.n_features_in_ == 3 + + +def test_predict_with_reordered_columns(make_df): + X = make_df(DATA) + y = make_series(make_df, TARGET) + + estimator = BaseTargetMeanEstimator(bins=2).fit(X, y) + y_pred = estimator._predict(X[["cat", "num"]]) + + assert y_pred.tolist() == pytest.approx(PREDICTIONS["equal_width"]) + + +def test_values_outside_training_range_take_the_outer_bins(make_df): + X = make_df({"num": DATA["num"]}) + y = make_series(make_df, TARGET) + + estimator = BaseTargetMeanEstimator(bins=2).fit(X, y) + y_pred = estimator._predict(make_df({"num": [-1000, 1000]})) + + assert y_pred.tolist() == pytest.approx([5.0, 10.0]) + + +def test_constant_numerical_variable(make_df): + X = make_df({"num": [1.0, 1.0, 1.0, 1.0]}) + y = make_series(make_df, [1, 2, 3, 4]) + + estimator = BaseTargetMeanEstimator(bins=2).fit(X, y) + + assert estimator.encoder_dict_ == {"num": {"(-inf, 1.0]": 2.5}} + assert estimator._predict(X).tolist() == pytest.approx([2.5] * 4) + + +def test_error_if_predict_df_has_unseen_category(make_df): + X = make_df(DATA) + y = make_series(make_df, TARGET) + estimator = BaseTargetMeanEstimator(bins=2).fit(X, y) + + msg = "During the encoding, NaN values were introduced in the feature(s) cat." + with pytest.raises(ValueError, match=re.escape(msg)): + estimator._predict(make_df({"num": [1, 2], "cat": ["a", "z"]})) + + +def test_error_if_predict_df_has_values_in_bins_empty_in_train(make_df): + # the bins between 2 and 8 have no observations in the train set + X = make_df({"num": [0, 1, 2, 9, 10]}) + y = make_series(make_df, [1, 2, 3, 4, 5]) + estimator = BaseTargetMeanEstimator(bins=5).fit(X, y) + + msg = "During the encoding, NaN values were introduced in the feature(s) num." + with pytest.raises(ValueError, match=re.escape(msg)): + estimator._predict(make_df({"num": [0, 5]})) + + +@pytest.mark.parametrize( + "variable, value", [("num", None), ("num", float("nan")), ("cat", None)] +) +def test_error_if_df_contains_na(make_df, variable, value): + data_na = {**DATA, variable: [value] + DATA[variable][1:]} + X = make_df(DATA) + y = make_series(make_df, TARGET) + + with pytest.raises(ValueError, match=re.escape(MSG_NA)): + BaseTargetMeanEstimator(bins=2).fit(make_df(data_na), y) + + estimator = BaseTargetMeanEstimator(bins=2).fit(X, y) + with pytest.raises(ValueError, match=re.escape(MSG_NA)): + estimator._predict(make_df(data_na)) + + +def test_error_if_df_contains_inf(make_df): + data_inf = {**DATA, "num": [INF] + DATA["num"][1:]} + X = make_df(DATA) + y = make_series(make_df, TARGET) + + with pytest.raises(ValueError, match=re.escape(MSG_INF)): + BaseTargetMeanEstimator(bins=2).fit(make_df(data_inf), y) + + estimator = BaseTargetMeanEstimator(bins=2).fit(X, y) + with pytest.raises(ValueError, match=re.escape(MSG_INF)): + estimator._predict(make_df(data_inf)) + + +def test_error_if_predict_df_has_different_number_of_columns(make_df): + X = make_df(DATA) + y = make_series(make_df, TARGET) + estimator = BaseTargetMeanEstimator(bins=2).fit(X, y) + + 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)): + estimator._predict(X[["num"]]) + + +def test_error_if_not_fitted(make_df): + msg = ( + "This BaseTargetMeanEstimator instance is not fitted yet. Call 'fit' with " + "appropriate arguments before using this estimator." + ) + with pytest.raises(NotFittedError, match=re.escape(msg)): + BaseTargetMeanEstimator()._predict(make_df(DATA)) + + +def test_integer_column_names(): + X = pd.DataFrame({0: DATA["num"], 1: DATA["cat"]}) + y = pd.Series(TARGET) + + estimator = BaseTargetMeanEstimator(bins=2).fit(X, y) + + assert estimator.variables_numerical_ == [0] + assert estimator.variables_categorical_ == [1] + assert estimator.binner_dict_ == {0: BINNER_DICT_NUM["equal_width"]} + assert estimator.encoder_dict_[0] == pytest.approx(ENC_DICT_NUM["equal_width"]) + assert estimator.encoder_dict_[1] == pytest.approx(ENC_DICT_CAT) + assert estimator._predict(X).tolist() == pytest.approx(PREDICTIONS["equal_width"]) + + +def test_pandas_index_is_ignored(): + index = [10, 3, 7, 0, 1, 8, 2, 9, 4, 6] + X = pd.DataFrame(DATA, index=index) + y = pd.Series(TARGET, index=index) + + estimator = BaseTargetMeanEstimator(bins=2).fit(X, y) + + assert estimator.encoder_dict_["cat"] == pytest.approx(ENC_DICT_CAT) + assert estimator._predict(X).tolist() == pytest.approx(PREDICTIONS["equal_width"]) + + +def test_category_dtype(): + # the unused category "d" gets no target mean + cat = pd.Categorical(DATA["cat"], categories=["a", "b", "c", "d"]) + X = pd.DataFrame({"cat": cat}) + y = pd.Series(TARGET) + + estimator = BaseTargetMeanEstimator().fit(X, y) + + assert estimator.encoder_dict_["cat"] == pytest.approx(ENC_DICT_CAT) + assert estimator._predict(X).tolist() == pytest.approx( + [13 / 3, 13 / 3, 4.0, 4.0, 4.0, 7.5, 7.5, 7.5, 7.5, 13 / 3] + ) diff --git a/tests/test_prediction/test_check_estimator_prediction.py b/tests/test_prediction/test_check_estimator_prediction.py index afe45db71..4a72aa71d 100644 --- a/tests/test_prediction/test_check_estimator_prediction.py +++ b/tests/test_prediction/test_check_estimator_prediction.py @@ -11,7 +11,6 @@ EqualFrequencyDiscretiser, EqualWidthDiscretiser, ) -from feature_engine.encoding import MeanEncoder from tests.estimator_checks.dataframe_for_checks import test_df from tests.estimator_checks.fit_functionality_checks import check_error_if_y_not_passed @@ -213,19 +212,9 @@ def test_attributes_upon_fitting(_strategy, _bins, estimator): assert transformer.strategy == _strategy if _strategy == "equal_width": - assert ( - type(transformer._pipeline.named_steps["discretiser"]) - is EqualWidthDiscretiser - ) + assert type(transformer._discretiser) is EqualWidthDiscretiser else: - assert ( - type(transformer._pipeline.named_steps["discretiser"]) - is EqualFrequencyDiscretiser - ) - - assert type(transformer._pipeline.named_steps["encoder_num"]) is MeanEncoder - - assert type(transformer._pipeline.named_steps["encoder_cat"]) is MeanEncoder + assert type(transformer._discretiser) is EqualFrequencyDiscretiser @pytest.mark.parametrize("estimator", _estimators) From 452985c47314187e6b046b55bb676c03a29f8563 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Sat, 19 Sep 2026 11:35:49 +0200 Subject: [PATCH 2/2] Migrate TargetMeanRegressor to narwhals, add polars support Drop the pandas import and type fit and predict for any narwhals-supported dataframe, with the target as a series, numpy array or list. Rewrite the tests to run on pandas and polars, and add a polars example to the docstring. Co-Authored-By: Claude Opus 5 --- .../_prediction/target_mean_regressor.py | 33 +- .../test_target_mean_regressor.py | 370 +++++++----------- 2 files changed, 161 insertions(+), 242 deletions(-) diff --git a/feature_engine/_prediction/target_mean_regressor.py b/feature_engine/_prediction/target_mean_regressor.py index 231d7c268..8f778b489 100644 --- a/feature_engine/_prediction/target_mean_regressor.py +++ b/feature_engine/_prediction/target_mean_regressor.py @@ -1,5 +1,7 @@ +from typing import List, Union + import numpy as np -import pandas as pd +from narwhals.typing import IntoDataFrame, IntoSeries from sklearn.base import RegressorMixin from sklearn.utils.multiclass import type_of_target @@ -15,7 +17,8 @@ class TargetMeanRegressor(RegressorMixin, BaseTargetMeanEstimator): bin for each variable. The final estimation is the average of the target mean values across variables. - The TargetMeanRegressor() takes both numerical and categorical variables as input. + The TargetMeanRegressor() works with pandas and polars dataframes, and takes both + numerical and categorical variables as input. For numerical variables, the values are first sorted into bins of equal-width or equal-frequency. Then, the mean target value is estimated for each bin. If the variables are categorical, the mean target value is estimated for each category. @@ -36,7 +39,7 @@ class TargetMeanRegressor(RegressorMixin, BaseTargetMeanEstimator): the values will be sorted. strategy: str, default='equal_width' - Whether the bins should of equal width ('equal_width') or equal frequency + Whether the bins should be of equal width ('equal_width') or equal frequency ('equal_frequency'). Attributes @@ -83,18 +86,32 @@ class TargetMeanRegressor(RegressorMixin, BaseTargetMeanEstimator): .. [1] Miller, et al. "Predicting customer behaviour: The University of Melbourne’s KDD Cup report". JMLR Workshop and Conference Proceeding. KDD 2009 http://proceedings.mlr.press/v7/miller09/miller09.pdf + + Examples + -------- + + >>> import polars as pl + >>> from feature_engine._prediction.target_mean_regressor import TargetMeanRegressor + >>> X = pl.DataFrame(dict(x1=[1, 2, 3, 4, 5, 6], x2=["a", "a", "b", "b", "b", "a"])) + >>> y = pl.Series([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + >>> tmr = TargetMeanRegressor(bins=2) + >>> tmr.fit(X, y) + >>> tmr.predict(X) + array([2.5, 2.5, 3. , 4.5, 4.5, 4. ]) + >>> tmr.score(X, y) + 0.6 """ - def fit(self, X: pd.DataFrame, y: pd.Series): + def fit(self, X: IntoDataFrame, y: Union[IntoSeries, np.ndarray, List]): """ Learn the mean target value per category or bin. Parameters ---------- - X : pandas dataframe of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The training input samples. - y : pandas series of shape = [n_samples,] + y: Series, numpy array or list of shape = [n_samples,] The target variable. """ @@ -106,13 +123,13 @@ def fit(self, X: pd.DataFrame, y: pd.Series): return super().fit(X, y) - def predict(self, X: pd.DataFrame) -> np.ndarray: + def predict(self, X: IntoDataFrame) -> np.ndarray: """ Predict using the average of the target mean value across variables. Parameters ---------- - X : pandas dataframe of shape = [n_samples, ] + X: dataframe of shape = [n_samples, n_features] The input samples. Returns diff --git a/tests/test_prediction/test_target_mean_regressor.py b/tests/test_prediction/test_target_mean_regressor.py index f32792279..be4848133 100644 --- a/tests/test_prediction/test_target_mean_regressor.py +++ b/tests/test_prediction/test_target_mean_regressor.py @@ -1,250 +1,152 @@ +import re + import numpy as np +import pandas as pd import pytest +from sklearn.exceptions import NotFittedError from feature_engine._prediction.target_mean_regressor import TargetMeanRegressor +from tests.backend_helpers import make_series + +DATA = { + "cat_var_A": ["A"] * 5 + ["B"] * 5 + ["C"] * 5 + ["D"] * 5, + "cat_var_B": ["A"] * 6 + ["B"] * 2 + ["C"] * 2 + ["B"] * 2 + ["C"] * 2 + ["D"] * 6, + "num_var_A": [1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 4, 4, 4, 4, 4], + "num_var_B": [1, 1, 1, 1, 1, 1, 2, 2, 3, 3, 2, 2, 3, 3, 4, 4, 4, 4, 4, 4], +} +TARGET = [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3] +# with 2 bins, the numerical variables are split at 2.5. +PREDICTIONS_ALL_VARIABLES = ( + [41 / 120] * 5 + + [71 / 120, 0.925, 0.925, 1.325, 1.325, 1.675, 1.675, 2.075, 2.075, 289 / 120] + + [319 / 120] * 5 +) +R2_ALL_VARIABLES = 5533 / 6000 +MSG_BINARY = ( + "Trying to fit a regression to a binary target is not " + "allowed by this transformer. " +) + + +# init parameters +@pytest.mark.parametrize("bins", [0, -1, 2.5, "3", None, [5]]) +def test_error_if_bins_not_positive_integer(bins): + msg = f"bins must be a positive integer. Got {bins} instead." + with pytest.raises(ValueError, match=re.escape(msg)): + TargetMeanRegressor(bins=bins) + + +@pytest.mark.parametrize( + "strategy", ["arbitrary", "Equal_width", "", 1, None, ["equal_width"]] +) +def test_error_if_strategy_not_permitted(strategy): + msg = ( + "strategy takes only values 'equal_width' or 'equal_frequency'. " + f"Got {strategy} instead." + ) + with pytest.raises(ValueError, match=re.escape(msg)): + TargetMeanRegressor(strategy=strategy) -def test_regressor_categorical_variables(df_regression): - - X, y = df_regression - - tr = TargetMeanRegressor(variables="cat_var_A") - tr.fit(X, y) - pred = tr.predict(X) - - exp_pred = np.array( - [ - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 2.0, - 2.0, - 2.0, - 2.0, - 2.0, - 3.0, - 3.0, - 3.0, - 3.0, - 3.0, - ] - ) +@pytest.mark.parametrize( + "bins, strategy", [(1, "equal_width"), (5, "equal_frequency"), (10, "equal_width")] +) +def test_init_param_assignment(bins, strategy): + estimator = TargetMeanRegressor(bins=bins, strategy=strategy) + assert estimator.bins == bins + assert estimator.strategy == strategy - assert np.array_equal(pred, exp_pred) - - tr = TargetMeanRegressor(variables="cat_var_B") - tr.fit(X, y) - - pred = tr.predict(X) - - exp_pred = np.array( - [ - 0.16666667, - 0.16666667, - 0.16666667, - 0.16666667, - 0.16666667, - 0.16666667, - 1.5, - 1.5, - 1.5, - 1.5, - 1.5, - 1.5, - 1.5, - 1.5, - 2.83333333, - 2.83333333, - 2.83333333, - 2.83333333, - 2.83333333, - 2.83333333, - ] - ) - assert np.allclose(pred, exp_pred) - - tr = TargetMeanRegressor(variables=["cat_var_A", "cat_var_B"]) - tr.fit(X, y) - - pred = tr.predict(X) - - exp_pred = np.array( - [ - 0.08333333, - 0.08333333, - 0.08333333, - 0.08333333, - 0.08333333, - 0.58333333, - 1.25, - 1.25, - 1.25, - 1.25, - 1.75, - 1.75, - 1.75, - 1.75, - 2.41666667, - 2.91666667, - 2.91666667, - 2.91666667, - 2.91666667, - 2.91666667, - ] - ) +# fit and predict +@pytest.mark.parametrize( + "variables, expected", + [ + ("cat_var_A", [0.0] * 5 + [1.0] * 5 + [2.0] * 5 + [3.0] * 5), + ("cat_var_B", [1 / 6] * 6 + [1.5] * 8 + [17 / 6] * 6), + ( + ["cat_var_A", "cat_var_B"], + [1 / 12] * 5 + + [7 / 12, 1.25, 1.25, 1.25, 1.25, 1.75, 1.75, 1.75, 1.75, 29 / 12] + + [35 / 12] * 5, + ), + ], +) +def test_predict_categorical_variables(make_df, variables, expected): + X = make_df(DATA) + y = make_series(make_df, TARGET) - assert np.allclose(pred, exp_pred) - - -def test_classifier_numerical_variables(df_regression): - - X, y = df_regression - - tr = TargetMeanRegressor(variables="num_var_A", bins=2) - tr.fit(X, y) - - pred = tr.predict(X) - - exp_pred = np.array( - [ - 0.5, - 0.5, - 0.5, - 0.5, - 0.5, - 0.5, - 0.5, - 0.5, - 0.5, - 0.5, - 2.5, - 2.5, - 2.5, - 2.5, - 2.5, - 2.5, - 2.5, - 2.5, - 2.5, - 2.5, - ] - ) + y_pred = TargetMeanRegressor(variables=variables).fit(X, y).predict(X) - assert np.array_equal(pred, exp_pred) - - tr = TargetMeanRegressor(variables="num_var_B", bins=2) - tr.fit(X, y) - - pred = tr.predict(X) - - exp_pred = np.array( - [ - 0.7, - 0.7, - 0.7, - 0.7, - 0.7, - 0.7, - 0.7, - 0.7, - 2.3, - 2.3, - 0.7, - 0.7, - 2.3, - 2.3, - 2.3, - 2.3, - 2.3, - 2.3, - 2.3, - 2.3, - ] - ) + assert isinstance(y_pred, np.ndarray) + assert y_pred.tolist() == pytest.approx(expected) + + +@pytest.mark.parametrize( + "variables, expected", + [ + ("num_var_A", [0.5] * 10 + [2.5] * 10), + ("num_var_B", [0.7] * 8 + [2.3] * 2 + [0.7] * 2 + [2.3] * 8), + ( + ["num_var_A", "num_var_B"], + [0.6] * 8 + [1.4, 1.4, 1.6, 1.6] + [2.4] * 8, + ), + ], +) +def test_predict_numerical_variables(make_df, variables, expected): + X = make_df(DATA) + y = make_series(make_df, TARGET) + + y_pred = TargetMeanRegressor(variables=variables, bins=2).fit(X, y).predict(X) + + assert isinstance(y_pred, np.ndarray) + assert y_pred.tolist() == pytest.approx(expected) - np.array_equal(pred, exp_pred) - - tr = TargetMeanRegressor(variables=["num_var_A", "num_var_B"], bins=2) - tr.fit(X, y) - - pred = tr.predict(X) - - exp_pred = np.array( - [ - 0.6, - 0.6, - 0.6, - 0.6, - 0.6, - 0.6, - 0.6, - 0.6, - 1.4, - 1.4, - 1.6, - 1.6, - 2.4, - 2.4, - 2.4, - 2.4, - 2.4, - 2.4, - 2.4, - 2.4, - ] - ) - assert np.array_equal(pred, exp_pred) - - -def test_classifier_all_variables(df_regression): - - X, y = df_regression - - tr = TargetMeanRegressor(bins=2) - tr.fit(X, y) - - pred = tr.predict(X) - - exp_pred = np.array( - [ - 0.34166667, - 0.34166667, - 0.34166667, - 0.34166667, - 0.34166667, - 0.59166667, - 0.925, - 0.925, - 1.325, - 1.325, - 1.675, - 1.675, - 2.075, - 2.075, - 2.40833333, - 2.65833333, - 2.65833333, - 2.65833333, - 2.65833333, - 2.65833333, - ] +def test_predict_and_score_all_variables(make_df): + X = make_df(DATA) + y = make_series(make_df, TARGET) + + estimator = TargetMeanRegressor(bins=2).fit(X, y) + + assert estimator.predict(X).tolist() == pytest.approx(PREDICTIONS_ALL_VARIABLES) + assert estimator.score(X, y) == pytest.approx(R2_ALL_VARIABLES) + + +@pytest.mark.parametrize("to_target", [list, np.array]) +def test_target_as_list_or_array(make_df, to_target): + X = make_df(DATA) + y = to_target(TARGET) + + estimator = TargetMeanRegressor(bins=2).fit(X, y) + + assert estimator.predict(X).tolist() == pytest.approx(PREDICTIONS_ALL_VARIABLES) + assert estimator.score(X, y) == pytest.approx(R2_ALL_VARIABLES) + + +@pytest.mark.parametrize("target", [[0, 1] * 10, [1.0, 2.0] * 10, ["a", "b"] * 10]) +def test_error_if_target_is_binary(make_df, target): + X = make_df(DATA) + y = make_series(make_df, target) + + with pytest.raises(ValueError, match=re.escape(MSG_BINARY)): + TargetMeanRegressor().fit(X, y) + + +def test_error_if_not_fitted(make_df): + msg = ( + "This TargetMeanRegressor instance is not fitted yet. Call 'fit' with " + "appropriate arguments before using this estimator." ) + with pytest.raises(NotFittedError, match=re.escape(msg)): + TargetMeanRegressor().predict(make_df(DATA)) + - assert np.allclose(pred, exp_pred) +def test_integer_column_names(): + X = pd.DataFrame(DATA) + X.columns = [0, 1, 2, 3] + y = pd.Series(TARGET) + estimator = TargetMeanRegressor(bins=2).fit(X, y) -def test_error_when_y_is_binary(df_regression): - X, y = df_regression - y = [1.0, 2.0] - tr = TargetMeanRegressor(bins=2) - with pytest.raises(ValueError): - tr.fit(X, y) + assert estimator.predict(X).tolist() == pytest.approx(PREDICTIONS_ALL_VARIABLES) + assert estimator.score(X, y) == pytest.approx(R2_ALL_VARIABLES)