From 05e88e81622e9b55ee3906820c7cf021e4789240 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Wed, 26 Aug 2026 00:55:57 +0200 Subject: [PATCH 1/4] Migrate ArbitraryDiscretiser to narwhals, add polars support fit() needed no changes: it already delegates entirely to the already-migrated FitFromDictMixin._fit_from_dict(). The pandas dependency was in transform()'s post-hoc NaN-introduced check, which used X[...].isnull().sum().sum() / .columns / .any() / .tolist() - pandas-only calls that broke outright on polars input coming back from the now-migrated BaseDiscretiser.transform(). Replaced it with a narwhals-based per-column check, branched on return_boundaries rather than dtype: labels (return_boundaries=True) use None for missing values, which narwhals' is_null() detects correctly on both backends. Codes (return_boundaries=False) are numeric, so a numpy float cast + np.isnan is used instead of is_null()/is_nan() directly. That numeric-cast branch isn't just style - narwhals' is_null() (and polars' own null semantics) do NOT see a boxed np.nan sitting inside a polars Object-dtype column (return_object=True's output dtype): verified with a direct repro, is_null().any() returns False on a polars Object series holding all-NaN values, silently swallowing the warning/error this method exists to raise. is_nan() isn't usable there either - narwhals raises "is_nan only supported for numeric dtype, not Object". The numpy-float-cast approach sidesteps both issues and was confirmed to raise/warn correctly across all pandas/polars x return_object x return_boundaries combinations. Benchmarked old (pandas-only) vs new (narwhals) transform() at 10k/50k/100k rows x 1/2/10 cols on pandas input: return_object=False lands at parity (0.9-1.05x, within noise); return_object=True is 1.15-1.3x slower (e.g. 100k rows x 10 cols: 36.2ms old vs 44.9ms new) since the per-variable numpy float-cast replaces one vectorized pandas isnull().sum().sum() call. This falls within the "minimal loss" band used to decide against a pandas/polars split elsewhere in this migration, so a single narwhals-driven path was kept - no is_pandas branch was added. narwhals-on-polars is faster than narwhals-on-pandas at every size tested, consistent with the base branch's own findings. Verified: tests/test_discretisation full suite (114 passed, same 5 pre-existing check_estimator failures as the unmodified base branch - reproduced there too, predates this change). Rewrote test_arbitrary_discretiser.py per AGENTS.md: one parametrized test per behavior over pd.DataFrame/pl.DataFrame (previously pandas-only), switched pytest.raises()/pytest.warns() to the match= form instead of capturing and asserting on the record. flake8 and mypy clean. Module imports with pandas blocked. sphinx -W build clean (only the pre-existing unrelated linkcode_resolve warning). Verified the existing docstring/rst examples against real output before touching: the "Name: x" vs "Name: count" and bare-fit()-repr doctest mismatches are the same pre-existing pandas-3.0 doc-staleness noted in the base branch commit (reproduced on the unmodified file too) - left alone, out of scope here. Added a "With polars" example to both the class docstring and ArbitraryDiscretiser.rst, output verified against a real run. Co-Authored-By: Claude Sonnet 5 --- .../discretisation/ArbitraryDiscretiser.rst | 36 ++++++ feature_engine/discretisation/arbitrary.py | 63 +++++++-- .../test_arbitrary_discretiser.py | 120 ++++++++---------- 3 files changed, 137 insertions(+), 82 deletions(-) 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/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 From d35ae955aa9c4b2ffd624b3149c69b10e4464f68 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Mon, 14 Sep 2026 22:42:14 +0200 Subject: [PATCH 2/4] Use shared backend test fixtures and helpers in ArbitraryDiscretiser tests Build the California housing input from the data_california fixture on the backend under test, check isinstance(X, make_df) plus to_dict() contents, and use the make_df fixture and pytest.raises(match=re.escape(msg)). Co-Authored-By: Claude Opus 5 --- .../test_arbitrary_discretiser.py | 41 +++++++------------ 1 file changed, 15 insertions(+), 26 deletions(-) diff --git a/tests/test_discretisation/test_arbitrary_discretiser.py b/tests/test_discretisation/test_arbitrary_discretiser.py index 6827c2c48..81d882b32 100644 --- a/tests/test_discretisation/test_arbitrary_discretiser.py +++ b/tests/test_discretisation/test_arbitrary_discretiser.py @@ -1,38 +1,29 @@ import re -import narwhals as nw import numpy as np import pandas as pd -import polars as pl import pytest -from sklearn.datasets import fetch_california_housing from feature_engine.discretisation import ArbitraryDiscretiser +from tests.backend_helpers import to_dict +BINS = [0, 20, 40, 60, np.inf] -@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) -def test_arbitrary_discretiser(make_df): - california_dataset = fetch_california_housing() - data_pd = pd.DataFrame( - california_dataset.data, columns=california_dataset.feature_names - ) - user_dict = {"HouseAge": [0, 20, 40, 60, np.inf]} + +def test_arbitrary_discretiser(make_df, data_california): + user_dict = {"HouseAge": BINS} # ground truth via pandas.cut - bins are user-supplied and fixed, so both # backends must reproduce this exact output. + house_age = pd.Series(data_california["HouseAge"]) expected_codes = pd.cut( - data_pd["HouseAge"], - bins=[0, 20, 40, 60, np.inf], - labels=False, - include_lowest=True, - ).to_numpy() + house_age, bins=BINS, labels=False, include_lowest=True + ).tolist() expected_labels = ( - pd.cut(data_pd["HouseAge"], bins=[0, 20, 40, 60, np.inf], include_lowest=True) - .astype(str) - .to_numpy() + pd.cut(house_age, bins=BINS, include_lowest=True).astype(str).tolist() ) - data = make_df(data_pd) + data = make_df(data_california) transformer = ArbitraryDiscretiser( binning_dict=user_dict, return_object=False, return_boundaries=False @@ -46,18 +37,17 @@ def test_arbitrary_discretiser(make_df): assert transformer.variables_ == ["HouseAge"] assert transformer.binner_dict_ == user_dict # transform params - result_codes = nw.from_native(X, eager_only=True).get_column("HouseAge").to_numpy() - assert np.array_equal(result_codes, expected_codes) + assert isinstance(X, make_df) + assert to_dict(X)["HouseAge"] == expected_codes transformer = ArbitraryDiscretiser( binning_dict=user_dict, return_object=False, return_boundaries=True ) X = transformer.fit_transform(data) - result_labels = nw.from_native(X, eager_only=True).get_column("HouseAge").to_numpy() - assert np.array_equal(result_labels, expected_labels) + assert isinstance(X, make_df) + assert to_dict(X)["HouseAge"] == expected_labels -@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]} @@ -70,7 +60,6 @@ def test_error_if_input_df_contains_na_in_transform(make_df): transformer.transform(data_na) -@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, @@ -115,5 +104,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, match=msg): + with pytest.raises(ValueError, match=re.escape(msg)): ArbitraryDiscretiser(binning_dict=binning_dict) From 78838a6cff3f8490006540aa722b8f73a7971869 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Tue, 15 Sep 2026 12:06:16 +0200 Subject: [PATCH 3/4] Use frame_to_dict after the shared helper rename in #1045 Co-Authored-By: Claude Opus 5 --- tests/test_discretisation/test_arbitrary_discretiser.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_discretisation/test_arbitrary_discretiser.py b/tests/test_discretisation/test_arbitrary_discretiser.py index 81d882b32..a0727de38 100644 --- a/tests/test_discretisation/test_arbitrary_discretiser.py +++ b/tests/test_discretisation/test_arbitrary_discretiser.py @@ -5,7 +5,7 @@ import pytest from feature_engine.discretisation import ArbitraryDiscretiser -from tests.backend_helpers import to_dict +from tests.backend_helpers import frame_to_dict BINS = [0, 20, 40, 60, np.inf] @@ -38,14 +38,14 @@ def test_arbitrary_discretiser(make_df, data_california): assert transformer.binner_dict_ == user_dict # transform params assert isinstance(X, make_df) - assert to_dict(X)["HouseAge"] == expected_codes + assert frame_to_dict(X)["HouseAge"] == expected_codes transformer = ArbitraryDiscretiser( binning_dict=user_dict, return_object=False, return_boundaries=True ) X = transformer.fit_transform(data) assert isinstance(X, make_df) - assert to_dict(X)["HouseAge"] == expected_labels + assert frame_to_dict(X)["HouseAge"] == expected_labels def test_error_if_input_df_contains_na_in_transform(make_df): From 4afe6a1f5f15bb54267a82b4f7f9e79651a04055 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Fri, 18 Sep 2026 13:23:46 +0200 Subject: [PATCH 4/4] Group init tests and match errors in ArbitraryDiscretiser and BaseDiscretiser tests, check errors type Co-Authored-By: Claude Opus 5 --- feature_engine/discretisation/arbitrary.py | 2 +- .../test_arbitrary_discretiser.py | 76 +++++++++++++------ .../test_base_discretizer.py | 42 ++++++---- 3 files changed, 81 insertions(+), 39 deletions(-) diff --git a/feature_engine/discretisation/arbitrary.py b/feature_engine/discretisation/arbitrary.py index 4a6c75fd8..e504f1cfe 100644 --- a/feature_engine/discretisation/arbitrary.py +++ b/feature_engine/discretisation/arbitrary.py @@ -150,7 +150,7 @@ def __init__( f"variable. Got {binning_dict} instead." ) - if errors not in ["ignore", "raise"]: + if not isinstance(errors, str) or errors not in ["ignore", "raise"]: raise ValueError( "errors only takes values 'ignore' and 'raise'. " f"Got {errors} instead." diff --git a/tests/test_discretisation/test_arbitrary_discretiser.py b/tests/test_discretisation/test_arbitrary_discretiser.py index a0727de38..79e2bf1fc 100644 --- a/tests/test_discretisation/test_arbitrary_discretiser.py +++ b/tests/test_discretisation/test_arbitrary_discretiser.py @@ -9,7 +9,58 @@ BINS = [0, 20, 40, 60, np.inf] +MSG_NA = ( + "Some of the variables in the dataset contain NaN. Check and " + "remove those before using this transformer." +) + +# init parameters +@pytest.mark.parametrize("binning_dict", ["HOLA", 1, False, None, [0, 10, 20]]) +def test_error_if_binning_dict_not_dict_type(binning_dict): + msg = ( + "binning_dict must be a dictionary with the interval limits per " + f"variable. Got {binning_dict} instead." + ) + with pytest.raises(ValueError, match=re.escape(msg)): + ArbitraryDiscretiser(binning_dict=binning_dict) + + +@pytest.mark.parametrize( + "errors", ["medialuna", "Ignore", "", 1, None, ["ignore"], ("raise",)] +) +def test_error_if_errors_not_permitted_value(errors): + msg = f"errors only takes values 'ignore' and 'raise'. Got {errors} instead." + with pytest.raises(ValueError, match=re.escape(msg)): + ArbitraryDiscretiser(binning_dict={"Age": BINS}, errors=errors) + + +@pytest.mark.parametrize( + "binning_dict, return_object, return_boundaries, precision, errors", + [ + ({"HouseAge": BINS}, False, False, 3, "ignore"), + ({"HouseAge": BINS, "MedInc": [0, 5, np.inf]}, True, False, 1, "raise"), + ({"Age": [0, 10, np.inf]}, False, True, 10, "raise"), + ], +) +def test_init_param_assignment( + binning_dict, return_object, return_boundaries, precision, errors +): + transformer = ArbitraryDiscretiser( + binning_dict=binning_dict, + return_object=return_object, + return_boundaries=return_boundaries, + precision=precision, + errors=errors, + ) + assert transformer.binning_dict == binning_dict + assert transformer.return_object is return_object + assert transformer.return_boundaries is return_boundaries + assert transformer.precision == precision + assert transformer.errors == errors + + +# fit and transform def test_arbitrary_discretiser(make_df, data_california): user_dict = {"HouseAge": BINS} @@ -30,9 +81,6 @@ def test_arbitrary_discretiser(make_df, data_california): ) X = transformer.fit_transform(data) - # init params - assert transformer.return_object is False - assert transformer.return_boundaries is False # fit params assert transformer.variables_ == ["HouseAge"] assert transformer.binner_dict_ == user_dict @@ -56,15 +104,13 @@ def test_error_if_input_df_contains_na_in_transform(make_df): transformer = ArbitraryDiscretiser(binning_dict=age_dict) transformer.fit(data) - with pytest.raises(ValueError, match="Some of the variables in the dataset"): + with pytest.raises(ValueError, match=re.escape(MSG_NA)): transformer.transform(data_na) @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. + # values outside the bin edges learned in fit become NaN, which warns or raises 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]}) @@ -90,19 +136,3 @@ def test_error_when_nan_introduced_during_transform(make_df, return_object): transformer.fit(train) with pytest.raises(ValueError, match=re.escape(msg)): transformer.transform(test) - - -def test_error_if_not_permitted_value_is_errors(): - age_dict = {"Age": [0, 10, 20, 30, np.inf]} - with pytest.raises(ValueError, match="errors only takes values"): - ArbitraryDiscretiser(binning_dict=age_dict, errors="medialuna") - - -@pytest.mark.parametrize("binning_dict", ["HOLA", 1, False]) -def test_error_if_binning_dict_not_dict_type(binning_dict): - msg = ( - "binning_dict must be a dictionary with the interval limits per " - f"variable. Got {binning_dict} instead." - ) - with pytest.raises(ValueError, match=re.escape(msg)): - ArbitraryDiscretiser(binning_dict=binning_dict) diff --git a/tests/test_discretisation/test_base_discretizer.py b/tests/test_discretisation/test_base_discretizer.py index 6ce784cfe..6cde4bd55 100644 --- a/tests/test_discretisation/test_base_discretizer.py +++ b/tests/test_discretisation/test_base_discretizer.py @@ -1,3 +1,5 @@ +import re + import numpy as np import pandas as pd import pytest @@ -8,36 +10,46 @@ BINS = [0, 20, 40, 60, np.inf] -# test init params -@pytest.mark.parametrize("param", [0.1, "hola", (True, False), {"a": True}, 2]) +# init parameters +@pytest.mark.parametrize("param", [0.1, "hola", (True, False), {"a": True}, 2, None]) 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)): BaseDiscretiser(return_object=param) -@pytest.mark.parametrize("param", [0.1, "hola", (True, False), {"a": True}, 2]) +@pytest.mark.parametrize("param", [0.1, "hola", (True, False), {"a": True}, 2, None]) 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)): BaseDiscretiser(return_boundaries=param) -@pytest.mark.parametrize("param", [0.1, "hola", (True, False), {"a": True}, 0, -1]) +@pytest.mark.parametrize( + "param", [0.1, "hola", (True, False), {"a": True}, 0, -1, None] +) 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)): BaseDiscretiser(precision=param) -@pytest.mark.parametrize("params", [(False, 1), (True, 10)]) -def test_correct_param_assignment_at_init(params): - param1, param2 = params - t = BaseDiscretiser( - return_object=param1, return_boundaries=param1, precision=param2 +@pytest.mark.parametrize( + "return_object, return_boundaries, precision", + [(False, False, 1), (True, False, 10), (False, True, 3)], +) +def test_init_param_assignment(return_object, return_boundaries, precision): + transformer = BaseDiscretiser( + return_object=return_object, + return_boundaries=return_boundaries, + precision=precision, ) - assert t.return_object is param1 - assert t.return_boundaries is param1 - assert t.precision == param2 + assert transformer.return_object is return_object + assert transformer.return_boundaries is return_boundaries + assert transformer.precision == precision +# fit and transform class MockClassFit(BaseDiscretiser): def fit(self, X): # bins are hard-coded rather than learnt, so this mock works unchanged