Repository: scikit-learn/scikit-learn
File: sklearn/inspection/_plot/tests/test_boundary_decision_display.py (Python)

Code before:
    clf = LogisticRegression().fit(X, y)

    disp = DecisionBoundaryDisplay.from_estimator(
        clf,
        X,
        response_method=response_method,
        plot_method=plot_method,
        multiclass_colors=multiclass_colors,
    )

    if multiclass_colors is None:
        if len(clf.classes_) <= 10:
            colors = mpl.pyplot.get_cmap("tab10", 10).colors[: len(clf.classes_)]
        else:
            cmap = mpl.pyplot.get_cmap("gist_rainbow", len(clf.classes_))
            colors = cmap(np.linspace(0, 1, len(clf.classes_)))
    elif multiclass_colors == "plasma":
        colors = mpl.pyplot.get_cmap(multiclass_colors, len(clf.classes_)).colors
    elif multiclass_colors == "Blues":
        cmap = mpl.pyplot.get_cmap(multiclass_colors, len(clf.classes_))
        colors = cmap(np.linspace(0, 1, len(clf.classes_)))
    else:
        colors = [mpl.colors.to_rgba(color) for color in multiclass_colors]

    # part of a non-regression test for issue #32866
    # to make sure the colormap used has enough distinct colors
    assert disp.n_classes <= len(np.unique(colors, axis=0))

    if plot_method == "pcolormesh" and response_method == "predict":
        # pcolormesh with predict uses ListedColormap
        cmap = mpl.colors.ListedColormap(colors)
        assert disp.surface_.cmap == cmap
    elif plot_method != "contour":
        cmaps = [
            mpl.colors.LinearSegmentedColormap.from_list(
                f"colormap_{class_idx}", [(1.0, 1.0, 1.0, 1.0), (r, g, b, 1.0)]
            )
            for class_idx, (r, g, b, _) in enumerate(colors)
        ]
        # part of a non-regression test for issue #32866
        # to make sure every class has its own surface
        assert len(disp.surface_) == disp.n_classes

        for idx, quad in enumerate(disp.surface_):
            assert quad.cmap == cmaps[idx]
    else:
        assert_allclose(disp.surface_.colors, colors)

    # non-regression test for issue #32866 with `contour` (currently still fails)
    # if hasattr(disp.surface_, "levels"):
    #    assert len(disp.surface_.levels) >= disp.n_classes

Review comment (suggestion):
```suggestion
    # Part of a non-regression test for issue #32866 to make sure the colormap 
    # used has enough distinct colors.
```

Code after:
    # Make sure the colormap has enough distinct colors.
    assert disp.n_classes == len(np.unique(colors, axis=0))

    if response_method == "predict":
        cmap = mpl.colors.ListedColormap(colors)
        assert disp.surface_.cmap == cmap
    elif plot_method != "contour":
        cmaps = [
            mpl.colors.LinearSegmentedColormap.from_list(
                f"colormap_{class_idx}", [(1.0, 1.0, 1.0, 1.0), (r, g, b, 1.0)]
            )
            for class_idx, (r, g, b, _) in enumerate(colors)
        ]
        # Make sure every class has its own surface.
        assert len(disp.surface_) == disp.n_classes

        for idx, quad in enumerate(disp.surface_):
            assert quad.cmap == cmaps[idx]
    else:
        assert_allclose(disp.surface_.colors, colors)

    # non-regression test for issue #32866 (currently still fails)
    # if hasattr(disp.surface_, "levels"):
    #    assert len(disp.surface_.levels) >= disp.n_classes


@pytest.mark.parametrize(
    "estimator, n_blobs, expected_n_classes",
    [
        (DecisionTreeClassifier(random_state=0), 7, 7),
        (DecisionTreeClassifier(random_state=0), 2, 2),
        (KMeans(n_clusters=7, random_state=0), 7, 7),
        (KMeans(n_clusters=2, random_state=0), 2, 2),
        (DecisionTreeRegressor(random_state=0), 7, 2),
        (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"""

    X, y = make_blobs(n_samples=150, centers=n_blobs, n_features=2, random_state=42)
    clf = estimator.fit(X, y)
    disp = DecisionBoundaryDisplay.from_estimator(clf, X, response_method="predict")
    assert disp.n_classes == expected_n_classes

    # Test that setting class_of_interest always converts to a binary problem.
    disp_coi = DecisionBoundaryDisplay.from_estimator(
        clf, X, class_of_interest=y[0], response_method="predict"
    )