arviz 0.23.1__py3-none-any.whl → 1.0.0rc0__py3-none-any.whl

arviz/plots/backends/matplotlib/traceplot.py

@@ -1,526 +0,0 @@
-"""Matplotlib traceplot."""
-
-import warnings
-from itertools import cycle
-
-from matplotlib import gridspec
-import matplotlib.pyplot as plt
-import numpy as np
-from matplotlib.lines import Line2D
-import matplotlib.ticker as mticker
-
-from ....stats.density_utils import get_bins
-from ...distplot import plot_dist
-from ...plot_utils import _scale_fig_size, format_coords_as_labels
-from ...rankplot import plot_rank
-from . import backend_kwarg_defaults, backend_show, dealiase_sel_kwargs, matplotlib_kwarg_dealiaser
-
-
-def plot_trace(
-    data,
-    var_names,  # pylint: disable=unused-argument
-    divergences,
-    kind,
-    figsize,
-    rug,
-    lines,
-    circ_var_names,
-    circ_var_units,
-    compact,
-    compact_prop,
-    combined,
-    chain_prop,
-    legend,
-    labeller,
-    plot_kwargs,
-    fill_kwargs,
-    rug_kwargs,
-    hist_kwargs,
-    trace_kwargs,
-    rank_kwargs,
-    plotters,
-    divergence_data,
-    axes,
-    backend_kwargs,
-    backend_config,  # pylint: disable=unused-argument
-    show,
-):
-    """Plot distribution (histogram or kernel density estimates) and sampled values.
-
-    If `divergences` data is available in `sample_stats`, will plot the location of divergences as
-    dashed vertical lines.
-
-    Parameters
-    ----------
-    data : obj
-        Any object that can be converted to an az.InferenceData object
-        Refer to documentation of az.convert_to_dataset for details
-    var_names : string, or list of strings
-        One or more variables to be plotted.
-    divergences : {"bottom", "top", None, False}
-        Plot location of divergences on the traceplots. Options are "bottom", "top", or False-y.
-    kind : {"trace", "rank_bar", "rank_vlines"}, optional
-        Choose between plotting sampled values per iteration and rank plots.
-    figsize : figure size tuple
-        If None, size is (12, variables * 2)
-    rug : bool
-        If True adds a rugplot. Defaults to False. Ignored for 2D KDE. Only affects continuous
-        variables.
-    lines : tuple or list
-        List of tuple of (var_name, {'coord': selection}, [line_positions]) to be overplotted as
-        vertical lines on the density and horizontal lines on the trace.
-    circ_var_names : string, or list of strings
-        List of circular variables to account for when plotting KDE.
-    circ_var_units : str
-        Whether the variables in `circ_var_names` are in "degrees" or "radians".
-    combined : bool
-        Flag for combining multiple chains into a single line. If False (default), chains will be
-        plotted separately.
-    legend : bool
-        Add a legend to the figure with the chain color code.
-    plot_kwargs : dict
-        Extra keyword arguments passed to `arviz.plot_dist`. Only affects continuous variables.
-    fill_kwargs : dict
-        Extra keyword arguments passed to `arviz.plot_dist`. Only affects continuous variables.
-    rug_kwargs : dict
-        Extra keyword arguments passed to `arviz.plot_dist`. Only affects continuous variables.
-    hist_kwargs : dict
-        Extra keyword arguments passed to `arviz.plot_dist`. Only affects discrete variables.
-    trace_kwargs : dict
-        Extra keyword arguments passed to `plt.plot`
-    rank_kwargs : dict
-        Extra keyword arguments passed to `arviz.plot_rank`
-    Returns
-    -------
-    axes : matplotlib axes
-
-
-    Examples
-    --------
-    Plot a subset variables
-
-    .. plot::
-        :context: close-figs
-
-        >>> import arviz as az
-        >>> data = az.load_arviz_data('non_centered_eight')
-        >>> coords = {'school': ['Choate', 'Lawrenceville']}
-        >>> az.plot_trace(data, var_names=('theta_t', 'theta'), coords=coords)
-
-    Show all dimensions of multidimensional variables in the same plot
-
-    .. plot::
-        :context: close-figs
-
-        >>> az.plot_trace(data, compact=True)
-
-    Combine all chains into one distribution
-
-    .. plot::
-        :context: close-figs
-
-        >>> az.plot_trace(data, var_names=('theta_t', 'theta'), coords=coords, combined=True)
-
-
-    Plot reference lines against distribution and trace
-
-    .. plot::
-        :context: close-figs
-
-        >>> lines = (('theta_t',{'school': "Choate"}, [-1]),)
-        >>> az.plot_trace(data, var_names=('theta_t', 'theta'), coords=coords, lines=lines)
-
-    """
-    # Set plot default backend kwargs
-    if backend_kwargs is None:
-        backend_kwargs = {}
-
-    if circ_var_names is None:
-        circ_var_names = []
-
-    backend_kwargs = {**backend_kwarg_defaults(), **backend_kwargs}
-
-    if lines is None:
-        lines = ()
-
-    num_chain_props = len(data.chain) + 1 if combined else len(data.chain)
-    if not compact:
-        chain_prop = "color" if chain_prop is None else chain_prop
-    else:
-        chain_prop = (
-            {
-                "linestyle": ("solid", "dotted", "dashed", "dashdot"),
-            }
-            if chain_prop is None
-            else chain_prop
-        )
-        compact_prop = "color" if compact_prop is None else compact_prop
-
-    if isinstance(chain_prop, str):
-        chain_prop = {chain_prop: plt.rcParams["axes.prop_cycle"].by_key()[chain_prop]}
-    if isinstance(chain_prop, tuple):
-        warnings.warn(
-            "chain_prop as a tuple will be deprecated in a future warning, use a dict instead",
-            FutureWarning,
-        )
-        chain_prop = {chain_prop[0]: chain_prop[1]}
-    chain_prop = {
-        prop_name: [prop for _, prop in zip(range(num_chain_props), cycle(props))]
-        for prop_name, props in chain_prop.items()
-    }
-
-    if isinstance(compact_prop, str):
-        compact_prop = {compact_prop: plt.rcParams["axes.prop_cycle"].by_key()[compact_prop]}
-    if isinstance(compact_prop, tuple):
-        warnings.warn(
-            "compact_prop as a tuple will be deprecated in a future warning, use a dict instead",
-            FutureWarning,
-        )
-        compact_prop = {compact_prop[0]: compact_prop[1]}
-
-    if figsize is None:
-        figsize = (12, len(plotters) * 2)
-
-    backend_kwargs.setdefault("figsize", figsize)
-
-    trace_kwargs = matplotlib_kwarg_dealiaser(trace_kwargs, "plot")
-    trace_kwargs.setdefault("alpha", 0.35)
-
-    hist_kwargs = matplotlib_kwarg_dealiaser(hist_kwargs, "hist")
-    hist_kwargs.setdefault("alpha", 0.35)
-
-    plot_kwargs = matplotlib_kwarg_dealiaser(plot_kwargs, "plot")
-    fill_kwargs = matplotlib_kwarg_dealiaser(fill_kwargs, "fill_between")
-    rug_kwargs = matplotlib_kwarg_dealiaser(rug_kwargs, "scatter")
-    rank_kwargs = matplotlib_kwarg_dealiaser(rank_kwargs, "bar")
-    if compact:
-        rank_kwargs.setdefault("bar_kwargs", {})
-        rank_kwargs["bar_kwargs"].setdefault("alpha", 0.2)
-
-    textsize = plot_kwargs.pop("textsize", 10)
-
-    figsize, _, titlesize, xt_labelsize, linewidth, _ = _scale_fig_size(
-        figsize, textsize, rows=len(plotters), cols=2
-    )
-
-    trace_kwargs.setdefault("linewidth", linewidth)
-    plot_kwargs.setdefault("linewidth", linewidth)
-
-    # Check the input for lines
-    if lines is not None:
-        all_var_names = set(plotter[0] for plotter in plotters)
-
-        invalid_var_names = set()
-        for line in lines:
-            if line[0] not in all_var_names:
-                invalid_var_names.add(line[0])
-        if invalid_var_names:
-            warnings.warn(
-                "A valid var_name should be provided, found {} expected from {}".format(
-                    invalid_var_names, all_var_names
-                )
-            )
-
-    if axes is None:
-        fig = plt.figure(**backend_kwargs)
-        spec = gridspec.GridSpec(ncols=2, nrows=len(plotters), figure=fig)
-
-    # pylint: disable=too-many-nested-blocks
-    for idx, (var_name, selection, isel, value) in enumerate(plotters):
-        for idy in range(2):
-            value = np.atleast_2d(value)
-
-            circular = var_name in circ_var_names and not idy
-            if var_name in circ_var_names and idy:
-                circ_units_trace = circ_var_units
-            else:
-                circ_units_trace = False
-
-            if axes is None:
-                ax = fig.add_subplot(spec[idx, idy], polar=circular)
-            else:
-                ax = axes[idx, idy]
-
-            if len(value.shape) == 2:
-                if compact_prop:
-                    aux_plot_kwargs = dealiase_sel_kwargs(plot_kwargs, compact_prop, 0)
-                    aux_trace_kwargs = dealiase_sel_kwargs(trace_kwargs, compact_prop, 0)
-                else:
-                    aux_plot_kwargs = plot_kwargs
-                    aux_trace_kwargs = trace_kwargs
-
-                ax = _plot_chains_mpl(
-                    ax,
-                    idy,
-                    value,
-                    data,
-                    chain_prop,
-                    combined,
-                    xt_labelsize,
-                    rug,
-                    kind,
-                    aux_trace_kwargs,
-                    hist_kwargs,
-                    aux_plot_kwargs,
-                    fill_kwargs,
-                    rug_kwargs,
-                    rank_kwargs,
-                    circular,
-                    circ_var_units,
-                    circ_units_trace,
-                )
-
-            else:
-                sub_data = data[var_name].sel(**selection)
-                legend_labels = format_coords_as_labels(sub_data, skip_dims=("chain", "draw"))
-                legend_title = ", ".join(
-                    [
-                        f"{coord_name}"
-                        for coord_name in sub_data.coords
-                        if coord_name not in {"chain", "draw"}
-                    ]
-                )
-                value = value.reshape((value.shape[0], value.shape[1], -1))
-                compact_prop_iter = {
-                    prop_name: [prop for _, prop in zip(range(value.shape[2]), cycle(props))]
-                    for prop_name, props in compact_prop.items()
-                }
-                handles = []
-                for sub_idx, label in zip(range(value.shape[2]), legend_labels):
-                    aux_plot_kwargs = dealiase_sel_kwargs(plot_kwargs, compact_prop_iter, sub_idx)
-                    aux_trace_kwargs = dealiase_sel_kwargs(trace_kwargs, compact_prop_iter, sub_idx)
-                    ax = _plot_chains_mpl(
-                        ax,
-                        idy,
-                        value[..., sub_idx],
-                        data,
-                        chain_prop,
-                        combined,
-                        xt_labelsize,
-                        rug,
-                        kind,
-                        aux_trace_kwargs,
-                        hist_kwargs,
-                        aux_plot_kwargs,
-                        fill_kwargs,
-                        rug_kwargs,
-                        rank_kwargs,
-                        circular,
-                        circ_var_units,
-                        circ_units_trace,
-                    )
-                    if legend:
-                        handles.append(
-                            Line2D(
-                                [],
-                                [],
-                                label=label,
-                                **dealiase_sel_kwargs(aux_plot_kwargs, chain_prop, 0),
-                            )
-                        )
-                if legend and idy == 0:
-                    ax.legend(handles=handles, title=legend_title)
-
-            if value[0].dtype.kind == "i" and idy == 0:
-                xticks = get_bins(value)
-                ax.set_xticks(xticks[:-1])
-            y = 1 / textsize
-            if not idy:
-                ax.set_yticks([])
-                if circular:
-                    y = 0.13 if selection else 0.12
-            ax.set_title(
-                labeller.make_label_vert(var_name, selection, isel),
-                fontsize=titlesize,
-                wrap=True,
-                y=textsize * y,
-            )
-            ax.tick_params(labelsize=xt_labelsize)
-
-            xlims = ax.get_xlim()
-            ylims = ax.get_ylim()
-
-            if divergences:
-                div_selection = {k: v for k, v in selection.items() if k in divergence_data.dims}
-                divs = divergence_data.sel(**div_selection).values
-                # if combined:
-                #     divs = divs.flatten()
-                divs = np.atleast_2d(divs)
-
-                for chain, chain_divs in enumerate(divs):
-                    div_draws = data.draw.values[chain_divs]
-                    div_idxs = np.arange(len(chain_divs))[chain_divs]
-                    if div_idxs.size > 0:
-                        if divergences == "top":
-                            ylocs = ylims[1]
-                        else:
-                            ylocs = ylims[0]
-                        values = value[chain, div_idxs]
-
-                        if circular:
-                            tick = [ax.get_rmin() + ax.get_rmax() * 0.60, ax.get_rmax()]
-                            for val in values:
-                                ax.plot(
-                                    [val, val],
-                                    tick,
-                                    color="black",
-                                    markeredgewidth=1.5,
-                                    markersize=30,
-                                    alpha=trace_kwargs["alpha"],
-                                    zorder=0.6,
-                                )
-                        else:
-                            if kind == "trace" and idy:
-                                ax.plot(
-                                    div_draws,
-                                    np.zeros_like(div_idxs) + ylocs,
-                                    marker="|",
-                                    color="black",
-                                    markeredgewidth=1.5,
-                                    markersize=30,
-                                    linestyle="None",
-                                    alpha=hist_kwargs["alpha"],
-                                    zorder=0.6,
-                                )
-                            elif not idy:
-                                ax.plot(
-                                    values,
-                                    np.zeros_like(values) + ylocs,
-                                    marker="|",
-                                    color="black",
-                                    markeredgewidth=1.5,
-                                    markersize=30,
-                                    linestyle="None",
-                                    alpha=trace_kwargs["alpha"],
-                                    zorder=0.6,
-                                )
-
-            for _, _, vlines in (j for j in lines if j[0] == var_name and j[1] == selection):
-                if isinstance(vlines, (float, int)):
-                    line_values = [vlines]
-                else:
-                    line_values = np.atleast_1d(vlines).ravel()
-                    if not np.issubdtype(line_values.dtype, np.number):
-                        raise ValueError(f"line-positions should be numeric, found {line_values}")
-                if idy:
-                    ax.hlines(
-                        line_values,
-                        xlims[0],
-                        xlims[1],
-                        colors="black",
-                        linewidth=1.5,
-                        alpha=trace_kwargs["alpha"],
-                    )
-
-                else:
-                    ax.vlines(
-                        line_values,
-                        ylims[0],
-                        ylims[1],
-                        colors="black",
-                        linewidth=1.5,
-                        alpha=trace_kwargs["alpha"],
-                    )
-
-        if kind == "trace" and idy:
-            ax.set_xlim(left=data.draw.min(), right=data.draw.max())
-
-    if legend:
-        legend_kwargs = trace_kwargs if combined else plot_kwargs
-        handles = [
-            Line2D(
-                [], [], label=chain_id, **dealiase_sel_kwargs(legend_kwargs, chain_prop, chain_id)
-            )
-            for chain_id in range(data.sizes["chain"])
-        ]
-        if combined:
-            handles.insert(
-                0,
-                Line2D(
-                    [], [], label="combined", **dealiase_sel_kwargs(plot_kwargs, chain_prop, -1)
-                ),
-            )
-        ax.figure.axes[1].legend(handles=handles, title="chain", loc="upper right")
-
-    if axes is None:
-        axes = np.array(ax.figure.axes).reshape(-1, 2)
-
-    if backend_show(show):
-        plt.show()
-
-    return axes
-
-
-def _plot_chains_mpl(
-    axes,
-    idy,
-    value,
-    data,
-    chain_prop,
-    combined,
-    xt_labelsize,
-    rug,
-    kind,
-    trace_kwargs,
-    hist_kwargs,
-    plot_kwargs,
-    fill_kwargs,
-    rug_kwargs,
-    rank_kwargs,
-    circular,
-    circ_var_units,
-    circ_units_trace,
-):
-    if not circular:
-        circ_var_units = False
-
-    for chain_idx, row in enumerate(value):
-        if kind == "trace":
-            aux_kwargs = dealiase_sel_kwargs(trace_kwargs, chain_prop, chain_idx)
-            if idy:
-                axes.plot(data.draw.values, row, **aux_kwargs)
-                if circ_units_trace == "degrees":
-                    y_tick_locs = axes.get_yticks()
-                    y_tick_labels = [i + 2 * 180 if i < 0 else i for i in np.rad2deg(y_tick_locs)]
-                    axes.yaxis.set_major_locator(mticker.FixedLocator(y_tick_locs))
-                    axes.set_yticklabels([f"{i:.0f}°" for i in y_tick_labels])
-
-        if not combined:
-            aux_kwargs = dealiase_sel_kwargs(plot_kwargs, chain_prop, chain_idx)
-            if not idy:
-                axes = plot_dist(
-                    values=row,
-                    textsize=xt_labelsize,
-                    rug=rug,
-                    ax=axes,
-                    hist_kwargs=hist_kwargs,
-                    plot_kwargs=aux_kwargs,
-                    fill_kwargs=fill_kwargs,
-                    rug_kwargs=rug_kwargs,
-                    backend="matplotlib",
-                    show=False,
-                    is_circular=circ_var_units,
-                )
-
-    if kind == "rank_bars" and idy:
-        axes = plot_rank(data=value, kind="bars", ax=axes, **rank_kwargs)
-    elif kind == "rank_vlines" and idy:
-        axes = plot_rank(data=value, kind="vlines", ax=axes, **rank_kwargs)
-
-    if combined:
-        aux_kwargs = dealiase_sel_kwargs(plot_kwargs, chain_prop, -1)
-        if not idy:
-            axes = plot_dist(
-                values=value.flatten(),
-                textsize=xt_labelsize,
-                rug=rug,
-                ax=axes,
-                hist_kwargs=hist_kwargs,
-                plot_kwargs=aux_kwargs,
-                fill_kwargs=fill_kwargs,
-                rug_kwargs=rug_kwargs,
-                backend="matplotlib",
-                show=False,
-                is_circular=circ_var_units,
-            )
-    return axes

arviz/plots/backends/matplotlib/tsplot.py

@@ -1,121 +0,0 @@
-"""Matplotlib plot time series figure."""
-
-import matplotlib.pyplot as plt
-import numpy as np
-
-from ...plot_utils import _scale_fig_size
-from . import create_axes_grid, backend_show, matplotlib_kwarg_dealiaser, backend_kwarg_defaults
-
-
-def plot_ts(
-    x_plotters,
-    y_plotters,
-    y_mean_plotters,
-    y_hat_plotters,
-    y_holdout_plotters,
-    x_holdout_plotters,
-    y_forecasts_plotters,
-    y_forecasts_mean_plotters,
-    num_samples,
-    backend_kwargs,
-    y_kwargs,
-    y_hat_plot_kwargs,
-    y_mean_plot_kwargs,
-    vline_kwargs,
-    length_plotters,
-    rows,
-    cols,
-    textsize,
-    figsize,
-    legend,
-    axes,
-    show,
-):
-    """Matplotlib time series."""
-    if backend_kwargs is None:
-        backend_kwargs = {}
-
-    backend_kwargs = {
-        **backend_kwarg_defaults(),
-        **backend_kwargs,
-    }
-
-    if figsize is None:
-        figsize = (12, rows * 5)
-
-    backend_kwargs.setdefault("figsize", figsize)
-    backend_kwargs.setdefault("squeeze", False)
-
-    figsize, _, _, xt_labelsize, _, _ = _scale_fig_size(figsize, textsize, rows, cols)
-
-    if axes is None:
-        _, axes = create_axes_grid(length_plotters, rows, cols, backend_kwargs=backend_kwargs)
-
-    y_kwargs = matplotlib_kwarg_dealiaser(y_kwargs, "plot")
-    y_kwargs.setdefault("color", "blue")
-    y_kwargs.setdefault("zorder", 10)
-
-    y_hat_plot_kwargs = matplotlib_kwarg_dealiaser(y_hat_plot_kwargs, "plot")
-    y_hat_plot_kwargs.setdefault("color", "grey")
-    y_hat_plot_kwargs.setdefault("alpha", 0.1)
-
-    y_mean_plot_kwargs = matplotlib_kwarg_dealiaser(y_mean_plot_kwargs, "plot")
-    y_mean_plot_kwargs.setdefault("color", "red")
-    y_mean_plot_kwargs.setdefault("linestyle", "dashed")
-
-    vline_kwargs = matplotlib_kwarg_dealiaser(vline_kwargs, "plot")
-    vline_kwargs.setdefault("color", "black")
-    vline_kwargs.setdefault("linestyle", "dashed")
-
-    for i, ax_i in enumerate(np.ravel(axes)[:length_plotters]):
-        y_var_name, _, _, y_plotters_i = y_plotters[i]
-        x_var_name, _, _, x_plotters_i = x_plotters[i]
-
-        ax_i.plot(x_plotters_i, y_plotters_i, **y_kwargs)
-        ax_i.plot([], label="Actual", **y_kwargs)
-        if y_hat_plotters is not None or y_forecasts_plotters is not None:
-            ax_i.plot([], label="Fitted", **y_mean_plot_kwargs)
-            ax_i.plot([], label="Uncertainty", **y_hat_plot_kwargs)
-
-        ax_i.set_xlabel(x_var_name)
-        ax_i.set_ylabel(y_var_name)
-
-        if y_hat_plotters is not None:
-            *_, y_hat_plotters_i = y_hat_plotters[i]
-            *_, x_hat_plotters_i = x_plotters[i]
-            for j in range(num_samples):
-                ax_i.plot(x_hat_plotters_i, y_hat_plotters_i[..., j], **y_hat_plot_kwargs)
-
-            *_, x_mean_plotters_i = x_plotters[i]
-            *_, y_mean_plotters_i = y_mean_plotters[i]
-            ax_i.plot(x_mean_plotters_i, y_mean_plotters_i, **y_mean_plot_kwargs)
-
-        if y_holdout_plotters is not None:
-            *_, y_holdout_plotters_i = y_holdout_plotters[i]
-            *_, x_holdout_plotters_i = x_holdout_plotters[i]
-
-            ax_i.plot(x_holdout_plotters_i, y_holdout_plotters_i, **y_kwargs)
-            ax_i.axvline(x_plotters_i[-1], **vline_kwargs)
-
-        if y_forecasts_plotters is not None:
-            *_, y_forecasts_plotters_i = y_forecasts_plotters[i]
-            *_, x_forecasts_plotters_i = x_holdout_plotters[i]
-            for j in range(num_samples):
-                ax_i.plot(
-                    x_forecasts_plotters_i, y_forecasts_plotters_i[..., j], **y_hat_plot_kwargs
-                )
-
-            *_, x_forecasts_mean_plotters_i = x_holdout_plotters[i]
-            *_, y_forecasts_mean_plotters_i = y_forecasts_mean_plotters[i]
-            ax_i.plot(
-                x_forecasts_mean_plotters_i, y_forecasts_mean_plotters_i, **y_mean_plot_kwargs
-            )
-            ax_i.axvline(x_plotters_i[-1], **vline_kwargs)
-
-        if legend:
-            ax_i.legend(fontsize=xt_labelsize, loc="upper left")
-
-    if backend_show(show):
-        plt.show()
-
-    return axes