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

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

    Introduced in https://github.com/scikit-learn/scikit-learn/pull/33015"""

    X, y = make_blobs(n_samples=150, centers=n_classes, n_features=2, random_state=42)
    clf = estimator.fit(X, y)

    disp = DecisionBoundaryDisplay.from_estimator(clf, X, response_method="predict")

    if is_regressor(estimator) or is_outlier_detector(estimator):
        n_classes = 2  # by convention for regressors and outlier detectors
    assert disp.n_classes == n_classes

Review comment (suggestion):
I would prefer to distinctly name and set the actual number of blobs in the data and the expected n_classes attribute value as distinct test parameters to better convey the expectations of this test.

```suggestion
@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), 7, 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
```

Note: I have checked the code locally. It's possible I made a mistake while editing via the GitHub UI.

Code after:
@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,
            multiclass_colors="plasma",
            cmap="Blues",
        )

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