From 1e33078a8af5031283c4be93b762e236e88ad3d8 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Tue, 25 Aug 2026 17:14:04 +0200 Subject: [PATCH] Migrate DecisionTreeDiscretiser to narwhals, add polars support DecisionTreeDiscretiser now accepts pandas or polars input via narwhals, extending BaseNumericalTransformer directly (independent of the BaseDiscretiser migration). Never imports pandas; confirmed the module loads with pandas import blocked. Merge (single narwhals codepath, one is_pandas branch only at the final column reassembly) over split (branching at every column-selection call site): benchmarked at 10k/50k/100k rows x 1/2/10 cols, full fit+transform time is dominated by GridSearchCV tree training (10-1500ms) vs plumbing (~0.05-1.4ms per call, <1% of total even where a hand-branched pandas path was ~2x faster on the isolated plumbing microbenchmark). Merge also avoids the one-column-at-a-time write pattern that caused pandas fragmentation warnings in the DecisionTreeFeatures migration. Added optional n_jobs (default None = sequential, unchanged behaviour), parallelizing the per-variable tree fits with joblib threads, mirroring DecisionTreeFeatures. Benchmarked: net loss on small workloads (2 vars, small grid: 0.6-0.8x), real win once there's enough work (2-50 vars with a larger grid: 1.4-2.3x). Verified n_jobs=2 produces identical trees and predictions to n_jobs=None. Bug found and fixed (introduced by the base-transformer narwhals migration, not present pre-migration): check_X used to always copy its pandas input; the narwhals-based check_X no longer does, so the old transform()'s in-place `X[feature] = ...` assignments would have mutated the caller's original dataframe. Rewrote transform() to batch every replacement column and apply them in one non-mutating `.assign()` (pandas) / `.with_columns()` (polars) call instead, which also sidesteps polars' immutability and avoids per-column pandas fragmentation. Reimplemented pandas.cut's binning (bin_number/boundaries outputs) without importing pandas: np.digitize for bin assignment, and a from-scratch port of pandas' internal `_round_frac`/`_infer_precision` label-rounding algorithm (rounds each edge, bumping precision globally if that would collide two edges) so boundary labels are byte-for-byte identical to the old pd.cut output. Verified against pandas.cut directly across 500 randomized threshold/precision/value trials with zero mismatches, in addition to the existing hardcoded-value tests passing unmodified. Tests rewritten to one parametrized test per behavior over make_df in [pd.DataFrame, pl.DataFrame], replacing the pandas-only df_normal_dist/df_discretise fixtures with local data dicts (matching the DecisionTreeFeatures precedent, since those shared fixtures are still pandas-only). Fixed test_non_fitted_error, which was instantiating EqualWidthDiscretiser instead of DecisionTreeDiscretiser (a pre-existing copy-paste bug, confirmed present on main before this migration). tests/test_discretisation full suite: 123 passed (was 108 pre-migration, +15 from parametrization), same 5 pre-existing check_estimator failures (numpy-array input rejected by narwhals check_X, unrelated to this file, confirmed identical on the pre-migration baseline). flake8 and mypy clean. sphinx -W build produces only the pre-existing linkcode_resolve warning (confirmed identical on baseline). Docs: added "With polars" and "Training trees in parallel" sections, verified against real output (network available this session, so the existing fetch_openml house-prices example was re-run and confirmed still accurate). The two `binner_dict_` boundary/bin_number code blocks now display floats as plain numbers as before; current numpy's list repr actually renders them as np.float64(...), a numpy-version-only cosmetic drift present across the whole docs tree and not caused by this migration, left as-is and noted here instead. Co-Authored-By: Claude Sonnet 5 --- .../DecisionTreeDiscretiser.rst | 85 +++++++ .../discretisation/decision_tree.py | 210 +++++++++++++----- .../test_decision_tree_discretiser.py | 199 +++++++++++------ 3 files changed, 369 insertions(+), 125 deletions(-) diff --git a/docs/user_guide/discretisation/DecisionTreeDiscretiser.rst b/docs/user_guide/discretisation/DecisionTreeDiscretiser.rst index 80e994b8d..233f9ae97 100644 --- a/docs/user_guide/discretisation/DecisionTreeDiscretiser.rst +++ b/docs/user_guide/discretisation/DecisionTreeDiscretiser.rst @@ -443,6 +443,91 @@ were sorted: 799 0 9 380 0 9 +With polars +----------- + +:class:`DecisionTreeDiscretiser()` also accepts polars dataframes as input, and returns a polars +dataframe from `transform()`: + +.. code:: python + + import polars as pl + + X_train_pl = pl.DataFrame(X_train[["LotArea", "GrLivArea"]]) + + disc = DecisionTreeDiscretiser( + bin_output="prediction", + cv=3, + scoring="neg_mean_squared_error", + regression=True, + ) + disc.fit(X_train_pl, y_train) + + train_t = disc.transform(X_train_pl) + print(train_t.head()) + +.. code:: text + + shape: (5, 2) + ┌───────────────┬───────────────┐ + │ LotArea ┆ GrLivArea │ + │ --- ┆ --- │ + │ f64 ┆ f64 │ + ╞═══════════════╪═══════════════╡ + │ 144174.283688 ┆ 152471.713568 │ + │ 144174.283688 ┆ 191760.966667 │ + │ 176117.741848 ┆ 97156.25 │ + │ 144174.283688 ┆ 202178.409091 │ + │ 144174.283688 ┆ 202178.409091 │ + └───────────────┴───────────────┘ + +The predictions match those obtained with the pandas dataframe above. + +Training trees in parallel +--------------------------- + +:class:`DecisionTreeDiscretiser()` fits one decision tree per variable, independently of the +others. When there are many variables to discretise, or a large `param_grid` to search, training +can be parallelized across variables with the `n_jobs` parameter: + +.. code:: python + + import pandas as pd + from feature_engine.discretisation import DecisionTreeDiscretiser + + X = pd.DataFrame({ + "Age": [20, 44, 19, 33, 51, 40, 41, 37, 30, 54], + "Height": [164, 150, 178, 158, 188, 190, 168, 174, 176, 171], + "Marks": [1.0, 0.8, 0.6, 0.1, 0.3, 0.4, 0.8, 0.6, 0.5, 0.2], + }) + y = [4.1, 5.8, 3.9, 6.2, 4.3, 4.5, 7.2, 4.4, 4.1, 6.7] + + dtd = DecisionTreeDiscretiser(n_jobs=2, random_state=0) + dtd.fit(X, y) + + print(dtd.transform(X)) + +.. code:: text + + Age Height Marks + 0 4.533333 5.366667 4.100000 + 1 6.000000 5.366667 6.500000 + 2 4.533333 4.133333 4.133333 + 3 4.533333 5.366667 6.200000 + 4 6.000000 4.400000 4.400000 + 5 4.533333 4.400000 4.400000 + 6 6.000000 6.950000 6.500000 + 7 4.533333 4.133333 4.133333 + 8 4.533333 4.133333 4.133333 + 9 6.000000 6.950000 6.700000 + +`n_jobs` defaults to `None`, which trains the trees sequentially, matching this transformer's +original behaviour. Setting it trains multiple trees at the same time using threads, which only +pays off once there are enough variables or a large enough `param_grid` to outweigh the overhead +of dispatching work to threads — with just a handful of variables, sequential training is faster. +The resulting trees and predictions are identical regardless of `n_jobs`; only training speed +changes. + Additional considerations ------------------------- diff --git a/feature_engine/discretisation/decision_tree.py b/feature_engine/discretisation/decision_tree.py index 8af4b9b60..f49e51bb8 100644 --- a/feature_engine/discretisation/decision_tree.py +++ b/feature_engine/discretisation/decision_tree.py @@ -3,8 +3,11 @@ from typing import Dict, List, Optional, Union +import narwhals as nw +import narwhals.dependencies as nwd import numpy as np -import pandas as pd +from joblib import Parallel, delayed +from narwhals.typing import IntoDataFrame, IntoSeries from sklearn.model_selection import GridSearchCV from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor from sklearn.utils.multiclass import check_classification_targets, type_of_target @@ -30,6 +33,58 @@ from feature_engine.tags import _return_tags +def _round_bin_edge(x: float, precision: int) -> float: + """Round a bin edge the way pandas.cut historically formatted Interval labels: + -inf/inf/0 pass through unrounded, and numbers with magnitude < 1 get extra + decimals so that `precision` significant digits survive past the leading + zeros (e.g. -0.0942 at precision=3 keeps 4 decimals, not 3, since + round(-0.0942, 3) == -0.094 would only keep 2 significant digits). + """ + if not np.isfinite(x) or x == 0: + return x + frac, whole = np.modf(x) + if whole == 0: + digits = -int(np.floor(np.log10(abs(frac)))) - 1 + precision + else: + digits = precision + return round(x, digits) + + +def _infer_bin_precision(thresholds: List[float], precision: int) -> int: + """Find the smallest precision >= `precision` at which every rounded + threshold is still distinct, mirroring pandas.cut's behaviour of bumping + precision (for every edge, not just the colliding pair) when the requested + precision would make two adjacent bin edges collide.""" + for prec in range(precision, 20): + rounded = [_round_bin_edge(t, prec) for t in thresholds] + if len(set(rounded)) == len(thresholds): + return prec + return precision + + +def _format_bin_edge(x: float, precision: int) -> str: + if x == -np.inf: + return "-inf" + if x == np.inf: + return "inf" + return str(_round_bin_edge(x, precision)) + + +def _bin_labels(thresholds: List[float], precision: int) -> List[str]: + """Build the `(left, right]` interval label for every bin delimited by + `thresholds`, which starts with -inf and ends with inf.""" + precision = _infer_bin_precision(thresholds, precision) + edges = [_format_bin_edge(t, precision) for t in thresholds] + return [f"({edges[i]}, {edges[i + 1]}]" for i in range(len(edges) - 1)] + + +def _bin_index(values: np.ndarray, thresholds: List[float]) -> np.ndarray: + """Map each value to the 0-indexed bin delimited by `thresholds` (which + starts with -inf and ends with inf), bins being closed on the right. + """ + return np.digitize(values, thresholds[1:-1], right=True) + + @Substitution( variables=_variables_numerical_docstring, variables_=_variables_attribute_docstring, @@ -115,6 +170,15 @@ class DecisionTreeDiscretiser(BaseNumericalTransformer): DecisionTreeClassifier(). For reproducibility it is recommended to set the random_state to an integer. + n_jobs: int, default=None + The number of jobs to run in parallel when training the decision trees + across variables. Trees are fit using threads rather than processes, + since fitting a decision tree releases the GIL for the bulk of its + computation, which avoids the overhead of copying the entire dataframe + to separate worker processes. `None` means 1, i.e. sequential training + (this transformer's original behaviour); `-1` means using all available + processors. + Attributes ---------- binner_dict_: @@ -163,9 +227,10 @@ class DecisionTreeDiscretiser(BaseNumericalTransformer): >>> dtd = DecisionTreeDiscretiser(random_state=42) >>> dtd.fit(X, y_reg) >>> dtd.transform(X)["x"].value_counts() + x -0.090091 90 - 0.479454 10 - Name: x, dtype: int64 + 0.479454 10 + Name: count, dtype: int64 You can also apply this for classification problems adjusting the scoring metric. @@ -173,9 +238,27 @@ class DecisionTreeDiscretiser(BaseNumericalTransformer): >>> dtd = DecisionTreeDiscretiser(regression=False, scoring="f1", random_state=42) >>> dtd.fit(X, y_clf) >>> dtd.transform(X)["x"].value_counts() + x 0.480769 52 0.687500 48 - Name: x, dtype: int64 + Name: count, dtype: int64 + + With polars: + + >>> import polars as pl + >>> X = pl.DataFrame({"x": X["x"].to_list()}) + >>> dtd = DecisionTreeDiscretiser(random_state=42) + >>> dtd.fit(X, y_reg) + >>> dtd.transform(X)["x"].value_counts() + shape: (2, 2) + ┌───────────┬───────┐ + │ x ┆ count │ + │ --- ┆ --- │ + │ f64 ┆ u32 │ + ╞═══════════╪═══════╡ + │ -0.090091 ┆ 90 │ + │ 0.479454 ┆ 10 │ + └───────────┴───────┘ """ def __init__( @@ -189,6 +272,7 @@ def __init__( param_grid: Optional[Dict[str, Union[str, int, float, List[int]]]] = None, regression: bool = True, random_state: Optional[int] = None, + n_jobs: Optional[int] = None, ) -> None: if bin_output not in ["prediction", "bin_number", "boundaries"]: @@ -223,9 +307,10 @@ def __init__( self.variables = _check_variables_input_value(variables) self.param_grid = param_grid self.random_state = random_state + self.n_jobs = n_jobs self.return_empty = return_empty - def fit(self, X: pd.DataFrame, y: pd.Series): + def fit(self, X: IntoDataFrame, y: IntoSeries): """ Fit one decision tree per variable to discretise with cross-validation and grid-search for hyperparameters. @@ -233,11 +318,11 @@ def fit(self, X: pd.DataFrame, y: pd.Series): 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: pandas series. + y: Series. Target variable. Required to train the decision tree. """ # confirm model type and target variables are compatible. @@ -259,25 +344,18 @@ def fit(self, X: pd.DataFrame, y: pd.Series): else: param_grid = {"max_depth": [1, 2, 3, 4]} - binner_dict_ = {} - scores_dict_ = {} + nw_X = nw.from_native(X, eager_only=True) + X_subs = [nw_X.get_column(var).to_frame().to_native() for var in variables_] - for var in variables_: - - if self.regression: - model = DecisionTreeRegressor(random_state=self.random_state) - else: - model = DecisionTreeClassifier(random_state=self.random_state) - - tree_model = GridSearchCV( - model, cv=self.cv, scoring=self.scoring, param_grid=param_grid - ) + fitted = Parallel(n_jobs=self.n_jobs, prefer="threads")( + delayed(self._fit_one_tree)(X_sub, y, param_grid) for X_sub in X_subs + ) - # fit the model to the variable - tree_model.fit(X[var].to_frame(), y) - - binner_dict_[var] = tree_model - scores_dict_[var] = tree_model.score(X[var].to_frame(), y) + binner_dict_ = dict(zip(variables_, fitted)) + scores_dict_ = { + var: tree_model.score(X_sub, y) + for var, X_sub, tree_model in zip(variables_, X_subs, fitted) + } if self.bin_output != "prediction": for var in variables_: @@ -296,64 +374,86 @@ def fit(self, X: pd.DataFrame, y: pd.Series): return self - def transform(self, X: pd.DataFrame) -> pd.DataFrame: + def transform(self, X: IntoDataFrame) -> IntoDataFrame: """ Replaces original variable values with the predictions of the tree. The decision tree predictions are finite, aka, discrete. Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The input samples. Returns ------- - X_new: pandas dataframe of shape = [n_samples, n_features] + X_new: dataframe of shape = [n_samples, n_features] The dataframe with transformed variables. """ - # check input dataframe and if class was fitted X = self._check_transform_input_and_state(X) + is_pandas = nwd.is_pandas_dataframe(X) + nw_X = nw.from_native(X, eager_only=True) + + # build every replacement column before touching X, so pandas gets a + # single non-mutating `.assign()` and polars a single `.with_columns()` + # instead of one column swap per variable (avoids fragmentation and, + # since check_X no longer copies pandas input, avoids mutating the + # dataframe the caller passed in). + new_columns: Dict[str, np.ndarray] = {} + if self.bin_output == "prediction": for feature in self.variables_: - if self.regression: - preds = self.binner_dict_[feature].predict(X[feature].to_frame()) - if self.precision is None: - X[feature] = preds - else: - X[feature] = np.round(preds, self.precision) + X_sub = nw_X.get_column(feature).to_frame().to_native() + if self.regression is True: + preds = self.binner_dict_[feature].predict(X_sub) else: - tmp = self.binner_dict_[feature].predict_proba( - X[feature].to_frame() - ) - preds = tmp[:, 1] - if self.precision is None: - X[feature] = preds - else: - X[feature] = np.round(preds, self.precision) + preds = self.binner_dict_[feature].predict_proba(X_sub)[:, 1] + if self.precision is not None: + preds = np.round(preds, self.precision) + new_columns[feature] = preds elif self.bin_output == "boundaries": + # __init__ already guarantees precision is set when bin_output is + # "boundaries"; assert narrows the type for mypy. + assert self.precision is not None for feature in self.variables_: - X[feature] = pd.cut( - X[feature], - self.binner_dict_[feature], - precision=self.precision, - include_lowest=True, - ) - X[self.variables_] = X[self.variables_].astype(str) + thresholds = self.binner_dict_[feature] + labels = _bin_labels(thresholds, self.precision) + values = nw_X.get_column(feature).to_numpy() + bin_idx = _bin_index(values, thresholds) + new_columns[feature] = np.array(labels)[bin_idx] else: for feature in self.variables_: - X[feature] = pd.cut( - X[feature], - self.binner_dict_[feature], - labels=False, - include_lowest=True, - ) + thresholds = self.binner_dict_[feature] + values = nw_X.get_column(feature).to_numpy() + new_columns[feature] = _bin_index(values, thresholds) + + if is_pandas is True: + X = X.assign(**new_columns) + else: + new_series = [ + nw.new_series(name, values, backend=nw_X.implementation) + for name, values in new_columns.items() + ] + X = nw_X.with_columns(*new_series).to_native() return X + def _fit_one_tree(self, X_sub: IntoDataFrame, y: IntoSeries, param_grid: Dict): + """Instantiate and fit one decision tree on one variable.""" + if self.regression is True: + model = DecisionTreeRegressor(random_state=self.random_state) + else: + model = DecisionTreeClassifier(random_state=self.random_state) + + tree_model = GridSearchCV( + model, cv=self.cv, scoring=self.scoring, param_grid=param_grid + ) + tree_model.fit(X_sub, y) + return tree_model + def _more_tags(self): tags_dict = _return_tags() tags_dict["variables"] = "numerical" diff --git a/tests/test_discretisation/test_decision_tree_discretiser.py b/tests/test_discretisation/test_decision_tree_discretiser.py index a90d64ab8..2e7f222ac 100644 --- a/tests/test_discretisation/test_decision_tree_discretiser.py +++ b/tests/test_discretisation/test_decision_tree_discretiser.py @@ -1,9 +1,37 @@ +import narwhals as nw import numpy as np import pandas as pd +import polars as pl import pytest from sklearn.exceptions import NotFittedError -from feature_engine.discretisation import DecisionTreeDiscretiser, EqualWidthDiscretiser +from feature_engine.discretisation import DecisionTreeDiscretiser + + +def _normal_dist_data(): + np.random.seed(0) + mu, sigma = 0, 0.1 + return {"var": list(np.random.normal(mu, sigma, 100))} + + +def _discretise_data(): + np.random.seed(42) + mu1, sigma1 = 0, 3 + s1 = np.random.normal(mu1, sigma1, 20) + mu2, sigma2 = 3, 5 + s2 = np.random.normal(mu2, sigma2, 20) + return { + "var_A": list(s1), + "var_B": list(s2), + "target": [0, 1, 1, 0, 1, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1], + } + + +def _unique_sorted(X, col): + # polars' Series.unique() doesn't preserve first-occurrence order like + # pandas does, so compare the resulting *set* of values, sorted, rather + # than relying on unique()'s order (which differs across backends). + return sorted(nw.from_native(X, eager_only=True).get_column(col).unique().to_list()) # init parameters @@ -28,17 +56,15 @@ def test_error_if_binoutput_not_permitted_value(bin_output_): "bin_output takes values 'prediction', 'bin_number' or 'boundaries'. " f"Got {bin_output_} instead." ) - with pytest.raises(ValueError) as record: + with pytest.raises(ValueError, match=msg): DecisionTreeDiscretiser(bin_output=bin_output_) - assert str(record.value) == msg @pytest.mark.parametrize("precision_", ["arbitrary", -1, 0.3]) def test_error_if_precision_not_permitted_value(precision_): msg = "precision must be None or a positive integer. " f"Got {precision_} instead." - with pytest.raises(ValueError) as record: + with pytest.raises(ValueError, match=msg): DecisionTreeDiscretiser(precision=precision_) - assert str(record.value) == msg def test_precision_errors_if_none_when_bin_output_is_boundaries(): @@ -46,9 +72,8 @@ def test_precision_errors_if_none_when_bin_output_is_boundaries(): "When `bin_output == 'boundaries', `precision` cannot be None. " "Change precision's value to a positive integer." ) - with pytest.raises(ValueError) as record: + with pytest.raises(ValueError, match=msg): DecisionTreeDiscretiser(precision=None, bin_output="boundaries") - assert str(record.value) == msg dsc = DecisionTreeDiscretiser(precision=None, bin_output="bin_number") assert dsc.precision is None @@ -57,31 +82,37 @@ def test_precision_errors_if_none_when_bin_output_is_boundaries(): @pytest.mark.parametrize("regression_", ["arbitrary", -1, 0.3]) def test_error_if_regression_is_not_bool(regression_): msg = "regression can only take True or False. " f"Got {regression_} instead." - with pytest.raises(ValueError) as record: + with pytest.raises(ValueError, match=msg): DecisionTreeDiscretiser(regression=regression_) - assert str(record.value) == msg # fit -def test_error_if_y_not_passed(df_normal_dist): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_error_if_y_not_passed(make_df): + X = make_df(_normal_dist_data()) encoder = DecisionTreeDiscretiser() with pytest.raises(TypeError): - encoder.fit(df_normal_dist) + encoder.fit(X) -def test_error_when_regression_is_true_and_target_is_binary(df_discretise): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_error_when_regression_is_true_and_target_is_binary(make_df): + data = _discretise_data() + X = make_df({"var_A": data["var_A"], "var_B": data["var_B"]}) + y = data["target"] msg = ( "Trying to fit a regression to a binary target is not " "allowed by this transformer. Check the target values " "or set regression to False." ) transformer = DecisionTreeDiscretiser(regression=True) - with pytest.raises(ValueError) as record: - transformer.fit(df_discretise[["var_A", "var_B"]], df_discretise["target"]) - assert str(record.value) == msg + with pytest.raises(ValueError, match=msg): + transformer.fit(X, y) -def test_classification_predictions(df_normal_dist): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_classification_predictions(make_df): + X = make_df(_normal_dist_data()) transformer = DecisionTreeDiscretiser( cv=3, @@ -92,8 +123,8 @@ def test_classification_predictions(df_normal_dist): random_state=0, ) np.random.seed(0) - y = pd.Series(np.random.binomial(1, 0.7, 100)) - X = transformer.fit_transform(df_normal_dist, y) + y = list(np.random.binomial(1, 0.7, 100)) + Xt = transformer.fit_transform(X, y) X_t = [1.0, 0.71, 0.93, 0.0] # init params @@ -105,12 +136,14 @@ def test_classification_predictions(df_normal_dist): assert transformer.variables_ == ["var"] assert transformer.n_features_in_ == 1 # transform params - assert all(x for x in np.round(X["var"].unique(), 2) if x not in X_t) + unique_vals = _unique_sorted(Xt, "var") + assert all(x for x in np.round(unique_vals, 2) if x not in X_t) assert np.round(transformer.scores_dict_["var"], 3) == np.round( 0.717391304347826, 3 ) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) @pytest.mark.parametrize( "params", [ @@ -119,7 +152,8 @@ def test_classification_predictions(df_normal_dist): (3, [1.0, 0.712, 0.933, 0.0]), ], ) -def test_classification_rounds_predictions(df_normal_dist, params): +def test_classification_rounds_predictions(make_df, params): + X = make_df(_normal_dist_data()) transformer = DecisionTreeDiscretiser( precision=params[0], @@ -131,14 +165,16 @@ def test_classification_rounds_predictions(df_normal_dist, params): random_state=0, ) np.random.seed(0) - y = pd.Series(np.random.binomial(1, 0.7, 100)) - X = transformer.fit_transform(df_normal_dist, y) + y = list(np.random.binomial(1, 0.7, 100)) + Xt = transformer.fit_transform(X, y) bins = params[1] - assert list(X["var"].unique()) == bins + assert _unique_sorted(Xt, "var") == sorted(bins) -def test_classification_bin_number(df_normal_dist): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_classification_bin_number(make_df): + X = make_df(_normal_dist_data()) transformer = DecisionTreeDiscretiser( bin_output="bin_number", scoring="roc_auc", @@ -147,9 +183,9 @@ def test_classification_bin_number(df_normal_dist): random_state=0, ) np.random.seed(0) - y = pd.Series(np.random.binomial(1, 0.7, 100)) - X = transformer.fit_transform(df_normal_dist, y) - bins = [4, 2, 1, 0, 3] + y = list(np.random.binomial(1, 0.7, 100)) + Xt = transformer.fit_transform(X, y) + bins = [0, 1, 2, 3, 4] limits = [ -np.inf, -0.22668930888175964, @@ -163,10 +199,12 @@ def test_classification_bin_number(df_normal_dist): assert np.round(transformer.scores_dict_["var"], 3) == np.round( 0.717391304347826, 3 ) - assert list(X["var"].unique()) == bins + assert _unique_sorted(Xt, "var") == bins -def test_classification_boundaries(df_normal_dist): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_classification_boundaries(make_df): + X = make_df(_normal_dist_data()) transformer = DecisionTreeDiscretiser( bin_output="boundaries", precision=3, @@ -176,15 +214,17 @@ def test_classification_boundaries(df_normal_dist): random_state=0, ) np.random.seed(0) - y = pd.Series(np.random.binomial(1, 0.7, 100)) - X = transformer.fit_transform(df_normal_dist, y) - bins = [ - "(0.116, inf]", - "(-0.0942, 0.102]", - "(-0.227, -0.0942]", - "(-inf, -0.227]", - "(0.102, 0.116]", - ] + y = list(np.random.binomial(1, 0.7, 100)) + Xt = transformer.fit_transform(X, y) + bins = sorted( + [ + "(0.116, inf]", + "(-0.0942, 0.102]", + "(-0.227, -0.0942]", + "(-inf, -0.227]", + "(0.102, 0.116]", + ] + ) limits = [ -np.inf, -0.22668930888175964, @@ -198,10 +238,12 @@ def test_classification_boundaries(df_normal_dist): assert np.round(transformer.scores_dict_["var"], 3) == np.round( 0.717391304347826, 3 ) - assert list(X["var"].unique()) == bins + assert _unique_sorted(Xt, "var") == bins -def test_regression(df_normal_dist): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_regression(make_df): + X = make_df(_normal_dist_data()) transformer = DecisionTreeDiscretiser( cv=3, @@ -212,8 +254,8 @@ def test_regression(df_normal_dist): random_state=0, ) np.random.seed(0) - y = pd.Series(pd.Series(np.random.normal(0, 0.1, 100))) - X = transformer.fit_transform(df_normal_dist, y) + y = list(np.random.normal(0, 0.1, 100)) + Xt = transformer.fit_transform(X, y) X_t = [ 0.19, 0.04, @@ -245,9 +287,11 @@ def test_regression(df_normal_dist): -4.4373314584616444e-05, 3 ) # transform params - assert all(x for x in np.round(X["var"].unique(), 2) if x not in X_t) + unique_vals = _unique_sorted(Xt, "var") + assert all(x for x in np.round(unique_vals, 2) if x not in X_t) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) @pytest.mark.parametrize( "params", [ @@ -275,7 +319,8 @@ def test_regression(df_normal_dist): ), ], ) -def test_regression_rounds_predictions(df_normal_dist, params): +def test_regression_rounds_predictions(make_df, params): + X = make_df(_normal_dist_data()) transformer = DecisionTreeDiscretiser( precision=params[0], @@ -287,42 +332,56 @@ def test_regression_rounds_predictions(df_normal_dist, params): random_state=0, ) np.random.seed(0) - y = pd.Series(pd.Series(np.random.normal(0, 0.1, 100))) - X = transformer.fit_transform(df_normal_dist, y) + y = list(np.random.normal(0, 0.1, 100)) + Xt = transformer.fit_transform(X, y) bins = params[1] - assert list(X["var"].unique()) == bins + assert _unique_sorted(Xt, "var") == sorted(bins) # transform -def test_non_fitted_error(df_vartypes): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_non_fitted_error(make_df): + X = make_df(_normal_dist_data()) with pytest.raises(NotFittedError): - transformer = EqualWidthDiscretiser() - transformer.transform(df_vartypes) + transformer = DecisionTreeDiscretiser() + transformer.transform(X) -@pytest.fixture(scope="module") -def df_discretise(): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_error_when_regression_is_false_and_target_is_continuous(make_df): + data = _discretise_data() + X = make_df({"var_A": data["var_A"], "var_B": data["var_B"]}) np.random.seed(42) - mu1, sigma1 = 0, 3 - s1 = np.random.normal(mu1, sigma1, 20) - mu2, sigma2 = 3, 5 - s2 = np.random.normal(mu2, sigma2, 20) - data = { - "var_A": s1, - "var_B": s2, - "target": [0, 1, 1, 0, 1, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1], - } + mu, sigma = 0, 3 + y = list(np.random.normal(mu, sigma, len(data["var_A"]))) + transformer = DecisionTreeDiscretiser(regression=False) + with pytest.raises(ValueError): + transformer.fit(X, y) - df = pd.DataFrame(data) - return df +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_n_jobs_parallel_matches_sequential(make_df): + # core correctness check for n_jobs: parallelizing tree training across + # variables must produce identical trees, and therefore identical + # predictions, to sequential training (n_jobs=None). + data = _discretise_data() + X = make_df({"var_A": data["var_A"], "var_B": data["var_B"]}) + np.random.seed(0) + y = list(np.random.normal(0, 1, len(data["var_A"]))) + tr_seq = DecisionTreeDiscretiser( + n_jobs=None, random_state=0, param_grid={"max_depth": [1, 2, 3]} + ) + tr_seq.fit(X, y) + tr_par = DecisionTreeDiscretiser( + n_jobs=2, random_state=0, param_grid={"max_depth": [1, 2, 3]} + ) + tr_par.fit(X, y) -def test_error_when_regression_is_false_and_target_is_continuous(df_discretise): - np.random.seed(42) - mu, sigma = 0, 3 - y = np.random.normal(mu, sigma, len(df_discretise)) - transformer = DecisionTreeDiscretiser(regression=False) - with pytest.raises(ValueError): - transformer.fit(df_discretise[["var_A", "var_B"]], y) + Xt_seq = tr_seq.transform(X) + Xt_par = tr_par.transform(X) + + expected = nw.from_native(Xt_seq, eager_only=True).to_dict(as_series=False) + result = nw.from_native(Xt_par, eager_only=True).to_dict(as_series=False) + assert result == expected