Skip to content

Migrate TargetMeanRegressor to narwhals, add polars support #1069

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
solegalli wants to merge 2 commits into
base: narwhals-migration
Choose a base branch
from
Open

Migrate TargetMeanRegressor to narwhals, add polars support #1069

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9055577
Migrate BaseTargetMeanEstimator to narwhals, add polars support
solegalli Sep 19, 2026
452985c
Migrate TargetMeanRegressor to narwhals, add polars support
solegalli Sep 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
297 changes: 136 additions & 161 deletions feature_engine/_prediction/base_predictor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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 (
Expand All @@ -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

Expand All @@ -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
Expand Down Expand Up @@ -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."
Expand All @@ -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()
Expand Down
Loading