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

Code before:
                else:
                    cmap = plt.get_cmap(cmap, self.n_classes)
                    if not hasattr(cmap, "colors"):
                        # For LinearSegmentedColormap
                        colors = cmap(np.linspace(0, 1, self.n_classes))
                    else:
                        colors = cmap.colors
            elif isinstance(self.multiclass_colors, list):
                colors = [mpl.colors.to_rgba(color) for color in self.multiclass_colors]
            else:
                raise ValueError("'multiclass_colors' must be a list or a str.")

            self.multiclass_colors_ = colors

            if plot_method == "contour":
                # Plot only argmax
                response = (
                    self.response.argmax(axis=2)
                    if self.response.ndim == 3
                    else self.response
                )
                self.surface_ = plot_func(
                    self.xx0, self.xx1, response, colors=colors, **kwargs
                )
            elif plot_method == "pcolormesh" and self.response.ndim == 2:  # predict
                cmap = mpl.colors.ListedColormap(colors)
                self.surface_ = plot_func(
                    self.xx0, self.xx1, self.response, cmap=cmap, **kwargs
                )
            else:
                multiclass_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)
                ]
                self.surface_ = []
                for class_idx, cmap in enumerate(multiclass_cmaps):
                    if self.response.ndim == 2:  # predict
                        response = np.ma.array(
                            (self.response == class_idx),
                            mask=(self.response != class_idx),
                        )

                    else:  # predict_proba or decision_function
                        response = np.ma.array(
                            self.response[:, :, class_idx],
                            mask=(self.response.argmax(axis=2) != class_idx),
                        )


Review comment (suggestion):
For colormaps with `colors` attr, `plt.get_cmap(cmap, self.n_classes)` will have a max number of colors.

e.g., `plt.get_cmap('tab10', 12)`

will only give 10 distinct colors - it will just start repeating colors.

In a separate PR, we should probably fix this and add a test.


Code after:
                        "in the multiclass case."
                    )
                    del kwargs[kwarg]

            if self.multiclass_colors is None or isinstance(
                self.multiclass_colors, str
            ):
                if self.multiclass_colors is None:
                    cmap = "tab10" if self.n_classes <= 10 else "gist_rainbow"
                else:
                    cmap = self.multiclass_colors

                # Special case for the tab10 and tab20 colormaps that encode a
                # discrete set of colors that are easily distinguishable
                # contrary to other colormaps that are continuous.
                if cmap == "tab10" and self.n_classes <= 10:
                    colors = plt.get_cmap("tab10", 10).colors[: self.n_classes]
                elif cmap == "tab20" and self.n_classes <= 20:
                    colors = plt.get_cmap("tab20", 20).colors[: self.n_classes]
                else:
                    cmap = plt.get_cmap(cmap, self.n_classes)
                    if not hasattr(cmap, "colors"):
                        # Get `LinearSegmentedColormap` for non-qualitative cmaps
                        colors = cmap(np.linspace(0, 1, self.n_classes))
                    else:
                        colors = cmap.colors
            elif isinstance(self.multiclass_colors, list):
                colors = [mpl.colors.to_rgba(color) for color in self.multiclass_colors]
            else:
                raise ValueError("'multiclass_colors' must be a list or a str.")

            self.multiclass_colors_ = colors

            if self.response.ndim == 2:  # predict
                # `pcolormesh` requires cmap, for the others it makes no difference
                cmap = mpl.colors.ListedColormap(colors)
                self.surface_ = plot_func(
                    self.xx0, self.xx1, self.response, cmap=cmap, **kwargs
                )

            # predict_proba and decision_function differ for plotting methods
            elif plot_method == "contour":
                # Plot only integer class values
                self.surface_ = plot_func(
                    self.xx0,
                    self.xx1,
                    self.response.argmax(axis=2),
                    colors=colors,
                    **kwargs,
                )
            else: