diff --git a/AGENTS.md b/AGENTS.md index b5c3378cc..dd3524880 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -33,6 +33,13 @@ object that's already an instance of that module's class. - Check container emptiness with `len(x) == 0`, never `if not x:`. - `isinstance(...)` checks and `in`/`not in` membership tests are already explicit — leave them as-is, this rule isn't about those. +- The explicit `is True`/`is False` comparison is for flow control + (`if`/`while` conditions) only — don't tack it onto a variable + assignment. When a function already returns a strict bool (e.g. + `nwd.is_pandas_dataframe(X)`), assign it directly: + `is_pandas = nwd.is_pandas_dataframe(X)`, not + `is_pandas = nwd.is_pandas_dataframe(X) is True`. The `if`/`while` site + that later consumes `is_pandas` still spells out `if is_pandas is True:`. ## Comments @@ -75,6 +82,16 @@ and easy to miss without an actual comparison. - `pytest.raises(ExceptionType, match=msg)`, never `with pytest.raises() as record: ... assert str(record.value) == msg`. +- Dataframe-agnostic means one test, both backends: parametrize each + behavior over `@pytest.mark.parametrize("make_df", [pd.DataFrame, + pl.DataFrame])` and assert the same input produces the same output + values on both. Never write a separate pandas-only test and a + separate polars-only test for the same behavior — that duplicates + the test and hides the point of being dataframe-agnostic, which is + that the same input gives the same output regardless of backend. + Keep a test single-backend only when the behavior itself is + backend-specific (e.g. integer column names, which polars doesn't + support; pandas nullable extension dtypes). ## API changes diff --git a/docs/user_guide/creation/DecisionTreeFeatures.rst b/docs/user_guide/creation/DecisionTreeFeatures.rst index 56c6a4056..ba59800d2 100644 --- a/docs/user_guide/creation/DecisionTreeFeatures.rst +++ b/docs/user_guide/creation/DecisionTreeFeatures.rst @@ -485,6 +485,53 @@ are not there: 2670 1.843904 15709 1.843904 +With polars +----------- + +:class:`DecisionTreeFeatures()` works in the same way with a polars dataframe: + +.. code:: python + + import polars as pl + from feature_engine.creation import DecisionTreeFeatures + + X = pl.DataFrame({ + "Age": [20, 44, 19, 33, 51, 40, 41, 37, 30, 54], + "Height": [164, 150, 178, 158, 188, 190, 168, 174, 176, 171], + }) + y = [4.1, 5.8, 3.9, 6.2, 4.3, 4.5, 7.2, 4.4, 4.1, 6.7] + + dtf = DecisionTreeFeatures(features_to_combine=2, drop_original=True) + dtf.fit(X, y) + + print(dtf.transform(X)) + +The resulting values match those found with pandas: + +.. code:: text + + shape: (10, 3) + ┌───────────┬──────────────┬─────────────────────────┐ + │ tree(Age) ┆ tree(Height) ┆ tree(['Age', 'Height']) │ + │ --- ┆ --- ┆ --- │ + │ f64 ┆ f64 ┆ f64 │ + ╞═══════════╪══════════════╪═════════════════════════╡ + │ 4.533333 ┆ 5.366667 ┆ 4.1 │ + │ 6.0 ┆ 5.366667 ┆ 6.475 │ + │ 4.533333 ┆ 4.133333 ┆ 4.0 │ + │ 4.533333 ┆ 5.366667 ┆ 6.475 │ + │ 6.0 ┆ 4.4 ┆ 4.4 │ + │ 4.533333 ┆ 4.4 ┆ 4.4 │ + │ 6.0 ┆ 6.95 ┆ 6.475 │ + │ 4.533333 ┆ 4.133333 ┆ 4.4 │ + │ 4.533333 ┆ 4.133333 ┆ 4.0 │ + │ 6.0 ┆ 6.95 ┆ 6.475 │ + └───────────┴──────────────┴─────────────────────────┘ + +`get_feature_names_out()`, classification, and every other parameter shown +above with pandas work identically with polars. + + Creating features for classification ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -498,6 +545,46 @@ identical. We just need to set the parameter `regression` to False. classification, on the other hand, the features will contain the prediction of the class. +Training trees in parallel +~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Each tree is trained on its own feature combination independently of the others, so +when there are many combinations (a large number of variables and/or a high +`features_to_combine`) or a large `param_grid` to search, training can be +parallelized across combinations with the `n_jobs` parameter: + +.. code:: python + + import pandas as pd + from feature_engine.creation import DecisionTreeFeatures + + 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] + + dtf = DecisionTreeFeatures(features_to_combine=3, n_jobs=2, random_state=0) + dtf.fit(X, y) + + print(dtf.transform(X).columns.tolist()) + +.. code:: text + + ['Age', 'Height', 'Marks', 'tree(Age)', 'tree(Height)', 'tree(Marks)', + "tree(['Age', 'Height'])", "tree(['Age', 'Marks'])", + "tree(['Height', 'Marks'])", "tree(['Age', 'Height', 'Marks'])"] + +`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 feature +combinations or a large enough `param_grid` to outweigh the overhead of +dispatching work to threads — with just a handful of combinations, sequential +training is faster. The resulting trees and predictions are identical +regardless of `n_jobs`; only training speed changes. + + Additional resources -------------------- diff --git a/feature_engine/creation/decision_tree_features.py b/feature_engine/creation/decision_tree_features.py index aa76d4bbd..135a62815 100644 --- a/feature_engine/creation/decision_tree_features.py +++ b/feature_engine/creation/decision_tree_features.py @@ -1,8 +1,11 @@ import itertools from typing import Any, Dict, Iterable, 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.base import BaseEstimator, TransformerMixin from sklearn.model_selection import GridSearchCV from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor @@ -131,6 +134,15 @@ class DecisionTreeFeatures(TransformerMixin, BaseEstimator, GetFeatureNamesOutMi 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 feature combinations. 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. + {missing_values} {drop_original} @@ -210,6 +222,36 @@ class DecisionTreeFeatures(TransformerMixin, BaseEstimator, GetFeatureNamesOutMi 7 4.24 8 4.24 9 6.00 + + With polars: + + >>> import polars as pl + >>> from feature_engine.creation import DecisionTreeFeatures + >>> X = pl.DataFrame({ + ... "Age": [20, 44, 19, 33, 51, 40, 41, 37, 30, 54], + ... "Height": [164, 150, 178, 158, 188, 190, 168, 174, 176, 171], + ... }) + >>> y = [4.1, 5.8, 3.9, 6.2, 4.3, 4.5, 7.2, 4.4, 4.1, 6.7] + >>> dtf = DecisionTreeFeatures(features_to_combine=1) + >>> dtf.fit(X, y) + >>> dtf.transform(X) + shape: (10, 4) + ┌─────┬────────┬───────────┬──────────────┐ + │ Age ┆ Height ┆ tree(Age) ┆ tree(Height) │ + │ --- ┆ --- ┆ --- ┆ --- │ + │ i64 ┆ i64 ┆ f64 ┆ f64 │ + ╞═════╪════════╪═══════════╪══════════════╡ + │ 20 ┆ 164 ┆ 4.533333 ┆ 5.366667 │ + │ 44 ┆ 150 ┆ 6.0 ┆ 5.366667 │ + │ 19 ┆ 178 ┆ 4.533333 ┆ 4.133333 │ + │ 33 ┆ 158 ┆ 4.533333 ┆ 5.366667 │ + │ 51 ┆ 188 ┆ 6.0 ┆ 4.4 │ + │ 40 ┆ 190 ┆ 4.533333 ┆ 4.4 │ + │ 41 ┆ 168 ┆ 6.0 ┆ 6.95 │ + │ 37 ┆ 174 ┆ 4.533333 ┆ 4.133333 │ + │ 30 ┆ 176 ┆ 4.533333 ┆ 4.133333 │ + │ 54 ┆ 171 ┆ 6.0 ┆ 6.95 │ + └─────┴────────┴───────────┴──────────────┘ """ def __init__( @@ -223,6 +265,7 @@ def __init__( param_grid: Optional[Dict[str, Union[str, int, float, List[int]]]] = None, regression: bool = True, random_state: int = 0, + n_jobs: Optional[int] = None, missing_values: str = "raise", drop_original: bool = False, ) -> None: @@ -251,21 +294,22 @@ def __init__( self.param_grid = param_grid self.regression = regression self.random_state = random_state + self.n_jobs = n_jobs self.missing_values = missing_values self.drop_original = drop_original - def fit(self, X: pd.DataFrame, y: pd.Series): + def fit(self, X: IntoDataFrame, y: IntoSeries): """ Fits decision trees based on the input variable combinations with cross-validation and grid-search for hyperparameters. 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 variables to transform. - y: pandas Series or np.array = [n_samples,] + y: Series or np.array = [n_samples,] The target variable that is used to train the decision tree. """ # confirm model type and target variables are compatible. @@ -302,39 +346,48 @@ def fit(self, X: pd.DataFrame, y: pd.Series): how_to_combine=self.features_to_combine, variables=variables_ ) - estimators_ = [] - for features in input_features: - estimator = self._make_decision_tree(param_grid=param_grid) + is_pandas = nwd.is_pandas_dataframe(X) + nw_X = nw.from_native(X, eager_only=True) + X_subs = [] + for features in input_features: # single feature models - if isinstance(features, str): - estimator.fit(X[features].to_frame(), y) + if isinstance(features, (str, int)): + X_sub = nw_X.get_column(features).to_frame().to_native() # multi feature models + elif is_pandas is True: + X_sub = X[features] else: - estimator.fit(X[features], y) + X_sub = nw_X.select(features).to_native() + X_subs.append(X_sub) - estimators_.append(estimator) + estimators_ = Parallel(n_jobs=self.n_jobs, prefer="threads")( + delayed(self._fit_one_tree)(X_sub, y, param_grid) for X_sub in X_subs + ) self.variables_ = variables_ self.input_features_ = input_features self.estimators_ = estimators_ - self.feature_names_in_ = X.columns.tolist() + if is_pandas is True: + self.feature_names_in_ = list(X.columns) + else: + self.feature_names_in_ = nw_X.columns self.n_features_in_ = X.shape[1] return self - def transform(self, X: pd.DataFrame) -> pd.DataFrame: + def transform(self, X: IntoDataFrame) -> IntoDataFrame: """ Create and add new variables. Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The data to be transformed. Returns ------- - X_new: pandas dataframe. + X_new: dataframe. Either the original dataframe plus the new features or a dataframe of only the new features. """ @@ -351,50 +404,64 @@ def transform(self, X: pd.DataFrame) -> pd.DataFrame: _check_contains_na(X, self.variables_) _check_contains_inf(X, self.variables_) - # reorder variables to match train set - X = X[self.feature_names_in_] - - # create new features and add them to the original dataframe - # if regression or multiclass, we return the output of predict() - if self.regression is True: - for features, estimator in zip(self.input_features_, self.estimators_): - if isinstance(features, str): - preds = estimator.predict(X[features].to_frame()) - if self.precision is not None: - preds = np.round(preds, self.precision) - X.loc[:, f"tree({features})"] = preds - else: - preds = estimator.predict(X[features]) - if self.precision is not None: - preds = np.round(preds, self.precision) - X.loc[:, f"tree({features})"] = preds + is_pandas = nwd.is_pandas_dataframe(X) + # reorder variables to match train set + if is_pandas is True: + X = X[self.feature_names_in_] + else: + X = ( + nw.from_native(X, eager_only=True) + .select(self.feature_names_in_) + .to_native() + ) + nw_X = nw.from_native(X, eager_only=True) + + def get_x_sub(features): + if isinstance(features, (str, int)): + return nw_X.get_column(features).to_frame().to_native() + if is_pandas is True: + return X[features] + return nw_X.select(features).to_native() + + new_series = [] + new_columns = {} + # if regression or multiclass, we return the output of predict(); # if binary classification, we return the probability - elif self._is_binary == "binary": - for features, estimator in zip(self.input_features_, self.estimators_): - if isinstance(features, str): - preds = estimator.predict_proba(X[features].to_frame()) - if self.precision is not None: - preds = np.round(preds, self.precision) - X.loc[:, f"tree({features})"] = preds[:, 1] - else: - preds = estimator.predict_proba(X[features]) - if self.precision is not None: - preds = np.round(preds, self.precision) - X.loc[:, f"tree({features})"] = preds[:, 1] + for features, estimator in zip(self.input_features_, self.estimators_): + X_sub = get_x_sub(features) + col_name = f"tree({features})" + + if self.regression is True: + preds = estimator.predict(X_sub) + if self.precision is not None: + preds = np.round(preds, self.precision) + elif self._is_binary == "binary": + preds = estimator.predict_proba(X_sub)[:, 1] + if self.precision is not None: + preds = np.round(preds, self.precision) + else: + preds = estimator.predict(X_sub) - # if multiclass, we return the output of predict() - else: - for features, estimator in zip(self.input_features_, self.estimators_): - if isinstance(features, str): - preds = estimator.predict(X[features].to_frame()) - X.loc[:, f"tree({features})"] = preds - else: - preds = estimator.predict(X[features]) - X.loc[:, f"tree({features})"] = preds + if is_pandas is True: + new_columns[col_name] = preds + else: + new_series.append( + nw.new_series(col_name, preds, backend=nw_X.implementation) + ) - if self.drop_original: - X.drop(columns=self.variables_, inplace=True) + if is_pandas is True: + # assign() still inserts columns one at a time internally, so it + # doesn't avoid fragmentation with many feature combinations; + # building one DataFrame and joining it does (single insertion). + X = X.join(type(X)(new_columns, index=X.index)) + if self.drop_original is True: + X = X.drop(columns=self.variables_) + else: + nw_X = nw.from_native(X, eager_only=True).with_columns(*new_series) + if self.drop_original is True: + nw_X = nw_X.drop(self.variables_) + X = nw_X.to_native() return X @@ -414,6 +481,12 @@ def _make_decision_tree(self, param_grid: Dict): return tree_model + def _fit_one_tree(self, X_sub: IntoDataFrame, y: IntoSeries, param_grid: Dict): + """Instantiate and fit one decision tree on one feature combination.""" + estimator = self._make_decision_tree(param_grid=param_grid) + estimator.fit(X_sub, y) + return estimator + def _create_variable_combinations( self, variables: List, diff --git a/tests/test_creation/test_decision_tree_features.py b/tests/test_creation/test_decision_tree_features.py index 4e8a93e8c..a97f68475 100644 --- a/tests/test_creation/test_decision_tree_features.py +++ b/tests/test_creation/test_decision_tree_features.py @@ -1,5 +1,9 @@ +import warnings + +import narwhals as nw import numpy as np import pandas as pd +import polars as pl import pytest from sklearn.model_selection import GridSearchCV from sklearn.pipeline import Pipeline @@ -8,44 +12,87 @@ from feature_engine.creation import DecisionTreeFeatures from tests.estimator_checks.fit_functionality_checks import check_return_empty - -@pytest.fixture(scope="module") -def df_creation(): - data = { - "Name": [ - "tom", - "nick", - "krish", - "megan", - "peter", - "jordan", - "fred", - "sam", - "alexa", - "brittany", - ], - "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], - } - - df = pd.DataFrame(data) - return df - - -@pytest.fixture(scope="module") -def regression_target(): - return pd.Series([4.1, 5.8, 3.9, 6.2, 4.3, 4.5, 7.2, 4.4, 4.1, 6.7]) - - -@pytest.fixture(scope="module") -def classification_target(): - return pd.Series([1, 1, 1, 0, 0, 1, 0, 1, 0, 0]) - - -@pytest.fixture(scope="module") -def multiclass_target(): - return pd.Series([1, 1, 2, 2, 0, 1, 0, 1, 0, 0]) +DATA = { + "Name": [ + "tom", + "nick", + "krish", + "megan", + "peter", + "jordan", + "fred", + "sam", + "alexa", + "brittany", + ], + "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], +} +REGRESSION_Y = [4.1, 5.8, 3.9, 6.2, 4.3, 4.5, 7.2, 4.4, 4.1, 6.7] +BINARY_Y = [1, 1, 1, 0, 0, 1, 0, 1, 0, 0] +MULTICLASS_Y = [1, 1, 2, 2, 0, 1, 0, 1, 0, 0] + +COMBOS = [ + "Age", + "Height", + "Marks", + ["Age", "Height"], + ["Age", "Marks"], + ["Height", "Marks"], + ["Age", "Height", "Marks"], +] + + +def _select(X, combo): + cols = combo if isinstance(combo, list) else [combo] + return nw.from_native(X, eager_only=True).select(cols).to_native() + + +def _expected_tree_predictions( + X, + y, + scoring, + random_state, + regression=True, + binary=False, + precision=None, + param_grid=None, +): + # Fits a fresh GridSearchCV per combo on the same backend as X, so this + # works as the reference for both pandas and polars input alike. + if param_grid is None: + param_grid = {"max_depth": [1, 2, 3, 4]} + if regression is True: + est = DecisionTreeRegressor(random_state=random_state) + else: + est = DecisionTreeClassifier(random_state=random_state) + tree = GridSearchCV(est, cv=3, scoring=scoring, param_grid=param_grid) + + expected = {} + for combo in COMBOS: + X_sub = _select(X, combo) + tree.fit(X_sub, y) + if regression is True: + preds = tree.predict(X_sub) + elif binary is True: + preds = tree.predict_proba(X_sub)[:, 1] + else: + preds = tree.predict(X_sub) + if precision is not None: + preds = np.round(preds, precision) + expected[f"tree({combo})"] = list(preds) + return expected + + +def assert_df_equal(X, expected: dict) -> None: + result = nw.from_native(X, eager_only=True).to_dict(as_series=False) + assert list(result.keys()) == list(expected.keys()) + for col, values in expected.items(): + if all(isinstance(v, (int, float, np.integer, np.floating)) for v in values): + assert result[col] == pytest.approx(values, abs=1e-6) + else: + assert result[col] == values @pytest.mark.parametrize("precision", ["string", 0.1, -1, np.nan]) @@ -204,390 +251,226 @@ def test_create_variable_combinations_when_tuple(input_features, expected): assert combos == expected -def test_feature_creation_regression(df_creation, regression_target): - X = df_creation.copy() - y = regression_target.copy() - +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_feature_creation_regression(make_df): + X = make_df(DATA) scoring = "neg_mean_squared_error" rs = 0 tr = DecisionTreeFeatures(scoring=scoring, random_state=rs) - Xt = tr.fit_transform(X, y) - - # get expected - est = DecisionTreeRegressor(random_state=rs) - tree = GridSearchCV( - est, - cv=3, - scoring=scoring, - param_grid={"max_depth": [1, 2, 3, 4]}, - ) - - combos = [ - "Age", - "Height", - "Marks", - ["Age", "Height"], - ["Age", "Marks"], - ["Height", "Marks"], - ["Age", "Height", "Marks"], - ] - var_names = [f"tree({item})" for item in combos] - - X_exp = df_creation.copy() - for i in range(len(combos)): - varn = var_names[i] - combon = combos[i] - if isinstance(combon, str): - tree.fit(X[combon].to_frame(), y) - X_exp[varn] = tree.predict(X[combon].to_frame()) - else: - tree.fit(X[combon], y) - X_exp[varn] = tree.predict(X[combon]) - - pd.testing.assert_frame_equal(Xt, X_exp) + Xt = tr.fit_transform(X, REGRESSION_Y) + expected = dict(DATA) + expected.update(_expected_tree_predictions(X, REGRESSION_Y, scoring, rs)) + assert_df_equal(Xt, expected) -def test_feature_creation_regression_and_precision(df_creation, regression_target): - X = df_creation.copy() - y = regression_target.copy() +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_feature_creation_regression_and_precision(make_df): + X = make_df(DATA) scoring = "neg_mean_squared_error" rs = 0 tr = DecisionTreeFeatures(scoring=scoring, random_state=rs, precision=1) - Xt = tr.fit_transform(X, y) - - # get expected - est = DecisionTreeRegressor(random_state=rs) - tree = GridSearchCV( - est, - cv=3, - scoring=scoring, - param_grid={"max_depth": [1, 2, 3, 4]}, - ) - - combos = [ - "Age", - "Height", - "Marks", - ["Age", "Height"], - ["Age", "Marks"], - ["Height", "Marks"], - ["Age", "Height", "Marks"], - ] - var_names = [f"tree({item})" for item in combos] - - X_exp = df_creation.copy() - for i in range(len(combos)): - varn = var_names[i] - combon = combos[i] - if isinstance(combon, str): - tree.fit(X[combon].to_frame(), y) - preds = tree.predict(X[combon].to_frame()) - X_exp[varn] = np.round(preds, 1) - else: - tree.fit(X[combon], y) - preds = tree.predict(X[combon]) - X_exp[varn] = np.round(preds, 1) - - pd.testing.assert_frame_equal(Xt, X_exp) + Xt = tr.fit_transform(X, REGRESSION_Y) + expected = dict(DATA) + expected.update( + _expected_tree_predictions(X, REGRESSION_Y, scoring, rs, precision=1) + ) + assert_df_equal(Xt, expected) -def test_feature_creation_regression_drop_original(df_creation, regression_target): - X = df_creation.copy() - y = regression_target.copy() +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_feature_creation_regression_drop_original(make_df): + X = make_df(DATA) scoring = "neg_mean_squared_error" rs = 0 tr = DecisionTreeFeatures(scoring=scoring, random_state=rs, drop_original=True) - Xt = tr.fit_transform(X, y) - - # get expected - est = DecisionTreeRegressor(random_state=rs) - tree = GridSearchCV( - est, - cv=3, - scoring=scoring, - param_grid={"max_depth": [1, 2, 3, 4]}, - ) - - combos = [ - "Age", - "Height", - "Marks", - ["Age", "Height"], - ["Age", "Marks"], - ["Height", "Marks"], - ["Age", "Height", "Marks"], - ] - var_names = [f"tree({item})" for item in combos] - - X_exp = df_creation.copy() - for i in range(len(combos)): - varn = var_names[i] - combon = combos[i] - if isinstance(combon, str): - tree.fit(X[combon].to_frame(), y) - X_exp[varn] = tree.predict(X[combon].to_frame()) - else: - tree.fit(X[combon], y) - X_exp[varn] = tree.predict(X[combon]) - X_exp.drop(["Age", "Height", "Marks"], axis=1, inplace=True) - - pd.testing.assert_frame_equal(Xt, X_exp) + Xt = tr.fit_transform(X, REGRESSION_Y) + expected = {"Name": DATA["Name"]} + expected.update(_expected_tree_predictions(X, REGRESSION_Y, scoring, rs)) + assert_df_equal(Xt, expected) -def test_feature_creation_binary_classif(df_creation, classification_target): - X = df_creation.copy() - y = classification_target.copy() +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_feature_creation_binary_classif(make_df): + X = make_df(DATA) scoring = "roc_auc" rs = 0 tr = DecisionTreeFeatures(scoring=scoring, random_state=rs, regression=False) - Xt = tr.fit_transform(X, y) - - # get expected - est = DecisionTreeClassifier(random_state=rs) - tree = GridSearchCV( - est, - cv=3, - scoring=scoring, - param_grid={"max_depth": [1, 2, 3, 4]}, - ) - - combos = [ - "Age", - "Height", - "Marks", - ["Age", "Height"], - ["Age", "Marks"], - ["Height", "Marks"], - ["Age", "Height", "Marks"], - ] - var_names = [f"tree({item})" for item in combos] - - X_exp = df_creation.copy() - for i in range(len(combos)): - varn = var_names[i] - combon = combos[i] - if isinstance(combon, str): - tree.fit(X[combon].to_frame(), y) - preds = tree.predict_proba(X[combon].to_frame()) - X_exp[varn] = preds[:, 1] - else: - tree.fit(X[combon], y) - preds = tree.predict_proba(X[combon]) - X_exp[varn] = preds[:, 1] - - pd.testing.assert_frame_equal(Xt, X_exp) + Xt = tr.fit_transform(X, BINARY_Y) + expected = dict(DATA) + expected.update( + _expected_tree_predictions( + X, BINARY_Y, scoring, rs, regression=False, binary=True + ) + ) + assert_df_equal(Xt, expected) -def test_feature_creation_binary_classif_w_precision( - df_creation, classification_target -): - X = df_creation.copy() - y = classification_target.copy() +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_feature_creation_binary_classif_w_precision(make_df): + X = make_df(DATA) scoring = "roc_auc" rs = 0 tr = DecisionTreeFeatures( scoring=scoring, random_state=rs, regression=False, precision=2 ) - Xt = tr.fit_transform(X, y) - - # get expected - est = DecisionTreeClassifier(random_state=rs) - tree = GridSearchCV( - est, - cv=3, - scoring=scoring, - param_grid={"max_depth": [1, 2, 3, 4]}, - ) - - combos = [ - "Age", - "Height", - "Marks", - ["Age", "Height"], - ["Age", "Marks"], - ["Height", "Marks"], - ["Age", "Height", "Marks"], - ] - var_names = [f"tree({item})" for item in combos] - - X_exp = df_creation.copy() - for i in range(len(combos)): - varn = var_names[i] - combon = combos[i] - if isinstance(combon, str): - tree.fit(X[combon].to_frame(), y) - preds = tree.predict_proba(X[combon].to_frame()) - X_exp[varn] = np.round(preds[:, 1], 2) - else: - tree.fit(X[combon], y) - preds = tree.predict_proba(X[combon]) - X_exp[varn] = np.round(preds[:, 1], 2) - - pd.testing.assert_frame_equal(Xt, X_exp) + Xt = tr.fit_transform(X, BINARY_Y) + expected = dict(DATA) + expected.update( + _expected_tree_predictions( + X, BINARY_Y, scoring, rs, regression=False, binary=True, precision=2 + ) + ) + assert_df_equal(Xt, expected) -def test_feature_creation_binary_multiclass(df_creation, multiclass_target): - X = df_creation.copy() - y = multiclass_target.copy() +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_feature_creation_binary_multiclass(make_df): + X = make_df(DATA) scoring = "roc_auc" rs = 0 tr = DecisionTreeFeatures(scoring=scoring, random_state=rs, regression=False) - Xt = tr.fit_transform(X, y) - - # get expected - est = DecisionTreeClassifier(random_state=rs) - tree = GridSearchCV( - est, - cv=3, - scoring=scoring, - param_grid={"max_depth": [1, 2, 3, 4]}, - ) - - combos = [ - "Age", - "Height", - "Marks", - ["Age", "Height"], - ["Age", "Marks"], - ["Height", "Marks"], - ["Age", "Height", "Marks"], - ] - var_names = [f"tree({item})" for item in combos] - - X_exp = df_creation.copy() - for i in range(len(combos)): - varn = var_names[i] - combon = combos[i] - if isinstance(combon, str): - tree.fit(X[combon].to_frame(), y) - preds = tree.predict(X[combon].to_frame()) - X_exp[varn] = preds - else: - tree.fit(X[combon], y) - preds = tree.predict(X[combon]) - X_exp[varn] = preds - - pd.testing.assert_frame_equal(Xt, X_exp) - - -def test_get_feature_names_out(df_creation, regression_target): - X = df_creation.copy() - y = regression_target.copy() + Xt = tr.fit_transform(X, MULTICLASS_Y) - tr = DecisionTreeFeatures( - variables=["Age", "Marks"], + expected = dict(DATA) + expected.update( + _expected_tree_predictions( + X, MULTICLASS_Y, scoring, rs, regression=False, binary=False + ) ) + assert_df_equal(Xt, expected) - Xt = tr.fit_transform(X, y) - feat_out = Xt.columns.to_list() - assert tr.get_feature_names_out() == feat_out - assert tr.get_feature_names_out(X.columns.to_list()) == feat_out +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_get_feature_names_out(make_df): + X = make_df(DATA) + tr = DecisionTreeFeatures(variables=["Age", "Marks"]) + Xt = tr.fit_transform(X, REGRESSION_Y) + feat_out = list(nw.from_native(Xt, eager_only=True).columns) + assert tr.get_feature_names_out() == feat_out + assert tr.get_feature_names_out(list(DATA.keys())) == feat_out -def test_get_feature_names_out_from_pipeline(df_creation, regression_target): - X = df_creation.copy() - y = regression_target.copy() - - # set up transformer - tr = DecisionTreeFeatures( - variables=["Age", "Marks"], - ) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_get_feature_names_out_from_pipeline(make_df): + X = make_df(DATA) + tr = DecisionTreeFeatures(variables=["Age", "Marks"]) pipe = Pipeline([("transformer", tr)]) - - Xt = pipe.fit_transform(X, y) - feat_out = Xt.columns.to_list() + Xt = pipe.fit_transform(X, REGRESSION_Y) + feat_out = list(nw.from_native(Xt, eager_only=True).columns) assert pipe.get_feature_names_out(input_features=None) == feat_out - assert pipe.get_feature_names_out(input_features=X.columns.to_list()) == feat_out + assert pipe.get_feature_names_out(input_features=list(DATA.keys())) == feat_out +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) @pytest.mark.parametrize("_input_features", ["hola", ["Age", "Marks"]]) -def test_get_feature_names_out_raises_error_when_wrong_param( - _input_features, df_creation, regression_target -): - X = df_creation.copy() - y = regression_target.copy() - - tr = DecisionTreeFeatures( - variables=["Age", "Marks"], - ) - tr.fit(X, y) - +def test_get_feature_names_out_raises_error_when_wrong_param(make_df, _input_features): + X = make_df(DATA) + tr = DecisionTreeFeatures(variables=["Age", "Marks"]) + tr.fit(X, REGRESSION_Y) with pytest.raises(ValueError): tr.get_feature_names_out(input_features=_input_features) -def test_error_when_regression_true_and_target_binary( - df_creation, classification_target -): - X = df_creation.copy() - y = classification_target.copy() +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_error_when_regression_true_and_target_binary(make_df): + X = make_df(DATA) tr = DecisionTreeFeatures(regression=True) 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." + "allowed by this transformer. Check the target values " + "or set regression to False." ) with pytest.raises(ValueError, match=msg): - tr.fit(X, y) + tr.fit(X, BINARY_Y) -def test_user_enter_param_grid(df_creation, classification_target): - X = df_creation.copy() - y = classification_target.copy() +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_user_enter_param_grid(make_df): + X = make_df(DATA) scoring = "roc_auc" rs = 0 grid = {"max_depth": [1, 2, 3, 4]} tr = DecisionTreeFeatures( scoring=scoring, random_state=rs, regression=False, param_grid=grid ) - Xt = tr.fit_transform(X, y) - - # get expected - est = DecisionTreeClassifier(random_state=rs) - tree = GridSearchCV( - est, - cv=3, - scoring=scoring, - param_grid={"max_depth": [1, 2, 3, 4]}, - ) - - combos = [ - "Age", - "Height", - "Marks", - ["Age", "Height"], - ["Age", "Marks"], - ["Height", "Marks"], - ["Age", "Height", "Marks"], - ] - var_names = [f"tree({item})" for item in combos] - - X_exp = df_creation.copy() - for i in range(len(combos)): - varn = var_names[i] - combon = combos[i] - if isinstance(combon, str): - tree.fit(X[combon].to_frame(), y) - preds = tree.predict_proba(X[combon].to_frame()) - X_exp[varn] = preds[:, 1] - else: - tree.fit(X[combon], y) - preds = tree.predict_proba(X[combon]) - X_exp[varn] = preds[:, 1] + Xt = tr.fit_transform(X, BINARY_Y) - pd.testing.assert_frame_equal(Xt, X_exp) + expected = dict(DATA) + expected.update( + _expected_tree_predictions( + X, BINARY_Y, scoring, rs, regression=False, binary=True, param_grid=grid + ) + ) + assert_df_equal(Xt, expected) def test_check_return_empty(): # DecisionTreeFeatures is not part of the check_feature_engine_estimator # pipeline (test_check_estimator_creation.py only feeds MathFeatures, # RelativeFeatures and CyclicalFeatures into it), so return_empty is - # tested directly here instead. + # tested directly here instead. check_return_empty is a shared, + # pandas-only estimator-check helper used across the library. check_return_empty(DecisionTreeFeatures(regression=False)) + + +@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 + # feature combinations must produce identical trees, and therefore + # identical predictions, to sequential training (n_jobs=None). + X = make_df(DATA) + tr_seq = DecisionTreeFeatures(n_jobs=None, random_state=0) + tr_seq.fit(X, REGRESSION_Y) + tr_par = DecisionTreeFeatures(n_jobs=2, random_state=0) + tr_par.fit(X, REGRESSION_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) + assert_df_equal(Xt_par, expected) + + +def test_transform_does_not_fragment_pandas_output(): + # regression test: transform() used to assign one new tree column at a + # time (X[col_name] = preds), which triggers pandas' "DataFrame is + # highly fragmented" PerformanceWarning once there are enough feature + # combinations - fixed by building all new columns in one DataFrame + # and joining once. Needs enough variables to cross pandas' internal + # fragmentation threshold (a handful of combos won't trigger it). + rng = np.random.RandomState(0) + n_vars = 9 + X = pd.DataFrame( + rng.rand(200, n_vars), columns=[f"v{i}" for i in range(n_vars)] + ) + y = rng.rand(200) + + tr = DecisionTreeFeatures( + features_to_combine=3, param_grid={"max_depth": [1, 2]}, random_state=0 + ) + tr.fit(X, y) + + with warnings.catch_warnings(): + warnings.simplefilter("error", pd.errors.PerformanceWarning) + tr.transform(X) + + +def test_single_int_named_feature_combo(): + # regression test: a single-variable combo with an integer column name + # used to crash (isinstance(features, str) missed the int case), since + # X[features] for a bare int returns a 1D Series, not the 2D input + # sklearn requires - fixed to check isinstance(features, (str, int)). + # Integer column names are pandas-only - polars requires string columns. + df = pd.DataFrame({0: [1.0, 2, 3, 4, 5, 6, 7, 8], 1: [2.0, 3, 4, 5, 6, 7, 8, 9]}) + y = [1.0, 2, 3, 4, 5, 6, 7, 8] + transformer = DecisionTreeFeatures(features_to_combine=1, random_state=0) + transformer.fit(df, y) + Xt = transformer.transform(df) + assert "tree(0)" in Xt.columns + assert "tree(1)" in Xt.columns