def plot_2d_classification_predictions_and_uncertainty(
    inputs: Optional[np.ndarray] = None,
    preds: Optional[np.ndarray] = None,
    grid_uncertainty: Optional[np.ndarray] = None,
    uncertainty_fn: Optional[Callable[[np.ndarray], np.ndarray]] = None,
    grid: Optional[np.ndarray] = None,
    xx1: Optional[np.ndarray] = None,
    xx2: Optional[np.ndarray] = None,
    ax: Optional[matplotlib.axes.Axes] = None,
    figsize: tuple = (5, 5),
    title: Optional[str] = None,
    uncertainty_cmap: str = "Blues",
    preds_color: Tuple[str, str] = ("C2", "C3"),
    base_inputs_color: str = "black",
    marker_size: int = 10,
    legend: bool = True,
    uncertainty_label: str = "uncertainty",
    fontsize: int = 12,
    colorbar: bool = False,
    remove_ticks: bool = False,
    show: bool = False,
    return_pcolor: bool = False,
    **save_options,
) -> Optional[Collection]:
    """
    Plot inputs, predictions and uncertainty.

    :param inputs: Optional[np.ndarray]
        Input data. This must be a two-dimensional array, with different inputs over the first axis. The number of
        dimensions over the second axis must be 2.
    :param preds: Optional[np.ndarray]
        Predictions. This must be a one-dimensional array of zeros and ones, with different inputs over the first
        axis.
    :param grid_uncertainty: Optional[np.ndarray]
        Scalar uncertainty value for each element of the grid. This must have the same dimensionality as `grid`.
    :param uncertainty_fn: Optional[Callable[[np.ndarray], np.ndarray]]
        A function taking as argument a grid of inputs, and returning a scalar uncertainty value for each element
        in the grid.
    :param grid: Optional[np.ndarray]
        The grid to display the plot over.
    :param xx1: Optional[np.ndarray]
        Mesh over the first axis.
    :param xx2: Optional[np.ndarray]
        Mesh over the second axis.
    :param ax: Optional[matplotlib.axes.Axes]
        Axis matplotlib object.
    :param figsize: tuple
        Figure size.
    :param title: Optional[str]
        Figure title.
    :param uncertainty_cmap: str
        Uncertainty color map.
    :param preds_color: Tuple[str, str]
        Color of predictions. The first element of the tuple is associated to the prediction 1; the second to 0.
    :param base_inputs_color: str
        Color of inputs if `preds` is not available.
    :param marker_size: int
        Size of scattered inputs.
    :param legend: bool
        Whether to include a legend.
    :param uncertainty_label: str
        Uncertainty label in the legend.
    :param fontsize: int
        Font size in the legend.
    :param colorbar: bool
        Whether to include a color bar.
    :param remove_ticks: bool
        Whether to remove ticks from axes.
    :param save_options: dict
        Options to save the file with `matplotlib.pyplot.savefig`. If no option is given, the file is not saved.
    :param show: bool
        Whether to show the plot.
    :param return_pcolor: bool
        Whether to return the output of `matplotlib.pyplot.pcolor`.

    :return Optional[Collection]
        If `return_pcolor`, it returns the output of `matplotlib.pyplot.pcolor`.
    """
    if (xx1 is None or xx2 is None) and (inputs is None):
        raise ValueError(
            """If `inputs` is not passed, please pass either both `xx1` and `xx2`, or `grid`.
        """
        )
    if grid_uncertainty is None and uncertainty_fn is None:
        raise ValueError(
            """ One between `grid_uncertainty` and `uncertainty_fn` must be passed. If both are passed,
        `uncertainty_fn` is ignored."""
        )
    if grid is None and uncertainty_fn is None:
        raise ValueError(
            """If `grid` is not passed, then `uncertainty_fn` must be passed."""
        )
    if grid is not None and (uncertainty_fn is None and grid_uncertainty is None):
        raise ValueError(
            """If `grid` is not passed, then at least one of `uncertainty_fn` and `grid_uncertainty`
        must be passed."""
        )
    if grid_uncertainty is not None and grid is None:
        raise ValueError(
            """If `grid_uncertainty` is passed, then `grid` must also be passed."""
        )
    if grid is None:
        if xx1 is None or xx2 is None:
            mins, maxs = np.min(inputs, 0), np.max(inputs, 0)
            if xx1 is None:
                xx1 = np.linspace(
                    mins[0] - 1.25 * (maxs[0] - mins[0]),
                    maxs[0] + 1.25 * (maxs[0] - mins[0]),
                    100 if xx2 is None else len(xx2),
                )
            if xx2 is None:
                xx2 = np.linspace(
                    mins[1] - 1.25 * (maxs[1] - mins[1]),
                    maxs[1] + 1.25 * (maxs[1] - mins[1]),
                    100 if xx1 is None else len(xx1),
                )
        grid = np.array([[_xx1, _xx2] for _xx1 in xx1 for _xx2 in xx2])
    if grid_uncertainty is None:
        grid_uncertainty = uncertainty_fn(grid)
    size1 = len(np.unique(grid[:, 0])) if xx1 is None else len(xx1)
    size2 = len(np.unique(grid[:, 1])) if xx2 is None else len(xx2)
    grid_uncertainty = grid_uncertainty.reshape(size1, size2)
    grid = grid.reshape(size1, size2, 2)
    if ax is None:
        fig, ax = plt.subplots(1, 1, figsize=figsize)
    if title:
        ax.set_title(title, fontsize=fontsize + 2)
    im = ax.pcolor(
        grid[:, :, 0],
        grid[:, :, 1],
        grid_uncertainty,
        cmap=uncertainty_cmap,
        label=uncertainty_label,
    )
    if inputs is not None:
        ax.scatter(
            inputs[:, 0],
            inputs[:, 1],
            s=marker_size,
            c=(
                [preds_color[0] if i == 1 else preds_color[1] for i in preds]
                if preds is not None
                else base_inputs_color
            ),
        )
    if colorbar:
        plt.colorbar(im, ax=ax.ravel().tolist() if hasattr(ax, "ravel") else ax)
    if legend:
        ax.legend(fontsize=fontsize)
    if remove_ticks:
        ax.xticks([], [])
        ax.yticks([], [])
    if len(save_options) > 0:
        plt.savefig(**save_options)
    if show:
        plt.show()

    if return_pcolor:
        return im

    Plot inputs, predictions and uncertainty.

    :param inputs: Optional[np.ndarray]
        Input data. This must be a two-dimensional array, with different inputs over the first axis. The number of
        dimensions over the second axis must be 2.
    :param preds: Optional[np.ndarray]
        Predictions. This must be a one-dimensional array of zeros and ones, with different inputs over the first
        axis.
    :param grid_uncertainty: Optional[np.ndarray]
        Scalar uncertainty value for each element of the grid. This must have the same dimensionality as `grid`.
    :param uncertainty_fn: Optional[Callable[[np.ndarray], np.ndarray]]
        A function taking as argument a grid of inputs, and returning a scalar uncertainty value for each element
        in the grid.
    :param grid: Optional[np.ndarray]
        The grid to display the plot over.
    :param xx1: Optional[np.ndarray]
        Mesh over the first axis.
    :param xx2: Optional[np.ndarray]
        Mesh over the second axis.
    :param ax: Optional[matplotlib.axes.Axes]
        Axis matplotlib object.
    :param figsize: tuple
        Figure size.
    :param title: Optional[str]
        Figure title.
    :param uncertainty_cmap: str
        Uncertainty color map.
    :param preds_color: Tuple[str, str]
        Color of predictions. The first element of the tuple is associated to the prediction 1; the second to 0.
    :param base_inputs_color: str
        Color of inputs if `preds` is not available.
    :param marker_size: int
        Size of scattered inputs.
    :param legend: bool
        Whether to include a legend.
    :param uncertainty_label: str
        Uncertainty label in the legend.
    :param fontsize: int
        Font size in the legend.
    :param colorbar: bool
        Whether to include a color bar.
    :param remove_ticks: bool
        Whether to remove ticks from axes.
    :param save_options: dict
        Options to save the file with `matplotlib.pyplot.savefig`. If no option is given, the file is not saved.
    :param show: bool
        Whether to show the plot.
    :param return_pcolor: bool
        Whether to return the output of `matplotlib.pyplot.pcolor`.

    :return Optional[Collection]
        If `return_pcolor`, it returns the output of `matplotlib.pyplot.pcolor`.
    
fortuna/plot.py