Skip to content

Commit cd7f30b

Browse files
solegalliclaude
andcommitted
Migrate the selection base classes and helpers to narwhals, add polars support
BaseSelector.transform() returns the retained features in the train set order, in the same library as the input (pandas X[features], narwhals select otherwise). BaseRecursiveSelector.fit() trains the estimators on native frames and returns (nw_X, y). The helpers in base_selection_functions no longer import pandas: correlations are computed with numpy (np.corrcoef, or matrix products for pairwise complete observations when there are missing values), and feature importances are pandas Series for pandas input and dicts otherwise. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
1 parent 7c22dfa commit cd7f30b

7 files changed

Lines changed: 924 additions & 393 deletions

File tree

‎feature_engine/selection/base_recursive_selector.py‎

Lines changed: 48 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,23 @@
11
from types import GeneratorType
2-
from typing import List, Union
2+
from typing import List, Tuple, Union
33

4-
import pandas as pd
4+
import narwhals as nw
5+
import numpy as np
6+
from narwhals.typing import IntoDataFrame, IntoSeries
57
from sklearn.inspection import permutation_importance
68
from sklearn.model_selection import cross_validate
79

810
from feature_engine._check_init_parameters.check_variables import (
911
_check_variables_input_value,
1012
)
1113
from feature_engine.dataframe_checks import check_X_y
12-
from feature_engine.selection.base_selection_functions import get_feature_importances
14+
from feature_engine.selection.base_selection_functions import (
15+
_importance_series,
16+
_select_numerical_variables,
17+
get_feature_importances,
18+
)
1319
from feature_engine.selection.base_selector import BaseSelector
1420
from feature_engine.tags import _return_tags
15-
from feature_engine.variable_handling import (
16-
check_numerical_variables,
17-
find_numerical_variables,
18-
retain_variables_if_in_df,
19-
)
2021

2122
Variables = Union[None, int, str, List[Union[str, int]]]
2223

@@ -81,10 +82,13 @@ class BaseRecursiveSelector(BaseSelector):
8182
Performance of the model trained using the original dataset.
8283
8384
feature_importances_:
84-
Pandas Series with the feature importance (comes from step 2)
85+
The feature importance (comes from step 2). A pandas Series with the
86+
features as index when X is a pandas dataframe, and a dictionary with the
87+
features as keys otherwise.
8588
8689
feature_importances_std_:
87-
Pandas Series with the standard deviation of the feature importance.
90+
The standard deviation of the feature importance, as a pandas Series or a
91+
dictionary, like `feature_importances_`.
8892
8993
features_to_drop_:
9094
List with the features to remove from the dataset.
@@ -116,7 +120,9 @@ def __init__(
116120
):
117121

118122
if not isinstance(threshold, (int, float)):
119-
raise ValueError("threshold can only be integer or float")
123+
raise ValueError(
124+
f"threshold must be an integer or a float. Got {threshold} instead."
125+
)
120126

121127
super().__init__(confirm_variables)
122128
self.variables = _check_variables_input_value(variables)
@@ -126,84 +132,78 @@ def __init__(
126132
self.cv = cv
127133
self.groups = groups
128134

129-
def fit(self, X: pd.DataFrame, y: pd.Series):
135+
def fit(self, X: IntoDataFrame, y: IntoSeries) -> Tuple[nw.DataFrame, IntoSeries]:
130136
"""
131137
Find initial model performance. Sort features by importance.
132138
133139
Parameters
134140
----------
135-
X: pandas dataframe of shape = [n_samples, n_features]
141+
X: dataframe of shape = [n_samples, n_features]
136142
The input dataframe
137143
138144
y: array-like of shape (n_samples)
139145
Target variable. Required to train the estimator.
140-
"""
141146
142-
# check input dataframe
143-
X, y = check_X_y(X, y)
147+
Returns
148+
-------
149+
nw_X: narwhals dataframe
150+
The input dataframe, as a narwhals dataframe.
144151
145-
if self.variables is None:
146-
self.variables_ = find_numerical_variables(X)
147-
else:
148-
if self.confirm_variables is True:
149-
variables_ = retain_variables_if_in_df(X, self.variables)
150-
self.variables_ = check_numerical_variables(X, variables_)
151-
else:
152-
self.variables_ = check_numerical_variables(X, self.variables)
152+
y: Series or numpy array
153+
The target, checked.
154+
"""
155+
nw_X, y = check_X_y(X, y)
156+
157+
self.variables_ = _select_numerical_variables(
158+
X, self.variables, self.confirm_variables
159+
)
153160

154161
self._cv = list(self.cv) if isinstance(self.cv, GeneratorType) else self.cv
155162

156163
# check that there are more than 1 variable to select from
157164
self._check_variable_number()
158165

159-
# save input features
160166
self._get_feature_names_in(X)
161167

168+
X_model = nw_X.select(nw.col(*self.variables_)).to_native()
169+
162170
# train model with all features and cross-validation
163171
model = cross_validate(
164172
estimator=self.estimator,
165-
X=X[self.variables_],
173+
X=X_model,
166174
y=y,
167175
cv=self._cv,
168176
groups=self.groups,
169177
scoring=self.scoring,
170178
return_estimator=True,
171179
)
172180

173-
# store initial model performance
174181
self.initial_model_performance_ = model["test_score"].mean()
175182

176-
# Initialize a dataframe that will contain the list of the feature/coeff
177-
# importance for each cross validation fold
178-
feature_importances_cv = pd.DataFrame()
179-
180-
# Populate the feature_importances_cv dataframe with columns containing
181-
# the feature importance values for each model returned by the cross
182-
# validation.
183-
# There are as many columns as folds.
184-
for i in range(len(model["estimator"])):
185-
m = model["estimator"][i]
186-
183+
# one row of feature importance per cross-validation fold
184+
importances = []
185+
for m in model["estimator"]:
187186
if hasattr(m, "feature_importances_") or hasattr(m, "coef_"):
188-
feature_importances_cv[i] = get_feature_importances(m)
187+
importances.append(get_feature_importances(m))
189188
else:
190189
r = permutation_importance(
191190
m,
192-
X[self.variables_],
191+
X_model,
193192
y,
194193
n_repeats=1,
195194
random_state=10,
196195
)
197-
feature_importances_cv[i] = r.importances_mean
196+
importances.append(r.importances_mean)
197+
importances_arr = np.array(importances)
198198

199-
# Add the variables as index to feature_importances_cv
200-
feature_importances_cv.index = self.variables_
201-
202-
# Aggregate the feature importance returned in each fold
203-
self.feature_importances_ = feature_importances_cv.mean(axis=1)
204-
self.feature_importances_std_ = feature_importances_cv.std(axis=1)
199+
self.feature_importances_ = _importance_series(
200+
X, self.variables_, importances_arr.mean(axis=0)
201+
)
202+
self.feature_importances_std_ = _importance_series(
203+
X, self.variables_, importances_arr.std(axis=0, ddof=1)
204+
)
205205

206-
return X, y
206+
return nw_X, y
207207

208208
def _more_tags(self):
209209
tags_dict = _return_tags()

0 commit comments

Comments
 (0)