diff --git a/docs/user_guide/discretisation/GeometricWidthDiscretiser.rst b/docs/user_guide/discretisation/GeometricWidthDiscretiser.rst index 74e150763..940746d9d 100644 --- a/docs/user_guide/discretisation/GeometricWidthDiscretiser.rst +++ b/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 ~~~~~~~~~~~~~~ diff --git a/feature_engine/discretisation/geometric_width.py b/feature_engine/discretisation/geometric_width.py index 709381c71..41aa9bd18 100644 --- a/feature_engine/discretisation/geometric_width.py +++ b/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 diff --git a/tests/test_discretisation/test_geometric_width_discretiser.py b/tests/test_discretisation/test_geometric_width_discretiser.py index 6a4b56c2d..82a70add2 100644 --- a/tests/test_discretisation/test_geometric_width_discretiser.py +++ b/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,14 +56,16 @@ 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) @@ -58,34 +73,43 @@ def test_fit_and_transform_methods(df_normal_dist): # 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)