Migrate BaseDiscretiser to narwhals, add polars support

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

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)