11from 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
57from sklearn .inspection import permutation_importance
68from sklearn .model_selection import cross_validate
79
810from feature_engine ._check_init_parameters .check_variables import (
911 _check_variables_input_value ,
1012)
1113from 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+ )
1319from feature_engine .selection .base_selector import BaseSelector
1420from 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
2122Variables = 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