DOC Enhanced example visualization to RFE (scikit-learn#28862)

plon-Susk7 · web-flow · commit 79e14a77565e · 2024-04-22T15:19:41.000+02:00
diff --git a/examples/feature_selection/plot_rfe_digits.py b/examples/feature_selection/plot_rfe_digits.py
@@ -3,8 +3,14 @@
 Recursive feature elimination
 =============================
 
-A recursive feature elimination example showing the relevance of pixels in
-a digit classification task.
+This example demonstrates how Recursive Feature Elimination
+(:class:`~sklearn.feature_selection.RFE`) can be used to determine the
+importance of individual pixels for classifying handwritten digits.
+:class:`~sklearn.feature_selection.RFE` recursively removes the least
+significant features, assigning ranks based on their importance, where higher
+`ranking_` values denote lower importance. The ranking is visualized using both
+shades of blue and pixel annotations for clarity. As expected, pixels positioned
+at the center of the image tend to be more predictive than those near the edges.
 
 .. note::
 
@@ -16,21 +22,33 @@
 
 from sklearn.datasets import load_digits
 from sklearn.feature_selection import RFE
-from sklearn.svm import SVC
+from sklearn.linear_model import LogisticRegression
+from sklearn.pipeline import Pipeline
+from sklearn.preprocessing import MinMaxScaler
 
 # Load the digits dataset
 digits = load_digits()
 X = digits.images.reshape((len(digits.images), -1))
 y = digits.target
 
-# Create the RFE object and rank each pixel
-svc = SVC(kernel="linear", C=1)
-rfe = RFE(estimator=svc, n_features_to_select=1, step=1)
-rfe.fit(X, y)
-ranking = rfe.ranking_.reshape(digits.images[0].shape)
+pipe = Pipeline(
+    [
+        ("scaler", MinMaxScaler()),
+        ("rfe", RFE(estimator=LogisticRegression(), n_features_to_select=1, step=1)),
+    ]
+)
+
+pipe.fit(X, y)
+ranking = pipe.named_steps["rfe"].ranking_.reshape(digits.images[0].shape)
 
 # Plot pixel ranking
 plt.matshow(ranking, cmap=plt.cm.Blues)
+
+# Add annotations for pixel numbers
+for i in range(ranking.shape[0]):
+    for j in range(ranking.shape[1]):
+        plt.text(j, i, str(ranking[i, j]), ha="center", va="center", color="black")
+
 plt.colorbar()
-plt.title("Ranking of pixels with RFE")
+plt.title("Ranking of pixels with RFE\n(Logistic Regression)")
 plt.show()
