diff --git a/docs/user_guide/encoding/DecisionTreeEncoder.rst b/docs/user_guide/encoding/DecisionTreeEncoder.rst index 9c5ab968d..5e220597c 100644 --- a/docs/user_guide/encoding/DecisionTreeEncoder.rst +++ b/docs/user_guide/encoding/DecisionTreeEncoder.rst @@ -438,6 +438,63 @@ In the following image we also see a monotonic relationship after the encoding: be some sort of relationship between the target and the categories that can be captured by the decision tree. Use with caution. +With polars +----------- + +:class:`DecisionTreeEncoder()` works the same way with a polars dataframe. Let's create a toy +dataset: + +.. code:: python + + import polars as pl + from feature_engine.encoding import DecisionTreeEncoder + + X = pl.DataFrame({ + "city": ["London", "Manchester", "Liverpool", "London", "Manchester", "Liverpool"], + "price": [500, 300, 250, 520, 310, 260], + }) + y = pl.Series("target", [1, 0, 0, 1, 0, 1]) + +Let's set up :class:`DecisionTreeEncoder()` to encode `city` with a classification tree, and fit +it to the data: + +.. code:: python + + encoder = DecisionTreeEncoder(variables=["city"], regression=False, cv=2) + encoder.fit(X, y) + + encoder.encoder_dict_ + +We see the resulting mappings from category to the tree's predictions: + +.. code:: python + + {'city': {'London': 1.0, 'Manchester': 0.25, 'Liverpool': 0.25}} + +Now let's transform the data: + +.. code:: python + + encoder.transform(X) + +We obtain a polars dataframe with the categories in `city` replaced by the tree's predictions: + +.. code:: text + + shape: (6, 2) + ┌──────┬───────┐ + │ city ┆ price │ + │ --- ┆ --- │ + │ f64 ┆ i64 │ + ╞══════╪═══════╡ + │ 1.0 ┆ 500 │ + │ 0.25 ┆ 300 │ + │ 0.25 ┆ 250 │ + │ 1.0 ┆ 520 │ + │ 0.25 ┆ 310 │ + │ 0.25 ┆ 260 │ + └──────┴───────┘ + Additional resources -------------------- diff --git a/feature_engine/encoding/decision_tree.py b/feature_engine/encoding/decision_tree.py index 9b66fe32d..0ff35d8a4 100644 --- a/feature_engine/encoding/decision_tree.py +++ b/feature_engine/encoding/decision_tree.py @@ -1,11 +1,15 @@ # Authors: Soledad Galli # License: BSD 3 clause -from typing import List, Optional, Union +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 sklearn.pipeline import Pipeline +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 from feature_engine._docstrings.fit_attributes import ( @@ -27,13 +31,11 @@ ) from feature_engine._docstrings.substitute import Substitution from feature_engine.dataframe_checks import _check_contains_na, check_X_y -from feature_engine.discretisation import DecisionTreeDiscretiser from feature_engine.encoding._helper_functions import check_parameter_unseen from feature_engine.encoding.base_encoder import ( CategoricalInitMixin, CategoricalMethodsMixin, ) -from feature_engine.encoding.ordinal import OrdinalEncoder from feature_engine.tags import _return_tags _unseen_docstring = ( @@ -139,6 +141,15 @@ class DecisionTreeEncoder(CategoricalMethodsMixin, CategoricalInitMixin): fill_value: float, default=None The value used to encode unseen categories. Only used when `unseen='encode'`. + 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 ---------- encoder_dict_: @@ -212,6 +223,27 @@ class DecisionTreeEncoder(CategoricalMethodsMixin, CategoricalInitMixin): 2 3 0.666667 3 4 0.500000 4 5 0.500000 + + With polars: + + >>> import polars as pl + >>> X = pl.DataFrame(dict(x1 = [1,2,3,4,5], x2 = ["b", "b", "b", "a", "a"])) + >>> y = [0, 1, 1, 1, 0] + >>> dte = DecisionTreeEncoder(regression=False, cv=2) + >>> dte.fit(X, y) + >>> dte.transform(X) + shape: (5, 2) + ┌─────┬──────────┐ + │ x1 ┆ x2 │ + │ --- ┆ --- │ + │ i64 ┆ f64 │ + ╞═════╪══════════╡ + │ 1 ┆ 0.666667 │ + │ 2 ┆ 0.666667 │ + │ 3 ┆ 0.666667 │ + │ 4 ┆ 0.5 │ + │ 5 ┆ 0.5 │ + └─────┴──────────┘ """ def __init__( @@ -228,6 +260,7 @@ def __init__( precision: Optional[int] = None, unseen: str = "ignore", fill_value: Optional[float] = None, + n_jobs: Optional[int] = None, ) -> None: if encoding_method not in ["ordered", "arbitrary"]: @@ -261,22 +294,23 @@ def __init__( self.precision = precision self.unseen = unseen self.fill_value = fill_value + self.n_jobs = n_jobs - def fit(self, X: pd.DataFrame, y: pd.Series): + def fit(self, X: IntoDataFrame, y: IntoSeries): """ Fit a decision tree per variable. Parameters ---------- - X : pandas dataframe of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The training input samples. Can be the entire dataframe, not just the categorical variables. - y : pandas series. + y: Series. The target variable. Required to train the decision tree and for ordered ordinal encoding. """ - X, y = check_X_y(X, y) + nw_X, y = check_X_y(X, y) # confirm model type and target variables are compatible. if self.regression is True: @@ -303,68 +337,112 @@ def fit(self, X: pd.DataFrame, y: pd.Series): self._get_feature_names_in(X) return self - encoder = OrdinalEncoder( - encoding_method=self.encoding_method, - variables=variables_, - missing_values="raise", - ignore_format=self.ignore_format, - ) - - tree = DecisionTreeDiscretiser( - cv=self.cv, - scoring=self.scoring, - variables=variables_, - param_grid=param_grid, - regression=self.regression, - random_state=self.random_state, - ) + # only needed for "ordered": pairs the target with X once so every + # variable's group_by below can reuse it, instead of rebuilding it + # per variable. + nw_Xy = None + target_name = "__feature_engine_decision_tree_target__" + if self.encoding_method == "ordered": + if nwd.is_into_series(y): + y_nw = nw.from_native(y, series_only=True).alias(target_name) + else: + y_nw = nw.new_series( + name=target_name, values=y, backend=nw_X.implementation + ) + nw_Xy = nw_X.with_columns(y_nw) - # pipeline for the encoder - pipe = Pipeline( - [ - ("encoder", encoder), - ("tree", tree), - ] + mappings = Parallel(n_jobs=self.n_jobs, prefer="threads")( + delayed(self._fit_one_variable)( + nw_X, nw_Xy, var, y, target_name, param_grid + ) + for var in variables_ ) - Xt = pipe.fit_transform(X, y) - - encoder_ = {} - if self.precision is None: - for var in variables_: - encoder_[var] = dict(zip(X[var], Xt[var])) - else: - for var in variables_: - encoder_[var] = dict(zip(X[var], np.round(Xt[var], self.precision))) - if self.unseen == "encode": self._unseen = self.fill_value - self.encoder_dict_ = encoder_ + self.encoder_dict_ = dict(zip(variables_, mappings)) self.variables_ = variables_ self._get_feature_names_in(X) return self - def transform(self, X: pd.DataFrame) -> pd.DataFrame: + def transform(self, X: IntoDataFrame) -> IntoDataFrame: """ Replace categorical variables by the predictions of the decision tree. 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]. Dataframe with variables encoded with decision tree predictions. """ - X = self._check_transform_input_and_state(X) + nw_X = self._check_transform_input_and_state(X) _check_contains_na(X, self.variables_) - X = self._encode(X) + X = self._encode(nw_X) return X + def _fit_one_variable( + self, nw_X, nw_Xy, var, y: IntoSeries, target_name: str, param_grid: Dict + ) -> Dict: + """Learn the category-to-prediction mapping for one variable: encode its + categories to ordinal integers, fit a decision tree on those integers, and + predict on each unique category to get its final mapped value.""" + if self.encoding_method == "ordered": + # sort by (mean, category): group_by's own order isn't guaranteed + # across backends, and this tie-break on the category itself + # reproduces pandas' groupby(sort=True) + stable sort_values + # behavior for categories with equal target means. + categories = ( + nw_Xy.group_by(var, drop_null_keys=True) + .agg(nw.col(target_name).mean()) + .sort([target_name, var]) + .get_column(var) + .to_list() + ) + else: + categories = nw_X.get_column(var).unique(maintain_order=True).to_list() + + ordinal_map = {k: i for i, k in enumerate(categories, 0)} + + X_sub = nw_X.get_column(var).replace_strict(ordinal_map).to_frame().to_native() + estimator = self._fit_one_tree(X_sub, y, param_grid) + + # predict directly on the (few) unique ordinal codes instead of the + # full column: the tree's prediction for a category depends only on + # its ordinal code, so this gives identical results far cheaper. + X_pred = nw.new_series( + var, list(range(len(categories))), backend=nw_X.implementation + ).to_frame().to_native() + + if self.regression is True: + preds = estimator.predict(X_pred) + else: + preds = estimator.predict_proba(X_pred)[:, 1] + + if self.precision is not None: + preds = np.round(preds, self.precision) + + return dict(zip(categories, preds)) + + def _fit_one_tree(self, X_sub: IntoDataFrame, y: IntoSeries, param_grid: Dict): + """Instantiate and fit one decision tree on one variable's ordinal-encoded + values.""" + 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 _assign_param_grid(self): if self.param_grid: param_grid = self.param_grid diff --git a/tests/test_encoding/test_decision_tree_encoder.py b/tests/test_encoding/test_decision_tree_encoder.py index fd4cef789..bf31b342c 100644 --- a/tests/test_encoding/test_decision_tree_encoder.py +++ b/tests/test_encoding/test_decision_tree_encoder.py @@ -1,7 +1,10 @@ +import math import re +import narwhals as nw import numpy as np import pandas as pd +import polars as pl import pytest from sklearn.exceptions import NotFittedError @@ -9,6 +12,39 @@ from feature_engine.encoding import DecisionTreeEncoder +def _to_backend(df: pd.DataFrame, make_df): + """Rebuild a pandas fixture dataframe on the requested backend. + + Swaps float NaN for None in string columns - polars (unlike pandas) + rejects a float NaN mixed into an otherwise-string column. + """ + data = {} + for col in df.columns: + values = df[col].tolist() + if any(isinstance(v, str) for v in values): + values = [ + None if isinstance(v, float) and math.isnan(v) else v for v in values + ] + data[col] = values + return make_df(data) + + +def _assert_values(X, expected: dict) -> None: + """NaN-aware, backend-agnostic comparison of a dataframe's contents.""" + result = nw.from_native(X, eager_only=True).to_dict(as_series=False) + assert list(result.keys()) == list(expected.keys()) + for col, exp_values in expected.items(): + got_values = result[col] + assert len(got_values) == len(exp_values) + for got, exp in zip(got_values, exp_values): + if isinstance(exp, float) and math.isnan(exp): + assert got is None or (isinstance(got, float) and math.isnan(got)) + elif isinstance(exp, float): + assert got == pytest.approx(exp) + else: + assert got == exp + + # init parameters @pytest.mark.parametrize("enc_method", ["count", False, 1]) def test_error_if_encoding_method_not_permitted_value(enc_method): @@ -69,9 +105,13 @@ def test_init_param_assignment( # fit attributes -def test_encoding_dictionary(df_enc): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_encoding_dictionary(df_enc, make_df): + X = _to_backend(df_enc[["var_A", "var_B"]], make_df) + y = df_enc["target"].tolist() + encoder = DecisionTreeEncoder(regression=False) - encoder.fit(df_enc[["var_A", "var_B"]], df_enc["target"]) + encoder.fit(X, y) # Tree: var_A <= 1.5 -> 0.25 else 0.5 # Tree: var_B <= 0.5 -> 0.2 else 0.4 @@ -82,9 +122,34 @@ def test_encoding_dictionary(df_enc): assert encoder.encoder_dict_ == expected_encodings -def test_precision(df_enc): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_ordered_encoding_dictionary(df_enc, make_df): + X = _to_backend(df_enc[["var_A", "var_B"]], make_df) + y = df_enc["target"].tolist() + + encoder = DecisionTreeEncoder(regression=False, encoding_method="ordered") + encoder.fit(X, y) + + # ordered ranks: var_A -> B(mean 0.2) < A(mean 0.333) < C(mean 0.5) + # var_B -> A(mean 0.2) < B(mean 0.333) < C(mean 0.5) + # so var_A's ordinal codes are B=0, A=1, C=2 (split at code <= 0.5, + # i.e. B alone); var_B's are A=0, B=1, C=2 (split at code <= 0.5, i.e. A + # alone). Same tree-split logic as the arbitrary-encoding case above, + # applied to the reordered codes. + expected_encodings = { + "var_A": {"B": 0.2, "A": 0.4, "C": 0.4}, + "var_B": {"A": 0.2, "B": 0.4, "C": 0.4}, + } + assert encoder.encoder_dict_ == expected_encodings + + +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_precision(df_enc, make_df): + X = _to_backend(df_enc[["var_A", "var_B"]], make_df) + y = df_enc["target"].tolist() + encoder = DecisionTreeEncoder(regression=False, precision=1) - encoder.fit(df_enc[["var_A", "var_B"]], df_enc["target"]) + encoder.fit(X, y) # Tree: var_A <= 1.5 -> 0.25 else 0.5 # Tree: var_B <= 0.5 -> 0.2 else 0.4 @@ -95,73 +160,100 @@ def test_precision(df_enc): assert encoder.encoder_dict_ == expected_encodings -def test_classification(df_enc): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_classification(df_enc, make_df): + X = _to_backend(df_enc[["var_A", "var_B"]], make_df) + y = df_enc["target"].tolist() + encoder = DecisionTreeEncoder(regression=False) - encoder.fit(df_enc[["var_A", "var_B"]], df_enc["target"]) - X = encoder.transform(df_enc[["var_A", "var_B"]]) + encoder.fit(X, y) + Xt = encoder.transform(X) - transf_df = df_enc.copy() - transf_df["var_A"] = [0.25] * 16 + [0.5] * 4 # Tree: var_A <= 1.5 -> 0.25 else 0.5 - transf_df["var_B"] = [0.2] * 10 + [0.4] * 10 # Tree: var_B <= 0.5 -> 0.2 else 0.4 - pd.testing.assert_frame_equal(X, transf_df[["var_A", "var_B"]]) + expected = { + "var_A": [0.25] * 16 + [0.5] * 4, # Tree: var_A <= 1.5 -> 0.25 else 0.5 + "var_B": [0.2] * 10 + [0.4] * 10, # Tree: var_B <= 0.5 -> 0.2 else 0.4 + } + _assert_values(Xt, expected) -def test_regression(df_enc): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_regression(df_enc, make_df): + X = _to_backend(df_enc[["var_A", "var_B"]], make_df) random = np.random.RandomState(42) y = random.normal(0, 0.1, len(df_enc)) encoder = DecisionTreeEncoder( regression=True, random_state=random, ) - encoder.fit(df_enc[["var_A", "var_B"]], y) - X = encoder.transform(df_enc[["var_A", "var_B"]]) + encoder.fit(X, y) + Xt = encoder.transform(X) + + expected = { + "var_A": ( + [0.034348] * 6 + [-0.024679] * 10 + [-0.075473] * 4 + ), # Tree: var_A <= 1.5 -> 0.25 else 0.5 + "var_B": [0.044806] * 10 + [-0.079066] * 10, + } + nw_Xt = nw.from_native(Xt, eager_only=True) + rounded = { + col: [round(v, 6) for v in nw_Xt.get_column(col).to_list()] + for col in ["var_A", "var_B"] + } + assert rounded == expected - transf_df = df_enc.copy() - transf_df["var_A"] = ( - [0.034348] * 6 + [-0.024679] * 10 + [-0.075473] * 4 - ) # Tree: var_A <= 1.5 -> 0.25 else 0.5 - transf_df["var_B"] = [0.044806] * 10 + [-0.079066] * 10 - pd.testing.assert_frame_equal(X.round(6), transf_df[["var_A", "var_B"]]) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_fit_raises_error_if_df_contains_na(df_enc_na, make_df): + X = _to_backend(df_enc_na[["var_A", "var_B"]], make_df) + y = df_enc_na["target"].tolist() -def test_fit_raises_error_if_df_contains_na(df_enc_na): - # test case 4: when dataset contains na, fit method encoder = DecisionTreeEncoder(regression=False) msg = ( "Some of the variables in the dataset contain NaN. Check and " "remove those before using this transformer." ) with pytest.raises(ValueError, match=msg): - encoder.fit(df_enc_na[["var_A", "var_B"]], df_enc_na["target"]) + encoder.fit(X, y) + +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_transform_raises_error_if_df_contains_na(df_enc, df_enc_na, make_df): + X = _to_backend(df_enc[["var_A", "var_B"]], make_df) + X_na = _to_backend(df_enc_na[["var_A", "var_B"]], make_df) + y = df_enc["target"].tolist() -def test_transform_raises_error_if_df_contains_na(df_enc, df_enc_na): - # test case 4: when dataset contains na, transform method encoder = DecisionTreeEncoder(regression=False) - encoder.fit(df_enc[["var_A", "var_B"]], df_enc["target"]) + encoder.fit(X, y) msg = ( "Some of the variables in the dataset contain NaN. Check and " "remove those before using this transformer." ) with pytest.raises(ValueError, match=msg): - encoder.transform(df_enc_na[["var_A", "var_B"]]) + encoder.transform(X_na) + +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_classification_ignore_format(df_enc_numeric, make_df): + X = _to_backend(df_enc_numeric[["var_A", "var_B"]], make_df) + y = df_enc_numeric["target"].tolist() -def test_classification_ignore_format(df_enc_numeric): encoder = DecisionTreeEncoder( regression=False, ignore_format=True, ) - encoder.fit(df_enc_numeric[["var_A", "var_B"]], df_enc_numeric["target"]) - X = encoder.transform(df_enc_numeric[["var_A", "var_B"]]) + encoder.fit(X, y) + Xt = encoder.transform(X) - transf_df = df_enc_numeric.copy() - transf_df["var_A"] = [0.25] * 16 + [0.5] * 4 # Tree: var_A <= 1.5 -> 0.25 else 0.5 - transf_df["var_B"] = [0.2] * 10 + [0.4] * 10 # Tree: var_B <= 0.5 -> 0.2 else 0.4 - pd.testing.assert_frame_equal(X, transf_df[["var_A", "var_B"]]) + expected = { + "var_A": [0.25] * 16 + [0.5] * 4, # Tree: var_A <= 1.5 -> 0.25 else 0.5 + "var_B": [0.2] * 10 + [0.4] * 10, # Tree: var_B <= 0.5 -> 0.2 else 0.4 + } + _assert_values(Xt, expected) -def test_regression_ignore_format(df_enc_numeric): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_regression_ignore_format(df_enc_numeric, make_df): + X = _to_backend(df_enc_numeric[["var_A", "var_B"]], make_df) random = np.random.RandomState(42) y = random.normal(0, 0.1, len(df_enc_numeric)) encoder = DecisionTreeEncoder( @@ -169,18 +261,25 @@ def test_regression_ignore_format(df_enc_numeric): random_state=random, ignore_format=True, ) - encoder.fit(df_enc_numeric[["var_A", "var_B"]], y) - X = encoder.transform(df_enc_numeric[["var_A", "var_B"]]) - - transf_df = df_enc_numeric.copy() - transf_df["var_A"] = ( - [0.034348] * 6 + [-0.024679] * 10 + [-0.075473] * 4 - ) # Tree: var_A <= 1.5 -> 0.25 else 0.5 - transf_df["var_B"] = [0.044806] * 10 + [-0.079066] * 10 - pd.testing.assert_frame_equal(X.round(6), transf_df[["var_A", "var_B"]]) + encoder.fit(X, y) + Xt = encoder.transform(X) + + expected = { + "var_A": ( + [0.034348] * 6 + [-0.024679] * 10 + [-0.075473] * 4 + ), # Tree: var_A <= 1.5 -> 0.25 else 0.5 + "var_B": [0.044806] * 10 + [-0.079066] * 10, + } + nw_Xt = nw.from_native(Xt, eager_only=True) + rounded = { + col: [round(v, 6) for v in nw_Xt.get_column(col).to_list()] + for col in ["var_A", "var_B"] + } + assert rounded == expected def test_variables_cast_as_category(df_enc_category_dtypes): + # pandas Categorical dtype has no direct polars equivalent - pandas-only. df = df_enc_category_dtypes.copy() encoder = DecisionTreeEncoder(regression=False) encoder.fit(df[["var_A", "var_B"]], df["target"]) @@ -193,7 +292,11 @@ def test_variables_cast_as_category(df_enc_category_dtypes): assert X["var_A"].dtypes == float -def test_error_when_regression_is_true_and_target_is_binary(df_enc): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_error_when_regression_is_true_and_target_is_binary(df_enc, make_df): + X = _to_backend(df_enc[["var_A", "var_B"]], make_df) + y = df_enc["target"].tolist() + encoder = DecisionTreeEncoder(regression=True) msg = ( "Trying to fit a regression to a binary target is not " @@ -201,16 +304,18 @@ def test_error_when_regression_is_true_and_target_is_binary(df_enc): "or set regression to False." ) with pytest.raises(ValueError, match=msg): - encoder.fit(df_enc[["var_A", "var_B"]], df_enc["target"]) + encoder.fit(X, y) -def test_error_when_regression_is_false_and_target_is_continuous(df_enc): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_error_when_regression_is_false_and_target_is_continuous(df_enc, make_df): + X = _to_backend(df_enc[["var_A", "var_B"]], make_df) random = np.random.RandomState(42) y = random.normal(0, 10, len(df_enc)) encoder = DecisionTreeEncoder(regression=False) # the error message comes from sklearn api - won't test with pytest.raises(ValueError): - encoder.fit(df_enc[["var_A", "var_B"]], y) + encoder.fit(X, y) @pytest.mark.parametrize( @@ -225,54 +330,60 @@ def test_assigns_param_grid(grid): assert encoder._assign_param_grid() == grid -def test_unseen_is_encode(df_enc): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_unseen_is_encode(df_enc, make_df): + X = _to_backend(df_enc[["var_A", "var_B"]], make_df) + y = df_enc["target"].tolist() + encoder = DecisionTreeEncoder(unseen="encode", regression=False, fill_value=-1) - encoder.fit(df_enc[["var_A", "var_B"]], df_enc["target"]) + encoder.fit(X, y) - X_unseen_input = pd.DataFrame( + X_unseen_input = make_df( { "var_A": ["A", "ZZZ", "YYY"], "var_B": ["C", "YYY", "ZZZ"], } ) - - X_unseen_output = pd.DataFrame( - { - "var_A": [0.25, -1, -1], - "var_B": [0.4, -1, -1], - } - ) + expected = { + "var_A": [0.25, -1, -1], + "var_B": [0.4, -1, -1], + } Xt = encoder.transform(X_unseen_input) - pd.testing.assert_frame_equal(Xt, X_unseen_output) + _assert_values(Xt, expected) -def test_unseen_is_ignore(df_enc): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_unseen_is_ignore(df_enc, make_df): + X = _to_backend(df_enc[["var_A", "var_B"]], make_df) + y = df_enc["target"].tolist() + encoder = DecisionTreeEncoder(unseen="ignore", regression=False) - encoder.fit(df_enc[["var_A", "var_B"]], df_enc["target"]) + encoder.fit(X, y) - X_unseen_input = pd.DataFrame( + X_unseen_input = make_df( { "var_A": ["A", "ZZZ", "YYY"], "var_B": ["C", "YYY", "ZZZ"], } ) - - X_unseen_output = pd.DataFrame( - { - "var_A": [0.25, np.nan, np.nan], - "var_B": [0.4, np.nan, np.nan], - } - ) + expected = { + "var_A": [0.25, np.nan, np.nan], + "var_B": [0.4, np.nan, np.nan], + } Xt = encoder.transform(X_unseen_input) - pd.testing.assert_frame_equal(Xt, X_unseen_output) + _assert_values(Xt, expected) + +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_fit_errors_if_new_cat_values_and_unseen_is_raise_param(df_enc, make_df): + X = _to_backend(df_enc[["var_A", "var_B"]], make_df) + y = df_enc["target"].tolist() -def test_fit_errors_if_new_cat_values_and_unseen_is_raise_param(df_enc): encoder = DecisionTreeEncoder(unseen="raise", regression=False) - encoder.fit(df_enc[["var_A", "var_B"]], df_enc["target"]) - X = pd.DataFrame( + encoder.fit(X, y) + X_unseen = make_df( { "var_A": ["A", "ZZZ", "YYY"], "var_B": ["C", "YYY", "ZZZ"], @@ -285,56 +396,65 @@ def test_fit_errors_if_new_cat_values_and_unseen_is_raise_param(df_enc): ) # new categories will raise an error with pytest.raises(ValueError, match=msg): - encoder.transform(X) + encoder.transform(X_unseen) -def test_inverse_transform_when_no_unseen(): - X = pd.DataFrame({"words": ["dog", "dog", "dog", "cat", "cat", "cat", "bird"]}) - y = pd.Series([0, 0, 1, 1, 1, 1, 0]) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_inverse_transform_when_no_unseen(make_df): + X = make_df({"words": ["dog", "dog", "dog", "cat", "cat", "cat", "bird"]}) + y = [0, 0, 1, 1, 1, 1, 0] enc = DecisionTreeEncoder(regression=False) enc.fit(X, y) dft = enc.transform(X) - pd.testing.assert_frame_equal(enc.inverse_transform(dft), X) + Xi = enc.inverse_transform(dft) + _assert_values(Xi, {"words": ["dog", "dog", "dog", "cat", "cat", "cat", "bird"]}) -def test_inverse_transform_when_ignore_unseen(): - X = pd.DataFrame({"words": ["dog", "dog", "dog", "cat", "cat", "cat", "bird"]}) - y = pd.Series([0, 0, 1, 1, 1, 1, 0]) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_inverse_transform_when_ignore_unseen(make_df): + X = make_df({"words": ["dog", "dog", "dog", "cat", "cat", "cat", "bird"]}) + y = [0, 0, 1, 1, 1, 1, 0] enc = DecisionTreeEncoder(regression=False, unseen="ignore") enc.fit(X, y) - df1 = pd.DataFrame({"words": ["dog", "dog", "dog", "cat", "cat", "cat", "frog"]}) - df2 = pd.DataFrame({"words": ["dog", "dog", "dog", "cat", "cat", "cat", np.nan]}) + df1 = make_df({"words": ["dog", "dog", "dog", "cat", "cat", "cat", "frog"]}) dft = enc.transform(df1) - pd.testing.assert_frame_equal(enc.inverse_transform(dft), df2) + Xi = enc.inverse_transform(dft) + _assert_values( + Xi, {"words": ["dog", "dog", "dog", "cat", "cat", "cat", np.nan]} + ) -def test_inverse_transform_when_encode_unseen(): - X = pd.DataFrame({"words": ["dog", "dog", "dog", "cat", "cat", "cat", "bird"]}) - y = pd.Series([0, 0, 1, 1, 1, 1, 0]) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_inverse_transform_when_encode_unseen(make_df): + X = make_df({"words": ["dog", "dog", "dog", "cat", "cat", "cat", "bird"]}) + y = [0, 0, 1, 1, 1, 1, 0] enc = DecisionTreeEncoder(regression=False, unseen="encode", fill_value=1000) enc.fit(X, y) - df1 = pd.DataFrame({"words": ["dog", "dog", "dog", "cat", "cat", "cat", "frog"]}) - df2 = pd.DataFrame({"words": ["dog", "dog", "dog", "cat", "cat", "cat", np.nan]}) + df1 = make_df({"words": ["dog", "dog", "dog", "cat", "cat", "cat", "frog"]}) dft = enc.transform(df1) - pd.testing.assert_frame_equal(enc.inverse_transform(dft), df2) + Xi = enc.inverse_transform(dft) + _assert_values( + Xi, {"words": ["dog", "dog", "dog", "cat", "cat", "cat", np.nan]} + ) -def test_inverse_transform_raises_non_fitted_error(): - X = pd.DataFrame({"words": ["dog", "dog", "dog", "cat", "cat", "cat", "bird"]}) - y = pd.Series([0, 0, 1, 1, 1, 1, 0]) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_inverse_transform_raises_non_fitted_error(make_df): + X = make_df({"words": ["dog", "dog", "dog", "cat", "cat", "cat", "bird"]}) + y = [0, 0, 1, 1, 1, 1, 0] enc = DecisionTreeEncoder() # Test when fit is not called prior to transform. with pytest.raises(NotFittedError): enc.inverse_transform(X) - X.loc[len(X) - 1] = np.nan + X_na = make_df({"words": ["dog", "dog", "dog", "cat", "cat", "cat", None]}) with pytest.raises(ValueError): - enc.fit(X, y) + enc.fit(X_na, y) # Test when fit is not called prior to transform. with pytest.raises(NotFittedError): - enc.inverse_transform(X) + enc.inverse_transform(X_na)