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/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..40e84f493 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,47 @@ 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_] - - # return the dataframe with the selected features - return X.drop(columns=self.features_to_drop_) - - def _get_feature_names_in(self, X): - """Get the names and number of features in the train set. The dataframe - used during fit.""" - - self.feature_names_in_ = X.columns.to_list() + # 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] + + # pandas is faster than narwhals. + 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): + """Get the names and number of features in the train set (the dataframe + used during fit).""" + + 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/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/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..64e43b88c 100644 --- a/tests/test_selection/test_base_selector.py +++ b/tests/test_selection/test_base_selector.py @@ -1,55 +1,164 @@ +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"]), + (["Name", "City", "Age", "Marks", "dob"], []), + ], +) +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"]]) 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]))