Migrate GeometricWidthDiscretiser.fit() to narwhals, add polars support

docs/user_guide/discretisation/GeometricWidthDiscretiser.rst

@@ -144,6 +144,66 @@ In the following output, we see the interval limits determined for each variable
         2212.974,
         inf]}
 
+With polars
+-----------
+
+:class:`GeometricWidthDiscretiser()` works in the same way with a polars dataframe:
+
+.. code:: python
+
+    import numpy as np
+    import polars as pl
+    from feature_engine.discretisation import GeometricWidthDiscretiser
+
+    np.random.seed(42)
+    df = pl.DataFrame({"x": np.random.randint(1, 100, 100).astype(float)})
+
+    disc = GeometricWidthDiscretiser(bins=10)
+    Xt = disc.fit_transform(df)
+
+    print(Xt["x"].value_counts().sort("x"))
+
+The resulting bin counts:
+
+.. code:: text
+
+    shape: (9, 2)
+    ┌─────┬───────┐
+    │ x   ┆ count │
+    │ --- ┆ ---   │
+    │ i64 ┆ u32   │
+    ╞═════╪═══════╡
+    │ 0   ┆ 6     │
+    │ 1   ┆ 3     │
+    │ 3   ┆ 3     │
+    │ 4   ┆ 1     │
+    │ 5   ┆ 5     │
+    │ 6   ┆ 9     │
+    │ 7   ┆ 8     │
+    │ 8   ┆ 25    │
+    │ 9   ┆ 40    │
+    └─────┴───────┘
+
+And the fitted bin edges, matching what we'd get fitting on the same values with pandas:
+
+.. code:: python
+
+    disc.binner_dict_
+
+.. code:: python
+
+    {'x': [-inf,
+        3.573433146226546,
+        4.475691865644366,
+        5.895335641248283,
+        8.129050213617685,
+        11.643650760992958,
+        17.173639757979174,
+        25.874707744105372,
+        39.565256521047,
+        61.106419756718246,
+        inf]}
+
 Interval width
 ~~~~~~~~~~~~~~
 

feature_engine/discretisation/geometric_width.py

@@ -1,7 +1,8 @@
 from typing import List, Optional, Union
 
+import narwhals as nw
 import numpy as np
-import pandas as pd
+from narwhals.typing import IntoDataFrame, IntoSeries
 
 from feature_engine._check_init_parameters.check_init_input_params import (
     _check_return_empty_is_bool,
@@ -159,14 +160,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 geometric 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
@@ -177,10 +178,12 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
         X, variables_ = self._fit_setup(X)
 
         # fit
+        nw_X = nw.from_native(X, eager_only=True)
         binner_dict_ = {}
 
         for var in variables_:
-            min_, max_ = X[var].min(), X[var].max()
+            col = nw_X.get_column(var)
+            min_, max_ = col.min(), col.max()
             increment = np.power(max_ - min_, 1.0 / self.bins)
             bins = np.r_[
                 -np.inf, min_ + np.power(increment, np.arange(1, self.bins)), np.inf

tests/test_discretisation/test_geometric_width_discretiser.py

@@ -1,38 +1,51 @@
+import re
+
+import narwhals as nw
 import numpy as np
 import pandas as pd
 import pytest
 from sklearn.exceptions import NotFittedError
 
 from feature_engine.discretisation import GeometricWidthDiscretiser
+from tests.backend_helpers import frame_to_dict
+
+MSG_NA = (
+    "Some of the variables in the dataset contain NaN. Check and "
+    "remove those before using this transformer."
+)
 
 
-# test init params
+# init parameters
 @pytest.mark.parametrize("param", [0.1, "hola", (True, False), {"a": True}, 2])
 def test_raises_error_when_return_object_not_bool(param):
-    with pytest.raises(ValueError):
+    msg = f"return_object must be True or False. Got {param} instead."
+    with pytest.raises(ValueError, match=re.escape(msg)):
         GeometricWidthDiscretiser(return_object=param)
 
 
 @pytest.mark.parametrize("param", [0.1, "hola", (True, False), {"a": True}, 2])
 def test_raises_error_when_return_boundaries_not_bool(param):
-    with pytest.raises(ValueError):
+    msg = f"return_boundaries must be True or False. Got {param} instead."
+    with pytest.raises(ValueError, match=re.escape(msg)):
         GeometricWidthDiscretiser(return_boundaries=param)
 
 
 @pytest.mark.parametrize("param", [0.1, "hola", (True, False), {"a": True}, 0, -1])
 def test_raises_error_when_precision_not_int(param):
-    with pytest.raises(ValueError):
+    msg = f"precision must be a positive integer. Got {param} instead."
+    with pytest.raises(ValueError, match=re.escape(msg)):
         GeometricWidthDiscretiser(precision=param)
 
 
-@pytest.mark.parametrize("param", [0.1, "hola", (True, False), {"a": True}])
+@pytest.mark.parametrize("param", [0.1, "hola", (True, False), {"a": True}, None])
 def test_raises_error_when_bins_not_int(param):
-    with pytest.raises(ValueError):
+    msg = f"bins must be an integer. Got {param} instead."
+    with pytest.raises(ValueError, match=re.escape(msg)):
         GeometricWidthDiscretiser(bins=param)
 
 
 @pytest.mark.parametrize("params", [(False, 1), (True, 10)])
-def test_correct_param_assignment_at_init(params):
+def test_init_param_assignment(params):
     param1, param2 = params
     t = GeometricWidthDiscretiser(
         return_object=param1, return_boundaries=param1, precision=param2, bins=param2
@@ -43,49 +56,60 @@ def test_correct_param_assignment_at_init(params):
     assert t.bins == param2
 
 
-def test_fit_and_transform_methods(df_normal_dist):
+# fit and transform
+def test_fit_and_transform_methods(make_df, data_normal_dist):
     transformer = GeometricWidthDiscretiser(
         bins=10, variables=None, return_object=False
     )
-    X = transformer.fit_transform(df_normal_dist)
+    X = transformer.fit_transform(make_df(data_normal_dist))
 
     # manual calculation
-    min_, max_ = df_normal_dist["var"].min(), df_normal_dist["var"].max()
+    arr = np.array(data_normal_dist["var"])
+    min_, max_ = arr.min(), arr.max()
     increment = np.power(max_ - min_, 1.0 / 10)
     bins = np.r_[-np.inf, min_ + np.power(increment, np.arange(1, 10)), np.inf]
     bins = np.sort(bins)
 
     # fit params
     assert (transformer.binner_dict_["var"] == bins).all()
 
-    # transform params
-    assert (
-        X["var"] == pd.cut(df_normal_dist["var"], bins=bins, precision=7).cat.codes
-    ).all()
+    # transform params - ground truth from pandas.cut on the same bins; values
+    # must match regardless of which backend the input dataframe uses.
+    expected = pd.cut(pd.Series(arr), bins=bins, precision=7).cat.codes.tolist()
+    assert isinstance(X, make_df)
+    assert frame_to_dict(X)["var"] == expected
 
 
-def test_automatically_find_variables_and_return_as_object(df_normal_dist):
+def test_automatically_find_variables_and_return_as_object(make_df, data_normal_dist):
     transformer = GeometricWidthDiscretiser(bins=10, variables=None, return_object=True)
-    X = transformer.fit_transform(df_normal_dist)
-    assert X["var"].dtypes == "O"
+    X = transformer.fit_transform(make_df(data_normal_dist))
+    assert isinstance(X, make_df)
+    assert nw.from_native(X, eager_only=True).schema["var"] == nw.Object
 
 
-def test_error_if_input_df_contains_na_in_fit(df_na):
-    # test case 3: when dataset contains na, fit method
+def test_error_if_input_df_contains_na_in_fit(make_df):
+    df_na = make_df({"Age": [20.0, 21.0, None, 23.0]})
     transformer = GeometricWidthDiscretiser()
-    with pytest.raises(ValueError):
+    with pytest.raises(ValueError, match=re.escape(MSG_NA)):
         transformer.fit(df_na)
 
 
-def test_error_if_input_df_contains_na_in_transform(df_vartypes, df_na):
-    # test case 4: when dataset contains na, transform method
+def test_error_if_input_df_contains_na_in_transform(make_df):
+    df = make_df({"Age": [20.0, 21.0, 19.0, 23.0]})
+    df_na = make_df({"Age": [20.0, 21.0, None, 23.0]})
+
     transformer = GeometricWidthDiscretiser()
-    transformer.fit(df_vartypes)
-    with pytest.raises(ValueError):
-        transformer.transform(df_na[["Name", "City", "Age", "Marks", "dob"]])
+    transformer.fit(df)
+    with pytest.raises(ValueError, match=re.escape(MSG_NA)):
+        transformer.transform(df_na)
 
 
-def test_non_fitted_error(df_vartypes):
+def test_non_fitted_error(make_df):
+    df = make_df({"Age": [20.0, 21.0, 19.0, 23.0]})
     transformer = GeometricWidthDiscretiser()
-    with pytest.raises(NotFittedError):
-        transformer.transform(df_vartypes)
+    msg = (
+        "This GeometricWidthDiscretiser instance is not fitted yet. Call 'fit' "
+        "with appropriate arguments before using this estimator."
+    )
+    with pytest.raises(NotFittedError, match=re.escape(msg)):
+        transformer.transform(df)