From cd7f30b2759817b3f0722bcbf8b51d436e1bfec0 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Sat, 19 Sep 2026 11:27:15 +0200 Subject: [PATCH 1/3] Migrate the selection base classes and helpers to narwhals, add polars support BaseSelector.transform() returns the retained features in the train set order, in the same library as the input (pandas X[features], narwhals select otherwise). BaseRecursiveSelector.fit() trains the estimators on native frames and returns (nw_X, y). The helpers in base_selection_functions no longer import pandas: correlations are computed with numpy (np.corrcoef, or matrix products for pairwise complete observations when there are missing values), and feature importances are pandas Series for pandas input and dicts otherwise. Co-Authored-By: Claude Opus 5 --- .../selection/base_recursive_selector.py | 96 ++-- .../selection/base_selection_functions.py | 212 ++++++-- feature_engine/selection/base_selector.py | 44 +- tests/test_selection/conftest.py | 17 + .../test_base_recursive_selector.py | 257 +++++++++ .../test_base_selection_functions.py | 513 ++++++++++-------- tests/test_selection/test_base_selector.py | 178 ++++-- 7 files changed, 924 insertions(+), 393 deletions(-) create mode 100644 tests/test_selection/test_base_recursive_selector.py diff --git a/feature_engine/selection/base_recursive_selector.py b/feature_engine/selection/base_recursive_selector.py index fe9113077..1161572aa 100644 --- a/feature_engine/selection/base_recursive_selector.py +++ b/feature_engine/selection/base_recursive_selector.py @@ -1,7 +1,9 @@ from types import GeneratorType -from typing import List, Union +from typing import List, Tuple, Union -import pandas as pd +import narwhals as nw +import numpy as np +from narwhals.typing import IntoDataFrame, IntoSeries from sklearn.inspection import permutation_importance from sklearn.model_selection import cross_validate @@ -9,14 +11,13 @@ _check_variables_input_value, ) from feature_engine.dataframe_checks import check_X_y -from feature_engine.selection.base_selection_functions import get_feature_importances +from feature_engine.selection.base_selection_functions import ( + _importance_series, + _select_numerical_variables, + get_feature_importances, +) from feature_engine.selection.base_selector import BaseSelector from feature_engine.tags import _return_tags -from feature_engine.variable_handling import ( - check_numerical_variables, - find_numerical_variables, - retain_variables_if_in_df, -) Variables = Union[None, int, str, List[Union[str, int]]] @@ -81,10 +82,13 @@ class BaseRecursiveSelector(BaseSelector): Performance of the model trained using the original dataset. feature_importances_: - Pandas Series with the feature importance (comes from step 2) + The feature importance (comes from step 2). A pandas Series with the + features as index when X is a pandas dataframe, and a dictionary with the + features as keys otherwise. feature_importances_std_: - Pandas Series with the standard deviation of the feature importance. + The standard deviation of the feature importance, as a pandas Series or a + dictionary, like `feature_importances_`. features_to_drop_: List with the features to remove from the dataset. @@ -116,7 +120,9 @@ def __init__( ): if not isinstance(threshold, (int, float)): - raise ValueError("threshold can only be integer or float") + raise ValueError( + f"threshold must be an integer or a float. Got {threshold} instead." + ) super().__init__(confirm_variables) self.variables = _check_variables_input_value(variables) @@ -126,43 +132,45 @@ def __init__( self.cv = cv self.groups = groups - def fit(self, X: pd.DataFrame, y: pd.Series): + def fit(self, X: IntoDataFrame, y: IntoSeries) -> Tuple[nw.DataFrame, IntoSeries]: """ Find initial model performance. Sort features by importance. Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The input dataframe y: array-like of shape (n_samples) Target variable. Required to train the estimator. - """ - # check input dataframe - X, y = check_X_y(X, y) + Returns + ------- + nw_X: narwhals dataframe + The input dataframe, as a narwhals dataframe. - if self.variables is None: - self.variables_ = find_numerical_variables(X) - else: - if self.confirm_variables is True: - variables_ = retain_variables_if_in_df(X, self.variables) - self.variables_ = check_numerical_variables(X, variables_) - else: - self.variables_ = check_numerical_variables(X, self.variables) + y: Series or numpy array + The target, checked. + """ + nw_X, y = check_X_y(X, y) + + self.variables_ = _select_numerical_variables( + X, self.variables, self.confirm_variables + ) self._cv = list(self.cv) if isinstance(self.cv, GeneratorType) else self.cv # check that there are more than 1 variable to select from self._check_variable_number() - # save input features self._get_feature_names_in(X) + X_model = nw_X.select(nw.col(*self.variables_)).to_native() + # train model with all features and cross-validation model = cross_validate( estimator=self.estimator, - X=X[self.variables_], + X=X_model, y=y, cv=self._cv, groups=self.groups, @@ -170,40 +178,32 @@ def fit(self, X: pd.DataFrame, y: pd.Series): return_estimator=True, ) - # store initial model performance self.initial_model_performance_ = model["test_score"].mean() - # Initialize a dataframe that will contain the list of the feature/coeff - # importance for each cross validation fold - feature_importances_cv = pd.DataFrame() - - # Populate the feature_importances_cv dataframe with columns containing - # the feature importance values for each model returned by the cross - # validation. - # There are as many columns as folds. - for i in range(len(model["estimator"])): - m = model["estimator"][i] - + # one row of feature importance per cross-validation fold + importances = [] + for m in model["estimator"]: if hasattr(m, "feature_importances_") or hasattr(m, "coef_"): - feature_importances_cv[i] = get_feature_importances(m) + importances.append(get_feature_importances(m)) else: r = permutation_importance( m, - X[self.variables_], + X_model, y, n_repeats=1, random_state=10, ) - feature_importances_cv[i] = r.importances_mean + importances.append(r.importances_mean) + importances_arr = np.array(importances) - # Add the variables as index to feature_importances_cv - feature_importances_cv.index = self.variables_ - - # Aggregate the feature importance returned in each fold - self.feature_importances_ = feature_importances_cv.mean(axis=1) - self.feature_importances_std_ = feature_importances_cv.std(axis=1) + self.feature_importances_ = _importance_series( + X, self.variables_, importances_arr.mean(axis=0) + ) + self.feature_importances_std_ = _importance_series( + X, self.variables_, importances_arr.std(axis=0, ddof=1) + ) - return X, y + return nw_X, y def _more_tags(self): tags_dict = _return_tags() diff --git a/feature_engine/selection/base_selection_functions.py b/feature_engine/selection/base_selection_functions.py index 96fc45970..4cbadb60b 100644 --- a/feature_engine/selection/base_selection_functions.py +++ b/feature_engine/selection/base_selection_functions.py @@ -1,8 +1,11 @@ -from typing import List, Union from types import GeneratorType +from typing import List, Union +import narwhals as nw +import narwhals.dependencies as nwd import numpy as np -import pandas as pd +from narwhals.typing import IntoDataFrame +from scipy.stats import kendalltau, rankdata from sklearn.model_selection import cross_validate from feature_engine.variable_handling import ( @@ -36,8 +39,19 @@ def get_feature_importances(estimator): return importances +def _importance_series(X: IntoDataFrame, features, values: np.ndarray): + """ + Return the importance of each feature as a pandas Series indexed by the + features when X is a pandas dataframe, or as a dictionary with the features as + keys otherwise. + """ + if nwd.is_pandas_dataframe(X) is True: + return nw.get_native_namespace(X).Series(values, index=features) + return dict(zip(features, values.tolist())) + + def _select_all_variables( - X: pd.DataFrame, + X: IntoDataFrame, variables: Variables, confirm_variables: bool, exclude_datetime: bool = False, @@ -62,7 +76,7 @@ def _select_all_variables( def _select_numerical_variables( - X: pd.DataFrame, + X: IntoDataFrame, variables: Variables, confirm_variables: bool, ): @@ -85,8 +99,113 @@ def _select_numerical_variables( return variables_ +def _corrcoef(values: np.ndarray) -> np.ndarray: + # constant columns return NaN, like pandas, instead of warning. + with np.errstate(divide="ignore", invalid="ignore"): + return np.corrcoef(values, rowvar=False) + + +def _pearson_pairwise_complete(values: np.ndarray, finite: np.ndarray) -> np.ndarray: + """ + Pearson correlation of every pair of columns, using the rows where both are + finite. The sums of all pairs come from matrix products, which is much faster + than looping over the pairs. + """ + mask: np.ndarray = finite.astype(float) + with np.errstate(divide="ignore", invalid="ignore"): + # centring first keeps the sums below numerically stable. + mean = np.where(finite, values, 0.0).sum(axis=0) / mask.sum(axis=0) + x = np.where(finite, values - mean, 0.0) + n_obs = mask.T @ mask + sum_x = x.T @ mask + sum_xx = (x * x).T @ mask + var = sum_xx - sum_x * sum_x / n_obs + corr = (x.T @ x - sum_x * sum_x.T / n_obs) / np.sqrt(var * var.T) + + # when the variance of the shared rows is tiny compared with the sums, the + # subtraction above loses precision: recompute those pairs one by one. + unstable = (var <= 1e-8 * sum_xx) | (var.T <= 1e-8 * sum_xx.T) + corr[unstable] = np.nan + for i, j in zip(*np.nonzero(np.triu(unstable & (n_obs > 1), 1))): + rows = finite[:, i] & finite[:, j] + corr[i, j] = _corrcoef(x[rows][:, [i, j]])[0, 1] + return corr + + +def _spearman_pairwise_complete(nw_X: nw.DataFrame) -> np.ndarray: + """ + Spearman correlation of every pair of columns, ranking each pair on the rows + where both are finite, like pandas.DataFrame.corr(). + """ + variables = nw_X.columns + exprs = [] + for i, var_i in enumerate(variables): + for j in range(i + 1, len(variables)): + var_j = variables[j] + rows = nw.col(var_i).is_finite() & nw.col(var_j).is_finite() + x = nw.when(rows).then(nw.col(var_i)).rank("average") + y = nw.when(rows).then(nw.col(var_j)).rank("average") + dx = x - x.mean() + dy = y - y.mean() + exprs.append( + ((dx * dy).sum() / ((dx * dx).sum() * (dy * dy).sum()).sqrt()).alias( + f"__{i}_{j}__" + ) + ) + n_vars = len(variables) + corr = np.full((n_vars, n_vars), np.nan) + corr[np.triu_indices(n_vars, 1)] = np.array(nw_X.select(exprs).row(0), dtype=float) + return corr + + +def _correlation_matrix(X: IntoDataFrame, variables: list, method) -> np.ndarray: + """ + Correlation matrix of the variables. Like pandas.DataFrame.corr(), each pair of + variables is compared on the rows where both have finite values. Only the + values above the diagonal are used. + """ + if nwd.is_pandas_dataframe(X) is True: + values = X[variables].to_numpy(dtype=float, na_value=np.nan) + else: + nw_X = nw.from_native(X, eager_only=True).select(nw.col(*variables)) + values = nw_X.to_numpy().astype(float) + finite = np.isfinite(values) + + # numpy is faster than pandas and narwhals. + if method == "pearson": + if finite.all(): + return _corrcoef(values) + return _pearson_pairwise_complete(values, finite) + + if method == "spearman" and finite.all(): + # scipy ranks faster than pandas, and polars faster than scipy. + if nwd.is_pandas_dataframe(X) is True: + return _corrcoef(rankdata(values, axis=0)) + return _corrcoef(nw_X.select(nw.all().rank("average")).to_numpy()) + + # pandas is faster than narwhals. + if nwd.is_pandas_dataframe(X) is True: + return X[variables].corr(method=method).to_numpy() + + if method == "spearman": + return _spearman_pairwise_complete(nw_X) + + # kendall and callables are computed pair by pair, as pandas does. + corr_func = (lambda a, b: kendalltau(a, b)[0]) if method == "kendall" else method + n_vars = len(variables) + corr = np.full((n_vars, n_vars), np.nan) + for i in range(n_vars): + for j in range(i + 1, n_vars): + rows = finite[:, i] & finite[:, j] + if rows.all(): + corr[i, j] = corr_func(values[:, i], values[:, j]) + elif rows.any(): + corr[i, j] = corr_func(values[rows, i], values[rows, j]) + return corr + + def find_correlated_features( - X: pd.DataFrame, + X: IntoDataFrame, variables: list[Union[str, int]], method: str, threshold: float, @@ -96,7 +215,7 @@ def find_correlated_features( Parameters ---------- - X : pandas dataframe of shape = [n_samples, n_features] + X : dataframe of shape = [n_samples, n_features] The training dataset. variables : list @@ -133,36 +252,32 @@ def find_correlated_features( correlated with the key. The key + the values should be the same as the set found in `correlated_feature_groups`. """ - # the correlation matrix - correlated_matrix = X[variables].corr(method=method).to_numpy() + correlated_matrix = _correlation_matrix(X, variables, method) # the correlated pairs correlated_mask = np.triu(np.abs(correlated_matrix), 1) > threshold - examined = set() + examined: np.ndarray = np.zeros(len(variables), dtype=bool) correlated_groups = list() features_to_drop = list() correlated_dict = {} for i, f_i in enumerate(variables): - if f_i not in examined: - examined.add(f_i) - temp_set = set([f_i]) - for j, f_j in enumerate(variables): - if f_j not in examined: - if correlated_mask[i, j] == 1: - examined.add(f_j) - features_to_drop.append(f_j) - temp_set.add(f_j) - if len(temp_set) > 1: - correlated_groups.append(temp_set) - correlated_dict[f_i] = temp_set.difference({f_i}) + if examined.item(i) is False: + examined[i] = True + correlated = np.flatnonzero(correlated_mask[i] & ~examined) + if len(correlated) > 0: + examined[correlated] = True + correlated_features = [variables[j] for j in correlated] + features_to_drop.extend(correlated_features) + correlated_groups.append({f_i, *correlated_features}) + correlated_dict[f_i] = set(correlated_features) return correlated_groups, features_to_drop, correlated_dict def single_feature_performance( - X: pd.DataFrame, - y: pd.Series, + X: IntoDataFrame, + y, variables: List[Union[str, int]], estimator, cv, @@ -174,7 +289,7 @@ def single_feature_performance( Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The input dataframe y: array-like of shape (n_samples) @@ -211,12 +326,13 @@ def single_feature_performance( feature_performance_std = {} cv = list(cv) if isinstance(cv, GeneratorType) else cv + nw_X = nw.from_native(X, eager_only=True) # train a model for every feature and store the performance for feature in variables: model = cross_validate( estimator, - X[feature].to_frame(), + nw_X.get_column(feature).to_frame().to_native(), y, cv=cv, groups=groups, @@ -230,8 +346,8 @@ def single_feature_performance( def find_feature_importance( - X: pd.DataFrame, - y: pd.Series, + X: IntoDataFrame, + y, estimator, cv, scoring, @@ -245,7 +361,7 @@ def find_feature_importance( Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The input dataframe y: array-like of shape (n_samples) @@ -268,14 +384,15 @@ def find_feature_importance( Returns ------- - feature_importance: pd.Series - A pandas Series with the feature name as index and its importance as value. The - importance is given by the coefficients of linear models or the impurity gain - from tree-based models. - - feature_importance_std: pd.Series - A pandas Series with the feature name as key and the standard deviation of the - feature importance as value. + feature_importance: pandas Series or dict + The importance of each feature, given by the coefficients of linear models or + the impurity gain from tree-based models. A pandas Series with the feature + names as index when X is a pandas dataframe, and a dictionary with the + feature names as keys otherwise. + + feature_importance_std: pandas Series or dict + The standard deviation of the importance of each feature, as a pandas Series + or a dictionary, like `feature_importance`. """ cv = list(cv) if isinstance(cv, GeneratorType) else cv @@ -289,19 +406,16 @@ def find_feature_importance( return_estimator=True, ) - # dataframe to store the feature importance for each cv fold - feature_importances_cv = pd.DataFrame() - - # Populate dataframe with columns containing the feature importance values - # for each cv fold. There are as many columns as folds. - for i in range(len(model["estimator"])): - m = model["estimator"][i] - feature_importances_cv[i] = get_feature_importances(m) + importances = np.array([get_feature_importances(m) for m in model["estimator"]]) - # add the variables as the index to feature_importances_cv - feature_importances_cv.index = X.columns + # pandas keeps the columns index, with its name and dtype. + if nwd.is_pandas_dataframe(X) is True: + features = X.columns + else: + features = nw.from_native(X, eager_only=True).columns - # aggregate the feature importance returned in each fold - feature_importances_ = feature_importances_cv.mean(axis=1) - feature_importances_std_ = feature_importances_cv.std(axis=1) + feature_importances_ = _importance_series(X, features, importances.mean(axis=0)) + feature_importances_std_ = _importance_series( + X, features, importances.std(axis=0, ddof=1) + ) return feature_importances_, feature_importances_std_ diff --git a/feature_engine/selection/base_selector.py b/feature_engine/selection/base_selector.py index cfa8f1c95..0e4ff5d97 100644 --- a/feature_engine/selection/base_selector.py +++ b/feature_engine/selection/base_selector.py @@ -1,5 +1,7 @@ +import narwhals as nw +import narwhals.dependencies as nwd import numpy as np -import pandas as pd +from narwhals.typing import IntoDataFrame from sklearn.base import BaseEstimator, TransformerMixin from sklearn.utils.validation import check_is_fitted @@ -41,41 +43,43 @@ def __init__( self.confirm_variables = confirm_variables - def transform(self, X: pd.DataFrame) -> pd.DataFrame: + def transform(self, X: IntoDataFrame) -> IntoDataFrame: """ Return dataframe with selected features. Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features]. + X: dataframe of shape = [n_samples, n_features]. The input dataframe. Returns ------- - X_new: pandas dataframe of shape = [n_samples, n_selected_features] - Pandas dataframe with the selected features. + X_new: dataframe of shape = [n_samples, n_selected_features] + The dataframe with the selected features, in the same library as the + input. """ - - # check if fit is performed prior to transform check_is_fitted(self) - - # check if input is a dataframe - X = check_X(X) - - # check if number of columns in test dataset matches to train dataset + nw_X = check_X(X) _check_X_matches_training_df(X, self.n_features_in_) - # reorder df to match train set - X = X[self.feature_names_in_] + # selecting in the train set order also restores the train column order. + features_to_drop = set(self.features_to_drop_) + features = [f for f in self.feature_names_in_ if f not in features_to_drop] - # return the dataframe with the selected features - return X.drop(columns=self.features_to_drop_) + # pandas is faster than narwhals. + if nwd.is_pandas_dataframe(X) is True: + return X[features] + else: + return nw_X.select(nw.col(*features)).to_native() - def _get_feature_names_in(self, X): - """Get the names and number of features in the train set. The dataframe - used during fit.""" + def _get_feature_names_in(self, X: IntoDataFrame): + """Get the names and number of features in the train set (the dataframe + used during fit).""" - self.feature_names_in_ = X.columns.to_list() + if nwd.is_pandas_dataframe(X) is True: + self.feature_names_in_ = list(X.columns) + else: + self.feature_names_in_ = nw.from_native(X, eager_only=True).columns self.n_features_in_ = X.shape[1] return self diff --git a/tests/test_selection/conftest.py b/tests/test_selection/conftest.py index e41d7ce4e..7006979c8 100644 --- a/tests/test_selection/conftest.py +++ b/tests/test_selection/conftest.py @@ -25,6 +25,23 @@ def df_test(): return X, y +@pytest.fixture(scope="module") +def data_classification(): + """The data of df_test as a plain dict, with the target under "target".""" + X, y = make_classification( + n_samples=1000, + n_features=12, + n_redundant=4, + n_clusters_per_class=1, + weights=[0.50], + class_sep=2, + random_state=1, + ) + data = {f"var_{i}": X[:, i].tolist() for i in range(12)} + data["target"] = y.tolist() + return data + + @pytest.fixture(scope="module") def df_test_with_groups(): # Parameters diff --git a/tests/test_selection/test_base_recursive_selector.py b/tests/test_selection/test_base_recursive_selector.py new file mode 100644 index 000000000..a4b0a3c6b --- /dev/null +++ b/tests/test_selection/test_base_recursive_selector.py @@ -0,0 +1,257 @@ +import re + +import narwhals as nw +import numpy as np +import pandas as pd +import pytest +from sklearn.ensemble import RandomForestClassifier +from sklearn.linear_model import LogisticRegression +from sklearn.model_selection import GroupKFold, StratifiedKFold +from sklearn.neighbors import KNeighborsClassifier + +from feature_engine.selection.base_recursive_selector import BaseRecursiveSelector +from tests.backend_helpers import make_series + +VARIABLES = ["var_0", "var_4", "var_7"] + + +def _split_target(make_df, data): + X = make_df({k: v for k, v in data.items() if k != "target"}) + return X, make_series(make_df, data["target"]) + + +# init parameters +@pytest.mark.parametrize("threshold", [None, [0.1], "a_string", {"a": 1}]) +def test_error_if_threshold_not_number(threshold): + msg = f"threshold must be an integer or a float. Got {threshold} instead." + with pytest.raises(ValueError, match=re.escape(msg)): + BaseRecursiveSelector(RandomForestClassifier(), threshold=threshold) + + +@pytest.mark.parametrize( + "estimator, scoring, cv, groups, threshold, confirm_variables", + [ + (RandomForestClassifier(), "roc_auc", 3, None, 0.01, False), + (LogisticRegression(), "accuracy", StratifiedKFold(), [1, 2], 1, True), + (KNeighborsClassifier(), "r2", GroupKFold(), None, -0.5, False), + ], +) +def test_init_param_assignment( + estimator, scoring, cv, groups, threshold, confirm_variables +): + selector = BaseRecursiveSelector( + estimator, + scoring=scoring, + cv=cv, + groups=groups, + threshold=threshold, + confirm_variables=confirm_variables, + ) + assert selector.estimator is estimator + assert selector.scoring == scoring + assert selector.cv is cv + assert selector.groups == groups + assert selector.threshold == threshold + assert selector.confirm_variables is confirm_variables + + +# fit +@pytest.mark.parametrize("cv_type", ["int", "splitter", "generator"]) +def test_fit_with_feature_importances(make_df, data_classification, cv_type): + X, y = _split_target(make_df, data_classification) + cv = {"int": 3, "splitter": StratifiedKFold(n_splits=3)}.get(cv_type) + if cv_type == "generator": + cv = StratifiedKFold(n_splits=3).split(X, y) + + selector = BaseRecursiveSelector( + RandomForestClassifier(n_estimators=5, random_state=1), + scoring="roc_auc", + cv=cv, + variables=VARIABLES, + ) + selector.fit(X, y) + + assert selector.variables_ == VARIABLES + assert selector.feature_names_in_ == [f"var_{i}" for i in range(12)] + assert selector.n_features_in_ == 12 + assert selector.initial_model_performance_ == pytest.approx(0.9947395643178775) + assert dict(selector.feature_importances_) == pytest.approx( + { + "var_0": 0.05016049596989658, + "var_4": 0.5704427408209541, + "var_7": 0.37939676320914945, + } + ) + assert dict(selector.feature_importances_std_) == pytest.approx( + { + "var_0": 0.02070511883333705, + "var_4": 0.031350384984021235, + "var_7": 0.03942210400299374, + } + ) + + +def test_fit_with_coefficients(make_df, data_classification): + X, y = _split_target(make_df, data_classification) + selector = BaseRecursiveSelector( + LogisticRegression(), scoring="roc_auc", cv=3, variables=VARIABLES + ) + selector.fit(X, y) + + assert selector.initial_model_performance_ == pytest.approx(0.996732746280939) + assert dict(selector.feature_importances_) == pytest.approx( + { + "var_0": 2.0421118575507067, + "var_4": 0.36861802531867144, + "var_7": 2.68269483714549, + } + ) + assert dict(selector.feature_importances_std_) == pytest.approx( + { + "var_0": 0.08627973281236784, + "var_4": 0.12490483002792199, + "var_7": 0.13862536091514688, + } + ) + + +def test_fit_with_permutation_importance(make_df, data_classification): + # KNN has no coef_ or feature_importances_, so importance comes from + # permutation_importance. + X, y = _split_target(make_df, data_classification) + selector = BaseRecursiveSelector( + KNeighborsClassifier(), scoring="accuracy", cv=3, variables=VARIABLES + ) + selector.fit(X, y) + + assert selector.initial_model_performance_ == pytest.approx(0.991997986009962) + assert dict(selector.feature_importances_) == pytest.approx( + { + "var_0": 0.0050000000000000044, + "var_4": 0.0753333333333333, + "var_7": 0.42766666666666664, + } + ) + assert dict(selector.feature_importances_std_) == pytest.approx( + { + "var_0": 0.0010000000000000009, + "var_4": 0.0005773502691896263, + "var_7": 0.0037859388972001223, + } + ) + + +def test_feature_importances_type(make_df, data_classification): + X, y = _split_target(make_df, data_classification) + selector = BaseRecursiveSelector( + LogisticRegression(), scoring="roc_auc", cv=3, variables=VARIABLES + ) + selector.fit(X, y) + + importance_type = pd.Series if make_df is pd.DataFrame else dict + assert isinstance(selector.feature_importances_, importance_type) + assert isinstance(selector.feature_importances_std_, importance_type) + assert list(selector.feature_importances_.keys()) == VARIABLES + + +def test_fit_returns_narwhals_frame_and_target(make_df, data_classification): + X, y = _split_target(make_df, data_classification) + selector = BaseRecursiveSelector( + LogisticRegression(), scoring="roc_auc", cv=3, variables=VARIABLES + ) + nw_X, y_ = selector.fit(X, y) + + assert isinstance(nw_X, nw.DataFrame) + assert nw_X.to_native() is X + assert isinstance(y_, type(y)) + + +def test_fit_finds_numerical_variables(make_df, data_classification): + X, y = _split_target( + make_df, + { + "var_0": data_classification["var_0"], + "cat": ["a", "b"] * 500, + "var_4": data_classification["var_4"], + "target": data_classification["target"], + }, + ) + selector = BaseRecursiveSelector(LogisticRegression(), scoring="roc_auc", cv=3) + selector.fit(X, y) + + assert selector.variables_ == ["var_0", "var_4"] + assert selector.feature_names_in_ == ["var_0", "cat", "var_4"] + assert list(selector.feature_importances_.keys()) == ["var_0", "var_4"] + + +def test_fit_with_confirm_variables(make_df, data_classification): + X, y = _split_target(make_df, data_classification) + selector = BaseRecursiveSelector( + LogisticRegression(), + scoring="roc_auc", + cv=3, + variables=["var_0", "var_4", "Hola"], + confirm_variables=True, + ) + selector.fit(X, y) + + assert selector.variables_ == ["var_0", "var_4"] + assert list(selector.feature_importances_.keys()) == ["var_0", "var_4"] + + +@pytest.mark.parametrize("target_type", [list, np.array]) +def test_fit_with_list_and_array_target(make_df, data_classification, target_type): + X, _ = _split_target(make_df, data_classification) + y = target_type(data_classification["target"]) + selector = BaseRecursiveSelector( + LogisticRegression(), scoring="roc_auc", cv=3, variables=VARIABLES + ) + selector.fit(X, y) + + assert selector.initial_model_performance_ == pytest.approx(0.996732746280939) + + +def test_error_if_only_one_variable(make_df, data_classification): + X, y = _split_target(make_df, data_classification) + selector = BaseRecursiveSelector( + LogisticRegression(), scoring="roc_auc", cv=3, variables=["var_0"] + ) + msg = ( + "The selector needs at least 2 or more variables to select from. " + "Got only 1 variable: ['var_0']." + ) + with pytest.raises(ValueError, match=re.escape(msg)): + selector.fit(X, y) + + +def test_feature_importances_are_pandas_series_indexed_by_variables( + data_classification, +): + X, y = _split_target(pd.DataFrame, data_classification) + selector = BaseRecursiveSelector( + LogisticRegression(), scoring="roc_auc", cv=3, variables=VARIABLES + ) + selector.fit(X, y) + + pd.testing.assert_series_equal( + selector.feature_importances_, + pd.Series( + [2.0421118575507067, 0.36861802531867144, 2.68269483714549], + index=VARIABLES, + ), + ) + + +def test_fit_with_integer_column_names(data_classification): + X, y = _split_target(pd.DataFrame, data_classification) + X.columns = list(range(12)) + selector = BaseRecursiveSelector( + LogisticRegression(), scoring="roc_auc", cv=3, variables=[0, 4, 7] + ) + selector.fit(X, y) + + assert selector.variables_ == [0, 4, 7] + assert selector.feature_names_in_ == list(range(12)) + assert dict(selector.feature_importances_) == pytest.approx( + {0: 2.0421118575507067, 4: 0.36861802531867144, 7: 2.68269483714549} + ) diff --git a/tests/test_selection/test_base_selection_functions.py b/tests/test_selection/test_base_selection_functions.py index b2345a53e..0c4cb0e03 100644 --- a/tests/test_selection/test_base_selection_functions.py +++ b/tests/test_selection/test_base_selection_functions.py @@ -1,333 +1,364 @@ +from datetime import datetime + +import numpy as np import pandas as pd import pytest from sklearn.ensemble import RandomForestClassifier -from sklearn.model_selection import StratifiedKFold, GroupKFold - +from sklearn.linear_model import Lasso, LogisticRegression +from sklearn.model_selection import GroupKFold, StratifiedKFold from feature_engine.selection.base_selection_functions import ( _select_all_variables, _select_numerical_variables, find_correlated_features, find_feature_importance, + get_feature_importances, single_feature_performance, ) - - -@pytest.fixture -def df(): - df = pd.DataFrame( - { - "Name": ["tom", "nick", "krish", "jack"], - "City": ["London", "Manchester", "Liverpool", "Bristol"], - "Age": [20, 21, 19, 18], - "Marks": [0.9, 0.8, 0.7, 0.6], - "date_range": pd.date_range("2020-02-24", periods=4, freq="min"), - "date_obj0": ["2020-02-24", "2020-02-25", "2020-02-26", "2020-02-27"], - } +from tests.backend_helpers import make_series + +DATA_VARTYPES = { + "Name": ["tom", "nick", "krish", "jack"], + "City": ["London", "Manchester", "Liverpool", "Bristol"], + "Age": [20, 21, 19, 18], + "Marks": [0.9, 0.8, 0.7, 0.6], + "date_range": [datetime(2020, 2, 24, 0, minute) for minute in range(4)], + "date_obj0": ["2020-02-24", "2020-02-25", "2020-02-26", "2020-02-27"], +} + +# a is uncorrelated, b is correlated with c and d, and c and d only with b. +DATA_CORR = { + "a": [1, -1, 0, 0, 0, 0, 2], + "b": [0, 0, 1, -1, 1, -1, 0], + "c": [0, 0, 1, -1, 0, 0, 1], + "d": [0, 0, 0, 0, 1, -1, 1], +} + +EXPECTED_MEAN = { + "var_0": 0.5813469607144305, + "var_1": 0.5325152703164752, + "var_2": 0.5023573007759755, + "var_3": 0.47596844810700234, + "var_4": 0.9696712897767115, + "var_5": 0.5078009005719849, + "var_6": 0.966096275433625, + "var_7": 0.9918595739378872, + "var_8": 0.521667767752105, + "var_9": 0.9476311088509884, + "var_10": 0.4871054926777818, + "var_11": 0.5180029642379039, +} + +EXPECTED_STD = { + "var_0": 0.0035430274728173775, + "var_1": 0.0046697767238672565, + "var_2": 0.023714708852568194, + "var_3": 0.04219857610624132, + "var_4": 0.010364079344188424, + "var_5": 0.03203946151605523, + "var_6": 0.0063709642968091335, + "var_7": 0.0014579159677989356, + "var_8": 0.027570153897628277, + "var_9": 0.014363240810578251, + "var_10": 0.020283618255582142, + "var_11": 0.02707242215734807, +} + + +EXPECTED_IMPORTANCE_MEAN = { + "var_0": 0.008110472647428566, + "var_1": 0.004425867029009318, + "var_2": 0.0014110527658847542, + "var_3": 0.0, + "var_4": 0.09519163119147249, + "var_5": 0.005151538162222261, + "var_6": 0.06819196935501609, + "var_7": 0.7958920351591532, + "var_8": 0.005514122712728161, + "var_9": 0.006699116878609683, + "var_10": 0.0006441114577744834, + "var_11": 0.008768082640701015, +} + +EXPECTED_IMPORTANCE_STD = { + "var_0": 0.0044896289532760465, + "var_1": 0.004400023300047043, + "var_2": 0.0012779435856766451, + "var_3": 0.0, + "var_4": 0.15831114958148926, + "var_5": 0.005860881861755391, + "var_6": 0.0666073858478959, + "var_7": 0.12339701768095801, + "var_8": 0.0037045342320885986, + "var_9": 0.0016309720229483885, + "var_10": 0.0011156337706026607, + "var_11": 0.005458498614789974, +} + + +def _split_target(make_df, data): + X = make_df({k: v for k, v in data.items() if k != "target"}) + return X, make_series(make_df, data["target"]) + + +def _pearson(x, y): + return np.corrcoef(x, y)[0, 1] + + +@pytest.fixture(scope="module") +def data_with_groups(): + rng = np.random.default_rng(1) + data = {f"var_{i}": rng.normal(size=100).tolist() for i in range(1, 6)} + data["target"] = rng.integers(0, 100, size=100).tolist() + groups = np.repeat(np.arange(1, 11), 10) + rng.shuffle(groups) + return data, groups.tolist() + + +@pytest.mark.parametrize( + "variables, confirm_variables, exclude_datetime, expected", + [ + (None, False, False, list(DATA_VARTYPES)), + (None, False, True, ["Name", "City", "Age", "Marks"]), + (["Name", "Age"], False, True, ["Name", "Age"]), + (["Name", "Age", "Hola"], True, True, ["Name", "Age"]), + ], +) +def test_select_all_variables( + make_df, variables, confirm_variables, exclude_datetime, expected +): + variables_ = _select_all_variables( + make_df(DATA_VARTYPES), variables, confirm_variables, exclude_datetime ) - df["Name"] = df["Name"].astype("category") - return df + assert variables_ == expected -def test_select_all_variables(df): - # select all variables - assert ( - _select_all_variables( - df, variables=None, confirm_variables=False, exclude_datetime=False - ) - == df.columns.to_list() +@pytest.mark.parametrize( + "variables, confirm_variables, expected", + [ + (None, False, ["Age", "Marks"]), + (["Marks"], False, ["Marks"]), + (["Marks", "Hola"], True, ["Marks"]), + ], +) +def test_select_numerical_variables(make_df, variables, confirm_variables, expected): + variables_ = _select_numerical_variables( + make_df(DATA_VARTYPES), variables, confirm_variables ) + assert variables_ == expected - # select all variables except datetime - assert _select_all_variables( - df, variables=None, confirm_variables=False, exclude_datetime=True - ) == ["Name", "City", "Age", "Marks"] - - # select subset of variables, without confirm - subset = ["Name", "City", "Age", "Marks"] - assert ( - _select_all_variables( - df, variables=subset, confirm_variables=False, exclude_datetime=True - ) - == subset - ) - # select subset of variables, with confirm - subset = ["Name", "City", "Age", "Marks", "Hola"] - assert ( - _select_all_variables( - df, variables=subset, confirm_variables=True, exclude_datetime=True - ) - == subset[:-1] +@pytest.mark.parametrize( + "variables, expected", + [ + (["a", "b", "c", "d"], ([{"b", "c", "d"}], ["c", "d"], {"b": {"c", "d"}})), + (["a", "c", "b", "d"], ([{"c", "b"}], ["b"], {"c": {"b"}})), + ], +) +def test_find_correlated_features(make_df, variables, expected): + X = make_df( + { + "a": [1, -1, 0, 0, 0, 0], + "b": [0, 0, 1, -1, 1, -1], + "c": [0, 0, 1, -1, 0, 0], + "d": [0, 0, 0, 0, 1, -1], + } ) + assert find_correlated_features(X, variables, "pearson", 0.7) == expected -def test_select_numerical_variables(df): - # select all numerical variables - assert _select_numerical_variables( - df, - variables=None, - confirm_variables=False, - ) == ["Age", "Marks"] - - # select subset of variables, without confirm - subset = ["Marks"] - assert ( - _select_numerical_variables( - df, - variables=subset, - confirm_variables=False, - ) - == subset - ) +@pytest.mark.parametrize("method", ["pearson", "spearman", "kendall", _pearson]) +def test_find_correlated_features_methods(make_df, method): + X = make_df(DATA_CORR) + groups, drop, dict_ = find_correlated_features(X, list(DATA_CORR), method, 0.5) + assert groups == [{"b", "c", "d"}] + assert drop == ["c", "d"] + assert dict_ == {"b": {"c", "d"}} - # select subset of variables, with confirm - subset = ["Marks", "Hola"] - assert ( - _select_numerical_variables( - df, - variables=subset, - confirm_variables=True, - ) - == subset[:-1] - ) +@pytest.mark.parametrize( + "method, expected", + [ + ("pearson", ([{"a", "c"}, {"b", "d"}], ["c", "d"], {"a": {"c"}, "b": {"d"}})), + ("spearman", ([{"b", "d"}], ["d"], {"b": {"d"}})), + ("kendall", ([{"b", "d"}], ["d"], {"b": {"d"}})), + (_pearson, ([{"a", "c"}, {"b", "d"}], ["c", "d"], {"a": {"c"}, "b": {"d"}})), + ], +) +def test_find_correlated_features_skips_missing_values(make_df, method, expected): + # each pair of variables is compared on the rows where neither is missing. + data = { + "a": [None, -1, 0, 0, 0, 0, 2], + "b": [0, 0, 1, -1, 1, -1, None], + "c": [0, 0, None, -1, 0, 0, 1], + "d": [0, 0, 0, 0, 1, -1, 1], + } + X = make_df(data) + assert find_correlated_features(X, list(data), method, 0.6) == expected -def test_find_correlated_features(): - # given a correlation-threshold of 0.7 - # a is uncorrelated, - # b is collinear to c and d, - # c and d are collinear only to b. - X = pd.DataFrame() - X["a"] = [1, -1, 0, 0, 0, 0] - X["b"] = [0, 0, 1, -1, 1, -1] - X["c"] = [0, 0, 1, -1, 0, 0] - X["d"] = [0, 0, 0, 0, 1, -1] +@pytest.mark.parametrize("method", ["pearson", "spearman", "kendall"]) +def test_find_correlated_features_ignores_constant_variables(make_df, method): + X = make_df({**DATA_CORR, "e": [1] * 7}) groups, drop, dict_ = find_correlated_features( - X, variables=["a", "b", "c", "d"], method="pearson", threshold=0.7 + X, ["e", *DATA_CORR], method, 0.5 ) - assert groups == [{"b", "c", "d"}] assert drop == ["c", "d"] assert dict_ == {"b": {"c", "d"}} - groups, drop, dict_ = find_correlated_features( - X, variables=["a", "c", "b", "d"], method="pearson", threshold=0.7 - ) - assert groups == [{"c", "b"}] - assert drop == ["b"] - assert dict_ == {"c": {"b"}} +def test_find_correlated_features_with_integer_column_names(): + X = pd.DataFrame({i: values for i, values in enumerate(DATA_CORR.values())}) + groups, drop, dict_ = find_correlated_features(X, [0, 1, 2, 3], "pearson", 0.5) + assert groups == [{1, 2, 3}] + assert drop == [2, 3] + assert dict_ == {1: {2, 3}} -def test_single_feature_performance(df_test): - X, y = df_test +def test_single_feature_performance(make_df, data_classification): + X, y = _split_target(make_df, data_classification) rf = RandomForestClassifier(n_estimators=5, random_state=1) - variables = X.columns.to_list() mean_, std_ = single_feature_performance( X=X, y=y, - variables=variables, + variables=list(EXPECTED_MEAN), estimator=rf, cv=3, scoring="roc_auc", ) - - expected_mean = { - "var_0": 0.5813469607144305, - "var_1": 0.5325152703164752, - "var_2": 0.5023573007759755, - "var_3": 0.47596844810700234, - "var_4": 0.9696712897767115, - "var_5": 0.5078009005719849, - "var_6": 0.966096275433625, - "var_7": 0.9918595739378872, - "var_8": 0.521667767752105, - "var_9": 0.9476311088509884, - "var_10": 0.4871054926777818, - "var_11": 0.5180029642379039, - } - expected_std = { - "var_0": 0.0035430274728173775, - "var_1": 0.0046697767238672565, - "var_2": 0.023714708852568194, - "var_3": 0.04219857610624132, - "var_4": 0.010364079344188424, - "var_5": 0.03203946151605523, - "var_6": 0.0063709642968091335, - "var_7": 0.0014579159677989356, - "var_8": 0.027570153897628277, - "var_9": 0.014363240810578251, - "var_10": 0.020283618255582142, - "var_11": 0.02707242215734807, - } - assert mean_ == expected_mean - assert std_ == expected_std + assert mean_ == pytest.approx(EXPECTED_MEAN) + assert std_ == pytest.approx(EXPECTED_STD) -def test_single_feature_performance_cv_generator(df_test): - X, y = df_test +@pytest.mark.parametrize("cv_type", ["splitter", "generator"]) +def test_single_feature_performance_with_cv_splitter( + make_df, data_classification, cv_type +): + X, y = _split_target(make_df, data_classification) rf = RandomForestClassifier(n_estimators=5, random_state=1) - variables = X.columns.to_list() cv = StratifiedKFold(n_splits=3) - for cv_ in [cv, cv.split(X, y)]: - mean_, _ = single_feature_performance( - X=X, - y=y, - variables=variables, - estimator=rf, - cv=cv_, - scoring="roc_auc", - ) - - expected_mean = { - "var_0": 0.5813469607144305, - "var_1": 0.5325152703164752, - "var_2": 0.5023573007759755, - "var_3": 0.47596844810700234, - "var_4": 0.9696712897767115, - "var_5": 0.5078009005719849, - "var_6": 0.966096275433625, - "var_7": 0.9918595739378872, - "var_8": 0.521667767752105, - "var_9": 0.9476311088509884, - "var_10": 0.4871054926777818, - "var_11": 0.5180029642379039, - } - assert mean_ == expected_mean + if cv_type == "generator": + cv = cv.split(X, y) + + mean_, _ = single_feature_performance( + X=X, y=y, variables=list(EXPECTED_MEAN), estimator=rf, cv=cv, scoring="roc_auc" + ) + assert mean_ == pytest.approx(EXPECTED_MEAN) -def test_single_feature_performance_with_groups(df_test_with_groups): - X, y, groups = df_test_with_groups +@pytest.mark.parametrize("target_type", [list, np.array]) +def test_single_feature_performance_with_list_and_array_target( + make_df, data_classification, target_type +): + X, _ = _split_target(make_df, data_classification) + y = target_type(data_classification["target"]) rf = RandomForestClassifier(n_estimators=5, random_state=1) - variables = X.columns.to_list() - scoring = "neg_mean_absolute_error" + + mean_, _ = single_feature_performance( + X=X, y=y, variables=["var_0", "var_4"], estimator=rf, cv=3, scoring="roc_auc" + ) + assert mean_ == pytest.approx( + {"var_0": EXPECTED_MEAN["var_0"], "var_4": EXPECTED_MEAN["var_4"]} + ) + + +def test_single_feature_performance_with_groups(make_df, data_with_groups): + data, groups = data_with_groups + X, y = _split_target(make_df, data) + rf = RandomForestClassifier(n_estimators=5, random_state=1) + variables = ["var_1", "var_2", "var_3", "var_4", "var_5"] cv = GroupKFold(n_splits=3) - cv_indices = cv.split(X=X, y=y, groups=groups) expected_mean_, expected_std_ = single_feature_performance( X=X, y=y, variables=variables, estimator=rf, - cv=cv_indices, - scoring=scoring, + cv=cv.split(X=X, y=y, groups=groups), + scoring="neg_mean_absolute_error", ) - mean_, std_ = single_feature_performance( X=X, y=y, variables=variables, estimator=rf, cv=cv, - scoring=scoring, + scoring="neg_mean_absolute_error", groups=groups, ) - assert mean_ == expected_mean_ assert std_ == expected_std_ -def test_find_feature_importance(df_test): - X, y = df_test +@pytest.mark.parametrize("cv_type", ["splitter", "generator"]) +def test_find_feature_importance(make_df, data_classification, cv_type): + X, y = _split_target(make_df, data_classification) rf = RandomForestClassifier(n_estimators=3, random_state=3) cv = StratifiedKFold(n_splits=3) - scoring = "recall" - - expected_mean = pd.Series( - data=[0.01, 0.0, 0.0, 0.0, 0.1, 0.01, 0.07, 0.8, 0.01, 0.01, 0.0, 0.01], - index=[ - "var_0", - "var_1", - "var_2", - "var_3", - "var_4", - "var_5", - "var_6", - "var_7", - "var_8", - "var_9", - "var_10", - "var_11", - ], - ) - expected_std = pd.Series( - data=[ - 0.0045, - 0.0044, - 0.0013, - 0.0, - 0.1583, - 0.0059, - 0.0666, - 0.1234, - 0.0037, - 0.0016, - 0.0011, - 0.0055, - ], - index=[ - "var_0", - "var_1", - "var_2", - "var_3", - "var_4", - "var_5", - "var_6", - "var_7", - "var_8", - "var_9", - "var_10", - "var_11", - ], - ) + if cv_type == "generator": + cv = cv.split(X, y) mean_, std_ = find_feature_importance( - X=X, - y=y, - estimator=rf, - cv=cv, - scoring=scoring, + X=X, y=y, estimator=rf, cv=cv, scoring="recall" ) - pd.testing.assert_series_equal(mean_.round(2), expected_mean) - pd.testing.assert_series_equal(std_.round(4), expected_std) - mean_, std_ = find_feature_importance( - X=X, - y=y, - estimator=rf, - cv=cv.split(X, y), - scoring=scoring, - ) - pd.testing.assert_series_equal(mean_.round(2), expected_mean) - pd.testing.assert_series_equal(std_.round(4), expected_std) + importance_type = pd.Series if make_df is pd.DataFrame else dict + assert isinstance(mean_, importance_type) + assert isinstance(std_, importance_type) + assert dict(mean_) == pytest.approx(EXPECTED_IMPORTANCE_MEAN) + assert dict(std_) == pytest.approx(EXPECTED_IMPORTANCE_STD) -def test_find_feature_importancewith_groups(df_test_with_groups): - X, y, groups = df_test_with_groups +def test_find_feature_importance_with_groups(make_df, data_with_groups): + data, groups = data_with_groups + X, y = _split_target(make_df, data) rf = RandomForestClassifier(n_estimators=3, random_state=1) cv = GroupKFold(n_splits=3) - scoring = "neg_mean_absolute_error" - cv_indices = cv.split(X=X, y=y, groups=groups) expected_mean_, expected_std_ = find_feature_importance( X=X, y=y, estimator=rf, - cv=cv_indices, - scoring=scoring, + cv=cv.split(X=X, y=y, groups=groups), + scoring="neg_mean_absolute_error", ) - mean_, std_ = find_feature_importance( X=X, y=y, estimator=rf, cv=cv, - scoring=scoring, - groups=groups + scoring="neg_mean_absolute_error", + groups=groups, ) + assert dict(mean_) == dict(expected_mean_) + assert dict(std_) == dict(expected_std_) - pd.testing.assert_series_equal(mean_, expected_mean_) - pd.testing.assert_series_equal(std_, expected_std_) + +def test_find_feature_importance_returns_series_indexed_by_columns(): + X = pd.DataFrame({"a": [1.0, 2, 3, 4, 5, 6], "b": [0.5, 0, 1, 1, 3, 2]}) + X.columns.name = "features" + y = pd.Series([1.0, 2, 3, 4, 5, 6]) + + mean_, std_ = find_feature_importance( + X=X, y=y, estimator=Lasso(alpha=0.01), cv=2, scoring="r2" + ) + assert list(mean_.index) == ["a", "b"] + assert mean_.index.name == "features" + assert std_.index.name == "features" + + +@pytest.mark.parametrize( + "estimator, expected", + [ + (LogisticRegression(), [0.728 ** (1 / 3), 1.0]), + (Lasso(), [0.5, 2.0]), + ], +) +def test_get_feature_importances(estimator, expected): + if isinstance(estimator, Lasso): + estimator.coef_ = np.array([-0.5, 2.0]) + else: + estimator.coef_ = np.array([[0.6, 0.0], [0.0, 1.0], [0.8, 0.0]]) + assert get_feature_importances(estimator) == pytest.approx(expected) diff --git a/tests/test_selection/test_base_selector.py b/tests/test_selection/test_base_selector.py index 54cc5bfc0..7ef31226b 100644 --- a/tests/test_selection/test_base_selector.py +++ b/tests/test_selection/test_base_selector.py @@ -1,55 +1,163 @@ +import re +from datetime import datetime + +import numpy as np +import pandas as pd import pytest -from pandas.testing import assert_frame_equal +from sklearn.exceptions import NotFittedError from feature_engine.selection.base_selector import BaseSelector +from tests.backend_helpers import frame_to_dict + +DATA = { + "Name": ["tom", "nick", "krish", None], + "City": ["London", "Manchester", "Liverpool", "Bristol"], + "Age": [20, 21, 19, 18], + "Marks": [0.9, None, 0.7, 0.6], + "dob": [datetime(2020, 2, 24, 0, minute) for minute in range(4)], +} + + +# init parameters +@pytest.mark.parametrize("confirm_variables", [None, "hola", [True], 1, 0.5]) +def test_error_if_confirm_variables_not_bool(confirm_variables): + msg = ( + "confirm_variables takes only values True and False. " + f"Got {confirm_variables} instead." + ) + with pytest.raises(ValueError, match=re.escape(msg)): + BaseSelector(confirm_variables=confirm_variables) -@pytest.mark.parametrize("val", [None, "hola", [True]]) -def test_confirm_variables_in_init(val): - with pytest.raises(ValueError): - BaseSelector(confirm_variables=val) +@pytest.mark.parametrize("confirm_variables", [True, False]) +def test_init_param_assignment(confirm_variables): + selector = BaseSelector(confirm_variables=confirm_variables) + assert selector.confirm_variables is confirm_variables -class MockClass(BaseSelector): - def __init__(self, variables=None, confirm_variables=False): - self.variables = variables - self.confirm_variables = confirm_variables +# fit and transform +class MockSelector(BaseSelector): + def __init__(self, features_to_drop=("Name", "Marks")): + self.features_to_drop = features_to_drop + self.confirm_variables = False def fit(self, X, y=None): - self.features_to_drop_ = ["Name", "Marks"] + self.features_to_drop_ = list(self.features_to_drop) self._get_feature_names_in(X) return self -def test_transform_method(df_vartypes): - transformer = MockClass() - transformer.fit(df_vartypes) - Xt = transformer.transform(df_vartypes) +def test_transform_drops_features(make_df): + X = make_df(DATA) + Xt = MockSelector().fit(X).transform(X) + + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == { + "City": ["London", "Manchester", "Liverpool", "Bristol"], + "Age": [20, 21, 19, 18], + "dob": DATA["dob"], + } + + +def test_transform_restores_train_column_order(make_df): + selector = MockSelector().fit(make_df(DATA)) + X = make_df({var: DATA[var] for var in ["dob", "Marks", "Age", "Name", "City"]}) + Xt = selector.transform(X) + + assert isinstance(Xt, make_df) + assert list(Xt.columns) == ["City", "Age", "dob"] + + +@pytest.mark.parametrize( + "features_to_drop, expected", + [ + ([], ["Name", "City", "Age", "Marks", "dob"]), + (["dob"], ["Name", "City", "Age", "Marks"]), + (["Age", "Name", "City", "dob"], ["Marks"]), + ], +) +def test_transform_returns_retained_features(make_df, features_to_drop, expected): + X = make_df(DATA) + Xt = MockSelector(features_to_drop).fit(X).transform(X) + + assert isinstance(Xt, make_df) + assert list(Xt.columns) == expected + + +def test_transform_does_not_modify_input(make_df): + X = make_df(DATA) + MockSelector().fit(X).transform(X) + assert frame_to_dict(X) == frame_to_dict(make_df(DATA)) + - # tests output of transform - assert_frame_equal(Xt, df_vartypes.drop(["Name", "Marks"], axis=1)) +def test_error_if_transform_df_has_different_number_of_columns(make_df): + selector = MockSelector().fit(make_df(DATA)) + msg = ( + "The number of columns in this dataset is different from the one used to " + "fit this transformer (when using the fit() method)." + ) + with pytest.raises(ValueError, match=re.escape(msg)): + selector.transform(make_df({"Age": DATA["Age"], "Marks": DATA["Marks"]})) + + +def test_error_if_transform_before_fit(make_df): + msg = ( + "This MockSelector instance is not fitted yet. Call 'fit' with " + "appropriate arguments before using this estimator." + ) + with pytest.raises(NotFittedError, match=re.escape(msg)): + MockSelector().transform(make_df(DATA)) - # tests this line: X = X[self.feature_names_in_] - assert_frame_equal( - transformer.transform(df_vartypes[["City", "Age", "Name", "Marks", "dob"]]), - Xt, + +def test_error_if_transform_input_not_dataframe(make_df): + selector = MockSelector().fit(make_df(DATA)) + msg = ( + "X must be a dataframe from a library supported by narwhals " + "(e.g. pandas, polars, PyArrow). Got instead." ) - # test error when there is a df shape missmatch - with pytest.raises(ValueError): - assert transformer.transform(df_vartypes[["Age", "Marks"]]) + with pytest.raises(TypeError, match=re.escape(msg)): + selector.transform(np.ones((4, 5))) + + +def test_get_feature_names_in(make_df): + selector = MockSelector() + selector._get_feature_names_in(make_df(DATA)) + assert selector.feature_names_in_ == ["Name", "City", "Age", "Marks", "dob"] + assert selector.n_features_in_ == 5 + + +def test_get_support(make_df): + selector = MockSelector().fit(make_df(DATA)) + assert selector.get_support() == [False, True, True, False, True] + assert list(selector.get_support(indices=True)) == [1, 2, 4] + + +def test_get_feature_names_out(make_df): + selector = MockSelector().fit(make_df(DATA)) + assert selector.get_feature_names_out() == ["City", "Age", "dob"] + + +def test_check_variable_number(): + selector = MockSelector() + selector.variables_ = ["Age"] + msg = ( + "The selector needs at least 2 or more variables to select from. " + "Got only 1 variable: ['Age']." + ) + with pytest.raises(ValueError, match=re.escape(msg)): + selector._check_variable_number() + +def test_transform_with_integer_column_names(): + X = pd.DataFrame({i: values for i, values in enumerate(DATA.values())}) + selector = MockSelector(features_to_drop=[0, 3]).fit(X) + Xt = selector.transform(X[[4, 3, 2, 1, 0]]) -def test_get_feature_names_in(df_vartypes): - tr = MockClass() - tr._get_feature_names_in(df_vartypes) - assert tr.n_features_in_ == df_vartypes.shape[1] - assert tr.feature_names_in_ == list(df_vartypes.columns) + assert selector.feature_names_in_ == [0, 1, 2, 3, 4] + pd.testing.assert_frame_equal(Xt, X[[1, 2, 4]]) -def test_get_support(df_vartypes): - tr = MockClass() - tr.fit(df_vartypes) - v_bool = [False, True, True, False, True] - v_ind = [1, 2, 4] - assert tr.get_support() == v_bool - assert list(tr.get_support(indices=True)) == v_ind +def test_transform_keeps_pandas_index(): + X = pd.DataFrame(DATA, index=[10, 11, 12, 13]) + Xt = MockSelector().fit(X).transform(X) + pd.testing.assert_frame_equal(Xt, X[["City", "Age", "dob"]]) From 8a105e39f563895c54bd4cf7e13a5ecc9ad25d9a Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Sat, 19 Sep 2026 12:16:30 +0200 Subject: [PATCH 2/3] Return an empty dataframe on polars when a selector drops every feature Co-Authored-By: Claude Opus 5 --- feature_engine/selection/base_selector.py | 4 ++++ tests/test_selection/test_base_selector.py | 1 + 2 files changed, 5 insertions(+) diff --git a/feature_engine/selection/base_selector.py b/feature_engine/selection/base_selector.py index 0e4ff5d97..40e84f493 100644 --- a/feature_engine/selection/base_selector.py +++ b/feature_engine/selection/base_selector.py @@ -70,6 +70,10 @@ def transform(self, X: IntoDataFrame) -> IntoDataFrame: if nwd.is_pandas_dataframe(X) is True: return X[features] else: + if len(features) == 0: + # nw.col() needs at least one name. polars frames without columns + # have no rows either. + return nw_X.select([]).to_native() return nw_X.select(nw.col(*features)).to_native() def _get_feature_names_in(self, X: IntoDataFrame): diff --git a/tests/test_selection/test_base_selector.py b/tests/test_selection/test_base_selector.py index 7ef31226b..64e43b88c 100644 --- a/tests/test_selection/test_base_selector.py +++ b/tests/test_selection/test_base_selector.py @@ -74,6 +74,7 @@ def test_transform_restores_train_column_order(make_df): ([], ["Name", "City", "Age", "Marks", "dob"]), (["dob"], ["Name", "City", "Age", "Marks"]), (["Age", "Name", "City", "dob"], ["Marks"]), + (["Name", "City", "Age", "Marks", "dob"], []), ], ) def test_transform_returns_retained_features(make_df, features_to_drop, expected): From 2969c7178ea24d11678bbeb248097ae4d9c68008 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Sat, 19 Sep 2026 12:42:50 +0200 Subject: [PATCH 3/3] Migrate DropHighPSIFeatures to narwhals, add polars support fit() no longer imports pandas and accepts any narwhals-supported dataframe. The numerical features are binned with numpy (the same intervals as EqualFrequencyDiscretiser and EqualWidthDiscretiser, learned from the basis set) and counted with np.bincount; the categories are counted with Series.factorize on pandas and a narwhals group_by otherwise. The dataframe is no longer copied into a basis and a test dataframe: a boolean mask selects the rows of each column. With pandas, split_col=None uses the index as before. Dataframes without an index, like polars, are split by the position of the rows. Also fixes cut_off=0 being ignored, the missing spaces in two error messages, and init accepting thresholds that are not numbers or 'auto'. Co-Authored-By: Claude Opus 5 --- .../selection/DropHighPSIFeatures.rst | 85 +- feature_engine/selection/drop_psi_features.py | 475 ++++--- .../test_drop_high_psi_features.py | 1266 ++++++++--------- 3 files changed, 966 insertions(+), 860 deletions(-) diff --git a/docs/user_guide/selection/DropHighPSIFeatures.rst b/docs/user_guide/selection/DropHighPSIFeatures.rst index b4ecf9bde..9229338f7 100644 --- a/docs/user_guide/selection/DropHighPSIFeatures.rst +++ b/docs/user_guide/selection/DropHighPSIFeatures.rst @@ -6,7 +6,7 @@ DropHighPSIFeatures =================== The :class:`DropHighPSIFeatures()` finds and removes features with changes in their -distribution, i.e. "unstable values", from a pandas dataframe. +distribution, i.e. "unstable values", from a pandas or polars dataframe. The stability of the distribution is computed using the **Population Stability Index (PSI)** and all features having a PSI value above a given threshold are removed. @@ -207,6 +207,13 @@ reference in the `split_col` parameter. If you don't, the split will be done bas values of the dataframe index. This might be a good option if the index contains meaningful values or if splitting just based on `split_frac`. +Polars dataframes don't have an index. If you don't indicate a column in `split_col`, +:class:`DropHighPSIFeatures()` uses the position of the rows instead, from 0 for the +first row to the number of rows minus 1 for the last one. This way, the first rows go +to the basis data set and the last rows to the test set. If the rows are in time order, +this compares the older observations with the most recent ones. The positions are also +what the `cut_off` refers to: with `cut_off=99`, the first 100 rows go to the basis set. + Python implementation ---------------------- @@ -255,7 +262,7 @@ with a shift in its distribution (*var_3* in this case): X = pd.DataFrame(X, columns=colnames) # Add a column with a shift. - X['var_3'][250:] = X['var_3'][250:] + 1 + X.loc[250:, 'var_3'] = X.loc[250:, 'var_3'] + 1 The default approach in :class:`DropHighPSIFeatures()` is to split the input dataframe `X` in two equally sized data sets. You can adjust the proportions by changing @@ -295,7 +302,7 @@ by the `transform` method. {'var_0': 0.07405459925568803, 'var_1': 0.09124093185820083, 'var_2': 0.16985790067687764, - 'var_3': 1.342485289730313, + 'var_3': 0.8324552340063687, 'var_4': 0.0743442762545251, 'var_5': 0.06809060587241555} @@ -942,6 +949,78 @@ noticeably between groups, whereas the distribution of *var_4* remains stable: .. image:: ../../images/PSI_distribution_case5.png +With polars +----------- + +:class:`DropHighPSIFeatures()` works in the same way with polars dataframes, and returns +a polars dataframe. Let's repeat Case 1 with a polars dataframe: + +.. code:: python + + import polars as pl + from sklearn.datasets import make_classification + from feature_engine.selection import DropHighPSIFeatures + + X, y = make_classification(n_samples=500, n_features=6, random_state=0) + + # Add a shift to the last 250 values of var_3. + X[250:, 3] = X[250:, 3] + 1 + X = pl.DataFrame(X, schema=["var_" + str(i) for i in range(6)]) + + transformer = DropHighPSIFeatures(split_frac=0.6) + X_transformed = transformer.fit_transform(X) + + print(transformer.psi_values_) + +We obtain the same PSI values as with pandas: + +.. code:: python + + {'var_0': 0.07405459925568803, + 'var_1': 0.09124093185820083, + 'var_2': 0.16985790067687764, + 'var_3': 0.8324552340063687, + 'var_4': 0.0743442762545251, + 'var_5': 0.06809060587241555} + +As the polars dataframe has no index, the data was split by the position of the rows. +The cut-off is the same as in Case 1: + +.. code:: python + + print(transformer.cut_off_) + +.. code:: python + + 299.4 + +This means that the first 300 rows, with positions 0 to 299, form the basis data set. +The transformed data is a polars dataframe without *var_3*: + +.. code:: python + + print(X_transformed.head()) + +.. code:: text + + shape: (5, 5) + ┌───────────┬───────────┬───────────┬───────────┬───────────┐ + │ var_0 ┆ var_1 ┆ var_2 ┆ var_4 ┆ var_5 │ + │ --- ┆ --- ┆ --- ┆ --- ┆ --- │ + │ f64 ┆ f64 ┆ f64 ┆ f64 ┆ f64 │ + ╞═══════════╪═══════════╪═══════════╪═══════════╪═══════════╡ + │ 0.476031 ┆ 0.494526 ┆ -0.565498 ┆ -1.596759 ┆ -0.283645 │ + │ -0.375222 ┆ -2.332233 ┆ 0.732424 ┆ -1.623769 ┆ 0.276038 │ + │ -0.561579 ┆ 0.719409 ┆ 1.608522 ┆ -0.702902 ┆ -0.224755 │ + │ 0.22584 ┆ -0.720572 ┆ 0.960557 ┆ -0.721353 ┆ 0.059809 │ + │ -0.446836 ┆ -0.66195 ┆ -0.061743 ┆ 1.390007 ┆ 0.293002 │ + └───────────┴───────────┴───────────┴───────────┴───────────┘ + +To split a polars dataframe by date, pass the name of the date column in `split_col`, +as in Case 3. If the column holds datetimes, pass a datetime as `cut_off`, for example +`datetime(1789, 7, 14)`, instead of a date. + + Additional resources -------------------- diff --git a/feature_engine/selection/drop_psi_features.py b/feature_engine/selection/drop_psi_features.py index ea94e8e6f..bb595e93f 100644 --- a/feature_engine/selection/drop_psi_features.py +++ b/feature_engine/selection/drop_psi_features.py @@ -1,10 +1,12 @@ import datetime -from typing import Dict, List, Union +import numbers +from typing import Dict, List, Optional, Union +import narwhals as nw +import narwhals.dependencies as nwd import numpy as np -import pandas as pd import scipy.stats as stats -from pandas.api.types import is_datetime64_any_dtype, is_numeric_dtype +from narwhals.typing import IntoDataFrame, IntoSeries from feature_engine._check_init_parameters.check_variables import ( _check_variables_input_value, @@ -27,10 +29,6 @@ _check_contains_na, check_X, ) -from feature_engine.discretisation import ( - EqualFrequencyDiscretiser, - EqualWidthDiscretiser, -) from feature_engine.selection.base_selector import BaseSelector from feature_engine.tags import _return_tags from feature_engine.variable_handling import ( @@ -103,7 +101,8 @@ class DropHighPSIFeatures(BaseSelector): First, you can indicate which variable should be used to guide the data split. This variable can be of any data type. If you do not enter a variable name, - DropHighPSIFeatures() will use the dataframe index. + DropHighPSIFeatures() will use the index of a pandas dataframe, or the row order of + dataframes that don't have an index, like polars dataframes. Next, you need to specify how that variable (or the index) should be used to split the data. You can specify a proportion of observations to be put in each data set, @@ -131,10 +130,12 @@ class DropHighPSIFeatures(BaseSelector): ---------- split_col: string or int, default=None. The variable that will be used to split the dataset into the basis and test - sets. If None, the dataframe index will be used. `split_col` can be a numerical, - categorical or datetime variable. If `split_col` is a categorical variable, and - the splitting criteria is given by `split_frac`, it will be assumed that the - labels of the variable are sorted alphabetically. + sets. If None, the index of a pandas dataframe will be used. Dataframes without + an index, like polars dataframes, are split by the position of the rows, + from 0 to n_samples - 1, so the first rows go to the basis set. `split_col` can + be a numerical, categorical or datetime variable. If `split_col` is a + categorical variable, and the splitting criteria is given by `split_frac`, it + will be assumed that the labels of the variable are sorted alphabetically. split_frac: float, default=0.5. The proportion of observations in each of the basis and test dataframes. If @@ -243,7 +244,8 @@ class DropHighPSIFeatures(BaseSelector): cut_off_: Value used to split the dataframe into basis and test. - This value is computed when not given as parameter. + This value is computed when not given as parameter. When `split_col` is None + and the dataframe has no index, it refers to the position of the rows. {feature_names_in_} @@ -295,6 +297,27 @@ class DropHighPSIFeatures(BaseSelector): 8 9 9 0 10 32 + 11 87 + 12 6 + 13 32 + 14 11 + 15 44 + 16 8 + 17 7 + 18 9 + 19 0 + + With a polars dataframe, which has no index, the rows are split by their + position: + + >>> import polars as pl + >>> X = pl.DataFrame(dict( + >>> x1 = [1,1,0,0,0,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0], + >>> x2 = [32,87,6,32,11,44,8,7,9,0,32,87,6,32,11,44,8,7,9,0], + >>> )) + >>> psi = DropHighPSIFeatures() + >>> psi.fit_transform(X).columns + ['x2'] """ def __init__( @@ -320,17 +343,16 @@ def __init__( f"{split_col} instead." ) - # split_frac and cut_off can't be None at the same time - if not split_frac and not cut_off: + if split_frac is None and cut_off is None: raise ValueError( - f"cut_off and split_frac cannot be both set to None " + f"cut_off and split_frac cannot be both set to None. " f"The current values are {split_frac, cut_off}. Please " f"specify a value for at least one of these parameters." ) - # check split_frac only if it will be used. - if split_frac and not cut_off: - if not (0 < split_frac < 1): + # split_frac is only used when cut_off is None. + if cut_off is None: + if not isinstance(split_frac, (float, int)) or not 0 < split_frac < 1: raise ValueError( f"split_frac must be a float between 0 and 1. Got {split_frac} " f"instead." @@ -344,17 +366,21 @@ def __init__( if not isinstance(switch, bool): raise ValueError(f"switch must be a boolean. Got {switch} instead.") - if (isinstance(threshold, str) and (threshold != "auto")) or ( - isinstance(threshold, (float, int)) and threshold < 0 + if not ( + (isinstance(threshold, str) and threshold == "auto") + or (isinstance(threshold, numbers.Real) and threshold >= 0) ): raise ValueError( f"threshold must be greater than 0 or 'auto'. Got {threshold} instead." ) if not isinstance(bins, int) or bins <= 1: - raise ValueError(f"bins must be an integer >= 1. Got {bins} instead.") + raise ValueError(f"bins must be an integer >= 2. Got {bins} instead.") - if strategy not in ["equal_width", "equal_frequency"]: + if not isinstance(strategy, str) or strategy not in [ + "equal_width", + "equal_frequency", + ]: raise ValueError( "strategy takes only values equal_width or equal_frequency. Got " f"{strategy} instead." @@ -366,7 +392,10 @@ def __init__( f"instead." ) - if missing_values not in ["raise", "ignore"]: + if not isinstance(missing_values, str) or missing_values not in [ + "raise", + "ignore", + ]: raise ValueError( f"missing_values takes only values 'raise' or 'ignore'. Got " f"{missing_values} instead." @@ -403,138 +432,108 @@ def __init__( self.missing_values = missing_values self.p_value = p_value - def fit(self, X: pd.DataFrame, y: pd.Series = None): + def fit(self, X: IntoDataFrame, y: Optional[IntoSeries] = None): """ Find features with high PSI values. Parameters ---------- - X : pandas dataframe of shape = [n_samples, n_features] + X : dataframe of shape = [n_samples, n_features] The training dataset. - y : pandas series. Default = None + y : Series. Default = None y is not needed in this transformer. You can pass y or None. """ - # check input dataframe - X = check_X(X) + nw_X = check_X(X) - # select variables to evaluate cat_variables_, num_variables_ = self._select_variables(X) - # check that split column is in the dataframe and remove from variable lists + # the split column is removed from the variables selected automatically cat_variables_, num_variables_ = self._check_split_column( - X, cat_variables_, num_variables_ + nw_X, cat_variables_, num_variables_ ) if self.missing_values == "raise": - # check if dataset contains na or inf _check_contains_na(X, num_variables_ + cat_variables_) - _check_contains_inf(X, num_variables_) + # the intervals can't be computed with inf values, even if NaN are ignored + _check_contains_inf(X, num_variables_) - # Split the dataframe into basis and test. - basis_df, test_df = self._split_dataframe(X) + is_basis = self._basis_mask(X, nw_X) + n_basis = int(is_basis.sum()) + n_test = is_basis.shape[0] - n_basis - # Check the shape of the returned dataframes for PSI calculations. - # The number of observations must be at least equal to the - # number of bins. - if min(basis_df.shape[0], test_df.shape[0]) < self.bins: + if min(n_basis, n_test) < self.bins: raise ValueError( "The number of rows in the basis and test datasets that will be used " f"in the PSI calculations must be at least larger than {self.bins}. " - "After splitting the original dataset based on the given cut_off or" - f"split_frac we have {basis_df.shape[0]} samples in the basis set, " - f"and {test_df.shape[0]} samples in the test set. " + "After splitting the original dataset based on the given cut_off or " + f"split_frac we have {n_basis} samples in the basis set, " + f"and {n_test} samples in the test set. " "Please adjust the value of the cut_off or split_frac." ) - # Switch basis and test dataframes if required. - if self.switch: - test_df, basis_df = basis_df, test_df - - # Set up parameters for numerical features - if len(num_variables_) > 0: - - # Set up the discretizer for numerical features - if self.strategy == "equal_width": - bucketer = EqualWidthDiscretiser(bins=self.bins) - else: - bucketer = EqualFrequencyDiscretiser(q=self.bins) + if self.switch is True: + is_basis = ~is_basis + n_basis, n_test = n_test, n_basis - # Set up the threshold for numerical features - if self.threshold == "auto": - threshold_num = self._calculate_auto_threshold( - basis_df.shape[0], - test_df.shape[0], - self.bins, - ) - else: - threshold_num = self.threshold - - # Set up the generic threshold for categorical features if used - if len(cat_variables_) > 0: - if self.threshold != "auto": - threshold_cat = self.threshold + if self.threshold == "auto": + threshold_num = self._calculate_auto_threshold(n_basis, n_test, self.bins) + else: + threshold_num = self.threshold - # Compute the PSI by looping over the features self.psi_values_: Dict = {} self.features_to_drop_ = [] - # Compute PSI for numerical features for feature in num_variables_: - # Discretize feature - basis_discrete = bucketer.fit_transform(basis_df[[feature]].dropna()) - test_discrete = bucketer.transform(test_df[[feature]].dropna()) - - # Determine percentage of observations per bin - basis_distrib, test_distrib = self._observation_frequency_per_bin( - basis_discrete, test_discrete + # pandas is faster than narwhals. + if nwd.is_pandas_dataframe(X) is True: + values = X[feature].to_numpy() + else: + values = nw_X.get_column(feature).to_numpy() + basis = values[is_basis] + test = values[~is_basis] + if self.missing_values == "ignore": + basis = basis[~np.isnan(basis)] + test = test[~np.isnan(test)] + self._check_observations(feature, basis, test) + + # the intervals are learned from the basis set only + limits = self._interval_limits(basis) + n_intervals = len(limits) + 1 + basis_counts = np.bincount( + np.searchsorted(limits, basis, side="left"), minlength=n_intervals ) - - # Calculate the PSI value - self.psi_values_[feature] = np.sum( - (test_distrib - basis_distrib) * np.log(test_distrib / basis_distrib) + test_counts = np.bincount( + np.searchsorted(limits, test, side="left"), minlength=n_intervals ) + self.psi_values_[feature] = self._psi(basis_counts, test_counts) - # Assess if feature should be dropped if self.psi_values_[feature] > threshold_num: self.features_to_drop_.append(feature) - # Compute the PSI for categorical features for feature in cat_variables_: - basis_discrete = basis_df[[feature]] - test_discrete = test_df[[feature]] - - # Determine percentage of observations per bin - basis_distrib, test_distrib = self._observation_frequency_per_bin( - basis_discrete, test_discrete - ) - - # Calculate the PSI value - self.psi_values_[feature] = np.sum( - (test_distrib - basis_distrib) * np.log(test_distrib / basis_distrib) + basis_counts, test_counts = self._category_counts( + X, nw_X, feature, is_basis ) + self.psi_values_[feature] = self._psi(basis_counts, test_counts) - # Determine the appropriate threshold for the categorical feature if self.threshold == "auto": - n_bins_cat = X[feature].nunique() + n_categories = np.count_nonzero(basis_counts + test_counts) threshold_cat = self._calculate_auto_threshold( - basis_df.shape[0], - test_df.shape[0], - n_bins_cat, + n_basis, n_test, n_categories ) + else: + threshold_cat = self.threshold - # Assess if feature should be dropped if self.psi_values_[feature] > threshold_cat: self.features_to_drop_.append(feature) - # store analyzed variables self.variables_ = num_variables_ + cat_variables_ - # save input features self._get_feature_names_in(X) return self - def _select_variables(self, X: pd.DataFrame): + def _select_variables(self, X: IntoDataFrame): """Based on the user input to the `variables` attribute in init, find the numerical and categorical variables for which the PSI should be calculated. @@ -568,7 +567,7 @@ def _select_variables(self, X: pd.DataFrame): def _check_split_column( self, - X: pd.DataFrame, + nw_X: nw.DataFrame, cat_variables: List[Union[str, int]], num_variables: List[Union[str, int]], ): @@ -578,12 +577,9 @@ def _check_split_column( It will get added if the variables are selected automatically. """ if self.split_col is not None: - # check that split_col is in the dataframe. - if self.split_col not in X.columns: + if self.split_col not in nw_X.columns: raise ValueError(f"{self.split_col} is not in the dataframe.") - # Remove the split_col from variables lists. Happens when variables are - # selected automatically by the transformer. if self.variables is None or self.variables == "all": if self.split_col in num_variables: num_variables.remove(self.split_col) @@ -592,117 +588,171 @@ def _check_split_column( return cat_variables, num_variables - def _observation_frequency_per_bin(self, basis, test): - """ - Obtain the fraction of observations per interval. - - Parameters - ---------- - basis : pd.DataFrame. - The basis Pandas DataFrame with discretised (i.e., binned) values. - - test: pd.DataFrame. - The test Pandas DataFrame with discretised (i.e., binned) values. - - Returns - ------- - distribution.basis: pd.Series. - Basis Pandas Series with percentage of observations per bin. - - distribution.meas: pd.Series. - Test Pandas Series with percentage of observations per bin. - """ - # Compute the feature distribution for basis and test - basis_distrib = basis.value_counts(normalize=True) - test_distrib = test.value_counts(normalize=True) - - # Align the two distributions by merging the buckets (bins). This ensures - # the number of bins is the same for the two distributions (in case of - # empty buckets). - distributions = ( - pd.DataFrame(basis_distrib) - .merge( - pd.DataFrame(test_distrib), - right_index=True, - left_index=True, - how="outer", + def _check_observations(self, feature, basis: np.ndarray, test: np.ndarray): + """Check that the basis and test sets have values of the feature, which can + be missing when NaN are ignored.""" + if len(basis) == 0 or len(test) == 0: + raise ValueError( + f"The variable {feature} has only missing values in the basis or in " + "the test set, so its PSI can't be computed. Got " + f"{len(basis)} values in the basis set and {len(test)} values in the " + "test set." ) - .fillna(self.min_pct_empty_bins) - .replace(to_replace=0, value=self.min_pct_empty_bins) - ) - distributions.columns = ["basis", "test"] - return distributions.basis, distributions.test - - def _split_dataframe(self, X: pd.DataFrame): + def _interval_limits(self, values: np.ndarray) -> np.ndarray: + """Inner limits of the intervals, the same ones EqualFrequencyDiscretiser + and EqualWidthDiscretiser find. The first and last intervals are open, so + values outside the basis range go to them.""" + if self.strategy == "equal_frequency": + quantiles = np.linspace(0, 1, self.bins + 1) + # round up quantiles that are not exact in base 2, as pandas.qcut does + np.putmask( + quantiles, + self.bins * quantiles != np.arange(self.bins + 1), + np.nextafter(quantiles, 1), + ) + limits = np.unique(np.quantile(values, quantiles, method="linear")) + else: + low, high = np.min(values), np.max(values) + # widen a constant range by 0.1%, as pandas.cut does + if low == high: + low = low - 0.001 * abs(low) if low != 0 else -0.001 + high = high + 0.001 * abs(high) if high != 0 else 0.001 + limits = np.unique(np.linspace(low, high, self.bins + 1)) + return limits[1:-1] + + def _category_counts( + self, X: IntoDataFrame, nw_X: nw.DataFrame, feature, is_basis: np.ndarray + ): + """Number of observations per category in the basis and test sets, with + the categories sorted. Missing values are not counted.""" + # pandas is faster than narwhals. + if nwd.is_pandas_dataframe(X) is True: + codes, categories = X[feature].factorize(sort=True) + is_value = codes >= 0 + basis_counts = np.bincount( + codes[is_basis & is_value], minlength=len(categories) + ) + test_counts = np.bincount( + codes[~is_basis & is_value], minlength=len(categories) + ) + else: + counts = ( + nw_X.select(nw.col(feature)) + .with_columns( + nw.new_series( + "__basis__", + is_basis, + backend=nw.get_native_namespace(nw_X), + ) + ) + .drop_nulls() + .group_by(feature) + .agg(nw.col("__basis__").sum(), nw.len().alias("__count__")) + .sort(feature) + ) + basis_counts = counts.get_column("__basis__").to_numpy() + test_counts = counts.get_column("__count__").to_numpy() - basis_counts + return basis_counts, test_counts + + def _psi(self, basis_counts: np.ndarray, test_counts: np.ndarray) -> float: + """PSI from the number of observations per interval or category.""" + # intervals or categories that are empty in both sets don't add to the PSI + observed = (basis_counts > 0) | (test_counts > 0) + basis = basis_counts[observed] / basis_counts.sum() + test = test_counts[observed] / test_counts.sum() + basis[basis == 0] = self.min_pct_empty_bins + test[test == 0] = self.min_pct_empty_bins + # min_pct_empty_bins=0 gives an inf PSI with empty intervals, without warnings + with np.errstate(divide="ignore"): + return np.sum((test - basis) * np.log(test / basis)) + + def _basis_mask(self, X: IntoDataFrame, nw_X: nw.DataFrame) -> np.ndarray: """ - Split dataframe according to a cut-off value and return two dataframes: the - basis dataframe contains all observations <= cut_off and the test dataframe the - observations > cut_off. - - If cut-off is a list, then the basis dataframe will contain all observations - which values are within the list, and the test dataframe all remaining - observations. + Find the observations of the basis dataset. - The cut-off value is associated to a specific column. + The basis dataset contains the observations whose value in `split_col` (or + in the index, or the row position when the dataframe has no index) is <= + cut_off, or is in cut_off when it is a list. The test dataset contains the + remaining observations. Parameters ---------- - X : pandas dataframe + X : dataframe + + nw_X : narwhals dataframe + X in narwhals format. Returns ------- - basis_df: pd.DataFrame - pandas dataframe with observations which value <= cut_off - - test_df: pd.DataFrame - pandas dataframe with observations which value > cut_off + is_basis: numpy array + Boolean array that is True for the observations of the basis dataset. """ - - # Identify the values according to which the split must be done. - if self.split_col is None: - reference = pd.Series(X.index) + # pandas is faster than narwhals, and has an index. + if nwd.is_pandas_dataframe(X) is True: + if self.split_col is None: + reference = X.index.to_series() + else: + reference = X[self.split_col] + n_missing = reference.isna().sum() else: - reference = X[self.split_col] + if self.split_col is None: + # without an index, the row order is the reference + reference = nw.new_series( + "__row__", + np.arange(nw_X.shape[0]), + backend=nw.get_native_namespace(nw_X), + ) + else: + reference = nw_X.get_column(self.split_col) + n_missing = reference.null_count() + if reference.dtype.is_float() is True: + n_missing += reference.is_nan().sum() - # Raise an error if there are missing values in the reference column. - if reference.isna().any(): + if n_missing > 0: raise ValueError( - f"There are {reference.isna().sum()} missing values in the reference" + f"There are {n_missing} missing values in the reference " "variable. Missing data are not allowed in the variable used to " "split the dataframe." ) - # If cut_off is not pre-defined, compute it. - if not self.cut_off: - self.cut_off_ = self._get_cut_off_value(reference) + if self.cut_off is None: + self.cut_off_ = self._get_cut_off_value(X, reference) else: self.cut_off_ = self.cut_off - # Split the original dataframe if isinstance(self.cut_off_, list): cut_off = self.cut_off_ - # isin with values castable to datetime (strings, dates) is - # deprecated in pandas; cast them to the reference dtype first. - if is_datetime64_any_dtype(reference): - cut_off = pd.to_datetime(cut_off) - is_within_cut_off = np.array(reference.isin(cut_off)) - + if nwd.is_pandas_dataframe(X) is True: + # isin with dates or strings is deprecated with datetime columns. + if reference.dtype.kind == "M": + cut_off = np.array(cut_off, dtype="datetime64[ns]") + is_basis = reference.isin(cut_off) + else: + # polars can't compare dates with datetimes. + if reference.dtype == nw.Datetime: + cut_off = ( + nw.new_series( + "__cut_off__", + cut_off, + backend=nw.get_native_namespace(nw_X), + ) + .cast(reference.dtype) + .to_list() + ) + is_basis = reference.is_in(cut_off) else: - is_within_cut_off = np.array(reference <= self.cut_off_) - - basis_df = X[is_within_cut_off] - test_df = X[~is_within_cut_off] + is_basis = reference <= self.cut_off_ - return basis_df, test_df + return is_basis.to_numpy() - def _get_cut_off_value(self, split_column): + def _get_cut_off_value(self, X: IntoDataFrame, split_column): """ Find the cut-off value to split the dataframe. It is implemented when the user does not enter a cut_off value as a parameter. It is calculated based on split_frac. - Finds the value in a pandas series at which we find the split_frac percentage + Finds the value in a series at which we find the split_frac percentage of observations. If the reference column is numerical, the cut-off value is determined using @@ -717,33 +767,44 @@ def _get_cut_off_value(self, split_column): Parameters ---------- - split_column: pd.Series. - Series for which the nth quantile will be computed. + X: dataframe + The dataframe to split. + + split_column: pandas or narwhals series. + Series for which the nth quantile will be computed: a pandas series if X + is a pandas dataframe, or a narwhals series otherwise. Returns ------- cut_off: (float, int, str, object). value for the cut-off. """ + if nwd.is_pandas_dataframe(X) is True: + if self.split_distinct is True: + split_column = split_column.drop_duplicates() - # In case split_distinct is used, extract series with unique values - if self.split_distinct: - split_column = pd.Series(split_column.unique()) - - # If the value is numerical, use numpy functionality - if is_numeric_dtype(split_column): - cut_off = np.quantile(split_column, self.split_frac) - - # Otherwise use value_counts combined with cumsum + if split_column.dtype.kind in "biufc": + cut_off = np.quantile(split_column, self.split_frac) + else: + cumulative = ( + split_column.value_counts(normalize=True).sort_index().cumsum() + ) + # the value whose cumulative fraction is the closest to split_frac + position = np.argmin(np.abs(cumulative.to_numpy() - self.split_frac)) + cut_off = cumulative.index.to_numpy()[position] else: - reference = pd.DataFrame( - split_column.value_counts(normalize=True).sort_index().cumsum() - ) + if self.split_distinct is True: + split_column = split_column.unique() - # Get the index (i.e. value) with the quantile that is the closest - # to the split_frac defined at initialization. - distance = abs(reference - self.split_frac) - cut_off = (distance.idxmin()).values[0] + if split_column.dtype.is_numeric() is True: + cut_off = np.quantile(split_column.to_numpy(), self.split_frac) + else: + proportions = split_column.value_counts( + name="__proportion__", normalize=True + ).sort(split_column.name) + cumulative = proportions.get_column("__proportion__").cum_sum() + position = np.argmin(np.abs(cumulative.to_numpy() - self.split_frac)) + cut_off = proportions.get_column(split_column.name).item(int(position)) return cut_off diff --git a/tests/test_selection/test_drop_high_psi_features.py b/tests/test_selection/test_drop_high_psi_features.py index 8cc296837..1cba57259 100644 --- a/tests/test_selection/test_drop_high_psi_features.py +++ b/tests/test_selection/test_drop_high_psi_features.py @@ -1,18 +1,41 @@ import math +import re from datetime import date, datetime +import narwhals as nw import numpy as np import pandas as pd import pytest from sklearn.datasets import make_classification +from sklearn.exceptions import NotFittedError from feature_engine.selection import DropHighPSIFeatures +from tests.backend_helpers import frame_to_dict + +# The expected PSI values were determined with the Probatus package: +# AutoDist(statistical_tests=["PSI"], binning_strategies="QuantileBucketer", +# bin_count=10).compute(data.iloc[0:500, :], data.iloc[500:, :]) +EXPECTED_PSI = { + "var_0": 0.043828484052281, + "var_1": 0.040929870747665395, + "var_2": 0.04330418495156895, + "var_3": 0.03773286532548153, + "var_4": 0.05047388515663041, + "var_5": 0.014717735595712466, + "drift_1": 8.283089355027482, + "drift_2": 8.283089355027482, + "cat_1": 0.0, + "drift_cat_1": 18.41883867587797, +} + +NUMERICAL = ["var_0", "var_1", "var_2", "var_3", "var_4", "var_5", "drift_1", "drift_2"] @pytest.fixture(scope="module") -def df(): - # create array with 4 correlated features and 2 independent ones - X, y = make_classification( +def data_drift(): + """6 features without drift, 2 numerical and 1 categorical feature that drift + between the first and the second half of the rows.""" + X, _ = make_classification( n_samples=1000, n_features=6, n_redundant=2, @@ -21,122 +44,205 @@ def df(): class_sep=2, random_state=1, ) + data = {f"var_{i}": X[:, i].tolist() for i in range(6)} + data["cat_1"] = ["A", "B"] * 500 + data["drift_1"] = list(range(1000)) + data["drift_2"] = [number / 2 for number in range(1000)] + data["drift_cat_1"] = ["A"] * 500 + ["B"] * 500 + return data + + +DATA_MIXED = { + "A": list(range(20)), + "B": [1, 2, 2, 1] * 5, + "C": ["A", "B", "D", "D"] * 5, + "time": [datetime(2019, 1, day + 1) for day in range(20)], +} - # transform array into pandas df - colnames = ["var_" + str(i) for i in range(6)] - X = pd.DataFrame(X, columns=colnames) - # Add a categorical column that does not drift - X["cat_1"] = ["A", "B"] * int(X.shape[0] / 2) +def basis_rows(transformer, X): + """Positions of the rows that go to the basis dataset.""" + is_basis = transformer._basis_mask(X, nw.from_native(X, eager_only=True)) + return np.flatnonzero(is_basis).tolist() - # Add drifted features that will be dropped during transformation. - X["drift_1"] = [number for number in range(X.shape[0])] - X["drift_2"] = [number / 2 for number in range(X.shape[0])] - X["drift_cat_1"] = ["A" for _ in range(int(X.shape[0] / 2))] + [ - "B" for _ in range(int(X.shape[0] / 2)) - ] - X["drift_cat_1"] = X["drift_cat_1"].astype("category") - return X +# init parameters +@pytest.mark.parametrize("split_col", [["hola"], 1.5, {"a": 1}]) +def test_error_if_split_col_not_allowed(split_col): + msg = f"split_col must be a string an integer or None. Got {split_col} instead." + with pytest.raises(ValueError, match=re.escape(msg)): + DropHighPSIFeatures(split_col=split_col) -@pytest.fixture(scope="module") -def df_mixed_types(): - df = pd.DataFrame( - { - "A": [it for it in range(0, 20)], - "B": [1, 2, 2, 1] * 5, - "C": ["A", "B", "D", "D"] * 5, - "time": [datetime(2019, 1, it + 1) for it in range(20)], - } +def test_error_if_split_frac_and_cut_off_are_none(): + msg = ( + "cut_off and split_frac cannot be both set to None. The current values are " + "(None, None). Please specify a value for at least one of these parameters." ) + with pytest.raises(ValueError, match=re.escape(msg)): + DropHighPSIFeatures(split_frac=None, cut_off=None) - return df +@pytest.mark.parametrize("split_frac", [0, 1, -0.5, 1.5, "0.5", [0.5]]) +def test_error_if_split_frac_not_allowed(split_frac): + msg = f"split_frac must be a float between 0 and 1. Got {split_frac} instead." + with pytest.raises(ValueError, match=re.escape(msg)): + DropHighPSIFeatures(split_frac=split_frac) -EXPECTED_PSI_NUM = { - "var_0": 0.043828484052281, - "var_1": 0.040929870747665395, - "var_2": 0.04330418495156895, - "var_3": 0.03773286532548153, - "var_4": 0.05047388515663041, - "var_5": 0.014717735595712466, - "drift_1": 8.283089355027482, - "drift_2": 8.283089355027482, -} -EXPECTED_PSI_ALL = { - "var_0": 0.043828484052281, - "var_1": 0.040929870747665395, - "var_2": 0.04330418495156895, - "var_3": 0.03773286532548153, - "var_4": 0.05047388515663041, - "var_5": 0.014717735595712466, - "drift_1": 8.283089355027482, - "drift_2": 8.283089355027482, - "cat_1": 0.0, - "drift_cat_1": 18.41883867587797, -} +@pytest.mark.parametrize("split_distinct", [1, "True", None, [True]]) +def test_error_if_split_distinct_not_bool(split_distinct): + msg = f"split_distinct must be a boolean. Got {split_distinct} instead." + with pytest.raises(ValueError, match=re.escape(msg)): + DropHighPSIFeatures(split_distinct=split_distinct) -EXPECTED_PSI_NUM_FEW = { - "var_2": 0.04330418495156895, - "var_3": 0.03773286532548153, - "drift_1": 8.283089355027482, - "drift_2": 8.283089355027482, -} -EXPECTED_PSI_MIXED = { - "var_0": 0.043828484052281, - "drift_1": 8.283089355027482, - "drift_cat_1": 18.41883867587797, -} +@pytest.mark.parametrize("switch", [1, "True", None, [True]]) +def test_error_if_switch_not_bool(switch): + msg = f"switch must be a boolean. Got {switch} instead." + with pytest.raises(ValueError, match=re.escape(msg)): + DropHighPSIFeatures(switch=switch) -EXPECTED_PSI_CAT_FEW = { - "cat_1": 0.0, - "drift_cat_1": 18.41883867587797, -} -EXPECTED_PSI_STRING = {"var_0": 0.043828484052281} - -_input_output = [ - (None, EXPECTED_PSI_NUM), - ("all", EXPECTED_PSI_ALL), - (["var_2", "var_3", "drift_1", "drift_2"], EXPECTED_PSI_NUM_FEW), - (["cat_1", "drift_cat_1"], EXPECTED_PSI_CAT_FEW), - (["var_0", "drift_1", "drift_cat_1"], EXPECTED_PSI_MIXED), - ("var_0", EXPECTED_PSI_STRING), -] - - -# ==== test main functionality of the class ==== -@pytest.mark.parametrize("variables, expected_psi", _input_output) -def test_fit_attributes(variables, expected_psi, df): - """Check the value of the fit attributes. - The expected PSI values used in the assertion were determined using - the Probatus package. - ``` - from probatus.stat_tests import AutoDist - psi_calculator = AutoDist(statistical_tests=["PSI"], - binning_strategies="QuantileBucketer", - bin_count=10) - train_df = data.iloc[0:500,:] - test_df = data.iloc[500:, :] - psi = psi_calculator.compute(train_df, test_df) - ``` - """ - transformer = DropHighPSIFeatures(variables=variables) - dft = transformer.fit_transform(df) - - assert transformer.variables_ == list(expected_psi.keys()) - assert transformer.psi_values_ == pytest.approx(expected_psi, 12) - assert transformer.features_to_drop_ == [ - var for var in expected_psi.keys() if "drift" in var - ] - pd.testing.assert_frame_equal(dft, df.drop(transformer.features_to_drop_, axis=1)) +@pytest.mark.parametrize("threshold", [-1, -0.1, "hola", None, [0.1]]) +def test_error_if_threshold_not_allowed(threshold): + msg = f"threshold must be greater than 0 or 'auto'. Got {threshold} instead." + with pytest.raises(ValueError, match=re.escape(msg)): + DropHighPSIFeatures(threshold=threshold) + + +@pytest.mark.parametrize("bins", [1, 0, -3, 2.5, "10", None]) +def test_error_if_bins_not_allowed(bins): + msg = f"bins must be an integer >= 2. Got {bins} instead." + with pytest.raises(ValueError, match=re.escape(msg)): + DropHighPSIFeatures(bins=bins) + + +@pytest.mark.parametrize("strategy", ["unknown", ["equal_width"], None, 1]) +def test_error_if_strategy_not_allowed(strategy): + msg = ( + "strategy takes only values equal_width or equal_frequency. " + f"Got {strategy} instead." + ) + with pytest.raises(ValueError, match=re.escape(msg)): + DropHighPSIFeatures(strategy=strategy) + + +@pytest.mark.parametrize("min_pct_empty_bins", [-1, -0.1, "unknown", None, [0.1]]) +def test_error_if_min_pct_empty_bins_not_allowed(min_pct_empty_bins): + msg = f"min_pct_empty_bins must be >= 0. Got {min_pct_empty_bins} instead." + with pytest.raises(ValueError, match=re.escape(msg)): + DropHighPSIFeatures(min_pct_empty_bins=min_pct_empty_bins) + + +@pytest.mark.parametrize("missing_values", ["hola", ["raise"], None, 1]) +def test_error_if_missing_values_not_allowed(missing_values): + msg = ( + "missing_values takes only values 'raise' or 'ignore'. " + f"Got {missing_values} instead." + ) + with pytest.raises(ValueError, match=re.escape(msg)): + DropHighPSIFeatures(missing_values=missing_values) + + +def test_error_if_split_col_in_variables(): + msg = ( + "hola cannot be used to split the data and be evaluated at the same time. " + "Either remove hola from the variables list or choose another splitting " + "criteria." + ) + with pytest.raises(ValueError, match=re.escape(msg)): + DropHighPSIFeatures(split_col="hola", variables=["hola", "chau"]) + + +@pytest.mark.parametrize("p_value", ["hola", -1.0, 10.0, 1, None]) +def test_error_if_p_value_not_allowed(p_value): + msg = f"p_value must be a float between 0 and 1. Got {p_value} instead." + with pytest.raises(ValueError, match=re.escape(msg)): + DropHighPSIFeatures(p_value=p_value) + + +@pytest.mark.parametrize( + "params", + [ + {}, + { + "split_col": "hola", + "split_frac": 0.6, + "split_distinct": True, + "cut_off": ["value_1", "value_2"], + "switch": True, + "threshold": 0.10, + "bins": 5, + "strategy": "equal_width", + "min_pct_empty_bins": 0.1, + "missing_values": "ignore", + "confirm_variables": True, + "p_value": 0.2, + }, + {"split_col": 0, "split_frac": None, "cut_off": 0, "threshold": "auto"}, + {"cut_off": date(2019, 1, 1), "threshold": 1, "min_pct_empty_bins": 0}, + ], +) +def test_init_param_assignment(params): + defaults = { + "split_col": None, + "split_frac": 0.5, + "split_distinct": False, + "cut_off": None, + "switch": False, + "threshold": 0.25, + "bins": 10, + "strategy": "equal_frequency", + "min_pct_empty_bins": 0.0001, + "missing_values": "raise", + "confirm_variables": False, + "p_value": 0.001, + } + transformer = DropHighPSIFeatures(**params) + for param, value in {**defaults, **params}.items(): + assert getattr(transformer, param) == value + + +# fit and transform +@pytest.mark.parametrize( + "variables, expected_variables", + [ + (None, NUMERICAL), + ("all", NUMERICAL + ["cat_1", "drift_cat_1"]), + ( + ["var_2", "var_3", "drift_1", "drift_2"], + ["var_2", "var_3", "drift_1", "drift_2"], + ), + (["cat_1", "drift_cat_1"], ["cat_1", "drift_cat_1"]), + (["var_0", "drift_cat_1", "drift_1"], ["var_0", "drift_1", "drift_cat_1"]), + ("var_0", ["var_0"]), + ], +) +@pytest.mark.parametrize("threshold", [0.25, "auto"]) +def test_fit_attributes_and_transform( + make_df, data_drift, variables, expected_variables, threshold +): + X = make_df(data_drift) + transformer = DropHighPSIFeatures(variables=variables, threshold=threshold) + Xt = transformer.fit_transform(X) + + features_to_drop = [var for var in expected_variables if "drift" in var] + assert transformer.variables_ == expected_variables + assert transformer.psi_values_ == pytest.approx( + {var: EXPECTED_PSI[var] for var in expected_variables} + ) + assert transformer.features_to_drop_ == features_to_drop + assert transformer.cut_off_ == 499.5 + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == { + var: values + for var, values in data_drift.items() + if var not in features_to_drop + } def test_auto_threshold_calculation(): - """Check the results of 'auto' threshold calculation""" transformer = DropHighPSIFeatures(threshold="auto", p_value=0.001, bins=10) assert math.isclose( transformer._calculate_auto_threshold(N=500, M=500, bins=10), @@ -154,634 +260,494 @@ def test_auto_threshold_calculation(): ) -@pytest.mark.parametrize("variables, expected_psi", _input_output) -def test_fit_attributes_with_autothreshold(variables, expected_psi, df): - """Check the value of the fit attributes. - The expected PSI values used in the assertion were determined using - the Probatus package. - ``` - from probatus.stat_tests import AutoDist - psi_calculator = AutoDist(statistical_tests=["PSI"], - binning_strategies="QuantileBucketer", - bin_count=10) - train_df = data.iloc[0:500,:] - test_df = data.iloc[500:, :] - psi = psi_calculator.compute(train_df, test_df) - ``` - """ - transformer = DropHighPSIFeatures(threshold="auto", variables=variables, bins=10) - transformer.fit(df) - - assert transformer.psi_values_ == pytest.approx(expected_psi, 12) +def test_auto_threshold_drops_features(make_df, data_drift): + # with p_value=0.2 the threshold for 10 intervals and 500 + 500 rows is 0.049. + transformer = DropHighPSIFeatures(threshold="auto", p_value=0.2, variables="all") + transformer.fit(make_df(data_drift)) assert transformer.features_to_drop_ == [ - var for var in expected_psi.keys() if "drift" in var + "var_4", + "drift_1", + "drift_2", + "drift_cat_1", ] -def test_calculation_when_strategy_equal_width(df): +def test_psi_with_equal_width_strategy(make_df, data_drift): transformer = DropHighPSIFeatures(strategy="equal_width") - transformer.fit(df) - - expected = { - "var_0": 0.014858665472468786, - "var_1": 0.04514737836588022, - "var_2": 0.03431479397506742, - "var_3": 0.04298209189840294, - "var_4": 0.02385796430263416, - "var_5": 0.046809664317794444, - "drift_1": 8.283089355027482, - "drift_2": 8.283089355027482, - } - - assert transformer.psi_values_ == pytest.approx(expected, 12) - assert transformer.features_to_drop_ == [ - var for var in expected.keys() if "drift" in var - ] - - -# ================ test init parameters ================= - -# Define two dictionaries with arguments: one with default values and -# one with arbitrary values. -default_dict = { - "split_col": None, - "split_frac": 0.5, - "split_distinct": False, - "cut_off": None, - "switch": False, - "threshold": 0.25, - "bins": 10, - "strategy": "equal_frequency", - "min_pct_empty_bins": 0.0001, - "missing_values": "raise", - "variables": None, - "p_value": 0.001, -} + transformer.fit(make_df(data_drift)) -args_dict = { - "split_col": "hola", - "split_frac": 0.6, - "split_distinct": True, - "cut_off": ["value_1", "value_2"], - "switch": True, - "threshold": 0.10, - "bins": 5, - "strategy": "equal_width", - "min_pct_empty_bins": 0.1, - "missing_values": "ignore", - "variables": ["chau", "adios"], - "p_value": 0.2, -} - -init_dict = [(None, default_dict), (args_dict, args_dict)] - - -@pytest.mark.parametrize("initialize, attribute_dict", init_dict) -def test_init_default_parameters(initialize, attribute_dict): - """Test the default param values are correctly assigned.""" - if initialize: - transformer = DropHighPSIFeatures(**attribute_dict) - else: - transformer = DropHighPSIFeatures() - - for attribute, value in attribute_dict.items(): - assert getattr(transformer, attribute) == value - - -def test_init_value_error_is_raised(): - with pytest.raises(ValueError): - DropHighPSIFeatures(split_col=["hola"]) - - with pytest.raises(ValueError): - DropHighPSIFeatures(split_col="hola", variables=["hola", "chau"]) - - with pytest.raises(ValueError): - DropHighPSIFeatures(split_frac=0) - - with pytest.raises(ValueError): - DropHighPSIFeatures(split_frac=1) - - with pytest.raises(ValueError): - DropHighPSIFeatures(split_frac=None, cut_off=None) - - with pytest.raises(ValueError): - DropHighPSIFeatures(split_distinct=1) - - with pytest.raises(ValueError): - DropHighPSIFeatures(bins=1) - - with pytest.raises(ValueError): - DropHighPSIFeatures(threshold=-1) - - with pytest.raises(ValueError): - DropHighPSIFeatures(switch=1) - - with pytest.raises(ValueError): - DropHighPSIFeatures(strategy="unknown") + assert transformer.psi_values_ == pytest.approx( + { + "var_0": 0.014858665472468786, + "var_1": 0.04514737836588022, + "var_2": 0.03431479397506742, + "var_3": 0.04298209189840294, + "var_4": 0.02385796430263416, + "var_5": 0.046809664317794444, + "drift_1": 8.283089355027482, + "drift_2": 8.283089355027482, + }, + ) + assert transformer.features_to_drop_ == ["drift_1", "drift_2"] - with pytest.raises(ValueError): - DropHighPSIFeatures(min_pct_empty_bins="unknown") - with pytest.raises(ValueError): - DropHighPSIFeatures(min_pct_empty_bins=-1) +def test_empty_bins_take_min_pct_empty_bins(make_df): + X = make_df({"x": [1, 2, 3, 4], "cat": ["A", "B", "A", "C"]}) + transformer = DropHighPSIFeatures( + variables="all", bins=2, min_pct_empty_bins=0.01 + ).fit(X) + # "B" is only in the basis set and "C" only in the test set. + psi_cat = (0.01 - 0.5) * math.log(0.01 / 0.5) + (0.5 - 0.01) * math.log(0.5 / 0.01) + # the basis set, [1, 2], has one value per interval; the test set, [3, 4], has + # both values in the second interval. + psi_x = (0.01 - 0.5) * math.log(0.01 / 0.5) + (1 - 0.5) * math.log(1 / 0.5) + assert transformer.psi_values_ == pytest.approx({"x": psi_x, "cat": psi_cat}) -@pytest.mark.parametrize("p_value", ["hola", -1, 10]) -def test_p_value_not_allowed(p_value): - with pytest.raises(ValueError): - DropHighPSIFeatures(p_value=p_value) +def test_constant_feature_has_psi_zero(make_df): + X = make_df({"x": [1.5] * 20, "y": list(range(20))}) + transformer = DropHighPSIFeatures(bins=3).fit(X) + assert transformer.psi_values_["x"] == 0 -# ================= test fit() functionality ================== +def test_split_col_not_included_in_variables(make_df, data_drift): + X = make_df(data_drift) -def test_split_col_not_included_in_variables(df): - """Check that the split column is not included among the features - to evaluate when these are selected automatically.""" - transformer = DropHighPSIFeatures(split_col="var_3", variables=None) - transformer.fit(df) - assert "var_3" not in transformer.variables_ - assert "var_3" not in transformer.psi_values_.keys() + transformer = DropHighPSIFeatures(split_col="var_3", variables=None).fit(X) + assert transformer.variables_ == [var for var in NUMERICAL if var != "var_3"] - transformer = DropHighPSIFeatures(split_col="var_3", variables="all") - transformer.fit(df) + transformer = DropHighPSIFeatures(split_col="var_3", variables="all").fit(X) assert "var_3" not in transformer.variables_ - assert "var_3" not in transformer.psi_values_.keys() - - transformer = DropHighPSIFeatures(split_col="cat_1", variables="all") - transformer.fit(df) - assert "cat_1" not in transformer.variables_ - assert "cat_1" not in transformer.psi_values_.keys() - + assert "var_3" not in transformer.psi_values_ -def test_error_split_col_not_in_df(df): - transformer = DropHighPSIFeatures(variables=None, split_col="var_0") - data = df.copy() - data = data.drop(["var_0"], axis=1) - msg = "var_0 is not in the dataframe." - with pytest.raises(ValueError) as record: - transformer.fit(data) - assert str(record.value) == msg + transformer = DropHighPSIFeatures(split_col="cat_1", variables="all").fit(X) + assert transformer.variables_ == NUMERICAL + ["drift_cat_1"] + assert "cat_1" not in transformer.psi_values_ -_input_output = [ - (["var_2", "var_3", "drift_1", "drift_2"], ["drift_1"], EXPECTED_PSI_NUM_FEW), - (["var_0", "drift_1", "drift_cat_1"], ["drift_cat_1"], EXPECTED_PSI_MIXED), -] +def test_error_if_split_col_not_in_df(make_df, data_drift): + X = make_df({var: data_drift[var] for var in ["var_1", "var_2"]}) + transformer = DropHighPSIFeatures(split_col="var_0") + with pytest.raises(ValueError, match=re.escape("var_0 is not in the dataframe.")): + transformer.fit(X) -@pytest.mark.parametrize("variables, variable, expected_psi", _input_output) -def test_confirm_variables(variables, variable, expected_psi, df): - data = df.copy() - data = data.drop(variable, axis=1) - del expected_psi[variable[0]] - +@pytest.mark.parametrize( + "variables, missing, expected_variables", + [ + (["var_2", "var_3", "drift_1", "drift_2"], "drift_1", ["var_2", "var_3"]), + (["var_0", "drift_1", "drift_cat_1"], "drift_cat_1", ["var_0", "drift_1"]), + ], +) +def test_confirm_variables(make_df, data_drift, variables, missing, expected_variables): + X = make_df({var: values for var, values in data_drift.items() if var != missing}) transformer = DropHighPSIFeatures(variables=variables, confirm_variables=True) - transformer.fit(data) + transformer.fit(X) - assert transformer.variables_ == list(expected_psi.keys()) - assert transformer.psi_values_ == pytest.approx(expected_psi, 12) - assert transformer.features_to_drop_ == [ - var for var in expected_psi.keys() if "drift" in var - ] + assert transformer.variables_ == expected_variables + ( + ["drift_2"] if "drift_2" in variables else [] + ) + assert transformer.psi_values_ == pytest.approx( + {var: EXPECTED_PSI[var] for var in transformer.variables_} + ) -def test_error_if_variables_is_none_and_no_numerical_in_df(df): - transformer = DropHighPSIFeatures(variables=None) - msg = "No numerical variables found in this dataframe. " - with pytest.raises(TypeError, match=msg): - transformer.fit(df[["cat_1", "drift_cat_1"]]) +def test_error_if_no_numerical_variables(make_df, data_drift): + X = make_df({var: data_drift[var] for var in ["cat_1", "drift_cat_1"]}) + msg = ( + "No numerical variables found in this dataframe. Check variable dtypes or " + "set return_empty to True to return an empty list instead." + ) + with pytest.raises(TypeError, match=re.escape(msg)): + DropHighPSIFeatures().fit(X) -def test_error_if_confirm_variables_returns_empty_list(df): +def test_error_if_confirm_variables_returns_empty_list(make_df, data_drift): + X = make_df({var: data_drift[var] for var in NUMERICAL}) transformer = DropHighPSIFeatures( variables=["cat_1", "drift_cat_1"], confirm_variables=True ) - data = df.copy() - data = data.drop(["cat_1", "drift_cat_1"], axis=1) msg = "None of the variables in the list are present in the dataframe." - with pytest.raises(ValueError) as record: - transformer.fit(data) - assert str(record.value) == msg - - -def test_error_if_na_in_split_col(df): - """Test an error is raised if the split column contains missing values.""" - data = df.copy() - data.loc[15, "var_3"] = np.nan - - transformer = DropHighPSIFeatures(split_col="var_3") - - with pytest.raises(ValueError): - transformer.fit(data) - - -def test_raise_error_if_na_in_df(df): - """Test an error is raised when missing values is set to raise.""" - data = df.copy() - data.loc[15, "var_3"] = np.nan - - transformer = DropHighPSIFeatures(missing_values="raise") - - with pytest.raises(ValueError): - transformer.fit(data) - - -def test_missing_value_ignored(df): - """Test if PSI are computed when missing values are present in the dataframe.""" - data = df.copy() - data.loc[15, "var_3"] = np.nan - - transformer = DropHighPSIFeatures(missing_values="ignore") - transformed = transformer.fit_transform(data) - - assert transformer.psi_values_ == pytest.approx(EXPECTED_PSI_NUM, 12) - assert transformer.features_to_drop_ == [ - var for var in EXPECTED_PSI_NUM if "drift" in var - ] - pd.testing.assert_frame_equal( - transformed, data.drop(transformer.features_to_drop_, axis=1) + with pytest.raises(ValueError, match=re.escape(msg)): + transformer.fit(X) + + +@pytest.mark.parametrize("split_col", ["var_3", "cat_1"]) +def test_error_if_na_in_split_col(make_df, data_drift, split_col): + data = {var: list(values) for var, values in data_drift.items()} + data[split_col][15] = None + data[split_col][17] = None + transformer = DropHighPSIFeatures(split_col=split_col, missing_values="ignore") + msg = ( + "There are 2 missing values in the reference variable. Missing data are not " + "allowed in the variable used to split the dataframe." ) - - -def test_raise_error_if_inf_in_df(df): - """Test an error is raised for inf when missing values is set to raise.""" - data = df.copy() - data.loc[15, "var_3"] = np.inf - - transformer = DropHighPSIFeatures(missing_values="raise") - - with pytest.raises(ValueError): - transformer.fit(data) - - -# ========= tests for _split_dataframe() fit ==== - -# tests for splits based on split_frac and numerical variables: - -quantile_test = [(0.5, 50), (0.33, 33), (0.17, 17), (0.81, 81)] - - -@pytest.mark.parametrize("split_frac, expected", quantile_test) -def test_calculation_quantile(split_frac, expected): - """Test the calculation of the quantiles using numerical values.""" - df = pd.DataFrame( - {"A": [it for it in range(0, 101)], "B": [it for it in range(0, 101)]} + with pytest.raises(ValueError, match=re.escape(msg)): + transformer.fit(make_df(data)) + + +@pytest.mark.parametrize("variable", ["var_3", "cat_1"]) +def test_error_if_na_in_variables(make_df, data_drift, variable): + data = {var: list(values) for var, values in data_drift.items()} + data[variable][15] = None + transformer = DropHighPSIFeatures(variables="all", missing_values="raise") + msg = ( + "Some of the variables in the dataset contain NaN. Check and remove those " + "before using this transformer." ) + with pytest.raises(ValueError, match=re.escape(msg)): + transformer.fit(make_df(data)) - test = DropHighPSIFeatures( - split_col="A", split_frac=split_frac, split_distinct=False - ) - test.fit_transform(df) - assert test.cut_off_ == expected +def test_missing_values_ignored(make_df, data_drift): + data = {var: list(values) for var, values in data_drift.items()} + data["var_3"][15] = None + data["cat_1"][15] = None + X = make_df(data) + transformer = DropHighPSIFeatures(variables="all", missing_values="ignore") + Xt = transformer.fit_transform(X) -quantile_test_skewed = [(50, 50), (1, 30), (10, 40), (7, 80)] + # a missing value in a feature does not remove the observation from the others + expected_psi = dict( + EXPECTED_PSI, var_3=0.03404846717863001, cat_1=4.016053504547722e-06 + ) + assert transformer.psi_values_ == pytest.approx(expected_psi) + assert transformer.features_to_drop_ == ["drift_1", "drift_2", "drift_cat_1"] + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == { + var: values for var, values in data.items() if "drift" not in var + } -@pytest.mark.parametrize("index, fraction", quantile_test_skewed) -def test_quatile_split_skewed_variables(index, fraction): - """Test the calculation of the quantiles using numerical and skewed variables.""" - df = pd.DataFrame( - { - "A": [index for it in range(0, fraction + 1)] - + [it for it in range(fraction + 1, 101)], - "B": [it for it in range(0, 101)], - } +@pytest.mark.parametrize("missing_values", ["raise", "ignore"]) +def test_error_if_inf_in_variables(make_df, data_drift, missing_values): + data = {var: list(values) for var, values in data_drift.items()} + data["var_3"][15] = np.inf + transformer = DropHighPSIFeatures(missing_values=missing_values) + msg = ( + "Some of the variables to transform contain inf values. Check and remove " + "those before using this transformer." ) - - test = DropHighPSIFeatures( - split_col="A", split_frac=fraction / 100, split_distinct=False + with pytest.raises(ValueError, match=re.escape(msg)): + transformer.fit(make_df(data)) + + +@pytest.mark.parametrize( + "rows, n_basis, n_test", [(range(500), 0, 500), (range(500, 1000), 500, 0)] +) +def test_error_if_only_missing_values_in_basis_or_test( + make_df, data_drift, rows, n_basis, n_test +): + data = {var: data_drift[var] for var in ["var_0", "var_1"]} + data["var_1"] = [None if row in rows else v for row, v in enumerate(data["var_1"])] + transformer = DropHighPSIFeatures(missing_values="ignore") + msg = ( + "The variable var_1 has only missing values in the basis or in the test set, " + f"so its PSI can't be computed. Got {n_basis} values in the basis set and " + f"{n_test} values in the test set." ) - test.fit_transform(df) - - assert test.cut_off_ == index - - -# tests for splits based on split_frac and categorical variables: - - -def test_calculation_distinct_value_categorical(): - """Test the calculation of the quantiles using distinct values when reference - variable is categorical.""" - df = pd.DataFrame( - {"A": [it for it in range(0, 30)], "C": ["A", "B", "C", "D", "D", "D"] * 5} + with pytest.raises(ValueError, match=re.escape(msg)): + transformer.fit(make_df(data)) + + +@pytest.mark.parametrize( + "cut_off, n_basis, n_test", [(5, 6, 994), (-1, 0, 1000), (5000, 1000, 0)] +) +def test_error_if_too_few_rows_in_basis_or_test( + make_df, data_drift, cut_off, n_basis, n_test +): + X = make_df({var: data_drift[var] for var in NUMERICAL}) + msg = ( + "The number of rows in the basis and test datasets that will be used in the " + "PSI calculations must be at least larger than 10. After splitting the " + "original dataset based on the given cut_off or split_frac we have " + f"{n_basis} samples in the basis set, and {n_test} samples in the test set. " + "Please adjust the value of the cut_off or split_frac." ) + with pytest.raises(ValueError, match=re.escape(msg)): + DropHighPSIFeatures(cut_off=cut_off).fit(X) - test = DropHighPSIFeatures(split_col="C", split_frac=0.5, split_distinct=False) - test.fit_transform(df) - assert test.cut_off_ == "C" +@pytest.mark.parametrize( + "split_frac, expected", [(0.5, 50), (0.33, 33), (0.17, 17), (0.81, 81)] +) +def test_cut_off_from_split_frac(make_df, split_frac, expected): + X = make_df({"A": list(range(101)), "B": list(range(101))}) + transformer = DropHighPSIFeatures(split_col="A", split_frac=split_frac) + transformer.fit(X) + assert transformer.cut_off_ == expected - test = DropHighPSIFeatures(split_col="C", split_frac=0.5, split_distinct=True) - test.fit_transform(df) - assert test.cut_off_ == "B" - df = pd.DataFrame( - {"A": [it for it in range(0, 30)], "C": ["A", "A", "A", "B", "C", "D"] * 5} +@pytest.mark.parametrize("value, fraction", [(50, 50), (1, 30), (10, 40), (7, 80)]) +def test_cut_off_from_split_frac_with_skewed_variable(make_df, value, fraction): + X = make_df( + { + "A": [value] * (fraction + 1) + list(range(fraction + 1, 101)), + "B": list(range(101)), + } ) - - test = DropHighPSIFeatures(split_col="C", split_frac=0.5, split_distinct=False) - - test.fit_transform(df) - assert test.cut_off_ == "A" - - test = DropHighPSIFeatures(split_col="C", split_frac=0.5, split_distinct=True) - test.fit_transform(df) - assert test.cut_off_ == "B" - - -numerical_split_distinct = [(True, [1, 2, 3], [4, 5, 6]), (False, [1], [2, 3, 4, 5, 6])] - - -@pytest.mark.parametrize("split_distinct, a_values, b_values", numerical_split_distinct) -def test_split_distinct_with_numerical_values(split_distinct, a_values, b_values): - """Test the split_distinct functionality with numerical variables.""" - # Define the testing dataframe - df = pd.DataFrame( + transformer = DropHighPSIFeatures(split_col="A", split_frac=fraction / 100) + transformer.fit(X) + assert transformer.cut_off_ == value + + +@pytest.mark.parametrize( + "values, split_distinct, expected", + [ + (["A", "B", "C", "D", "D", "D"], False, "C"), + (["A", "B", "C", "D", "D", "D"], True, "B"), + (["A", "A", "A", "B", "C", "D"], False, "A"), + (["A", "A", "A", "B", "C", "D"], True, "B"), + ], +) +def test_cut_off_from_split_frac_with_categorical_variable( + make_df, values, split_distinct, expected +): + X = make_df({"A": list(range(30)), "C": values * 5}) + transformer = DropHighPSIFeatures(split_col="C", split_distinct=split_distinct) + transformer.fit(X) + assert transformer.cut_off_ == expected + + +def test_psi_and_cut_off_with_different_split_col_types(make_df): + X = make_df(DATA_MIXED) + expected = { + "A": ({"B": 0.0, "C": 0.1621860432432657}, 9.5), + "B": ({"A": 3.0375978817052403, "C": 8.515489752777954}, 1.5), + "C": ({"A": 2.27819841127893, "B": 0.0}, "B"), + "time": ( + {"A": 8.283089355027482, "B": 0.0, "C": 0.1621860432432657}, + datetime(2019, 1, 10), + ), + None: ({"A": 8.283089355027482, "B": 0.0, "C": 0.1621860432432657}, 9.5), + } + for split_col, (psi_values, cut_off) in expected.items(): + transformer = DropHighPSIFeatures(split_col=split_col, variables="all") + transformer.fit(X) + assert transformer.psi_values_ == pytest.approx(psi_values) + assert transformer.cut_off_ == cut_off + + +@pytest.mark.parametrize( + "split_distinct, expected", + [(True, [0, 1, 2, 3, 4, 7, 8]), (False, [0, 1, 4, 7, 8])], +) +def test_split_distinct_with_numerical_values(make_df, split_distinct, expected): + X = make_df( { "ID": [1, 1, 2, 3, 1, 4, 5, 1, 1, 6], "numerical": [1, 1, 1, 4, 1, 4, 3, 7, 1, 3], } ) - a_expected = df[df.ID.isin(a_values)] - b_expected = df[df.ID.isin(b_values)] - # Run the split_dataframe method to extract the input of the PSI calculation. transformer = DropHighPSIFeatures(split_col="ID", split_distinct=split_distinct) - a, b = transformer._split_dataframe(df) - # Test if the functionality provides the expected results. - pd.testing.assert_frame_equal(a, a_expected) - pd.testing.assert_frame_equal(b, b_expected) - - -def test_calculation_df_split_with_different_variable_types(df_mixed_types): - """Test the split of the dataframe using different type of variables.""" - results = {} - cut_offs = {} - for split_col in df_mixed_types.columns: - test = DropHighPSIFeatures(split_frac=0.5, split_col=split_col, variables="all") - test.fit_transform(df_mixed_types) - results[split_col] = test.psi_values_ - cut_offs[split_col] = test.cut_off_ - - assert results["A"] == pytest.approx({"B": 0.0, "C": 0.1621860432432657}, 12) - assert results["B"] == pytest.approx( - {"A": 3.0375978817052403, "C": 8.515489752777954}, 12 - ) - assert results["C"] == pytest.approx({"A": 2.27819841127893, "B": 0.0}, 12) - assert results["time"] == pytest.approx( - {"A": 8.283089355027482, "B": 0.0, "C": 0.1621860432432657}, 12 - ) - - expected_cut_offs = { - "A": 9.5, - "B": 1.5, - "C": "B", - "time": np.datetime64(datetime(2019, 1, 10)), - } - - assert cut_offs == expected_cut_offs - - # Test when no dataframe with mixed data types when no split_col is provided. - test = DropHighPSIFeatures(split_frac=0.5, variables="all") - test.fit_transform(df_mixed_types) - assert test.psi_values_ == pytest.approx( - {"A": 8.283089355027482, "B": 0.0, "C": 0.1621860432432657}, 12 - ) - - -# =========== tests for user entered cut_off values =========== - - -type_test = [ - ("A", 14, [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14]), - ("B", 1, [0, 3, 4, 7, 8, 11, 12, 15, 16, 19]), - ("C", ["B"], [1, 5, 9, 13, 17]), - ("C", "B", [0, 1, 4, 5, 8, 9, 12, 13, 16, 17]), - ("time", datetime(2019, 1, 4), [0, 1, 2, 3]), -] - - -@pytest.mark.parametrize("col, cut_off, expected", type_test) -def test_split_using_cut_off(col, cut_off, expected, df_mixed_types): - """Test the cut off for different data types.""" - test = DropHighPSIFeatures(split_col=col, cut_off=cut_off) - a, b = test._split_dataframe(df_mixed_types) - - pd.testing.assert_frame_equal(a, df_mixed_types.loc[expected]) - pd.testing.assert_frame_equal( - b, df_mixed_types.loc[~df_mixed_types.index.isin(expected)] - ) - - -split_distinct_test = [ - ("A", [number for number in range(0, 100)]), - ( - "B", - ( - [0, 1, 2, 10, 11, 12, 20, 21, 22, 30, 31, 32, 40] - + [41, 42, 50, 51, 52, 60, 61, 62, 70, 71, 72, 80, 81] - + [82, 90, 91, 92, 100, 101, 102, 110, 111, 112, 120, 121, 122] - + [130, 131, 132, 140, 141, 142, 150, 151, 152, 160, 161, 162, 170] - + [171, 172, 180, 181, 182, 190, 191, 192] - ), - ), - ( - "C", - ( - [0, 1, 2, 10, 11, 12, 20, 21, 22, 30, 31, 32, 40] - + [41, 42, 50, 51, 52, 60, 61, 62, 70, 71, 72, 80, 81] - + [82, 90, 91, 92, 100, 101, 102, 110, 111, 112, 120, 121, 122] - + [130, 131, 132, 140, 141, 142, 150, 151, 152, 160, 161, 162, 170] - + [171, 172, 180, 181, 182, 190, 191, 192] - ), - ), - ( - "time", - ( - [0, 1, 2, 5, 6, 7, 10, 11, 12, 15, 16, 17, 20, 21, 22, 25, 26] - + [27, 30, 31, 32, 35, 36, 37, 40, 41, 42, 45, 46, 47, 50, 51, 52, 55] - + [56, 57, 60, 61, 62, 65, 66, 67, 70, 71, 72, 75, 76, 77, 80, 81, 82] - + [85, 86, 87, 90, 91, 92, 95, 96, 97] - ), - ), -] - - -@pytest.mark.parametrize("col, expected_index", split_distinct_test) -def test_split_distinct(col, expected_index): - """Test the cut off for different data types. - For columns B, C and time we have 6 distinct values, 5 appearing 20 times and - 1 appearing 100 times. A 50% split based on the number of values will result - in 2 groups of 3. One has 60 appearances (in total) and the other has 140. - """ - data = pd.DataFrame( + assert basis_rows(transformer, X) == expected + + +@pytest.mark.parametrize( + "split_col, cut_off, expected", + [ + ("A", 14, list(range(15))), + ("B", 1, [0, 3, 4, 7, 8, 11, 12, 15, 16, 19]), + ("C", ["B"], [1, 5, 9, 13, 17]), + ("C", "B", [0, 1, 4, 5, 8, 9, 12, 13, 16, 17]), + ("time", datetime(2019, 1, 4), [0, 1, 2, 3]), + ("A", [1, 2, 10, 11, 16], [1, 2, 10, 11, 16]), + ("B", [2], [1, 2, 5, 6, 9, 10, 13, 14, 17, 18]), + ("C", ["B", "D"], [1, 2, 3, 5, 6, 7, 9, 10, 11, 13, 14, 15, 17, 18, 19]), + ("time", [date(2019, 1, day) for day in [1, 2, 5, 18]], [0, 1, 4, 17]), + ("time", [datetime(2019, 1, day) for day in [1, 2, 5, 18]], [0, 1, 4, 17]), + (None, 5, [0, 1, 2, 3, 4, 5]), + (None, [3, 7, 11], [3, 7, 11]), + ], +) +def test_split_with_cut_off(make_df, split_col, cut_off, expected): + X = make_df(DATA_MIXED) + transformer = DropHighPSIFeatures(split_col=split_col, cut_off=cut_off) + assert basis_rows(transformer, X) == expected + + +def test_cut_off_zero_is_used(make_df): + X = make_df({"A": [0, 1, 2, 3] * 10, "B": list(range(40))}) + transformer = DropHighPSIFeatures(split_col="A", cut_off=0, bins=5).fit(X) + assert transformer.cut_off_ == 0 + assert basis_rows(transformer, X) == list(range(0, 40, 4)) + + +@pytest.mark.parametrize("split_col", ["A", "B", "C", "time"]) +def test_split_distinct(make_df, split_col): + # 6 distinct values, 5 appear 20 times and 1 appears 100 times: the basis gets + # the rows of the first 3 distinct values. + X = make_df( { - "A": [it for it in range(0, 200)], + "A": list(range(200)), "B": [1, 2, 3, 4, 5, 6, 6, 6, 6, 6] * 20, "C": ["A", "B", "C", "D", "E", "F", "F", "F", "F", "F"] * 20, - "time": [date(2019, 1, it + 1) for it in range(5)] * 20 + "time": [date(2019, 1, day + 1) for day in range(5)] * 20 + [date(2019, 1, 31)] * 100, } ) - test = DropHighPSIFeatures(split_col=col, split_distinct=True) - a, b = test._split_dataframe(data) - - pd.testing.assert_frame_equal(a, data.loc[expected_index]) - pd.testing.assert_frame_equal(b, data.loc[~data.index.isin(expected_index)]) - - -cut_off_list_test = [ - ("A", [1, 2, 10, 11, 16]), - ("B", [2]), - ("C", ["B", "D"]), - ("time", [date(2019, 1, day) for day in [1, 2, 5, 7, 12, 15, 18]]), -] - - -@pytest.mark.parametrize("col, cut_off_list", cut_off_list_test) -def test_split_by_list(df_mixed_types, col, cut_off_list): - """Test elements a correctly selected when cut_off is a list.""" - test = DropHighPSIFeatures(split_col=col, cut_off=cut_off_list, bins=3) - a, b = test._split_dataframe(df_mixed_types) - - if pd.api.types.is_datetime64_any_dtype(df_mixed_types[col]): - cut_off_list = pd.to_datetime(cut_off_list) - - pd.testing.assert_frame_equal( - a, df_mixed_types[df_mixed_types[col].isin(cut_off_list)] - ) - pd.testing.assert_frame_equal( - b, df_mixed_types[~df_mixed_types[col].isin(cut_off_list)] + if split_col == "A": + expected = list(range(100)) + elif split_col == "time": + expected = [row for row in range(100) if row % 5 < 3] + else: + expected = [row for row in range(200) if row % 10 < 3] + transformer = DropHighPSIFeatures(split_col=split_col, split_distinct=True) + assert basis_rows(transformer, X) == expected + + +@pytest.mark.parametrize( + "split_frac, cut_off, expected_cut_off, n_basis", + [(0.5, None, 499.5, 500), (0.6, None, 599.4, 600), (0.5, 250, 250, 251)], +) +def test_split_by_row_order_when_split_col_is_none( + make_df, data_drift, split_frac, cut_off, expected_cut_off, n_basis +): + # pandas uses the index, here the default index is the position of the rows. + X = make_df(data_drift) + transformer = DropHighPSIFeatures(split_frac=split_frac, cut_off=cut_off).fit(X) + assert transformer.cut_off_ == pytest.approx(expected_cut_off) + assert basis_rows(transformer, X) == list(range(n_basis)) + + +@pytest.mark.parametrize( + "split_frac, cut_off, expected_cut_off, expected", + [ + (0.5, None, 499.5, set(range(500))), + (0.6, None, 599.4, set(range(600))), + (0.5, 250, 250, set(range(251))), + ], +) +def test_split_by_index_of_shuffled_pandas_df( + data_drift, split_frac, cut_off, expected_cut_off, expected +): + X = pd.DataFrame(data_drift).sample(frac=1, random_state=0) + transformer = DropHighPSIFeatures(split_frac=split_frac, cut_off=cut_off) + is_basis = transformer._basis_mask(X, nw.from_native(X)) + assert transformer.cut_off_ == pytest.approx(expected_cut_off) + assert set(X.index[is_basis]) == expected + + +def test_split_by_datetime_index_of_pandas_df(data_drift): + X = pd.DataFrame(data_drift) + X.index = pd.date_range("2020-01-01", periods=1000, freq="h") + transformer = DropHighPSIFeatures(cut_off=pd.Timestamp("2020-01-20")).fit(X) + assert transformer.psi_values_ == pytest.approx( + { + "var_0": 0.026166596537012143, + "var_1": 0.05683862097207795, + "var_2": 0.04828999318701926, + "var_3": 0.03704211767222722, + "var_4": 0.07053922062598025, + "var_5": 0.011565170354399454, + "drift_1": 8.270551209760598, + "drift_2": 8.270551209760598, + }, ) -# Tests for split on index on shuffled dataframe -def test_split_shuffled_df_default(df): - """Test the default parameters when the index is shuffled.""" - # Shuffle the dataframe - df_shuffled = df.sample(frac=1) - test = DropHighPSIFeatures() - base, test = test._split_dataframe(df_shuffled) - - # The base dataframe should contain indexes from 0 to 499 - set(base.index) == {_ for _ in range(0, 500)} - - # The test dataframe should contain indexes from 500 to 999 - set(test.index) == {_ for _ in range(500, 999)} - - -def test_split_shuffled_df_split_frac(df): - """Test split_frac when the index is shuffled.""" - # Shuffle the dataframe - df_shuffled = df.sample(frac=1) - test = DropHighPSIFeatures(split_frac=0.6) - base, test = test._split_dataframe(df_shuffled) - - # The base dataframe should contain indexes from 0 to 599 - set(base.index) == {_ for _ in range(0, 600)} - - # The test dataframe should contain indexes from 600 to 999 - set(test.index) == {_ for _ in range(600, 999)} - - -def test_split_shuffled_df_cut_off(df): - """Test the cut_off when the index is shuffled.""" - # Shuffle the dataframe - df_shuffled = df.sample(frac=1) - test = DropHighPSIFeatures(cut_off=250) - base, test = test._split_dataframe(df_shuffled) - - # The base dataframe should contain indexes from 0 to 250 - set(base.index) == {_ for _ in range(0, 251)} - - # The test dataframe should contain indexes from 251 to 999 - set(test.index) == {_ for _ in range(251, 999)} - - -# ===== end of tests for _split_dataframe() ======= - -# ==== more tests for fit functionality ============ - +def test_switch(make_df): + data_a = { + "a": [1.0, 2, 3, 1], + "b": [1.0, 2, 3, 4], + "c": [1, 2, 3, 4], + "d": [1.7, 4.7, 6.6, 7.8], + } + data_b = { + "a": [4.0, 3, 5, 1], + "b": [11.0, 1, 2, 4], + "c": [4, 2, 2, 4], + "d": [4.7, 4.7, 7.6, 7.8], + } + X_order = make_df({var: data_a[var] + data_b[var] for var in data_a}) + X_reverse = make_df({var: data_b[var] + data_a[var] for var in data_a}) -def test_switch(): - """Test the functionality to switch the basis.""" + case = DropHighPSIFeatures(bins=3, switch=False, min_pct_empty_bins=0.001) + case.fit(X_order) + switch_case = DropHighPSIFeatures(bins=3, switch=True, min_pct_empty_bins=0.001) + switch_case.fit(X_reverse) - df_a = pd.DataFrame( - { - "a": [1.0, 2, 3, 1], - "b": [1.0, 2, 3, 4], - "c": [1, 2, 3, 4], - "d": [1.7, 4.7, 6.6, 7.8], - } - ) - - df_b = pd.DataFrame( + assert case.psi_values_ == switch_case.psi_values_ + assert case.psi_values_ == pytest.approx( { - "a": [4.0, 3, 5, 1], - "b": [11.0, 1, 2, 4], - "c": [4, 2, 2, 4], - "d": [4.7, 4.7, 7.6, 7.8], + "a": 1.0986122886681098, + "b": 1.5481305636876856, + "c": 1.5481305636876856, + "d": 1.5481305636876856, } ) - df_order = pd.concat([df_a, df_b]).reset_index(drop=True) - df_reverse = pd.concat([df_b, df_a]).reset_index(drop=True) - case = DropHighPSIFeatures( - split_frac=0.5, bins=3, switch=False, min_pct_empty_bins=0.001 - ) - case.fit(df_order) +def test_transform_restores_train_column_order(make_df, data_drift): + X = make_df(data_drift) + transformer = DropHighPSIFeatures().fit(X) + X_reordered = make_df({var: data_drift[var] for var in reversed(data_drift)}) + Xt = transformer.transform(X_reordered) - switch_case = DropHighPSIFeatures( - split_frac=0.5, bins=3, switch=True, min_pct_empty_bins=0.001 - ) - switch_case.fit(df_reverse) - - assert case.psi_values_ == switch_case.psi_values_ + assert isinstance(Xt, make_df) + assert list(Xt.columns) == [ + var for var in data_drift if var not in ["drift_1", "drift_2"] + ] -def test_observation_frequency_per_bin(): - """Test empty bins are populated by a tiny amount.""" - a = pd.DataFrame({"A": [1, 2, 4]}) - b = pd.DataFrame({"A": [1, 2, 3]}) - transformer = DropHighPSIFeatures() - a_bins, b_bins = transformer._observation_frequency_per_bin(a, b) +def test_error_if_transform_df_has_different_number_of_columns(make_df, data_drift): + transformer = DropHighPSIFeatures().fit(make_df(data_drift)) + X = make_df({**data_drift, "A": [1] * 1000}) + msg = ( + "The number of columns in this dataset is different from the one used to " + "fit this transformer (when using the fit() method)." + ) + with pytest.raises(ValueError, match=re.escape(msg)): + transformer.transform(X) - expected_a_bins = pd.Series([0.3333333, 0.333333, 0.0001, 0.333333]) - expected_b_bins = pd.Series([0.3333333, 0.333333, 0.333333, 0.0001]) - pd.testing.assert_series_equal( - a_bins.reset_index(drop=True), expected_a_bins, check_names=False - ) - pd.testing.assert_series_equal( - b_bins.reset_index(drop=True), expected_b_bins, check_names=False +def test_error_if_not_fitted(make_df, data_drift): + msg = ( + "This DropHighPSIFeatures instance is not fitted yet. Call 'fit' with " + "appropriate arguments before using this estimator." ) + with pytest.raises(NotFittedError, match=re.escape(msg)): + DropHighPSIFeatures().transform(make_df(data_drift)) -def test_transform_feature_to_drop_not_present(df): - """Test transform when the feature to drop in not in the dataframe.""" - test = DropHighPSIFeatures() - test.fit(df) +def test_input_df_is_not_modified(make_df, data_drift): + X = make_df(data_drift) + DropHighPSIFeatures(variables="all", switch=True).fit_transform(X) + assert frame_to_dict(X) == data_drift - # Define new dataframe with additional column - data = df.copy() - data["A"] = [1] * 1000 - # Remove one of the feature to drop - data = data.drop("drift_1", axis=1) - with pytest.raises(KeyError): - test.transform(data) +def test_category_dtype(data_drift): + X = pd.DataFrame(data_drift) + X["drift_cat_1"] = X["drift_cat_1"].astype("category") + X["cat_1"] = pd.Categorical(X["cat_1"], categories=["C", "B", "A"]) + transformer = DropHighPSIFeatures(variables=["cat_1", "drift_cat_1"]) + Xt = transformer.fit_transform(X) + assert transformer.psi_values_ == pytest.approx( + {"cat_1": 0.0, "drift_cat_1": 18.41883867587797} + ) + pd.testing.assert_frame_equal(Xt, X.drop(columns="drift_cat_1")) -def test_transform_different_number_of_columns(df): - """Test transform on df with different number of features to train set.""" - test = DropHighPSIFeatures() - test.fit(df) - # Define new dataframe with additional column - data = df.copy() - data["A"] = [1] * 1000 +def test_integer_column_names(data_drift): + X = pd.DataFrame(data_drift) + X.columns = list(range(X.shape[1])) + transformer = DropHighPSIFeatures(split_col=6, variables="all") + Xt = transformer.fit_transform(X) - with pytest.raises(ValueError): - test.transform(data) + # the column cat_1, which alternates A and B, puts even rows in the basis set. + assert transformer.cut_off_ == "A" + assert transformer.variables_ == [0, 1, 2, 3, 4, 5, 7, 8, 9] + assert transformer.features_to_drop_ == [] + pd.testing.assert_frame_equal(Xt, X) + + transformer = DropHighPSIFeatures(variables=[0, 7, 9]).fit(X) + assert transformer.psi_values_ == pytest.approx( + { + 0: EXPECTED_PSI["var_0"], + 7: EXPECTED_PSI["drift_1"], + 9: EXPECTED_PSI["drift_cat_1"], + }, + ) + pd.testing.assert_frame_equal(transformer.transform(X), X.drop(columns=[7, 9]))