Migrate EqualWidthDiscretiser to narwhals, add polars support

docs/user_guide/discretisation/EqualWidthDiscretiser.rst

@@ -48,9 +48,10 @@ potentially impact the model's performance in this scenario.
 EqualWidthDiscretiser
 ---------------------
 
-Feture-engine's :class:`EqualWidthDiscretiser()` applies equal width discretisation to numerical variables. It uses
-the `pandas.cut()` function under the hood to find the interval limits and then sort the continuous variables into
-the bins.
+Feture-engine's :class:`EqualWidthDiscretiser()` applies equal width discretisation to numerical variables. It finds
+the interval limits from each variable's minimum and maximum value, then sorts the continuous variables into the
+bins. It works with pandas, polars, and any other dataframe library supported by
+`narwhals <https://narwhals-dev.github.io/narwhals/>`_.
 
 You can specify the variables to be discretised by passing their names in a list when you set up the transformer. Alternatively,
 :class:`EqualWidthDiscretiser()` will automatically infer the data types and compute the interval limits for all numeric
@@ -271,7 +272,7 @@ If we want to output the intervals limits instead of integers, we can set `retur
 .. code:: python
 
     # Set up the discretisation transformer
-    disc = EqualFrequencyDiscretiser(
+    disc = EqualWidthDiscretiser(
         bins=10,
         variables=['LotArea','GrLivArea'],
         return_boundaries=True)
@@ -301,6 +302,48 @@ While we can't use these
 variables to train machine learning models, as opposed to the variables discretised into integers, they are very useful
 in this format for data analysis, and we can use any feature-engine encoder for further processing.
 
+With polars
+~~~~~~~~~~~
+
+:class:`EqualWidthDiscretiser()` works in the same way with a polars dataframe:
+
+.. code:: python
+
+    import polars as pl
+    from feature_engine.discretisation import EqualWidthDiscretiser
+
+    df = pl.DataFrame({
+        "x": [10400, 3675, 8640, 11670, 10667, 6120, 9500, 14000, 7200, 5300],
+    })
+
+    disc = EqualWidthDiscretiser(bins=5)
+
+    print(disc.fit_transform(df))
+
+The resulting values match those found with pandas:
+
+.. code:: text
+
+    shape: (10, 1)
+    ┌─────┐
+    │ x   │
+    │ --- │
+    │ i64 │
+    ╞═════╡
+    │ 3   │
+    │ 0   │
+    │ 2   │
+    │ 3   │
+    │ 3   │
+    │ 1   │
+    │ 2   │
+    │ 4   │
+    │ 1   │
+    │ 0   │
+    └─────┘
+
+`return_object`, `return_boundaries`, and `binner_dict_` work identically to the pandas examples above.
+
 See Also
 --------
 

feature_engine/discretisation/base_discretiser.py

@@ -1,7 +1,11 @@
 # Authors: Morgan Sell <morganpsell@gmail.com>
 # License: BSD 3 clause
 
-import pandas as pd
+from typing import List
+
+import narwhals as nw
+import numpy as np
+from narwhals.typing import IntoDataFrame
 
 from feature_engine._base_transformers.base_numerical import BaseNumericalTransformer
 
@@ -41,45 +45,133 @@ def __init__(
         self.return_boundaries = return_boundaries
         self.precision = precision
 
-    def transform(self, X: pd.DataFrame) -> pd.DataFrame:
+    def transform(self, X: IntoDataFrame) -> IntoDataFrame:
         """Sort the variable values into the intervals.
 
         Parameters
         ----------
-        X: pandas dataframe of shape = [n_samples, n_features]
+        X: dataframe of shape = [n_samples, n_features]
             The data to transform.
 
         Returns
         -------
-        X_new: pandas dataframe of shape = [n_samples, n_features]
+        X_new: dataframe of shape = [n_samples, n_features]
             The transformed data with the discrete variables.
         """
 
         # check input dataframe and if class was fitted
         X = self._check_transform_input_and_state(X)
 
-        # transform variables
+        # bin edges are already fixed by fit(), so sorting values into them is a
+        # plain numpy searchsorted - vectorizable identically for every backend,
+        # no pandas/polars-specific path needed.
+        nw_X = nw.from_native(X, eager_only=True)
+        native_namespace = nw_X.__native_namespace__()
+
         if self.return_boundaries is True:
-            for feature in self.variables_:
-                X[feature] = pd.cut(
-                    X[feature],
-                    self.binner_dict_[feature],
-                    precision=self.precision,
-                    include_lowest=True,
+            new_columns = [
+                nw.new_series(
+                    feature,
+                    _bin_labels(
+                        nw_X.get_column(feature).to_numpy(),
+                        self.binner_dict_[feature],
+                        self.precision,
+                    ),
+                    backend=native_namespace,
                 )
-            X[self.variables_] = X[self.variables_].astype(str)
-
+                for feature in self.variables_
+            ]
         else:
-            for feature in self.variables_:
-                X[feature] = pd.cut(
-                    X[feature],
-                    self.binner_dict_[feature],
-                    labels=False,
-                    include_lowest=True,
+            # nw.Object mirrors the pandas "O" dtype astype() used to produce,
+            # and is what feature-engine's categorical encoders detect on
+            # every narwhals-supported backend (see variable_handling).
+            dtype = nw.Object if self.return_object is True else None
+            new_columns = [
+                nw.new_series(
+                    feature,
+                    _bin_codes(
+                        nw_X.get_column(feature).to_numpy(),
+                        self.binner_dict_[feature],
+                        self.return_object,
+                    ),
+                    dtype=dtype,
+                    backend=native_namespace,
                 )
+                for feature in self.variables_
+            ]
 
-            # return object
-            if self.return_object:
-                X[self.variables_] = X[self.variables_].astype("O")
+        X = nw_X.with_columns(*new_columns).to_native()
 
         return X
+
+
+def _digitize(values: np.ndarray, bins_arr: np.ndarray):
+    """0-based bin index per value, right-closed intervals with the lowest edge
+    included - mirrors pandas.cut(bins=bins, include_lowest=True), which is
+    itself built on this same bins.searchsorted() call. Values outside the
+    bin range, and NaNs, are flagged via na_mask rather than given a code.
+    """
+    ids = np.asarray(np.searchsorted(bins_arr, values, side="left"))
+    ids[values == bins_arr[0]] = 1
+    na_mask: np.ndarray = np.isnan(values) | (ids == len(bins_arr)) | (ids == 0)
+    return ids - 1, na_mask
+
+
+def _bin_codes(values: np.ndarray, bins: List[float], return_object: bool):
+    bins_arr: np.ndarray = np.asarray(bins, dtype=float)
+    codes, na_mask = _digitize(values, bins_arr)
+
+    # match pandas.cut(labels=False): int codes, upcast to float only when a
+    # NaN placeholder is actually needed.
+    if na_mask.any():
+        codes = codes.astype(np.float64)
+        codes[na_mask] = np.nan
+    if return_object is True:
+        codes = codes.astype(object)
+
+    return codes
+
+
+def _bin_labels(values: np.ndarray, bins: List[float], precision: int):
+    bins_arr: np.ndarray = np.asarray(bins, dtype=float)
+    codes, na_mask = _digitize(values, bins_arr)
+
+    labels = np.asarray(_format_bin_labels(bins_arr, precision), dtype=object)
+    out: np.ndarray = np.empty(len(values), dtype=object)
+    out[~na_mask] = labels[codes[~na_mask]]
+    out[na_mask] = None
+
+    return out
+
+
+def _format_bin_labels(bins_arr: np.ndarray, precision: int) -> List[str]:
+    """"(lower, upper]" text per bin, replicating pandas.cut's own label
+    formatting: widen precision until break values are unique, then shrink
+    the lowest edge so include_lowest values still read as inside the first
+    interval.
+    """
+    precision = _infer_precision(precision, bins_arr)
+    breaks = [_round_frac(b, precision) for b in bins_arr]
+    breaks[0] = breaks[0] - 10 ** (-precision)
+    return [f"({breaks[i]}, {breaks[i + 1]}]" for i in range(len(breaks) - 1)]
+
+
+def _round_frac(x: float, precision: int) -> float:
+    if not np.isfinite(x) or x == 0:
+        return float(x)
+    frac, whole = np.modf(x)
+    if whole == 0:
+        digits = -int(np.floor(np.log10(abs(frac)))) - 1 + precision
+    else:
+        digits = precision
+    return float(np.around(x, digits))
+
+
+def _infer_precision(base_precision: int, bins_arr: np.ndarray) -> int:
+    # widen precision until every rounded break is unique - otherwise two
+    # adjacent bins could render with identical label text.
+    for precision in range(base_precision, 20):
+        levels = [_round_frac(b, precision) for b in bins_arr]
+        if len(set(levels)) == len(bins_arr):
+            return precision
+    return base_precision

feature_engine/discretisation/equal_width.py

@@ -3,7 +3,9 @@
 
 from typing import List, Optional, Union
 
-import pandas as pd
+import narwhals as nw
+import numpy as np
+from narwhals.typing import IntoDataFrame, IntoSeries
 
 from feature_engine._check_init_parameters.check_init_input_params import (
     _check_return_empty_is_bool,
@@ -164,14 +166,14 @@ def __init__(
         self.return_empty = return_empty
         self.bins = bins
 
-    def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
+    def fit(self, X: IntoDataFrame, y: Optional[IntoSeries] = None):
         """
         Learn the boundaries of the equal width intervals / bins for each
         variable.
 
         Parameters
         ----------
-        X: pandas dataframe of shape = [n_samples, n_features]
+        X: dataframe of shape = [n_samples, n_features]
             The training dataset. Can be the entire dataframe, not just the variables
             to be transformed.
         y: None
@@ -184,23 +186,39 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
         # fit
         binner_dict_ = {}
 
-        for var in variables_:
-            tmp, bins = pd.cut(
-                x=X[var],
-                bins=self.bins,
-                retbins=True,
-                duplicates="drop",
-                include_lowest=True,
-            )
-
-            # Prepend/Append infinities
-            bins = list(bins)
-            bins[0] = float("-inf")
-            bins[len(bins) - 1] = float("inf")
-            binner_dict_[var] = bins
+        if len(variables_) > 0:
+            # one narwhals call for every variable at once, instead of a
+            # get_column() round-trip per variable.
+            arr = nw.from_native(X, eager_only=True).select(variables_).to_numpy()
+            mins = arr.min(axis=0)
+            maxs = arr.max(axis=0)
+            for var, mn, mx in zip(variables_, mins, maxs):
+                binner_dict_[var] = _equal_width_edges(mn, mx, self.bins)
 
         self.binner_dict_ = binner_dict_
         self.variables_ = variables_
         self._get_feature_names_in(X)
 
         return self
+
+
+def _equal_width_edges(mn: float, mx: float, bins: int) -> List[float]:
+    """Bin-edge computation matching pandas.cut(bins=int, duplicates="drop"):
+    widen a constant [mn, mx] by 0.1% so linspace still produces positive-
+    width bins, then collapse duplicate edges the same way. The outer edges
+    are then clipped to +-inf, same as the pre-migration code did to the
+    retbins output, so transform() never needs an out-of-range branch.
+    """
+    if mn == mx:
+        mn = mn - 0.001 * abs(mn) if mn != 0 else -0.001
+        mx = mx + 0.001 * abs(mx) if mx != 0 else 0.001
+
+    edges = np.linspace(mn, mx, bins + 1)
+    unique_edges = np.unique(edges)
+    if len(unique_edges) < len(edges) and len(edges) != 2:
+        edges = unique_edges
+
+    edges_: List[float] = edges.tolist()
+    edges_[0] = float("-inf")
+    edges_[-1] = float("inf")
+    return edges_

tests/test_discretisation/test_base_discretizer.py

@@ -1,5 +1,6 @@
 import numpy as np
 import pandas as pd
+import polars as pl
 import pytest
 from sklearn.datasets import fetch_california_housing
 
@@ -38,42 +39,42 @@ def test_correct_param_assignment_at_init(params):
 
 class MockClassFit(BaseDiscretiser):
     def fit(self, X):
-        california_dataset = fetch_california_housing()
-        data = pd.DataFrame(
-            california_dataset.data, columns=california_dataset.feature_names
-        )
+        # bins are hard-coded rather than learnt, so this mock works unchanged
+        # on both pandas and polars input.
         self.variables_ = ["HouseAge"]
         self.binner_dict_ = {"HouseAge": [0, 20, 40, 60, np.inf]}
-        self.n_features_in_ = data.shape[1]
-        self.feature_names_in_ = california_dataset.feature_names
+        self.n_features_in_ = X.shape[1]
+        self.feature_names_in_ = list(X.columns)
         return self
 
 
-def test_transform():
+@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame])
+def test_transform(make_df):
     california_dataset = fetch_california_housing()
-    data = pd.DataFrame(
+    data_pd = pd.DataFrame(
         california_dataset.data, columns=california_dataset.feature_names
     )
 
-    data_t1 = data.copy()
-    data_t2 = data.copy()
-
-    # HouseAge is the median house age in the block group.
-    data_t1["HouseAge"] = pd.cut(
-        data["HouseAge"], bins=[0, 20, 40, 60, np.inf], include_lowest=True
-    )
-    data_t1["HouseAge"] = data_t1["HouseAge"].astype(str)
-    data_t2["HouseAge"] = pd.cut(
-        data["HouseAge"],
+    # ground truth via pandas.cut: bins are fixed by MockClassFit, so both
+    # backends must reproduce this exact output.
+    expected_codes = pd.cut(
+        data_pd["HouseAge"],
         bins=[0, 20, 40, 60, np.inf],
         labels=False,
         include_lowest=True,
+    ).to_numpy()
+    expected_labels = (
+        pd.cut(data_pd["HouseAge"], bins=[0, 20, 40, 60, np.inf], include_lowest=True)
+        .astype(str)
+        .to_numpy()
     )
 
+    data = make_df(data_pd)
+
     transformer = MockClassFit(return_boundaries=False)
     X = transformer.fit_transform(data)
-    pd.testing.assert_frame_equal(X, data_t2)
+    assert np.array_equal(X["HouseAge"].to_numpy(), expected_codes)
 
     transformer = MockClassFit(return_object=False, return_boundaries=True)
     X = transformer.fit_transform(data)
-    pd.testing.assert_frame_equal(X, data_t1)
+    assert np.array_equal(X["HouseAge"].to_numpy(), expected_labels)