Forest plot to compare credible intervals from a number of distributions.

    Generates a forest plot of 100*(credible_interval)% credible intervals from
    a trace or list of traces.

    Parameters
    ----------
    data : obj or list[obj]
        Any object that can be converted to an az.InferenceData object
        Refer to documentation of az.convert_to_dataset for details
    kind : str
        Choose kind of plot for main axis. Supports "forestplot" or "ridgeplot"
    model_names : list[str], optional
        List with names for the models in the list of data. Useful when
        plotting more that one dataset
    var_names: list[str], optional
        List of variables to plot (defaults to None, which results in all
        variables plotted)
    combined : bool
        Flag for combining multiple chains into a single chain. If False (default),
        chains will be plotted separately.
    credible_interval : float, optional
        Credible interval to plot. Defaults to 0.94.
    rope: tuple or dictionary of tuples
        Lower and upper values of the Region Of Practical Equivalence. If a list with one
        interval only is provided, the ROPE will be displayed across the y-axis. If more than one
        interval is provided the length of the list should match the number of variables.
    quartiles : bool, optional
        Flag for plotting the interquartile range, in addition to the credible_interval intervals.
        Defaults to True
    r_hat : bool, optional
        Flag for plotting Split R-hat statistics. Requires 2 or more chains. Defaults to False
    ess : bool, optional
        Flag for plotting the effective sample size. Requires 2 or more chains. Defaults to False
    colors : list or string, optional
        list with valid matplotlib colors, one color per model. Alternative a string can be passed.
        If the string is `cycle`, it will automatically chose a color per model from the
        matplotlibs cycle. If a single color is passed, eg 'k', 'C2', 'red' this color will be used
        for all models. Defauls to 'cycle'.
    textsize: float
        Text size scaling factor for labels, titles and lines. If None it will be autoscaled based
        on figsize.
    linewidth : int
        Line width throughout. If None it will be autoscaled based on figsize.
    markersize : int
        Markersize throughout. If None it will be autoscaled based on figsize.
    ridgeplot_alpha : float
        Transparency for ridgeplot fill.  If 0, border is colored by model, otherwise
        a black outline is used.
    ridgeplot_overlap : float
        Overlap height for ridgeplots.
    figsize : tuple
        Figure size. If None it will be defined automatically.

    Returns
    -------
    gridspec : matplotlib GridSpec

    Examples
    --------
    Forestpĺot

    .. plot::
        :context: close-figs

        >>> import arviz as az
        >>> non_centered_data = az.load_arviz_data('non_centered_eight')
        >>> fig, axes = az.plot_forest(non_centered_data,
        >>>                            kind='forestplot',
        >>>                            var_names=['theta'],
        >>>                            combined=True,
        >>>                            ridgeplot_overlap=3,
        >>>                            figsize=(9, 7))
        >>> axes[0].set_title('Estimated theta for 8 schools model')

    Ridgeplot

    .. plot::
        :context: close-figs

        >>> fig, axes = az.plot_forest(non_centered_data,
        >>>                            kind='ridgeplot',
        >>>                            var_names=['theta'],
        >>>                            combined=True,
        >>>                            ridgeplot_overlap=3,
        >>>                            colors='white',
        >>>                            figsize=(9, 7))
        >>> axes[0].set_title('Estimated theta for 8 schools model')

def plot_forest(
    data,
    kind="forestplot",
    model_names=None,
    var_names=None,
    combined=False,
    credible_interval=0.94,
    rope=None,
    quartiles=True,
    ess=False,
    r_hat=False,
    colors="cycle",
    textsize=None,
    linewidth=None,
    markersize=None,
    ridgeplot_alpha=None,
    ridgeplot_overlap=2,
    figsize=None,
):
    """Forest plot to compare credible intervals from a number of distributions.

    Generates a forest plot of 100*(credible_interval)% credible intervals from
    a trace or list of traces.

    Parameters
    ----------
    data : obj or list[obj]
        Any object that can be converted to an az.InferenceData object
        Refer to documentation of az.convert_to_dataset for details
    kind : str
        Choose kind of plot for main axis. Supports "forestplot" or "ridgeplot"
    model_names : list[str], optional
        List with names for the models in the list of data. Useful when
        plotting more that one dataset
    var_names: list[str], optional
        List of variables to plot (defaults to None, which results in all
        variables plotted)
    combined : bool
        Flag for combining multiple chains into a single chain. If False (default),
        chains will be plotted separately.
    credible_interval : float, optional
        Credible interval to plot. Defaults to 0.94.
    rope: tuple or dictionary of tuples
        Lower and upper values of the Region Of Practical Equivalence. If a list with one
        interval only is provided, the ROPE will be displayed across the y-axis. If more than one
        interval is provided the length of the list should match the number of variables.
    quartiles : bool, optional
        Flag for plotting the interquartile range, in addition to the credible_interval intervals.
        Defaults to True
    r_hat : bool, optional
        Flag for plotting Split R-hat statistics. Requires 2 or more chains. Defaults to False
    ess : bool, optional
        Flag for plotting the effective sample size. Requires 2 or more chains. Defaults to False
    colors : list or string, optional
        list with valid matplotlib colors, one color per model. Alternative a string can be passed.
        If the string is `cycle`, it will automatically chose a color per model from the
        matplotlibs cycle. If a single color is passed, eg 'k', 'C2', 'red' this color will be used
        for all models. Defauls to 'cycle'.
    textsize: float
        Text size scaling factor for labels, titles and lines. If None it will be autoscaled based
        on figsize.
    linewidth : int
        Line width throughout. If None it will be autoscaled based on figsize.
    markersize : int
        Markersize throughout. If None it will be autoscaled based on figsize.
    ridgeplot_alpha : float
        Transparency for ridgeplot fill.  If 0, border is colored by model, otherwise
        a black outline is used.
    ridgeplot_overlap : float
        Overlap height for ridgeplots.
    figsize : tuple
        Figure size. If None it will be defined automatically.

    Returns
    -------
    gridspec : matplotlib GridSpec

    Examples
    --------
    Forestpĺot

    .. plot::
        :context: close-figs

        >>> import arviz as az
        >>> non_centered_data = az.load_arviz_data('non_centered_eight')
        >>> fig, axes = az.plot_forest(non_centered_data,
        >>>                            kind='forestplot',
        >>>                            var_names=['theta'],
        >>>                            combined=True,
        >>>                            ridgeplot_overlap=3,
        >>>                            figsize=(9, 7))
        >>> axes[0].set_title('Estimated theta for 8 schools model')

    Ridgeplot

    .. plot::
        :context: close-figs

        >>> fig, axes = az.plot_forest(non_centered_data,
        >>>                            kind='ridgeplot',
        >>>                            var_names=['theta'],
        >>>                            combined=True,
        >>>                            ridgeplot_overlap=3,
        >>>                            colors='white',
        >>>                            figsize=(9, 7))
        >>> axes[0].set_title('Estimated theta for 8 schools model')
    """
    if not isinstance(data, (list, tuple)):
        data = [data]

    datasets = [convert_to_dataset(datum) for datum in reversed(data)]

    var_names = _var_names(var_names, datasets)

    ncols, width_ratios = 1, [3]

    if ess:
        ncols += 1
        width_ratios.append(1)

    if r_hat:
        ncols += 1
        width_ratios.append(1)

    plot_handler = PlotHandler(
        datasets, var_names=var_names, model_names=model_names, combined=combined, colors=colors
    )

    if figsize is None:
        figsize = (min(12, sum(width_ratios) * 2), plot_handler.fig_height())

    (figsize, _, titlesize, xt_labelsize, auto_linewidth, auto_markersize) = _scale_fig_size(
        figsize, textsize, 1.1, 1
    )

    if linewidth is None:
        linewidth = auto_linewidth

    if markersize is None:
        markersize = auto_markersize

    fig, axes = plt.subplots(
        nrows=1,
        ncols=ncols,
        figsize=figsize,
        gridspec_kw={"width_ratios": width_ratios},
        sharey=True,
        constrained_layout=True,
    )

    axes = np.atleast_1d(axes)
    if kind == "forestplot":
        plot_handler.forestplot(
            credible_interval,
            quartiles,
            xt_labelsize,
            titlesize,
            linewidth,
            markersize,
            axes[0],
            rope,
        )
    elif kind == "ridgeplot":
        plot_handler.ridgeplot(ridgeplot_overlap, linewidth, ridgeplot_alpha, axes[0])
    else:
        raise TypeError(
            "Argument 'kind' must be one of 'forestplot' or "
            "'ridgeplot' (you provided {})".format(kind)
        )

    idx = 1
    if ess:
        plot_handler.plot_neff(axes[idx], xt_labelsize, titlesize, markersize)
        idx += 1

    if r_hat:
        plot_handler.plot_rhat(axes[idx], xt_labelsize, titlesize, markersize)
        idx += 1

    for ax in axes:
        ax.grid(False)
        # Remove ticklines on y-axes
        for ticks in ax.yaxis.get_major_ticks():
            ticks.tick1On = False
            ticks.tick2On = False

        for loc, spine in ax.spines.items():
            if loc in ["left", "right"]:
                spine.set_visible(False)

        if len(plot_handler.data) > 1:
            plot_handler.make_bands(ax)

    labels, ticks = plot_handler.labels_and_ticks()
    axes[0].set_yticks(ticks)
    axes[0].set_yticklabels(labels)
    all_plotters = list(plot_handler.plotters.values())
    y_max = plot_handler.y_max() - all_plotters[-1].group_offset
    if kind == "ridgeplot":  # space at the top
        y_max += ridgeplot_overlap
    axes[0].set_ylim(-all_plotters[0].group_offset, y_max)

    return fig, axes