Migrate DropCorrelatedFeatures to narwhals, add polars support

docs/user_guide/selection/DropCorrelatedFeatures.rst

@@ -6,12 +6,14 @@ DropCorrelatedFeatures
 ======================
 
 The :class:`DropCorrelatedFeatures()` finds and removes correlated variables from a dataframe.
-Correlation is calculated with `pandas.corr()`. All correlation methods supported by `pandas.corr()`
-can be used in the selection, including Spearman, Kendall, or Spearman. You can also pass a
-bespoke correlation function, provided it returns a value between -1 and 1.
+It supports the correlation methods of `pandas.corr()`: Pearson, Spearman, and Kendall. You can
+also pass a bespoke correlation function, provided it returns a value between -1 and 1. As in
+`pandas.corr()`, each pair of variables is compared using the rows where both have values.
+
+:class:`DropCorrelatedFeatures()` works with pandas and polars dataframes.
 
 Features are removed on first found first removed basis, without any further insight. That is,
-the first feature will be retained an all subsequent features that are correlated with this, will
+the first feature will be retained and all subsequent features that are correlated with this, will
 be removed.
 
 The transformer will examine all numerical variables automatically. Note that you could pass a
@@ -95,7 +97,7 @@ to `var_0` and will therefore be removed.
 
    {'var_0': {'var_8'}, 'var_4': {'var_6', 'var_7', 'var_9'}}
 
-Similarly, `var_4` is a key and will be retained, whereas the variables 6, 7 and 8 were
+Similarly, `var_4` is a key and will be retained, whereas the variables 6, 7 and 9 were
 found correlated to `var_4` and will therefore be removed.
 
 The features that will be removed from the dataset are stored in a different attribute
@@ -137,6 +139,50 @@ Below we see the resulting dataframe:
     4 -0.186530
 
 
+With polars
+-----------
+
+:class:`DropCorrelatedFeatures()` works in the same way with a polars dataframe, and returns
+a polars dataframe:
+
+.. code:: python
+
+    import polars as pl
+    from sklearn.datasets import make_classification
+    from feature_engine.selection import DropCorrelatedFeatures
+
+    X, y = make_classification(n_samples=1000,
+                               n_features=12,
+                               n_redundant=4,
+                               n_clusters_per_class=1,
+                               weights=[0.50],
+                               class_sep=2,
+                               random_state=1)
+
+    X = pl.DataFrame(X, schema=['var_' + str(i) for i in range(12)])
+
+    tr = DropCorrelatedFeatures(threshold=0.8)
+    Xt = tr.fit_transform(X)
+
+    print(tr.features_to_drop_)
+
+The same features are found correlated as with pandas:
+
+.. code:: python
+
+    ['var_8', 'var_6', 'var_7', 'var_9']
+
+The transformed data is a polars dataframe without those features:
+
+.. code:: python
+
+    print(Xt.columns)
+
+.. code:: python
+
+    ['var_0', 'var_1', 'var_2', 'var_3', 'var_4', 'var_5', 'var_10', 'var_11']
+
+
 Additional resources
 --------------------
 

feature_engine/selection/base_recursive_selector.py

@@ -1,22 +1,23 @@
 from types import GeneratorType
-from typing import List, Union
+from typing import List, Tuple, Union
 
-import pandas as pd
+import narwhals as nw
+import numpy as np
+from narwhals.typing import IntoDataFrame, IntoSeries
 from sklearn.inspection import permutation_importance
 from sklearn.model_selection import cross_validate
 
 from feature_engine._check_init_parameters.check_variables import (
     _check_variables_input_value,
 )
 from feature_engine.dataframe_checks import check_X_y
-from feature_engine.selection.base_selection_functions import get_feature_importances
+from feature_engine.selection.base_selection_functions import (
+    _importance_series,
+    _select_numerical_variables,
+    get_feature_importances,
+)
 from feature_engine.selection.base_selector import BaseSelector
 from feature_engine.tags import _return_tags
-from feature_engine.variable_handling import (
-    check_numerical_variables,
-    find_numerical_variables,
-    retain_variables_if_in_df,
-)
 
 Variables = Union[None, int, str, List[Union[str, int]]]
 
@@ -81,10 +82,13 @@ class BaseRecursiveSelector(BaseSelector):
         Performance of the model trained using the original dataset.
 
     feature_importances_:
-        Pandas Series with the feature importance (comes from step 2)
+        The feature importance (comes from step 2). A pandas Series with the
+        features as index when X is a pandas dataframe, and a dictionary with the
+        features as keys otherwise.
 
     feature_importances_std_:
-        Pandas Series with the standard deviation of the feature importance.
+        The standard deviation of the feature importance, as a pandas Series or a
+        dictionary, like `feature_importances_`.
 
     features_to_drop_:
         List with the features to remove from the dataset.
@@ -116,7 +120,9 @@ def __init__(
     ):
 
         if not isinstance(threshold, (int, float)):
-            raise ValueError("threshold can only be integer or float")
+            raise ValueError(
+                f"threshold must be an integer or a float. Got {threshold} instead."
+            )
 
         super().__init__(confirm_variables)
         self.variables = _check_variables_input_value(variables)
@@ -126,84 +132,78 @@ def __init__(
         self.cv = cv
         self.groups = groups
 
-    def fit(self, X: pd.DataFrame, y: pd.Series):
+    def fit(self, X: IntoDataFrame, y: IntoSeries) -> Tuple[nw.DataFrame, IntoSeries]:
         """
         Find initial model performance. Sort features by importance.
 
         Parameters
         ----------
-        X: pandas dataframe of shape = [n_samples, n_features]
+        X: dataframe of shape = [n_samples, n_features]
            The input dataframe
 
         y: array-like of shape (n_samples)
            Target variable. Required to train the estimator.
-        """
 
-        # check input dataframe
-        X, y = check_X_y(X, y)
+        Returns
+        -------
+        nw_X: narwhals dataframe
+            The input dataframe, as a narwhals dataframe.
 
-        if self.variables is None:
-            self.variables_ = find_numerical_variables(X)
-        else:
-            if self.confirm_variables is True:
-                variables_ = retain_variables_if_in_df(X, self.variables)
-                self.variables_ = check_numerical_variables(X, variables_)
-            else:
-                self.variables_ = check_numerical_variables(X, self.variables)
+        y: Series or numpy array
+            The target, checked.
+        """
+        nw_X, y = check_X_y(X, y)
+
+        self.variables_ = _select_numerical_variables(
+            X, self.variables, self.confirm_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()
 
-        # save input features
         self._get_feature_names_in(X)
 
+        X_model = nw_X.select(nw.col(*self.variables_)).to_native()
+
         # train model with all features and cross-validation
         model = cross_validate(
             estimator=self.estimator,
-            X=X[self.variables_],
+            X=X_model,
             y=y,
             cv=self._cv,
             groups=self.groups,
             scoring=self.scoring,
             return_estimator=True,
         )
 
-        # store initial model performance
         self.initial_model_performance_ = model["test_score"].mean()
 
-        # Initialize a dataframe that will contain the list of the feature/coeff
-        # importance for each cross validation fold
-        feature_importances_cv = pd.DataFrame()
-
-        # Populate the feature_importances_cv dataframe with columns containing
-        # the feature importance values for each model returned by the cross
-        # validation.
-        # There are as many columns as folds.
-        for i in range(len(model["estimator"])):
-            m = model["estimator"][i]
-
+        # one row of feature importance per cross-validation fold
+        importances = []
+        for m in model["estimator"]:
             if hasattr(m, "feature_importances_") or hasattr(m, "coef_"):
-                feature_importances_cv[i] = get_feature_importances(m)
+                importances.append(get_feature_importances(m))
             else:
                 r = permutation_importance(
                     m,
-                    X[self.variables_],
+                    X_model,
                     y,
                     n_repeats=1,
                     random_state=10,
                 )
-                feature_importances_cv[i] = r.importances_mean
+                importances.append(r.importances_mean)
+        importances_arr = np.array(importances)
 
-        # Add the variables as index to feature_importances_cv
-        feature_importances_cv.index = self.variables_
-
-        # Aggregate the feature importance returned in each fold
-        self.feature_importances_ = feature_importances_cv.mean(axis=1)
-        self.feature_importances_std_ = feature_importances_cv.std(axis=1)
+        self.feature_importances_ = _importance_series(
+            X, self.variables_, importances_arr.mean(axis=0)
+        )
+        self.feature_importances_std_ = _importance_series(
+            X, self.variables_, importances_arr.std(axis=0, ddof=1)
+        )
 
-        return X, y
+        return nw_X, y
 
     def _more_tags(self):
         tags_dict = _return_tags()