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

Code before:
    response = getattr(clf, response_method)(grid).reshape(*disp.response.shape)
    assert_allclose(response, disp.response)

    assert len(disp.surface_) == len(clf.classes_)
    # Get which class has highest response and check it is plotted
    highest_class = np.argmax(response, axis=2)
    for idx, quadmesh in enumerate(disp.surface_):
        # Note quadmesh mask is True (i.e. masked) when `idx` is NOT the highest class
        assert_array_equal(
            highest_class != idx,
            quadmesh.get_array().mask.reshape(*highest_class.shape),
        )


@pytest.mark.parametrize(
    "multiclass_colors, n_classes",
    [
        (None, 3),
        (None, 15),
        ("plasma", 3),
        ("plasma", 15),
        ("Blues", 3),
        ("Blues", 15),
        (["red", "green", "blue"], 3),
        (
            [
                "red",
                "green",
                "blue",
                "yellow",
                "orange",
                "purple",
                "pink",
                "brown",
                "gray",
                "cyan",
                "magenta",
                "lime",
                "navy",
                "teal",
                "olive",
            ],
            15,
        ),
    ],
)
@pytest.mark.parametrize(
    "response_method", ["decision_function", "predict_proba", "predict"]
)
@pytest.mark.parametrize("plot_method", ["contourf", "contour", "pcolormesh"])
def test_multiclass_colors_cmap(

Review comment (suggestion):
I had to combine colors and number of classes here because the list version has to match the exact number of classes.

Code after:
        response_method=response_method,
    )

    grid = np.concatenate([disp.xx0.reshape(-1, 1), disp.xx1.reshape(-1, 1)], axis=1)
    response = getattr(clf, response_method)(grid).reshape(*disp.response.shape)
    assert_allclose(response, disp.response)

    assert len(disp.surface_) == len(clf.classes_)
    # Get which class has highest response and check it is plotted
    highest_class = np.argmax(response, axis=2)
    for idx, quadmesh in enumerate(disp.surface_):
        # Note quadmesh mask is True (i.e. masked) when `idx` is NOT the highest class
        assert_array_equal(
            highest_class != idx,
            quadmesh.get_array().mask.reshape(*highest_class.shape),
        )


@pytest.mark.parametrize(
    "multiclass_colors, n_classes",
    [
        (None, 3),
        (None, 11),
        ("plasma", 3),
        ("Blues", 3),
        (["red", "green", "blue"], 3),
    ],
)
@pytest.mark.parametrize(
    "response_method", ["decision_function", "predict_proba", "predict"]
)
@pytest.mark.parametrize("plot_method", ["contourf", "contour", "pcolormesh"])
def test_multiclass_colors_cmap(
    pyplot,
    n_classes,
    response_method,
    plot_method,
    multiclass_colors,
):
    """Check correct cmap used for all `multiclass_colors` inputs."""
    import matplotlib as mpl

    if parse_version(mpl.__version__) < parse_version("3.5"):
        pytest.skip(
            "Matplotlib >= 3.5 is needed for `==` to check equivalence of colormaps"
        )

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

    disp = DecisionBoundaryDisplay.from_estimator(