# AUTOGENERATED! DO NOT EDIT! File to edit: 00_inspect.ipynb (unless otherwise specified).

__all__ = ['COLORS', 'get_inspector', 'identify_type', 'ModelType', 'generate_model_html']

# Cell
# Meant to be colorblind-friendly
COLORS = {"blue": "#377eb8", "orange": "#ff7f00", "green": "#4daf4a", "pink": "#f781bf"}

# Cell
from enum import auto, Enum
from typing import Callable, Iterable, List, Optional, Sequence, Union
import warnings

from fastcore.basics import basic_repr, GetAttr, store_attr
from fastcore.test import test_fig_exists
from IPython.display import HTML
from matplotlib.axes import Axes
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
import sklearn
from sklearn.base import ClassifierMixin, clone, RegressorMixin
from sklearn.dummy import DummyClassifier, DummyRegressor
from sklearn.inspection import permutation_importance
from sklearn.linear_model._base import LinearModel, LinearClassifierMixin
from sklearn.model_selection._search import BaseSearchCV
from sklearn.preprocessing import MinMaxScaler
from sklearn.tree import BaseDecisionTree, plot_tree
from sklearn.utils import check_X_y
from sklearn.utils.validation import check_is_fitted
import waterfall_chart

from .delegate import delegates
from .explore import show_correlation
from .tune import (
    calculate_metrics_by_thresh_binary,
    calculate_metrics_by_thresh_multi,
    confusion_matrix,
)

# Cell
class _Inspector(GetAttr):
    """Model inspector base class

    Users should use `get_inspector` to generate appropriate
    `_Inspector` objects rather than instantiating this class or its
    subclasses directly.

    Subclasses of this class are based on distinctions among model types
    only, including the distinction between binary and multiclass
    classification. Plotting functionality that depends on how many
    columns are in the feature DataFrame is delegated to a separate
    `_Plotter` class hierarchy. We use custom `__getattr__` and
    `__dir__` methods to expose the `_Plotter` methods through the
    `_Inspector` object, so that the user does not need to interact with
    the `_Plotter` directly.
    """

    def __init__(self, model, X, y):
        check_is_fitted(model)
        check_X_y(X, y)
        if not isinstance(model, (DummyClassifier, DummyRegressor)):
            model._check_n_features(X, reset=False)

        store_attr()
        self._plotter = self._get_plotter_class()(self.model, self.X, self.y)
        self.default = self._plotter

    __repr__ = basic_repr(["model"])

    def _get_plotter_class(self):
        result = _Plotter
        num_features = len(self.X.columns)
        model_type = identify_type(self.model, self.y)
        if model_type is ModelType.REGRESSION:
            if num_features == 1:
                result = _1dPlotter
            elif num_features == 2:
                if hasattr(self.model, "estimators_"):
                    result = _Reg2dMultiPlotter
                else:
                    result = _Reg2dPlotter
        elif model_type is ModelType.BINARY:
            if num_features == 1:
                result = _Bin1dPlotter
            elif num_features == 2:
                if hasattr(self.model, "estimators_"):
                    result = _Bin2dMultiPlotter
                else:
                    result = _Bin2dPlotter
        else:
            if num_features == 1:
                result = _1dPlotter
            elif num_features == 2:
                if hasattr(self.model, "estimators_"):
                    result = _Multi2dMultiPlotter
                else:
                    result = _Multi2dPlotter
        return result

    @delegates(sklearn.inspection.permutation_importance)
    def permutation_importance(
        self,
        **kwargs,
    ) -> pd.Series:
        """Calculate permutation importance"""
        if kwargs is None:
            kwargs = {}
        kwargs = {**{"n_jobs": -1}, **kwargs}

        importances = permutation_importance(self.model, self.X, self.y, **kwargs)[
            "importances_mean"
        ]
        return pd.Series(importances, index=self.X.columns)

    def plot_permutation_importance(
        self,
        ax: Optional[Axes] = None,
        sort: bool = True,
        importance_kwargs: Optional[dict] = None,
        plot_kwargs: Optional[dict] = None,
    ) -> Axes:
        """Plot a correlation matrix for `self.X` and `self.y`

        Parameters:
        - `sort`: Sort features by decreasing importance
        - `ax`: Matplotlib `Axes` object. Plot will be added to this object
        if provided; otherwise a new `Axes` object will be generated.
        - `importance_kwargs`: kwargs to pass to
        `sklearn.inspection.permutation_importance`
        - `plot_kwargs`: kwargs to pass to `pd.Series.plot.barh`
        """
        if importance_kwargs is None:
            importance_kwargs = {}
        importance = self.permutation_importance(**importance_kwargs)
        if sort:
            importance = importance.sort_values()

        if plot_kwargs is None:
            plot_kwargs = {}
        ax = importance.plot.barh(**plot_kwargs)
        ax.set(title="Feature importances")
        return ax

    @delegates(show_correlation)
    def show_correlation(self, **kwargs) -> Axes:
        """Show a correlation matrix for `self.X` and `self.y`"""
        return show_correlation(
            df=pd.concat((self.X, self.y), axis="columns"),
            **kwargs,
        )

# Cell
def get_inspector(model, X, y) -> _Inspector:
    """Get an appropriate inspector for your model and data

    Parameters:
    - `model`: Fitted sklearn model
    - `X`: Feature matrix with same shape and column meanings as
    features `model` was trained on
    - `y`: Target series with same length as `X` and same meaning as
    target `model` was trained on
    """
    model_type = identify_type(model, y)
    if isinstance(model, LinearModel):
        return _LinRegInspector(model, X, y)
    elif isinstance(model, LinearClassifierMixin):
        return (
            _LinBinInspector(model, X, y)
            if model_type == ModelType.BINARY
            else _LinMultiInspector(model, X, y)
        )
    elif isinstance(model, BaseDecisionTree):
        return (
            _TreeRegInspector(model, X, y)
            if model_type == ModelType.REGRESSION
            else _TreeBinInspector(model, X, y)
            if model_type == ModelType.BINARY
            else _TreeMultiInspector(model, X, y)
        )
    elif isinstance(model, BaseSearchCV):
        return _SearchInspector(model, X, y)
    elif model_type == ModelType.BINARY:
        return _BinClasInspector(model, X, y)
    elif model_type == ModelType.MULTICLASS:
        return _MultiClasInspector(model, X, y)
    elif model_type == ModelType.REGRESSION:
        return _RegInspector(model, X, y)
    else:
        raise NotImplementedError("Model not supported")

# Cell
# export
def identify_type(model, y):
    if isinstance(model, RegressorMixin):
        return ModelType.REGRESSION
    elif isinstance(model, ClassifierMixin):
        return ModelType.BINARY if len(y.unique()) == 2 else ModelType.MULTICLASS
    else:
        return ModelType.OTHER

# Cell
# export
class ModelType(Enum):
    REGRESSION = auto()
    BINARY = auto()
    MULTICLASS = auto()
    OTHER = auto()

# Cell
class _BinClasInspector(_Inspector):
    def calculate_metrics_by_thresh(
        self,
        metrics: Union[Callable, Sequence[Callable]],
    ) -> pd.DataFrame:
        """Calculate classification metrics as a function of threshold

        Assumes that `self.model` has a `.predict_proba()` method.

        Parameters:
        - `metrics`: Callables that take `y_true`, `y_pred` as
        positional arguments and return a number. Must have a `__name__`
        attribute and must be able to handle `np.nan` values.

        Returns: DataFrame with one column "thresh" indicating the
        thresholds used, which is 0 and the sorted set of values that
        occur in `y_prob`, and an additional column for each input
        metric giving the value of that metric at that threshold.
        """
        return calculate_metrics_by_thresh_binary(
            y_true=self.y,
            y_prob=self.model.predict_proba(self.X),
            metrics=metrics,
        )

    @delegates(sklearn.metrics.confusion_matrix)
    def confusion_matrix(
        self,
        thresh: float = 0.5,
        **kwargs,
    ) -> pd.DataFrame:
        """Get confusion matrix

        Assumes that `self.model` has a `.predict_proba()` method. Uses
        `self.y` as ground-truth values,
        `self.model.predict_proba(self.X)[:, 1] > thresh` as
        predictions.

        Parameters:
        - `thresh`: Probability threshold for counting a prediction as
        positive
        """
        return confusion_matrix(
            y_true=self.y,
            y_pred=self.model.predict_proba(self.X)[:, 1] > thresh,
            **kwargs,
        )

# Cell
class _MultiClasInspector(_Inspector):
    def calculate_metrics_by_thresh(
        self,
        metrics: Union[Callable, Sequence[Callable]],
    ) -> pd.DataFrame:
        """Calculate classification metrics as a function of threshold

        Assumes that `self.model` has a `.predict_proba()` method. Uses
        `self.y` as ground-truth values,
        `self.model.predict_proba(self.X)[:, 1] > thresh` as
        predictions.

        Parameters:
        - `metrics`: Callables that take `y_true`, `y_pred` as
        positional arguments and return a number. Must have a `__name__`
        attribute and must be able to handle `np.nan` values.

        Returns: DataFrame with one column "thresh" indicating the
        thresholds used, which is 0 and the sorted set of values that
        occur in `y_prob`, and an additional column for each input
        metric giving the value of that metric at that threshold.
        """
        return calculate_metrics_by_thresh_multi(
            y_true=self.y,
            y_prob=self.model.predict_proba(self.X),
            metrics=metrics,
        )

    @delegates(pd.DataFrame().style.background_gradient)
    def confusion_matrix(
        self,
        **kwargs,
    ) -> pd.DataFrame:
        """Get confusion matrix

        Uses `self.y` as ground-truth values,
        `self.model.predict(self.X)` as predictions.
        """

        return confusion_matrix(
            y_true=self.y,
            y_pred=self.model.predict(self.X),
            **kwargs,
        )

# Cell
class _RegInspector(_Inspector):
    def plot_pred_vs_act(
        self,
        ax: Optional[Axes] = None,
        scatter_kwargs: Optional[dict] = None,
        line_kwargs: Optional[dict] = None,
    ) -> Axes:
        """Plot predicted vs. actual values

        Parameters:
        - `ax`: Matplotlib `Axes` object. Plot will be added to this object
        if provided; otherwise a new `Axes` object will be generated.
        - `scatter_kwargs`: kwargs to pass to `plt.scatter`
        - `line_kwargs`: kwargs to pass to `plt.plot`
        """
        if ax is None:
            _, ax = plt.subplots()

        if scatter_kwargs is None:
            scatter_kwargs = {}
        scatter_kwargs = {**{"alpha": 0.3}, **scatter_kwargs}
        if "c" not in scatter_kwargs and "color" not in scatter_kwargs:
            scatter_kwargs["c"] = "k"
        y_pred = self.model.predict(self.X)
        ax.scatter(self.y, y_pred, **scatter_kwargs)

        if line_kwargs is None:
            line_kwargs = {}
        line_kwargs = {
            **{"label": "predicted=actual", "linestyle": "dashed"},
            **line_kwargs,
        }
        ax.plot(
            [self.y.min(), self.y.max()],
            [self.y.min(), self.y.max()],
            **line_kwargs,
        )

        ax.set(xlabel="Actual", ylabel="Predicted")
        ax.legend()
        return ax

    def plot_residuals(
        self,
        axes: Optional[np.array] = None,
        scatter_kwargs: Optional[dict] = None,
        line_kwargs: Optional[dict] = None,
        hist_kwargs: Optional[dict] = None,
    ) -> Axes:
        """Plot residuals

        Parameters:
        - `axes`: 1D array of two Matplotlib `Axes` objects. Plot will
        be added to these objects if provided; otherwise a new array of
        `Axes` objects will be generated.
        - `scatter_kwargs`: kwargs to pass to `plt.scatter`
        - `line_kwargs`: kwargs to pass to `plt.plot` for line at y=0
        - `hist_kwargs`: kwargs to pass to `plt.hist` for histogram of
        residuals
        """
        if axes is None:
            _, axes = plt.subplots(
                1, 2, gridspec_kw={"width_ratios": [4, 1]}, sharey=True
            )

        if scatter_kwargs is None:
            scatter_kwargs = {}
        scatter_kwargs = {**{"alpha": 0.3}, **scatter_kwargs}
        if "c" not in scatter_kwargs and "color" not in scatter_kwargs:
            scatter_kwargs["c"] = "k"
        axes[0].scatter(
            x=self.y.index, y=self.y - self.model.predict(self.X), **scatter_kwargs
        )

        if line_kwargs is None:
            line_kwargs = {}
        line_kwargs = {**{"linestyle": "dashed"}, **line_kwargs}
        axes[0].plot([self.y.index.min(), self.y.index.max()], [0, 0], **line_kwargs)
        axes[0].set(
            ylabel="actual - predicted", xlabel=self.y.index.name, title="Residuals"
        )

        if hist_kwargs is None:
            hist_kwargs = {}
        hist_kwargs = {**{"orientation": "horizontal", "color": "k"}, **hist_kwargs}
        axes[1].hist(self.y - self.model.predict(self.X), **hist_kwargs)
        return axes

# Cell
# export
def generate_model_html(
    intercept: float,
    coefs: Sequence[float],
    feature_names: Iterable[str],
    target_name: str,
    intercept_formatter: str = ".2f",
    coef_formatter: str = ".2f",
):
    if len(coefs) != len(feature_names):
        raise ValueError("len(coefs) != len(feature_cols)")
    model_string = f"""
        <span style='color:{COLORS["pink"]}'>{target_name}</span>
        = <span style='color:{COLORS["orange"]}'>{intercept:{intercept_formatter}}</span>
    """  # noqa: E501
    for coef, feature_col in zip(coefs, feature_names):
        model_string += f"""
            <span style='color:{COLORS["green"]}'>{"+" if coef >= 0 else "-"} {abs(coef):{coef_formatter}}</span>
            * <span style='color:{COLORS["blue"]}'>{feature_col}</span>
        """  # noqa: E501
    return model_string

# Cell
def _plot_waterfall(
    X,
    y,
    item: Union[pd.Series, np.array],
    intercept: float,
    coefs: Sequence[float],
    y_lab="",
    bar_num_formatter: str = ".1f",
    tick_num_formatter: str = ".2f",
    sorted_value=True,
    threshold=0.01,
    blue_color=COLORS["blue"],
    green_color=COLORS["green"],
    red_color=COLORS["orange"],
    **waterfall_kwargs,
):
    index = ["int"] + [
        f"{name}: {val:{tick_num_formatter}}" for name, val in zip(X.columns, item)
    ]
    vals = [intercept] + list(np.array(coefs) * item)
    waterfall_chart.plot(
        index=index,
        data=vals,
        x_lab="Feature name and value",
        y_lab=y_lab,
        formatting=f"{{:,{bar_num_formatter}}}",
        net_label=y.name,
        sorted_value=sorted_value,
        threshold=threshold,
        blue_color=blue_color,
        green_color=green_color,
        red_color=red_color,
        **waterfall_kwargs,
    )
    return plt.gca()

# Cell
class _LinRegInspector(_RegInspector):
    """Linear regression model inspector"""

    def plot_coefs_vs_hparam(self, hparam: str, vals: Sequence[float]):
        """Plot coefficient values against a hyperparameter

        Parameters:
        - `hparam`: Name of hyperparameter; must be an attribute of
        `self.model`
        - `vals`: Values of that hyperparameter to use
        """
        current_val = getattr(self.model, hparam)
        model = clone(self.model)
        setattr(model, hparam, vals[-1])
        model.fit(self.X, self.y)
        column_order = model.coef_.argsort()[::-1]
        X = self.X.iloc[:, column_order]

        coefs = []
        for val in vals:
            setattr(model, hparam, val)
            coefs.append(model.fit(X, self.y).coef_)

        fig, ax = plt.subplots()
        ax.plot(vals, coefs)
        ax.plot([vals.min(), vals.max()], [0, 0], linestyle="dotted", c="k")
        ax.axvline(current_val, c="k", label="current value")
        ax.set(xlabel=hparam, ylabel="coefficient value")
        ax.legend(X.columns, bbox_to_anchor=(1.05, 1.0), loc="upper left")
        return ax

    @delegates(waterfall_chart.plot)
    def plot_waterfall(
        self,
        item: Union[pd.Series, np.array],
        bar_num_formatter: str = ".1f",
        tick_num_formatter: str = ".2f",
        sorted_value=True,
        threshold=0.01,
        blue_color=COLORS["blue"],
        green_color=COLORS["green"],
        red_color=COLORS["orange"],
        **kwargs,
    ):
        """Make a waterfall chart showing how each feature contributes
        to the prediction for the input item.

        Parameters:
        - `item`: Input item, with the same shape and value meanings as
        a single row from `self.X`
        - `bar_num_formatter`: Bar label format specifier
        - `tick_num_formatter`: Tick label format specifier

        Additional keyword arguments will be passed to
        `waterfall_chart.plot`
        """
        return _plot_waterfall(
            X=self.X,
            y=self.y,
            item=item,
            intercept=self.model.intercept_,
            coefs=self.model.coef_,
            y_lab="Contribution to prediction",
            bar_num_formatter=bar_num_formatter,
            tick_num_formatter=tick_num_formatter,
            sorted_value=sorted_value,
            threshold=threshold,
            blue_color=blue_color,
            green_color=green_color,
            red_color=red_color,
            **kwargs,
        )

    def show_model(
        self,
        intercept_formatter: str = ".2f",
        coef_formatter: str = ".2f",
    ):
        """Show model equation

        Parameters:
        - `intercept_formatter`: Intercept format specifier
        - `coef_formatter`: Intercept format specifier
        """
        return HTML(
            generate_model_html(
                intercept=self.model.intercept_,
                coefs=self.model.coef_,
                feature_names=self.X.columns,
                target_name=self.y.name,
                intercept_formatter=intercept_formatter,
                coef_formatter=coef_formatter,
            )
        )

# Cell
class _LinBinInspector(_BinClasInspector):
    """Linear binary classification model inspector"""

    def plot_coefs_vs_hparam(self, hparam: str, vals: Sequence[float]) -> np.array:
        """Plot coefficient values against a hyperparameter

        Parameters:
        - `hparam`: Name of hyperparameter; must be an attribute of
        `self.model`
        - `vals`: Values of that hyperparameter to use
        """
        current_val = getattr(self.model, hparam)
        model = clone(self.model)
        setattr(model, hparam, vals[-1])
        column_order = model.fit(self.X, self.y).coef_[0].argsort()[::-1]
        X = self.X.iloc[:, column_order]

        coef_arrays = []
        for val in vals:
            setattr(model, hparam, val)
            coef_arrays.append(model.fit(self.X, self.y).coef_)

        fig, ax = plt.subplots()
        ax.plot(vals, [coefs[0] for coefs in coef_arrays])
        ax.plot([vals.min(), vals.max()], [0, 0], linestyle="dotted", c="k")
        ax.set(xlabel=hparam, ylabel="Coefficient Value")
        ax.axvline(current_val, c="k", label="current value")
        ax.legend(X.columns, bbox_to_anchor=(1.05, 1.0), loc="upper left")
        return ax

    def show_model(
        self,
        intercept_formatter: str = ".2f",
        coef_formatter: str = ".2f",
    ):
        """Show model equation

        Parameters:
        - `intercept_formatter`: Intercept format specifier
        - `coef_formatter`: Intercept format specifier
        """
        return HTML(
            generate_model_html(
                intercept=self.model.intercept_[0],
                coefs=self.model.coef_[0],
                feature_names=self.X.columns,
                target_name=f"log-odds({self.y.name})",
                intercept_formatter=intercept_formatter,
                coef_formatter=coef_formatter,
            )
        )

    @delegates(waterfall_chart.plot)
    def plot_waterfall(
        self,
        item: Union[pd.Series, np.array],
        bar_num_formatter: str = ".1f",
        tick_num_formatter: str = ".2f",
        sorted_value=True,
        threshold=0.01,
        blue_color=COLORS["blue"],
        green_color=COLORS["green"],
        red_color=COLORS["orange"],
        **kwargs,
    ):
        """Make a waterfall chart showing how each feature contributes
        to the prediction for the input item for a binary classification
        model.

        Parameters:
        - `item`: Input item, with the same shape and value meanings as
        a single row from `self.X`
        - `bar_num_formatter`: Bar label format specifier
        - `tick_num_formatter`: Tick label format specifier
        - ``waterfall_kwargs`: kwargs to pass to `waterfall_chart.plot`
        """
        return _plot_waterfall(
            X=self.X,
            y=self.y,
            item=item,
            intercept=self.model.intercept_[0],
            coefs=self.model.coef_[0],
            y_lab="Contribution to predicted log-odds",
            bar_num_formatter=bar_num_formatter,
            tick_num_formatter=tick_num_formatter,
            sorted_value=sorted_value,
            threshold=threshold,
            blue_color=blue_color,
            green_color=green_color,
            red_color=red_color,
            **kwargs,
        )

# Cell
class _LinMultiInspector(_MultiClasInspector):
    """Linear multiclass classification model inspector"""

    def plot_coefs_vs_hparam(self, hparam: str, vals: Sequence[float]) -> np.array:
        """Plot coefficient values against a hyperparameter

        Parameters:
        - `hparam`: Name of hyperparameter; must be an attribute of
        `self.model`
        - `vals`: Values of that hyperparameter to use
        """
        current_val = getattr(self.model, hparam)
        model = clone(self.model)
        setattr(model, hparam, vals[-1])
        column_order = model.fit(self.X, self.y).coef_[0].argsort()[::-1]
        X = self.X.iloc[:, column_order]

        coef_arrays = []
        for val in vals:
            setattr(model, hparam, val)
            coef_arrays.append(model.fit(self.X, self.y).coef_)

        num_target_vals = len(set(self.y))

        fig, axes = plt.subplots(
            num_target_vals, 1, sharex=True, sharey=True, constrained_layout=True
        )
        for target_val_num in range(num_target_vals):
            axes[target_val_num].plot(
                vals, [coefs[target_val_num] for coefs in coef_arrays]
            )
            axes[target_val_num].set_title(f"y={sorted(set(self.y))[target_val_num]}")
        axes[0].set(xlabel=hparam, ylabel="Coefficient Value")
        for ax in axes:
            ax.axvline(current_val, c="k", label="current value")
            ax.plot([vals.min(), vals.max()], [0, 0], linestyle="dotted", c="k")
        axes[0].legend(X.columns, bbox_to_anchor=(1.05, 1.0), loc="upper left")
        return axes

    def show_model(
        self,
        intercept_formatter: str = ".2f",
        coef_formatter: str = ".2f",
    ):
        """Show model equation

        Parameters:
        - `intercept_formatter`: Intercept format specifier
        - `coef_formatter`: Intercept format specifier
        """
        model_string = ""
        for target_name, coefs, intercept in zip(
            np.unique(self.y), self.model.coef_, self.model.intercept_
        ):
            model_string += f"""
                    <p>
                        {generate_model_html(
                                intercept=intercept,
                                coefs=coefs,
                                feature_names=self.X.columns,
                                target_name=f"log-odds({self.y.name} = {target_name})",
                                intercept_formatter=intercept_formatter,
                                coef_formatter=coef_formatter,
                            )
                        }
                    </p>
                """
        return HTML(model_string)

# Cell
class _TreeMixin(_Inspector):
    """Mixin for decision tree model inspectors"""

    @delegates(plot_tree)
    def show_model(self, ax: Optional[Axes] = None, **kwargs):
        """Show decision tree"""
        if ax is None:
            # these dimensions seem to work well empirically
            max_size = 50
            depth = (
                kwargs["max_depth"] + 1
                if "max_depth" in kwargs
                else self.model.get_depth()
            )
            fig_height = min(depth * 2.2, max_size)
            width = (
                2 * kwargs["max_depth"]
                if "max_depth" in kwargs
                else self.model.get_n_leaves()
            )
            fig_width = min(width * 3.5, max_size)
            _, ax = plt.subplots(figsize=(fig_width, fig_height))
        kwargs = {"filled": True, "fontsize": 12, **kwargs}
        return plot_tree(
            self.model,
            feature_names=self.X.columns,
            class_names=np.unique(self.y).astype(str),
            ax=ax,
            **kwargs,
        )[0].axes

# Cell
class _TreeBinInspector(_BinClasInspector, _TreeMixin):
    """Inspector for binary decision tree models"""

    pass

# Cell
class _TreeMultiInspector(_MultiClasInspector, _TreeMixin):
    """Inspector for multiclass decision tree models"""

    pass

# Cell
class _TreeRegInspector(_RegInspector, _TreeMixin):
    """Inspector for decision tree regressors models"""

    pass

# Cell
class _SearchInspector(_Inspector):
    """Inspector for `BaseSearchCV` models with one hyperparameter"""

    def plot_scores_vs_hparam(
        self,
        hparam: Optional[str] = None,
        score_cols: Optional[Union[str, List[str]]] = None,
        ax=None,
    ):
        """Plot model scores against values of one hyperparameter

        Parameters:
        - `hparam`: Name of the hyperparameter to plot against. Must be
        provided if there are multiple hyperparameters. Any other
        hyperparameters will be fixed at the value they have in
        `self.model.best_params_`.
        - `score_cols`: Name of score columns to plot. By default will
        be the mean test and (if present) train score for the primary
        scoring metric.
        - `ax`: Matplotlib `Axes` object. Plot will be added to this object
        if provided; otherwise a new `Axes` object will be generated.
        """

        def _get_hparam():
            hparams = list(self.model.param_grid.keys())
            if len(hparams) == 1:
                return hparams[0]
            else:
                raise ValueError(
                    "Must provide `hparam` if there are multiple possibilities"
                )

        def _filter_to_best_values_of_other_hparams(results, hparam):
            return results.query(_get_query_string(hparam))

        def _get_query_string(hparam):
            return " & ".join(
                f"param_{param} == {repr(val)}"
                for param, val in self.model.best_params_.items()
                if param != hparam
            )

        def _get_default_score_cols(results):
            score_cols = []
            if self.model.refit is True:
                score_cols.append("mean_test_score")
            else:
                score_cols.append(f"mean_test_{self.model.refit}")
            train_cols = [
                col.replace("test", "train")
                for col in score_cols
                if col.replace("test", "train") in results
            ]
            return score_cols + train_cols

        def _plot_scores(results, hparam, score_cols, ax):
            if isinstance(self.model.best_params_[hparam], str):
                results.plot.bar(x=f"param_{hparam}", y=score_cols, ax=ax)

            else:
                for col in score_cols:
                    results.plot(x=f"param_{hparam}", y=col, ax=ax)
                ax.axvline(
                    self.model.best_params_[hparam],
                    c="k",
                    label=f"{hparam} = {self.model.best_params_[hparam]}",
                )

        if hparam is None:
            hparam = _get_hparam()

        results = pd.DataFrame(self.model.cv_results_)
        if len(self.model.param_grid) > 1:
            results = _filter_to_best_values_of_other_hparams(results, hparam)

        if score_cols is None:
            score_cols = _get_default_score_cols(results)

        if ax is None:
            _, ax = plt.subplots()

        _plot_scores(results, hparam, score_cols, ax)
        ax.set_title(_get_query_string(hparam).replace("==", "="))
        ax.legend()
        return ax

    @delegates(pd.DataFrame().style.background_gradient)
    def show_score_vs_hparam_pair(self, hparams=None, score_col=None, **kwargs):
        """Show model scores against a pair of hyperparameters

        Parameters:
        - `hparams`: Name of the hyperparameters to plot against.
        The first two hyperparameters in `self.model.param_grid` will be
        used by default. Any other hyperparameters will be fixed at the
        value they have in `self.model.best_params_`.
        - `score_col`: Name of score column to plot. By default will
        be the mean test score for the primary scoring metric.
        """
        score_col = (
            score_col
            if score_col is not None
            else "mean_test_score"
            if self.model.refit is True
            else f"mean_test_{self.model.refit}"
        )

        all_hparams = list(self.model.best_params_.keys())
        assert len(all_hparams) > 1, "Method requires at least two hyperparameters"
        if hparams is None:
            if len(all_hparams) == 2:
                hparams = all_hparams
            else:
                raise ValueError(
                    """
                        `hparams` must be specified unless there are
                        exactly two hyperparameters
                    """
                )

        results = pd.DataFrame(self.model.cv_results_)
        return (
            results.pivot_table(
                index=f"param_{hparams[0]}", columns=f"param_{hparams[1]}"
            )
            .loc[:, score_col]
            .style.background_gradient(**kwargs)
        )

# Cell
class _Plotter:
    def __init__(self, model, X, y):
        store_attr()

    __repr__ = basic_repr(["model"])

# Cell
class _1dPlotter(_Plotter):
    def plot(
        self,
        plot_data: bool = True,
        ax: Optional[Axes] = None,
        line_kwargs: Optional[dict] = None,
        scatter_kwargs: Optional[dict] = None,
    ) -> Axes:
        """Plot predictions from a regression or multiclass model with a
        single input as a line

        Parameters:
        - `plot_data`: Make a scatter plot of the data
        - `line_kwargs`: kwargs to pass to `ax.plot` for plotting
        predictions
        - `scatter_kwargs`: kwargs to pass to `ax.scatter` for plotting data
        - `ax`: Matplotlib `Axes` object. Plot will be added to this object
        if provided; otherwise a new `Axes` object will be generated.
        """

        def _plot_preds(ax, line_kwargs):
            X_sorted = self.X.sort_values(self.X.columns[0])
            ax.plot(
                X_sorted.iloc[:, 0],
                self.model.predict(X_sorted),
                **line_kwargs,
            )
            return ax

        if ax is None:
            _, ax = plt.subplots()
        if plot_data:
            if scatter_kwargs is None:
                scatter_kwargs = {}
            scatter_kwargs = {**{"alpha": 0.3}, **scatter_kwargs}
            if "c" not in scatter_kwargs and "color" not in scatter_kwargs:
                scatter_kwargs["c"] = "k"
            ax.scatter(self.X.iloc[:, 0], self.y, **scatter_kwargs)

        if line_kwargs is None:
            line_kwargs = {}
        line_kwargs = {**{"label": "predictions"}, **line_kwargs}
        ax = _plot_preds(ax, line_kwargs)

        ax.set(xlabel=self.X.columns[0], ylabel=self.y.name)
        ax.legend()
        return ax

# Cell
class _Bin1dPlotter(_Plotter):
    def plot(
        self,
        thresh: float = 0.5,
        plot_data: bool = True,
        ax: Optional[Axes] = None,
        prob_line_kwargs: Optional[dict] = None,
        thresh_line_kwargs: Optional[dict] = None,
        scatter_kwargs: Optional[dict] = None,
        scatter_kwargs_correct: Optional[dict] = None,
        scatter_kwargs_incorrect: Optional[dict] = None,
    ) -> Axes:
        """Plot predictions from a binary classification model that
        provides probabilities and has a single input

        Parameters:
        - `thresh`: Probability threshold for counting a prediction as
        positive
        - `plot_data`: Make a scatter plot of the data
        - `ax`: Matplotlib `Axes` object. Plot will be added to this
        object if provided; otherwise a new `Axes` object will be
        generated.
        - `prob_line_kwargs`: kwargs to pass to `ax.plot` for plotting
        model probabilities
        - `thresh_line_kwargs`: kwargs to pass to `ax.plot` for plotting
        threshold
        - `scatter_kwargs`: kwargs to pass to `ax.scatter` for plotting
        all data points
        - `scatter_kwargs_correct`: kwargs to pass to `ax.scatter` for
        plotting data points that the model predicted correctly.
        Overrides `scatter_kwargs`.
        - `scatter_kwargs_incorrect`: kwargs to pass to `ax.scatter` for
        plotting data points that the model predicted incorrectly.
        Overrides `scatter_kwargs`.
        """

        def _set_scatter_kwargs(
            scatter_kwargs, scatter_kwargs_correct, scatter_kwargs_incorrect
        ):
            if scatter_kwargs is None:
                scatter_kwargs = {}
            scatter_kwargs = {**{"alpha": 0.3}, **scatter_kwargs}

            if scatter_kwargs_correct is None:
                scatter_kwargs_correct = {}
            scatter_kwargs_correct = {
                **{"label": "correct"},
                **scatter_kwargs,
                **scatter_kwargs_correct,
            }
            if scatter_kwargs_incorrect is None:
                scatter_kwargs_incorrect = {}
            scatter_kwargs_incorrect = {
                **{"label": "incorrect"},
                **scatter_kwargs,
                **scatter_kwargs_correct,
            }
            if (
                "c" not in scatter_kwargs_correct
                and "color" not in scatter_kwargs_correct
            ):
                scatter_kwargs_correct["c"] = "b"
            if (
                "c" not in scatter_kwargs_incorrect
                and "color" not in scatter_kwargs_incorrect
            ):
                scatter_kwargs_incorrect["c"] = "orange"

            return scatter_kwargs_correct, scatter_kwargs_incorrect

        def _plot_probs(ax):
            num_points = 100
            X = np.linspace(self.X.min(), self.X.max(), num_points)
            ax.plot(
                X,
                self.model.predict_proba(X)[:, 1],
                **prob_line_kwargs,
            )
            return ax

        if ax is None:
            _, ax = plt.subplots()

        y_vals = np.unique(self.y)
        y_numeric = pd.Series(np.where(self.y == y_vals[0], 0, 1), index=self.y.index)

        if plot_data:
            scatter_kwargs_correct, scatter_kwargs_incorrect = _set_scatter_kwargs(
                scatter_kwargs=scatter_kwargs,
                scatter_kwargs_correct=scatter_kwargs_correct,
                scatter_kwargs_incorrect=scatter_kwargs_incorrect,
            )
            is_correct = y_numeric == (self.model.predict_proba(self.X)[:, 1] > thresh)
            ax.scatter(
                self.X.loc[is_correct].iloc[:, 0],
                y_numeric.loc[is_correct],
                **scatter_kwargs_correct,
            )
            ax.scatter(
                self.X.loc[~is_correct].iloc[:, 0],
                y_numeric.loc[~is_correct],
                **scatter_kwargs_incorrect,
            )

        if prob_line_kwargs is None:
            prob_line_kwargs = {}
        prob_line_kwargs = {**{"label": "probability"}, **prob_line_kwargs}
        ax = _plot_probs(ax)

        if thresh_line_kwargs is None:
            thresh_line_kwargs = {}
        thresh_line_kwargs = {
            **{"label": f"threshold={thresh:.2f}"},
            **thresh_line_kwargs,
        }
        if "c" not in thresh_line_kwargs and "color" not in thresh_line_kwargs:
            thresh_line_kwargs["c"] = "k"

        ax.plot(
            self.X.iloc[:, 0],
            thresh * np.ones(self.X.shape),
            **thresh_line_kwargs,
        )

        ax.set(xlabel=self.X.columns[0], ylabel=self.y.name)
        ax.legend()
        return ax

# Cell
class _2dPlotter(_Plotter):
    def _plot_data(self, ax, y=None, **scatter_kwargs):
        if y is None:
            y = self.y
        X_normalized = MinMaxScaler().fit_transform(self.X) * 99
        ax.scatter(
            X_normalized[:, 0] + 0.5,
            X_normalized[:, 1].max() - X_normalized[:, 1] + 0.5,
            c=y,
            **scatter_kwargs,
        )
        ax.set(xlabel=self.X.columns[0], ylabel=self.X.columns[1])
        return ax

    def _create_grid(self, num_points=20):
        x0_grid = np.linspace(
            self.X.iloc[:, 0].min(), self.X.iloc[:, 0].max(), num_points
        )
        x1_grid = np.linspace(
            self.X.iloc[:, 1].min(), self.X.iloc[:, 1].max(), num_points
        )
        return np.meshgrid(x0_grid, x1_grid)

    def _get_grid_preds(self, x0_grid, x1_grid):
        return self.model.predict(
            np.hstack((x0_grid.reshape(-1, 1), x1_grid.reshape(-1, 1)))
        ).reshape(x0_grid.shape)

# Cell
class _Reg2dPlotter(_2dPlotter):
    def plot(
        self,
        plot_data: bool = True,
        tick_formatter: Optional[str] = ".2f",
        ax: Axes = None,
        heatmap_kwargs: Optional[dict] = None,
        scatter_kwargs: Optional[dict] = None,
    ):
        """Plot predictions from a model with two inputs as a heatmap.

        Parameters:
        - `plot_data`: Make a scatter plot of the data
        - `tick_formatter`: Tick label format specifier
        - `ax`: Matplotlib `Axes` object. Plot will be added to this object
        if provided; otherwise a new `Axes` object will be generated.
        - `heatmap_kwargs`: kwargs to pass to `sns.heatmap` for plotting
        predictions
        - `scatter_kwargs`: kwargs to pass to `ax.scatter` for plotting data
        """

        def _plot_preds(ax, **heatmap_kwargs):
            x_grid = np.linspace(self.X.iloc[:, 0].min(), self.X.iloc[:, 0].max(), 100)
            y_grid = np.linspace(self.X.iloc[:, 1].max(), self.X.iloc[:, 1].min(), 100)

            preds = self.model.predict(
                np.transpose(
                    [np.tile(x_grid, len(y_grid)), np.repeat(y_grid, len(x_grid))]
                )
            ).reshape(len(y_grid), len(x_grid))
            preds = pd.DataFrame(preds, columns=x_grid, index=y_grid)
            return sns.heatmap(
                preds,
                vmin=self.y.min(),
                vmax=self.y.max(),
                ax=ax,
                **heatmap_kwargs,
            )

        if ax is None:
            _, ax = plt.subplots()
        if heatmap_kwargs is None:
            heatmap_kwargs = {}
        heatmap_kwargs = {**{"cmap": "viridis"}, **heatmap_kwargs}
        if scatter_kwargs is None:
            scatter_kwargs = {}
        scatter_kwargs = {
            **{"cmap": "viridis", "edgecolor": "k", "zorder": 999},
            **scatter_kwargs,
        }

        if plot_data:
            ax = self._plot_data(ax=ax, **scatter_kwargs)
        ax = _plot_preds(ax=ax, **heatmap_kwargs)
        if tick_formatter is not None:
            _format_ticks(ax=ax, formatter=tick_formatter)
        ax.set(
            xlabel=self.X.columns[0],
            ylabel=self.X.columns[1],
            title=self.y.name,
        )
        return ax

    def plot3d(
        self,
        plot_data: bool = True,
        ax: Axes = None,
        surf_kwargs: Optional[dict] = None,
        scatter_kwargs: Optional[dict] = None,
    ):
        """Plot data and predictions in 3d

        Best viewed with a tool such as https://github.com/matplotlib/ipympl
        that supports rotating the output

        Parameters:
        - `plot_data`: Make a scatter plot of the data
        - `ax`: Matplotlib `Axes` object. Plot will be added to this object
        if provided; otherwise a new `Axes` object will be generated.
        - `surf_kwargs`: kwargs to pass to `ax.plot_surface` for plotting
        predictions
        - `scatter_kwargs`: kwargs to pass to `ax.scatter` for plotting data
        """

        def _plot_preds(ax):

            x0_grid, x1_grid = self._create_grid()
            ax.plot_surface(
                x0_grid,
                x1_grid,
                self._get_grid_preds(x0_grid, x1_grid),
                rstride=1,
                cstride=1,
                vmin=self.y.min(),
                vmax=self.y.max(),
                **surf_kwargs,
            )
            return ax

        if ax is None:
            fig = plt.figure(figsize=(8, 8))
            ax = fig.add_subplot(111, projection="3d")
        if surf_kwargs is None:
            surf_kwargs = {}
        surf_kwargs = {**{"alpha": 0.4, "cmap": "viridis"}, **surf_kwargs}
        if scatter_kwargs is None:
            scatter_kwargs = {}
        scatter_kwargs = {**{"cmap": "viridis"}, **scatter_kwargs}
        if plot_data:
            ax.scatter(
                self.X.iloc[:, 0],
                self.X.iloc[:, 1],
                self.y,
                c=self.y,
                **scatter_kwargs,
            )
        ax = _plot_preds(ax)
        ax.set(
            xlabel=self.X.columns[0],
            ylabel=self.X.columns[1],
            zlabel=self.y.name,
        )
        return ax

# Cell
class _2dMultiPlotterMixin:
    def plot_components(self, axes=None, **kwargs):
        """Plot components estimators

        Parameters:
        - `axes`: NumPy array of Matplotlib `Axes` objects
        - `kwargs`: kwargs to pass to `self.plot`
        """
        if axes is None:
            _, axes = plt.subplots(
                3, 3, constrained_layout=True, figsize=(8, 6), sharex=True, sharey=True
            )
        if kwargs is None:
            kwargs = {}
        kwargs = {**{"plot_data": False}, **kwargs}

        for row_num, row in enumerate(axes):
            for col_num, col in enumerate(row):
                estimator_num = col_num + len(row) * row_num
                if estimator_num < len(self.model.estimators_):
                    get_inspector(
                        self.model.estimators_[estimator_num], self.X, self.y
                    ).plot(ax=axes[row_num, col_num], **kwargs)
                    axes[row_num, col_num].set(title=None)
        return axes

    def plot_components3d(self, axes=None, **kwargs):
        """Plot components estimators in 3D

        Parameters:
        - `axes`: NumPy array of Matplotlib `Axes` objects
        - `kwargs`: kwargs to pass to `self.plot3d`
        """
        if axes is None:
            fig = plt.figure(constrained_layout=True, figsize=(8, 6))
            ax1 = fig.add_subplot(331, projection="3d")
            for ax_ix in range(2, 10):
                fig.add_subplot(
                    330 + ax_ix, projection="3d", sharez=ax1, sharex=ax1, sharey=ax1
                )
            axes = np.array(fig.axes).reshape(3, 3)
        if kwargs is None:
            kwargs = {}
        kwargs = {**{"plot_data": False}, **kwargs}

        for row_num, row in enumerate(axes):
            for col_num, col in enumerate(row):
                estimator_num = col_num + len(row) * row_num
                if estimator_num < len(self.model.estimators_):
                    get_inspector(
                        self.model.estimators_[estimator_num], self.X, self.y
                    ).plot3d(ax=axes[row_num, col_num], **kwargs)
                    axes[row_num, col_num].set(title=None)
        return axes

# Cell
class _Reg2dMultiPlotter(_Reg2dPlotter, _2dMultiPlotterMixin):
    pass

# Cell
class _Bin2dPlotter(_2dPlotter):
    def plot(
        self,
        plot_data: bool = True,
        tick_formatter: Optional[str] = ".2f",
        ax: Optional[Axes] = None,
        heatmap_kwargs: Optional[dict] = None,
        scatter_kwargs: Optional[dict] = None,
    ):
        """Plot data and predictions

        Parameters:
        - `plot_data`: Make a scatter plot of the data
        - `tick_formatter`: Tick label format specifier
        - `ax`: Matplotlib `Axes` object. Plot will be added to this object
        if provided; otherwise a new `Axes` object will be generated.
        - `heatmap_kwargs`: kwargs to pass to `sns.heatmap` for plotting
        predictions
        - `scatter_kwargs`: kwargs to pass to `ax.scatter` for plotting data
        """

        def _plot_preds(y_vals, label_to_num, ax, **scatter_kwargs):
            num_points = 100
            x_grid = np.linspace(
                self.X.iloc[:, 0].min(), self.X.iloc[:, 0].max(), num_points
            )
            y_grid = np.linspace(
                self.X.iloc[:, 1].max(), self.X.iloc[:, 1].min(), num_points
            )

            preds = self.model.predict_proba(
                np.transpose(
                    [np.tile(x_grid, len(y_grid)), np.repeat(y_grid, len(x_grid))]
                )
            )[:, 1].reshape(len(y_grid), len(x_grid))
            preds = pd.DataFrame(preds, columns=x_grid, index=y_grid)
            sns.heatmap(preds, **heatmap_kwargs, ax=ax)
            return ax

        def _wash_out(ax):
            rectangle = plt.Rectangle((0, 0), 100, 100, fc="w", alpha=0.5)
            ax.add_patch(rectangle)
            return ax

        if ax is None:
            _, ax = plt.subplots()

        y_vals = np.unique(self.y)
        label_to_num = {label: num for label, num in zip(y_vals, range(len(y_vals)))}

        if heatmap_kwargs is None:
            heatmap_kwargs = {}
        heatmap_kwargs = {
            **{"cmap": sns.diverging_palette(10, 220, n=99), "vmin": 0, "vmax": 1},
            **heatmap_kwargs,
        }
        _plot_preds(y_vals, label_to_num, ax=ax, **heatmap_kwargs)

        if plot_data:
            if scatter_kwargs is None:
                scatter_kwargs = {}
            scatter_kwargs = {
                **{
                    "cmap": ax.collections[0].colorbar.cmap,
                    "edgecolor": "k",
                    "zorder": 999,
                },
                **scatter_kwargs,
            }
            self._plot_data(y=self.y.map(label_to_num), ax=ax, **scatter_kwargs)
        _format_ticks(ax=ax, formatter=tick_formatter)
        return ax

    def plot3d(
        self,
        thresh=0.5,
        plot_prob: bool = True,
        plot_thresh: bool = True,
        plot_data: bool = True,
        ax: Axes = None,
        prob_surf_kwargs: Optional[dict] = None,
        thresh_surf_kwargs: Optional[dict] = None,
        scatter_kwargs: Optional[dict] = None,
    ):
        """Plot data and predictions in 3D

        Best viewed with a tool such as https://github.com/matplotlib/ipympl
        that supports rotating the output

        Parameters:
        - `thresh`: Probability threshold for counting a prediction as
        positive
        - `plot_prob`: Whether to plot the model probabilities
        - `plot_thresh`: Whether to plot a classification threshold
        - `plot_data`: Whether to plot the data
        - `ax`: Matplotlib `Axes` object. Plot will be added to this object
        if provided; otherwise a new `Axes` object will be generated.
        - `prob_surf_kwargs`: kwargs to pass to the model probability
        surface
        - `thresh_surf_kwargs`: kwargs to pass to the threshold surface
        - `scatter_kwargs`: kwargs to pass to the scatter plot of the data
        """

        def _get_grid_probs():
            return self.model.predict_proba(
                np.hstack((x0_grid.reshape(-1, 1), x1_grid.reshape(-1, 1)))
            )[:, 1].reshape(x0_grid.shape)

        if ax is None:
            fig = plt.figure(figsize=(8, 8))
            ax = fig.add_subplot(111, projection="3d")
        if prob_surf_kwargs is None:
            prob_surf_kwargs = {}
        prob_surf_kwargs = {**{"alpha": 0.4, "cmap": "bwr_r"}, **prob_surf_kwargs}
        if thresh_surf_kwargs is None:
            thresh_surf_kwargs = {}
        thresh_surf_kwargs = {**{"alpha": 0.4, "color": "k"}, **thresh_surf_kwargs}
        if scatter_kwargs is None:
            scatter_kwargs = {}

        x0_grid, x1_grid = self._create_grid()

        if plot_prob:
            ax.plot_surface(
                x0_grid,
                x1_grid,
                _get_grid_probs(),
                rstride=1,
                cstride=1,
                **prob_surf_kwargs,
            )
        if plot_data:
            y_vals = np.unique(self.y)
            y_numeric = pd.Series(
                np.where(self.y == y_vals[0], 0, 1), index=self.y.index
            )
            y_pred = pd.Series(
                self.model.predict_proba(self.X)[:, 1] > thresh, index=self.y.index
            )
            ax.scatter(
                self.X.loc[y_numeric == 1].iloc[:, 0],
                self.X.loc[y_numeric == 1].iloc[:, 1],
                y_pred.loc[y_numeric == 1],
                c="b",
                label="positive",
                **scatter_kwargs,
            )
            ax.scatter(
                self.X.loc[y_numeric == 0].iloc[:, 0],
                self.X.loc[y_numeric == 0].iloc[:, 1],
                y_pred.loc[y_numeric == 0],
                c="r",
                label="negative",
                **scatter_kwargs,
            )
            ax.legend()
        if plot_thresh:
            ax.plot_surface(
                x0_grid,
                x1_grid,
                thresh * np.ones((len(x0_grid), len(x1_grid))),
                rstride=1,
                cstride=1,
                **thresh_surf_kwargs,
            )
        ax.set(
            xlabel=self.X.columns[0],
            ylabel=self.X.columns[1],
            zlabel=f"{self.y.name} prediction",
        )
        return ax

# Cell
class _Bin2dMultiPlotter(_Bin2dPlotter, _2dMultiPlotterMixin):
    pass

# Cell
class _Multi2dPlotter(_2dPlotter):
    def plot(
        self,
        plot_data: bool = True,
        tick_formatter: Optional[str] = ".2f",
        ax: Optional[Axes] = None,
        heatmap_kwargs: Optional[dict] = None,
        scatter_kwargs: Optional[dict] = None,
    ):
        """Plot data and predictions

        Parameters:
        - `plot_data`: Make a scatter plot of the data
        - `tick_formatter`: Tick label format specifier
        - `ax`: Matplotlib `Axes` object. Plot will be added to this object
        if provided; otherwise a new `Axes` object will be generated.
        - `heatmap_kwargs`: kwargs to pass to `sns.heatmap` for plotting
        predictions
        - `scatter_kwargs`: kwargs to pass to `ax.scatter` for plotting data
        """

        def _plot_preds(y_vals, label_to_num, ax, **scatter_kwargs):
            num_points = 100
            x_grid = np.linspace(
                self.X.iloc[:, 0].min(), self.X.iloc[:, 0].max(), num_points
            )
            y_grid = np.linspace(
                self.X.iloc[:, 1].max(), self.X.iloc[:, 1].min(), num_points
            )

            preds = self.model.predict(
                np.transpose(
                    [np.tile(x_grid, len(y_grid)), np.repeat(y_grid, len(x_grid))]
                )
            ).reshape(len(y_grid), len(x_grid))
            preds = pd.DataFrame(preds, columns=x_grid, index=y_grid)
            for col in preds:
                preds.loc[:, col] = preds.loc[:, col].map(label_to_num)
            sns.heatmap(preds.astype(int), **heatmap_kwargs, ax=ax)
            return ax

        def _set_colorbar(y_vals, ax):
            colorbar = ax.collections[0].colorbar
            r = colorbar.vmax - colorbar.vmin
            colorbar.set_ticks(
                [
                    colorbar.vmin + r / len(y_vals) * (0.5 + i)
                    for i in range(len(y_vals))
                ]
            )
            colorbar.set_ticklabels(y_vals)
            return colorbar

        def _wash_out(ax):
            rectangle = plt.Rectangle((0, 0), 100, 100, fc="w", alpha=0.5)
            ax.add_patch(rectangle)
            return ax

        if ax is None:
            _, ax = plt.subplots()

        y_vals = np.unique(self.y)
        label_to_num = {label: num for label, num in zip(y_vals, range(len(y_vals)))}

        if heatmap_kwargs is None:
            heatmap_kwargs = {}
        heatmap_kwargs = {
            **{"cmap": sns.color_palette(None, len(y_vals))},
            **heatmap_kwargs,
        }
        _plot_preds(y_vals, label_to_num, ax=ax, **heatmap_kwargs)
        _wash_out(ax)
        colorbar = _set_colorbar(y_vals=y_vals, ax=ax)

        if plot_data:
            if scatter_kwargs is None:
                scatter_kwargs = {}
            scatter_kwargs = {
                **{"cmap": colorbar.cmap, "edgecolor": "k", "zorder": 999},
                **scatter_kwargs,
            }
            self._plot_data(y=self.y.map(label_to_num), ax=ax, **scatter_kwargs)
        _format_ticks(ax=ax, formatter=tick_formatter)
        return ax

    def plot3d(
        self,
        plot_data: bool = True,
        ax: Axes = None,
        surf_kwargs: Optional[dict] = None,
        scatter_kwargs: Optional[dict] = None,
    ):
        """Plot data and predictions in 3D

        Best viewed with a tool such as https://github.com/matplotlib/ipympl
        that supports rotating the output

        Parameters:
        - `plot_data`: Make a scatter plot of the data
        - `ax`: Matplotlib `Axes` object. Plot will be added to this object
        if provided; otherwise a new `Axes` object will be generated.
        - `surf_kwargs`: kwargs to pass to `ax.plot_surface` for plotting
        predictions
        - `scatter_kwargs`: kwargs to pass to `ax.scatter` for plotting data
        """
        if ax is None:
            fig = plt.figure(figsize=(8, 8))
            ax = fig.add_subplot(111, projection="3d")
        if surf_kwargs is None:
            surf_kwargs = {}
        surf_kwargs = {**{"alpha": 0.4, "cmap": "viridis"}, **surf_kwargs}
        if scatter_kwargs is None:
            scatter_kwargs = {}

        y_vals = np.unique(self.y)
        label_to_num = {label: num for label, num in zip(y_vals, range(len(y_vals)))}
        y_int = self.y.map(label_to_num)

        y_pred_int = pd.Series(self.model.predict(self.X), index=self.y.index).map(
            label_to_num
        )
        x0_grid, x1_grid = self._create_grid(num_points=20)
        grid_preds = pd.DataFrame(self._get_grid_preds(x0_grid, x1_grid)).applymap(
            lambda x: label_to_num[x]
        )

        for val in y_vals:
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                ax.plot_surface(
                    x0_grid,
                    x1_grid,
                    grid_preds[grid_preds == val],
                    rstride=1,
                    cstride=1,
                    alpha=0.3,
                )
            if plot_data:
                ax.scatter(
                    self.X.iloc[:, 0].loc[y_int == val],
                    self.X.iloc[:, 1].loc[y_int == val],
                    y_pred_int.loc[y_int == val],
                    **scatter_kwargs,
                )

        return ax

# Cell
def _format_ticks(ax, formatter):
    labels = [item.get_text() for item in ax.get_xticklabels()]
    ax.set_xticklabels([f"{float(label):{formatter}}" for label in labels])

    labels = [item.get_text() for item in ax.get_yticklabels()]
    ax.set_yticklabels([f"{float(label):{formatter}}" for label in labels])