Repository: feature-engine/feature_engine

Issue:
expand probe features selection to some how account for the dispersion when selecting features
at the moment, it only compares mean importance of features to mean of the probes. We should be able to account for the dispersion somehow. Maybe also allow the comparison to max of probes instead of mean of probes, or mean plus std.


Discussion:

Patch:
diff --git a/docs/user_guide/selection/ProbeFeatureSelection.rst b/docs/user_guide/selection/ProbeFeatureSelection.rst
index e01fef6da..882d0ce76 100644
--- a/docs/user_guide/selection/ProbeFeatureSelection.rst
+++ b/docs/user_guide/selection/ProbeFeatureSelection.rst
@@ -43,7 +43,14 @@ Selecting features
 
 After assigning a value of feature importance to each feature, including the probes,
 :class:`ProbeFeatureSelection()` will select those variables whose importance is greater
-than the mean importance of all probes.
+than:
+
+- the mean importance of all probes
+- the maximum importance of all probes
+- the mean plus 3 times the standard deviation of the importance of the probes
+
+The threshold for feature selection can be controlled through the parameter `threshold`
+when setting up the transformer.
 
 Feature selection process
 -------------------------
@@ -54,7 +61,7 @@ strategy:
 1. Add 1 or more random features to the dataset
 2. Train a machine learning model using all features including the random ones
 3. Derive feature importance from the fitted model
-4. Take the average importance of the random features
+4. Take the average (or maximum or mean+std) importance of the random features
 5. Select features whose importance is greater than the importance of the random variables (step 4)
 
 
@@ -64,7 +71,7 @@ strategy:
 1. Add 1 or more random features to the dataset
 2. Train a machine learning per feature and per probe
 3. Determine the feature importance as the performance of the single feature model
-4. Take the average importance of the random features
+4. Take the average (or maximum or mean+std) importance of the random features
 5. Select features whose importance is greater than the importance of the random variables (step 4)
 
 Rationale of probe feature selection
@@ -506,11 +513,57 @@ If most variables are continuous, introduce features with normal and uniform dis
 If you have one hot encoded features or sparse matrices, binary features might be a better
 option.
 
+Changing the probe importance threshold
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+We can make the selection process more aggressive by using the maximum of the probe features
+or the mean plus 3 times the standard deviation as threshold to select features.
+
+In the following example, we'll use the same random forest and the same probe features,
+but this time, we'll select features whose importance is greater than the mean plus 3 times
+the standard deviation of the probes:
+
+.. code:: python
+
+    sel = ProbeFeatureSelection(
+        estimator=RandomForestClassifier(),
+        variables=None,
+        scoring="precision",
+        n_probes=1,
+        distribution="all",
+        threshold = "mean_plus_std",
+        cv=5,
+        random_state=150,
+        confirm_variables=False
+    )
+
+    sel.fit(X_train, y_train)
+
+We now inspect the variables that will be removed:
+
+.. code:: python
+
+    sel.features_to_drop_
+
+We see that now, several variables will be removed from the dataset:
+
+.. code:: python
+
+    ['mean smoothness',
+     'mean symmetry',
+     'mean fractal dimension',
+     'texture error',
+     'smoothness error',
+     'compactness error',
+     'concave points error',
+     'symmetry error',
+     'fractal dimension error']
+
 Using the individual feature strategy
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-We will now repeat the process, but we will train a random forest per feature instead, and
-use the roc-auc as a measure of feature importance:
+We will now select features by training a random forest per feature and using the roc-auc
+obtained from that model as a measure of feature importance:
 
 .. code:: python
 
@@ -557,6 +610,26 @@ many of the features return roc-auc that are not significantly different from th
 (error bars overlaps). So, even if the transformer would not drop those features, we
 could decide to discard them after analysis of this plot.
 
+Alternatively, we can set the threshold to be more aggressive and drop features whose
+importance is smaller than the mean plus three times the standard deviation of the
+importance of the probes, as follows:
+
+
+.. code:: python
+
+    sel = ProbeFeatureSelection(
+        estimator=RandomForestClassifier(n_estimators=5, random_state=1),
+        variables=None,
+        collective=False,
+        scoring="roc_auc",
+        n_probes=1,
+        distribution="all",
+        threshold = "mean_plus_std",
+        cv=5,
+        random_state=150,
+        confirm_variables=False
+    ).fit(X_train, y_train)
+
 
 Additional resources
 --------------------
diff --git a/feature_engine/selection/probe_feature_selection.py b/feature_engine/selection/probe_feature_selection.py
index 35a5f037a..ec112b3e4 100644
--- a/feature_engine/selection/probe_feature_selection.py
+++ b/feature_engine/selection/probe_feature_selection.py
@@ -70,8 +70,8 @@ class ProbeFeatureSelection(BaseSelector):
 
     Finally, ProbeFeatureSelection() selects the features whose importance is greater
     than those of the probes. In the case of there being more than one probe feature,
-    ProbeFeatureSelection() takes the average feature importance of all the probe
-    features.
+    ProbeFeatureSelection() can take the average, maximum, or mean plus 3 std feature
+    importance of all the probe features as threshold for the feature selection.
 
     The variables whose feature importance is smaller than the feature importance of
     the probe feature(s) are dropped from the dataset.
@@ -108,6 +108,14 @@ class ProbeFeatureSelection(BaseSelector):
         to `n_categories`. If `distribution` is 'poisson', then samples are taken from
         `np.random.poisson(n_categories, n_obs)`.
 
+    threshold: str, default='mean'
+        Indicates how to combine the importance of the probe features as threshold for
+        the feature selection. If 'mean', then features are selected if their
+        importance is greater than the mean of the probes. If 'max', then features are
+        selected if their importance is greater than the maximun importance of all
+        probes. If 'mean_plus_std', then features are selected if their importance is
+        greater than the mean plus three times the standard deviation of the probes.
+
     {cv}
 
     {groups}
@@ -180,6 +188,7 @@ def __init__(
         n_probes: int = 1,
         distribution: Union[str, list] = "normal",
         n_categories: int = 10,
+        threshold: str = "mean",
         cv=5,
         groups=None,
         random_state: int = 0,
@@ -222,6 +231,16 @@ def __init__(
                 f"n_categories must be a positive integer. Got {n_categories} instead."
             )
 
+        if not isinstance(threshold, str) or threshold not in [
+            "mean",
+            "max",
+            "mean_plus_std",
+        ]:
+            raise ValueError(
+                "threshold takes values 'mean', 'max' or 'mean_plus_std'. "
+                f"Got {threshold} instead."
+            )
+
         super().__init__(confirm_variables)
         self.estimator = estimator
         self.variables = variables
@@ -232,6 +251,7 @@ def __init__(
         self.cv = cv
         self.groups = groups
         self.n_probes = n_probes
+        self.threshold = threshold
         self.random_state = random_state
 
     def fit(self, X: pd.DataFrame, y: pd.Series):
@@ -341,21 +361,35 @@ def _get_features_to_drop(self):
         feature importance of all the probe features.
         """
 
-        # if more than 1 probe feature, calculate average feature importance
+        # if more than 1 probe feature, calculate threshold based on
+        # probe feature importance.
         if self.probe_features_.shape[1] > 1:
-            probe_features_avg_importance = self.feature_importances_[
-                self.probe_features_.columns
-            ].values.mean()
+            if self.threshold == "mean":
+                threshold = self.feature_importances_[
+                    self.probe_features_.columns
+                ].values.mean()
+            elif self.threshold == "max":
+                threshold = self.feature_importances_[
+                    self.probe_features_.columns
+                ].values.max()
+            else:
+                threshold = (
+                    self.feature_importances_[
+                        self.probe_features_.columns
+                    ].values.mean()
+                    + 3
+                    * self.feature_importances_[
+                        self.probe_features_.columns
+                    ].values.std()
+                )
 
         else:
-            probe_features_avg_importance = self.feature_importances_[
-                self.probe_features_.columns
-            ].values
+            threshold = self.feature_importances_[self.probe_features_.columns].values
 
         features_to_drop = []
 
         for var in self.variables_:
-            if self.feature_importances_[var] < probe_features_avg_importance:
+            if self.feature_importances_[var] < threshold:
                 features_to_drop.append(var)
 
         return features_to_drop


Test patch:
diff --git a/tests/test_selection/test_probe_feature_selection.py b/tests/test_selection/test_probe_feature_selection.py
index a1e3a2207..58d489122 100644
--- a/tests/test_selection/test_probe_feature_selection.py
+++ b/tests/test_selection/test_probe_feature_selection.py
@@ -96,7 +96,7 @@ def test_n_categories_raises_error(n_cat):
 @pytest.mark.parametrize("n_cat", [[10], [10, 1], {10}, {10, 1}])
 def test_n_categories_raises_error_with_collections(n_cat):
     # I had problems with collections and string matching so
-    # I splitted the test
+    # I split the test
     with pytest.raises(ValueError):
         ProbeFeatureSelection(
             estimator=DecisionTreeRegressor(),
@@ -104,6 +104,28 @@ def test_n_categories_raises_error_with_collections(n_cat):
         )
 
 
+@pytest.mark.parametrize("thresh", ["mean", "max", "mean_plus_std"])
+def test_threshold_init_param(thresh):
+    sel = ProbeFeatureSelection(
+        estimator=DecisionTreeRegressor(),
+        threshold=thresh,
+    )
+    assert sel.threshold == thresh
+
+
+@pytest.mark.parametrize("thresh", [1, "string"])
+def test_threshold_raises_error(thresh):
+    msg = (
+        "threshold takes values 'mean', 'max' or 'mean_plus_std'. "
+        f"Got {thresh} instead."
+    )
+    with pytest.raises(ValueError, match=msg):
+        ProbeFeatureSelection(
+            estimator=DecisionTreeRegressor(),
+            threshold=thresh,
+        )
+
+
 def test_fit_transform_functionality(df_test):
     X, y = df_test
 
@@ -317,9 +339,11 @@ def test_generate_features_n_categories(n_cats, max_discrete, max_poisson):
     assert probe_features["poisson_probe_0"].max() == max_poisson
 
 
-def test_get_features_to_drop():
-    # 1 probe
-    sel = ProbeFeatureSelection(estimator=LogisticRegression(), n_probes=1)
+@pytest.mark.parametrize("thresh", ["mean", "max", "mean_plus_std"])
+def test_get_features_to_drop_with_one_probe(thresh):
+    sel = ProbeFeatureSelection(
+        estimator=LogisticRegression(), n_probes=1, threshold=thresh
+    )
     sel.feature_importances_ = pd.Series(
         [11, 12, 9, 10], index=["var1", "var2", "var3", "probe"]
     )
@@ -327,17 +351,30 @@ def test_get_features_to_drop():
     sel.variables_ = ["var1", "var2", "var3"]
     assert sel._get_features_to_drop() == ["var3"]
 
-    # 2 probes
-    sel = ProbeFeatureSelection(estimator=LogisticRegression(), n_probes=2)
+
+_thresholds = [
+    ("mean", ["var4"]),
+    ("max", ["var3", "var4"]),
+    ("mean_plus_std", ["var1", "var3", "var4"]),
+]
+
+
+@pytest.mark.parametrize("thresh, vars_to_drop", _thresholds)
+def test_get_features_to_drop_with_many_probes(thresh, vars_to_drop):
+    sel = ProbeFeatureSelection(
+        estimator=LogisticRegression(),
+        n_probes=2,
+        threshold=thresh,
+    )
     sel.feature_importances_ = pd.Series(
-        [11, 12, 10, 8.7, 10, 8],
+        [11, 20, 9.9, 8.7, 10, 8],
         index=["var1", "var2", "var3", "var4", "probe1", "probe2"],
     )
     sel.probe_features_ = pd.DataFrame(
         {"probe1": [1, 1, 1, 1, 1], "probe2": [1, 1, 1, 1, 1]}
     )
     sel.variables_ = ["var1", "var2", "var3", "var4"]
-    assert sel._get_features_to_drop() == ["var4"]
+    assert sel._get_features_to_drop() == vars_to_drop
 
 
 def test_cv_generator(df_test):