fix cv generator recursive selectors (#805)

Author: solegalli

feature_engine/selection/base_recursive_selector.py

@@ -1,3 +1,4 @@
+from types import GeneratorType
 from typing import List, Union
 
 import pandas as pd
@@ -144,6 +145,8 @@ def fit(self, X: pd.DataFrame, y: pd.Series):
             else:
                 self.variables_ = check_numerical_variables(X, self.variables)
 
+        self._cv = list(self.cv) if isinstance(self.cv, GeneratorType) else self.cv
+
         # check that there are more than 1 variable to select from
         self._check_variable_number()
 
@@ -155,7 +158,7 @@ def fit(self, X: pd.DataFrame, y: pd.Series):
             self.estimator,
             X[self.variables_],
             y,
-            cv=self.cv,
+            cv=self._cv,
             scoring=self.scoring,
             return_estimator=True,
         )

feature_engine/selection/recursive_feature_addition.py

@@ -170,7 +170,7 @@ def fit(self, X: pd.DataFrame, y: pd.Series):
             self.estimator,
             X[first_most_important_feature].to_frame(),
             y,
-            cv=self.cv,
+            cv=self._cv,
             scoring=self.scoring,
             return_estimator=True,
         )
@@ -197,7 +197,7 @@ def fit(self, X: pd.DataFrame, y: pd.Series):
                 self.estimator,
                 X[_selected_features + [feature]],
                 y,
-                cv=self.cv,
+                cv=self._cv,
                 scoring=self.scoring,
                 return_estimator=True,
             )

feature_engine/selection/recursive_feature_elimination.py

@@ -189,7 +189,7 @@ def fit(self, X: pd.DataFrame, y: pd.Series):
                 self.estimator,
                 X_tmp.drop(columns=feature),
                 y,
-                cv=self.cv,
+                cv=self._cv,
                 scoring=self.scoring,
                 return_estimator=False,
             )
@@ -216,7 +216,7 @@ def fit(self, X: pd.DataFrame, y: pd.Series):
                     self.estimator,
                     X_tmp,
                     y,
-                    cv=self.cv,
+                    cv=self._cv,
                     return_estimator=False,
                     scoring=self.scoring,
                 )

tests/test_selection/test_recursive_feature_addition.py

@@ -2,6 +2,7 @@
 import pytest
 from sklearn.ensemble import RandomForestClassifier
 from sklearn.linear_model import Lasso, LogisticRegression, LinearRegression
+from sklearn.model_selection import KFold
 from sklearn.tree import DecisionTreeRegressor
 
 from feature_engine.selection import RecursiveFeatureAddition
@@ -254,3 +255,23 @@ def test_feature_importance(load_diabetes_dataset):
 
     assert round(sel.feature_importances_, 2).to_list() == imps
     assert round(sel.feature_importances_std_, 2).to_list() == imps_std
+
+
+def test_cv_generator(load_diabetes_dataset):
+    X, y = load_diabetes_dataset
+    linear_model = LinearRegression()
+    cv = KFold(n_splits=3)
+    sel = RecursiveFeatureAddition(estimator=linear_model, scoring="r2", cv=3).fit(X, y)
+    expected = sel.transform(X)
+
+    sel = RecursiveFeatureAddition(estimator=linear_model, scoring="r2", cv=cv).fit(
+        X, y
+    )
+    test1 = sel.transform(X)
+    pd.testing.assert_frame_equal(expected, test1)
+
+    sel = RecursiveFeatureAddition(
+        estimator=linear_model, scoring="r2", cv=cv.split(X, y)
+    ).fit(X, y)
+    test2 = sel.transform(X)
+    pd.testing.assert_frame_equal(expected, test2)

tests/test_selection/test_recursive_feature_elimination.py

@@ -2,6 +2,7 @@
 import pytest
 from sklearn.ensemble import RandomForestClassifier
 from sklearn.linear_model import Lasso, LinearRegression, LogisticRegression
+from sklearn.model_selection import KFold
 from sklearn.tree import DecisionTreeRegressor
 
 from feature_engine.selection import RecursiveFeatureElimination
@@ -276,3 +277,25 @@ def test_feature_importance(load_diabetes_dataset):
 
     assert round(sel.feature_importances_, 2).to_list() == imps
     assert round(sel.feature_importances_std_, 2).to_list() == imps_std
+
+
+def test_cv_generator(load_diabetes_dataset):
+    X, y = load_diabetes_dataset
+    linear_model = LinearRegression()
+    cv = KFold(n_splits=3)
+    sel = RecursiveFeatureElimination(estimator=linear_model, scoring="r2", cv=3).fit(
+        X, y
+    )
+    expected = sel.transform(X)
+
+    sel = RecursiveFeatureElimination(estimator=linear_model, scoring="r2", cv=cv).fit(
+        X, y
+    )
+    test1 = sel.transform(X)
+    pd.testing.assert_frame_equal(expected, test1)
+
+    sel = RecursiveFeatureElimination(
+        estimator=linear_model, scoring="r2", cv=cv.split(X, y)
+    ).fit(X, y)
+    test2 = sel.transform(X)
+    pd.testing.assert_frame_equal(expected, test2)