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

Code before:
            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


@pytest.mark.parametrize(
    "estimator, n_classes",
    [
        (DecisionTreeClassifier(random_state=0), 7),
        (DecisionTreeClassifier(random_state=0), 2),
        (KMeans(n_clusters=7, random_state=0), 7),
        (KMeans(n_clusters=2, random_state=0), 2),
        (DecisionTreeRegressor(random_state=0), 7),
        (IsolationForest(random_state=0), 7),
    ],
)
def test_n_classes_attribute(pyplot, estimator, n_classes):
    """Check that `n_classes` is set correctly.

Review comment (suggestion):
```suggestion
        # Part of a non-regression test for issue #32866 to make sure every class
        # has its own surface.
```

Code after:

        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"
    )
    assert disp_coi.n_classes == 2


def test_cmap_and_colors_logic(pyplot):
    """Check the handling logic for `cmap` and `colors`."""
    X, y = load_iris_2d_scaled()
    clf = LogisticRegression().fit(X, y)

    with pytest.warns(
        UserWarning,
        match="'cmap' is ignored in favor of 'multiclass_colors'",
    ):
        DecisionBoundaryDisplay.from_estimator(
            clf,
            X,