Migrate DropHighPSIFeatures to narwhals, add polars support

docs/user_guide/selection/DropHighPSIFeatures.rst

@@ -6,7 +6,7 @@ DropHighPSIFeatures
 ===================
 
 The :class:`DropHighPSIFeatures()` finds and removes features with changes in their
-distribution, i.e. "unstable values", from a pandas dataframe.
+distribution, i.e. "unstable values", from a pandas or polars dataframe.
 The stability of the distribution is computed using the **Population Stability
 Index (PSI)** and all features having a PSI value above a given threshold are removed.
 
@@ -207,6 +207,13 @@ reference in the `split_col` parameter. If you don't, the split will be done bas
 values of the dataframe index. This might be a good option if the index contains meaningful
 values or if splitting just based on `split_frac`.
 
+Polars dataframes don't have an index. If you don't indicate a column in `split_col`,
+:class:`DropHighPSIFeatures()` uses the position of the rows instead, from 0 for the
+first row to the number of rows minus 1 for the last one. This way, the first rows go
+to the basis data set and the last rows to the test set. If the rows are in time order,
+this compares the older observations with the most recent ones. The positions are also
+what the `cut_off` refers to: with `cut_off=99`, the first 100 rows go to the basis set.
+
 
 Python implementation
 ----------------------
@@ -255,7 +262,7 @@ with a shift in its distribution (*var_3* in this case):
     X = pd.DataFrame(X, columns=colnames)
 
     # Add a column with a shift.
-    X['var_3'][250:] = X['var_3'][250:] + 1
+    X.loc[250:, 'var_3'] = X.loc[250:, 'var_3'] + 1
 
 The default approach in :class:`DropHighPSIFeatures()` is to split the
 input dataframe `X` in two equally sized data sets. You can adjust the proportions by changing
@@ -295,7 +302,7 @@ by the `transform` method.
     {'var_0': 0.07405459925568803,
     'var_1': 0.09124093185820083,
     'var_2': 0.16985790067687764,
-    'var_3': 1.342485289730313,
+    'var_3': 0.8324552340063687,
     'var_4': 0.0743442762545251,
     'var_5': 0.06809060587241555}
 
@@ -942,6 +949,78 @@ noticeably between groups, whereas the distribution of *var_4* remains stable:
 .. image:: ../../images/PSI_distribution_case5.png
 
 
+With polars
+-----------
+
+:class:`DropHighPSIFeatures()` works in the same way with polars dataframes, and returns
+a polars dataframe. Let's repeat Case 1 with a polars dataframe:
+
+.. code:: python
+
+    import polars as pl
+    from sklearn.datasets import make_classification
+    from feature_engine.selection import DropHighPSIFeatures
+
+    X, y = make_classification(n_samples=500, n_features=6, random_state=0)
+
+    # Add a shift to the last 250 values of var_3.
+    X[250:, 3] = X[250:, 3] + 1
+    X = pl.DataFrame(X, schema=["var_" + str(i) for i in range(6)])
+
+    transformer = DropHighPSIFeatures(split_frac=0.6)
+    X_transformed = transformer.fit_transform(X)
+
+    print(transformer.psi_values_)
+
+We obtain the same PSI values as with pandas:
+
+.. code:: python
+
+    {'var_0': 0.07405459925568803,
+    'var_1': 0.09124093185820083,
+    'var_2': 0.16985790067687764,
+    'var_3': 0.8324552340063687,
+    'var_4': 0.0743442762545251,
+    'var_5': 0.06809060587241555}
+
+As the polars dataframe has no index, the data was split by the position of the rows.
+The cut-off is the same as in Case 1:
+
+.. code:: python
+
+    print(transformer.cut_off_)
+
+.. code:: python
+
+    299.4
+
+This means that the first 300 rows, with positions 0 to 299, form the basis data set.
+The transformed data is a polars dataframe without *var_3*:
+
+.. code:: python
+
+    print(X_transformed.head())
+
+.. code:: text
+
+    shape: (5, 5)
+    ┌───────────┬───────────┬───────────┬───────────┬───────────┐
+    │ var_0     ┆ var_1     ┆ var_2     ┆ var_4     ┆ var_5     │
+    │ ---       ┆ ---       ┆ ---       ┆ ---       ┆ ---       │
+    │ f64       ┆ f64       ┆ f64       ┆ f64       ┆ f64       │
+    ╞═══════════╪═══════════╪═══════════╪═══════════╪═══════════╡
+    │ 0.476031  ┆ 0.494526  ┆ -0.565498 ┆ -1.596759 ┆ -0.283645 │
+    │ -0.375222 ┆ -2.332233 ┆ 0.732424  ┆ -1.623769 ┆ 0.276038  │
+    │ -0.561579 ┆ 0.719409  ┆ 1.608522  ┆ -0.702902 ┆ -0.224755 │
+    │ 0.22584   ┆ -0.720572 ┆ 0.960557  ┆ -0.721353 ┆ 0.059809  │
+    │ -0.446836 ┆ -0.66195  ┆ -0.061743 ┆ 1.390007  ┆ 0.293002  │
+    └───────────┴───────────┴───────────┴───────────┴───────────┘
+
+To split a polars dataframe by date, pass the name of the date column in `split_col`,
+as in Case 3. If the column holds datetimes, pass a datetime as `cut_off`, for example
+`datetime(1789, 7, 14)`, instead of a date.
+
+
 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()