diff --git a/docs/user_guide/discretisation/ArbitraryDiscretiser.rst b/docs/user_guide/discretisation/ArbitraryDiscretiser.rst index b4d81e604..9c42a2cd1 100644 --- a/docs/user_guide/discretisation/ArbitraryDiscretiser.rst +++ b/docs/user_guide/discretisation/ArbitraryDiscretiser.rst @@ -110,6 +110,42 @@ obtain monotonic relationships between the variable and the target, you can do s seamlessly by setting `return_object` to True. You can find an example of discretisation followed by encoding to obtain monotonic releationships `here `_. +With polars +----------- + +:class:`ArbitraryDiscretiser()` also works with polars dataframes. + +.. code:: python + + import polars as pl + import numpy as np + from feature_engine.discretisation import ArbitraryDiscretiser + + X = pl.DataFrame({ + "MedInc": [1.5, 3.0, 5.0, 8.0, 0.8], + }) + + user_dict = {"MedInc": [0, 2, 4, 6, np.inf]} + + transformer = ArbitraryDiscretiser(binning_dict=user_dict, return_boundaries=False) + X_t = transformer.fit_transform(X) + print(X_t) + +.. code:: text + + shape: (5, 1) + ┌────────┐ + │ MedInc │ + │ --- │ + │ i64 │ + ╞════════╡ + │ 0 │ + │ 1 │ + │ 2 │ + │ 3 │ + │ 0 │ + └────────┘ + Additional resources -------------------- diff --git a/feature_engine/discretisation/arbitrary.py b/feature_engine/discretisation/arbitrary.py index 5776cd71d..4a6c75fd8 100644 --- a/feature_engine/discretisation/arbitrary.py +++ b/feature_engine/discretisation/arbitrary.py @@ -4,7 +4,9 @@ import warnings from typing import Dict, List, Optional, Union -import pandas as pd +import narwhals as nw +import numpy as np +from narwhals.typing import IntoDataFrame, IntoSeries from feature_engine._base_transformers.mixins import FitFromDictMixin from feature_engine._docstrings.fit_attributes import ( @@ -110,6 +112,27 @@ class ArbitraryDiscretiser(BaseDiscretiser, FitFromDictMixin): 3 25 1 17 Name: x, dtype: int64 + + With polars: + + >>> import polars as pl + >>> from feature_engine.discretisation import ArbitraryDiscretiser + >>> X = pl.DataFrame({"x": [10, 30, 60, 90]}) + >>> bins = dict(x=[0, 25, 50, 75, 100]) + >>> ad = ArbitraryDiscretiser(binning_dict=bins) + >>> ad.fit(X) + >>> ad.transform(X) + shape: (4, 1) + ┌─────┐ + │ x │ + │ --- │ + │ i64 │ + ╞═════╡ + │ 0 │ + │ 1 │ + │ 2 │ + │ 3 │ + └─────┘ """ def __init__( @@ -138,13 +161,13 @@ def __init__( self.binning_dict = binning_dict self.errors = errors - def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): + def fit(self, X: IntoDataFrame, y: Optional[IntoSeries] = None): """ This transformer does not learn any parameter. 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. @@ -161,30 +184,44 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): return self - 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. """ X = super().transform(X) - # check if NaN values were introduced by the discretisation procedure. - if X[self.variables_].isnull().sum().sum() > 0: - - # obtain the name(s) of the columns with null values - nan_columns = ( - X[self.variables_].columns[X[self.variables_].isnull().any()].tolist() - ) + # check if NaN values were introduced by the discretisation procedure. + nw_X = nw.from_native(X, eager_only=True) + if self.return_boundaries is True: + # missing labels are set to None by _bin_labels() + nan_columns = [ + var + for var in self.variables_ + if nw_X.get_column(var).is_null().any() + ] + else: + # codes are numeric; when return_object=True they're boxed as + # python floats in an Object column, where polars' is_null()/ + # is_nan() can't see NaN - a numpy float cast is reliable on + # both backends. + nan_columns = [ + var + for var in self.variables_ + if np.isnan(nw_X.get_column(var).to_numpy().astype(float)).any() + ] + + if len(nan_columns) > 0: if len(nan_columns) > 1: nan_columns_str = ", ".join(nan_columns) else: diff --git a/feature_engine/discretisation/base_discretiser.py b/feature_engine/discretisation/base_discretiser.py index 6c61d05d3..8bce3021f 100644 --- a/feature_engine/discretisation/base_discretiser.py +++ b/feature_engine/discretisation/base_discretiser.py @@ -1,7 +1,11 @@ # Authors: Morgan Sell # 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 diff --git a/tests/test_discretisation/test_arbitrary_discretiser.py b/tests/test_discretisation/test_arbitrary_discretiser.py index f1b2db712..6827c2c48 100644 --- a/tests/test_discretisation/test_arbitrary_discretiser.py +++ b/tests/test_discretisation/test_arbitrary_discretiser.py @@ -1,35 +1,39 @@ +import re + +import narwhals as nw import numpy as np import pandas as pd +import polars as pl import pytest -from numpy.random import default_rng -from scipy.stats import skewnorm from sklearn.datasets import fetch_california_housing from feature_engine.discretisation import ArbitraryDiscretiser -def test_arbitrary_discretiser(): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_arbitrary_discretiser(make_df): california_dataset = fetch_california_housing() - data = pd.DataFrame( + data_pd = pd.DataFrame( california_dataset.data, columns=california_dataset.feature_names ) user_dict = {"HouseAge": [0, 20, 40, 60, np.inf]} - 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 user-supplied and fixed, 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 = ArbitraryDiscretiser( binning_dict=user_dict, return_object=False, return_boundaries=False ) @@ -42,53 +46,38 @@ def test_arbitrary_discretiser(): assert transformer.variables_ == ["HouseAge"] assert transformer.binner_dict_ == user_dict # transform params - pd.testing.assert_frame_equal(X, data_t2) + result_codes = nw.from_native(X, eager_only=True).get_column("HouseAge").to_numpy() + assert np.array_equal(result_codes, expected_codes) transformer = ArbitraryDiscretiser( binning_dict=user_dict, return_object=False, return_boundaries=True ) X = transformer.fit_transform(data) - pd.testing.assert_frame_equal(X, data_t1) + result_labels = nw.from_native(X, eager_only=True).get_column("HouseAge").to_numpy() + assert np.array_equal(result_labels, expected_labels) -def test_error_if_input_df_contains_na_in_transform(df_vartypes, df_na): - # test case 1: when dataset contains na, transform method +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_error_if_input_df_contains_na_in_transform(make_df): + # test case 1: when dataset contains na, transform method raises age_dict = {"Age": [0, 10, 20, 30, np.inf]} + data = make_df({"Age": [20.0, 21.0, 19.0, 18.0]}) + data_na = make_df({"Age": [20.0, 21.0, None, 18.0]}) - with pytest.raises(ValueError): - transformer = ArbitraryDiscretiser(binning_dict=age_dict) - transformer.fit(df_vartypes) - transformer.transform(df_na[["Name", "City", "Age", "Marks", "dob"]]) - - -def test_error_when_nan_introduced_during_transform(): - # test error when NA are introduced during the discretisation. - rng = default_rng() - - # create dataframe with 2 variables, 1 normal and 1 skewed - random = skewnorm.rvs(a=-50, loc=4, size=100) - random = random - min(random) # Shift so the minimum value is equal to zero. - - train = pd.concat( - [ - pd.Series(rng.standard_normal(100)), - pd.Series(random), - ], - axis=1, - ) - - train.columns = ["var_a", "var_b"] + transformer = ArbitraryDiscretiser(binning_dict=age_dict) + transformer.fit(data) + with pytest.raises(ValueError, match="Some of the variables in the dataset"): + transformer.transform(data_na) - # create a dataframe with 2 variables normally distributed - test = pd.concat( - [ - pd.Series(rng.standard_normal(100)), - pd.Series(rng.standard_normal(100)), - ], - axis=1, - ) - test.columns = ["var_a", "var_b"] +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +@pytest.mark.parametrize("return_object", [False, True]) +def test_error_when_nan_introduced_during_transform(make_df, return_object): + # test warning/error when NA are introduced during the discretisation, + # i.e. when a value in the data to transform falls outside the bin edges + # fitted on the training data. + train = make_df({"var_a": [-4.0, -1.0, 1.0, 4.0], "var_b": [1.0, 2.0, 3.0, 4.0]}) + test = make_df({"var_a": [-4.0, -1.0, 1.0, 4.0], "var_b": [10.0, 20.0, 30.0, 40.0]}) msg = ( "During the discretisation, NaN values were introduced " @@ -98,29 +87,25 @@ def test_error_when_nan_introduced_during_transform(): limits_dict = {"var_a": [-5, -2, 0, 2, 5], "var_b": [0, 2, 5]} # check for warning when errors equals 'ignore' - with pytest.warns(UserWarning) as record: - transformer = ArbitraryDiscretiser(binning_dict=limits_dict, errors="ignore") - transformer.fit(train) + transformer = ArbitraryDiscretiser( + binning_dict=limits_dict, return_object=return_object, errors="ignore" + ) + transformer.fit(train) + with pytest.warns(UserWarning, match=re.escape(msg)): transformer.transform(test) - # check that only one warning was returned - assert len(record) == 1 - # check that message matches - assert record[0].message.args[0] == msg - # check for error when errors equals 'raise' - with pytest.raises(ValueError) as record: - transformer = ArbitraryDiscretiser(binning_dict=limits_dict, errors="raise") - transformer.fit(train) + transformer = ArbitraryDiscretiser( + binning_dict=limits_dict, return_object=return_object, errors="raise" + ) + transformer.fit(train) + with pytest.raises(ValueError, match=re.escape(msg)): transformer.transform(test) - # check that error message matches - assert str(record.value) == msg - def test_error_if_not_permitted_value_is_errors(): age_dict = {"Age": [0, 10, 20, 30, np.inf]} - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="errors only takes values"): ArbitraryDiscretiser(binning_dict=age_dict, errors="medialuna") @@ -130,8 +115,5 @@ def test_error_if_binning_dict_not_dict_type(binning_dict): "binning_dict must be a dictionary with the interval limits per " f"variable. Got {binning_dict} instead." ) - with pytest.raises(ValueError) as record: + with pytest.raises(ValueError, match=msg): ArbitraryDiscretiser(binning_dict=binning_dict) - - # check that error message matches - assert str(record.value) == msg diff --git a/tests/test_discretisation/test_base_discretizer.py b/tests/test_discretisation/test_base_discretizer.py index fc8110ff1..f852bf164 100644 --- a/tests/test_discretisation/test_base_discretizer.py +++ b/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)