From deb11db4ea6f5c5b02e71baa5020f4e85a7ef914 Mon Sep 17 00:00:00 2001 From: Swapnil Dwivedi Date: Fri, 6 Nov 2026 13:04:55 +0530 Subject: [PATCH] Use fitted bins for numerical PSI thresholds --- .../selection/DropHighPSIFeatures.rst | 4 +++- feature_engine/selection/drop_psi_features.py | 20 +++++++++---------- .../test_drop_high_psi_features.py | 18 +++++++++++++++++ 3 files changed, 31 insertions(+), 11 deletions(-) diff --git a/docs/user_guide/selection/DropHighPSIFeatures.rst b/docs/user_guide/selection/DropHighPSIFeatures.rst index b4ecf9bde..e9d065990 100644 --- a/docs/user_guide/selection/DropHighPSIFeatures.rst +++ b/docs/user_guide/selection/DropHighPSIFeatures.rst @@ -73,7 +73,9 @@ to the PSI value. The most commonly used thresholds are: - Below 10%, the variable has not experienced a significant shift. - Above 25%, the variable has experienced a major shift. - Between those two values, the shift is intermediate. -- 'auto': the threshold will be calculated based on the size of the base and target datasets and the number of bins. +- 'auto': the threshold is calculated from the sizes of the basis and test datasets and + the number of bins fitted for each feature. Repeated quantiles can lead to fewer bins + than requested when using the equal-frequency strategy. When 'auto', the threshold is calculated using the chi2 approximation, proposed by B. Yurdakul: diff --git a/feature_engine/selection/drop_psi_features.py b/feature_engine/selection/drop_psi_features.py index ea94e8e6f..473156270 100644 --- a/feature_engine/selection/drop_psi_features.py +++ b/feature_engine/selection/drop_psi_features.py @@ -460,16 +460,6 @@ def fit(self, X: pd.DataFrame, y: pd.Series = None): else: bucketer = EqualFrequencyDiscretiser(q=self.bins) - # 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": @@ -495,6 +485,16 @@ def fit(self, X: pd.DataFrame, y: pd.Series = None): (test_distrib - basis_distrib) * np.log(test_distrib / basis_distrib) ) + if self.threshold == "auto": + # Equal-frequency discretization can drop duplicate boundaries. + threshold_num = self._calculate_auto_threshold( + basis_df.shape[0], + test_df.shape[0], + len(bucketer.binner_dict_[feature]) - 1, + ) + else: + threshold_num = self.threshold + # Assess if feature should be dropped if self.psi_values_[feature] > threshold_num: self.features_to_drop_.append(feature) diff --git a/tests/test_selection/test_drop_high_psi_features.py b/tests/test_selection/test_drop_high_psi_features.py index 8cc296837..29c9ca4bb 100644 --- a/tests/test_selection/test_drop_high_psi_features.py +++ b/tests/test_selection/test_drop_high_psi_features.py @@ -178,6 +178,24 @@ def test_fit_attributes_with_autothreshold(variables, expected_psi, df): ] +@pytest.mark.parametrize("bins", [2, 10]) +def test_auto_threshold_uses_realized_quantile_bins(bins): + X = pd.DataFrame( + { + "score": [0] * 50 + [1] * 50 + [0] * 27 + [1] * 73, + "stable": list(range(100)) * 2, + } + ) + transformer = DropHighPSIFeatures(bins=bins, threshold="auto") + result = transformer.fit_transform(X) + + # Repeated quantiles leave two bins. The PSI for score is above the + # one-degree-of-freedom threshold, but below the nine-degree threshold. + assert transformer.psi_values_["score"] == pytest.approx(0.23 * np.log(0.73 / 0.27)) + assert transformer.features_to_drop_ == ["score"] + assert list(result.columns) == ["stable"] + + def test_calculation_when_strategy_equal_width(df): transformer = DropHighPSIFeatures(strategy="equal_width") transformer.fit(df)