FIX n_classes in DecisionBoundaryDisplay with custom estimators (scikit-learn#33202)

Co-authored-by: Lucy Liu <jliu176@gmail.com>

Author: AnneBeyer

doc/whats_new/upcoming_changes/sklearn.inspection/33202.fix.rst

@@ -0,0 +1,4 @@
+- In :class:`inspection.DecisionBoundaryDisplay`, `n_classes` is now inferred more
+  robustly from the estimator. If it fails for custom estimators, a comprehensive error
+  message is shown.
+  By :user:`Anne Beyer <AnneBeyer>`.

examples/cluster/plot_inductive_clustering.py

@@ -25,7 +25,7 @@
 
 import matplotlib.pyplot as plt
 
-from sklearn.base import BaseEstimator, clone
+from sklearn.base import BaseEstimator, ClusterMixin, clone
 from sklearn.cluster import AgglomerativeClustering
 from sklearn.datasets import make_blobs
 from sklearn.ensemble import RandomForestClassifier
@@ -50,7 +50,7 @@ def _classifier_has(attr):
     )
 
 
-class InductiveClusterer(BaseEstimator):
+class InductiveClusterer(ClusterMixin, BaseEstimator):
     def __init__(self, clusterer, classifier):
         self.clusterer = clusterer
         self.classifier = classifier
@@ -60,6 +60,7 @@ def fit(self, X, y=None):
         self.classifier_ = clone(self.classifier)
         y = self.clusterer_.fit_predict(X)
         self.classifier_.fit(X, y)
+        self.labels_ = y
         return self
 
     @available_if(_classifier_has("predict"))
@@ -122,7 +123,12 @@ def plot_scatter(X, color, alpha=0.5):
 
 # Plotting decision regions
 DecisionBoundaryDisplay.from_estimator(
-    inductive_learner, X, response_method="predict", alpha=0.4, ax=ax
+    inductive_learner,
+    X,
+    response_method="predict",
+    multiclass_colors="viridis",
+    alpha=0.4,
+    ax=ax,
 )
 plt.title("Classify unknown instances")
 

sklearn/inspection/_plot/decision_boundary.py

@@ -12,6 +12,7 @@
 from sklearn.utils._optional_dependencies import check_matplotlib_support
 from sklearn.utils._response import _get_response_values
 from sklearn.utils._set_output import _get_adapter_from_container
+from sklearn.utils.multiclass import type_of_target
 from sklearn.utils.validation import (
     _is_arraylike_not_scalar,
     _num_features,
@@ -576,6 +577,31 @@ def from_estimator(
             encoder.classes_ = estimator.classes_
             response = encoder.transform(response)
 
+        # infer n_classes from the estimator
+        if (
+            class_of_interest is not None
+            or is_regressor(estimator)
+            or is_outlier_detector(estimator)
+        ):
+            n_classes = 2
+        elif is_classifier(estimator) and hasattr(estimator, "classes_"):
+            n_classes = len(estimator.classes_)
+        elif is_clusterer(estimator) and hasattr(estimator, "labels_"):
+            n_classes = len(np.unique(estimator.labels_))
+        else:
+            target_type = type_of_target(response)
+            if target_type in ("binary", "continuous"):
+                n_classes = 2
+            elif target_type == "multiclass":
+                n_classes = len(np.unique(response))
+            else:
+                raise ValueError(
+                    "Number of classes or labels cannot be inferred from "
+                    f"{estimator.__class__.__name__}. Please make sure your estimator "
+                    "follows scikit-learn's estimator API as described here: "
+                    "https://scikit-learn.org/stable/developers/develop.html#rolling-your-own-estimator"
+                )
+
         if response.ndim == 1:
             response = response.reshape(*xx0.shape)
         else:
@@ -591,17 +617,6 @@ def from_estimator(
             else:
                 response = response.reshape(*xx0.shape, response.shape[-1])
 
-        if (
-            class_of_interest is not None
-            or is_regressor(estimator)
-            or is_outlier_detector(estimator)
-        ):
-            n_classes = 2
-        elif is_classifier(estimator):
-            n_classes = len(estimator.classes_)
-        elif is_clusterer(estimator):
-            n_classes = len(np.unique(estimator.labels_))
-
         if xlabel is None:
             xlabel = X.columns[0] if hasattr(X, "columns") else ""
 

sklearn/inspection/_plot/tests/test_boundary_decision_display.py

@@ -19,7 +19,8 @@
 from sklearn.inspection import DecisionBoundaryDisplay
 from sklearn.inspection._plot.decision_boundary import _check_boundary_response_method
 from sklearn.linear_model import LogisticRegression
-from sklearn.preprocessing import scale
+from sklearn.pipeline import Pipeline
+from sklearn.preprocessing import StandardScaler, scale
 from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor
 from sklearn.utils._testing import (
     _convert_container,
@@ -703,6 +704,34 @@ def test_multiclass_colors_cmap(
     #    assert len(disp.surface_.levels) >= disp.n_classes
 
 
+# estimator classes for non-regression test cases for issue #33194
+class CustomBinaryEstimator(BaseEstimator):
+    def fit(self, X, y):
+        self.fitted_ = True
+        return self
+
+    def predict(self, X):
+        return np.arange(X.shape[0]) % 2
+
+
+class CustomMulticlassEstimator(BaseEstimator):
+    def fit(self, X, y):
+        self.fitted_ = True
+        return self
+
+    def predict(self, X):
+        return np.arange(X.shape[0]) % 7
+
+
+class CustomContinuousEstimator(BaseEstimator):
+    def fit(self, X, y):
+        self.fitted_ = True
+        return self
+
+    def predict(self, X):
+        return np.arange(X.shape[0]) * 0.5
+
+
 @pytest.mark.parametrize(
     "estimator, n_blobs, expected_n_classes",
     [
@@ -712,12 +741,57 @@ def test_multiclass_colors_cmap(
         (KMeans(n_clusters=2, random_state=0), 2, 2),
         (DecisionTreeRegressor(random_state=0), 7, 2),
         (IsolationForest(random_state=0), 7, 2),
+        (CustomBinaryEstimator(), 2, 2),
+        (CustomMulticlassEstimator(), 7, 7),
+        (CustomContinuousEstimator(), 7, 2),
+        (
+            Pipeline(
+                [
+                    ("scale", StandardScaler()),
+                    ("dt", DecisionTreeClassifier(random_state=0)),
+                ]
+            ),
+            7,
+            7,
+        ),
+        # non-regression test case for issue #33194
+        (
+            Pipeline(
+                [
+                    ("scale", StandardScaler()),
+                    ("kmeans", KMeans(n_clusters=7, random_state=0)),
+                ]
+            ),
+            7,
+            7,
+        ),
+        (
+            Pipeline(
+                [
+                    ("scale", StandardScaler()),
+                    ("reg", DecisionTreeRegressor(random_state=0)),
+                ]
+            ),
+            7,
+            2,
+        ),
+        (
+            Pipeline(
+                [
+                    ("scale", StandardScaler()),
+                    ("kmeans", IsolationForest(random_state=0)),
+                ]
+            ),
+            7,
+            2,
+        ),
     ],
 )
 def test_n_classes_attribute(pyplot, estimator, n_blobs, expected_n_classes):
     """Check that `n_classes` is set correctly.
 
-    Introduced in https://github.com/scikit-learn/scikit-learn/pull/33015"""
+    Introduced in https://github.com/scikit-learn/scikit-learn/pull/33015.
+    """
 
     X, y = make_blobs(n_samples=150, centers=n_blobs, n_features=2, random_state=42)
     clf = estimator.fit(X, y)
@@ -731,6 +805,27 @@ def test_n_classes_attribute(pyplot, estimator, n_blobs, expected_n_classes):
     assert disp_coi.n_classes == 2
 
 
+def test_n_classes_raises_if_not_inferrable(pyplot):
+    """Check behaviour if `n_classes` can't be inferred.
+
+    Non-regression test for issue #33194.
+    """
+
+    class CustomUnknownEstimator(BaseEstimator):
+        def fit(self, X, y):
+            self.fitted_ = True
+            return self
+
+        def predict(self, X):
+            return np.array(0)
+
+    X, y = load_iris_2d_scaled()
+    est = CustomUnknownEstimator().fit(X, y)
+    msg = "Number of classes or labels cannot be inferred from CustomUnknownEstimator"
+    with pytest.raises(ValueError, match=msg):
+        DecisionBoundaryDisplay.from_estimator(est, X, response_method="predict")
+
+
 def test_cmap_and_colors_logic(pyplot):
     """Check the handling logic for `cmap` and `colors`."""
     X, y = load_iris_2d_scaled()