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

arviz/plots/essplot.py

@@ -1,319 +0,0 @@
-"""Plot quantile or local effective sample sizes."""
-
-import numpy as np
-import xarray as xr
-
-from ..data import convert_to_dataset
-from ..labels import BaseLabeller
-from ..rcparams import rcParams
-from ..sel_utils import xarray_var_iter
-from ..stats import ess
-from ..utils import _var_names, get_coords
-from .plot_utils import default_grid, filter_plotters_list, get_plotting_function
-
-
-def plot_ess(
-    idata,
-    var_names=None,
-    filter_vars=None,
-    kind="local",
-    relative=False,
-    coords=None,
-    figsize=None,
-    grid=None,
-    textsize=None,
-    rug=False,
-    rug_kind="diverging",
-    n_points=20,
-    extra_methods=False,
-    min_ess=400,
-    labeller=None,
-    ax=None,
-    extra_kwargs=None,
-    text_kwargs=None,
-    hline_kwargs=None,
-    rug_kwargs=None,
-    backend=None,
-    backend_kwargs=None,
-    show=None,
-    **kwargs,
-):
-    r"""Generate quantile, local, or evolution ESS plots.
-
-    The local and the quantile ESS plots are recommended for checking
-    that there are enough samples for all the explored regions of the
-    parameter space. Checking local and quantile ESS is particularly
-    relevant when working with HDI intervals as opposed to ESS bulk,
-    which is suitable for point estimates.
-
-    Parameters
-    ----------
-    idata : InferenceData
-        Any object that can be converted to an :class:`arviz.InferenceData` object
-        Refer to documentation of :func:`arviz.convert_to_dataset` for details.
-    var_names : list of str, optional
-        Variables to be plotted. Prefix the variables by ``~`` when you want to exclude
-        them from the plot. See :ref:`this section <common_var_names>` for usage examples.
-    filter_vars : {None, "like", "regex"}, default None
-        If `None` (default), interpret `var_names` as the real variables names. If "like",
-        interpret `var_names` as substrings of the real variables names. If "regex",
-        interpret `var_names` as regular expressions on the real variables names. See
-        :ref:`this section <common_filter_vars>` for usage examples.
-    kind : {"local", "quantile", "evolution"}, default "local"
-        Specify the kind of plot:
-
-        * The ``kind="local"`` argument generates the ESS' local efficiency for
-          estimating quantiles of a desired posterior.
-        * The ``kind="quantile"`` argument generates the ESS' local efficiency
-          for estimating small-interval probability of a desired posterior.
-        * The ``kind="evolution"`` argument generates the estimated ESS'
-          with incrised number of iterations of a desired posterior.
-
-    relative : bool, default False
-        Show relative ess in plot ``ress = ess / N``.
-    coords : dict, optional
-        Coordinates of `var_names` to be plotted. Passed to :meth:`xarray.Dataset.sel`.
-        See :ref:`this section <common_coords>` for usage examples.
-    grid : tuple, optional
-        Number of rows and columns. By default, the rows and columns are
-        automatically inferred. See :ref:`this section <common_grid>` for usage examples.
-    figsize : (float, float), optional
-        Figure size. If ``None`` it will be defined automatically.
-    textsize : float, optional
-        Text size scaling factor for labels, titles and lines. If ``None`` it will be autoscaled
-        based on `figsize`.
-    rug : bool, default False
-        Add a `rug plot <https://en.wikipedia.org/wiki/Rug_plot>`_ for a specific subset of values.
-    rug_kind : str, default "diverging"
-        Variable in sample stats to use as rug mask. Must be a boolean variable.
-    n_points : int, default 20
-        Number of points for which to plot their quantile/local ess or number of subsets
-        in the evolution plot.
-    extra_methods : bool, default False
-        Plot mean and sd ESS as horizontal lines. Not taken into account if ``kind = 'evolution'``.
-    min_ess : int, default 400
-        Minimum number of ESS desired. If ``relative=True`` the line is plotted at
-        ``min_ess / n_samples`` for local and quantile kinds and as a curve following
-        the ``min_ess / n`` dependency in evolution kind.
-    labeller : Labeller, optional
-        Class providing the method ``make_label_vert`` to generate the labels in the plot titles.
-        Read the :ref:`label_guide` for more details and usage examples.
-    ax : 2D array-like of matplotlib_axes or bokeh_figure, optional
-        A 2D array of locations into which to plot the densities. If not supplied, ArviZ will create
-        its own array of plot areas (and return it).
-    extra_kwargs : dict, optional
-        If evolution plot, `extra_kwargs` is used to plot ess tail and differentiate it
-        from ess bulk. Otherwise, passed to extra methods lines.
-    text_kwargs : dict, optional
-        Only taken into account when ``extra_methods=True``. kwargs passed to ax.annotate
-        for extra methods lines labels. It accepts the additional
-        key ``x`` to set ``xy=(text_kwargs["x"], mcse)``
-    hline_kwargs : dict, optional
-        kwargs passed to :func:`~matplotlib.axes.Axes.axhline` or to :class:`~bokeh.models.Span`
-        depending on the backend for the horizontal minimum ESS line.
-        For relative ess evolution plots the kwargs are passed to
-        :func:`~matplotlib.axes.Axes.plot` or to :class:`~bokeh.plotting.figure.line`
-    rug_kwargs : dict
-        kwargs passed to rug plot.
-    backend : {"matplotlib", "bokeh"}, default "matplotlib"
-        Select plotting backend.
-    backend_kwargs : dict, optional
-        These are kwargs specific to the backend being used, passed to
-        :func:`matplotlib.pyplot.subplots` or :class:`bokeh.plotting.figure`.
-        For additional documentation check the plotting method of the backend.
-    show : bool, optional
-        Call backend show function.
-    **kwargs
-        Passed as-is to :meth:`mpl:matplotlib.axes.Axes.hist` or
-        :meth:`mpl:matplotlib.axes.Axes.plot` function depending on the
-        value of `kind`.
-
-    Returns
-    -------
-    axes : matplotlib_axes or bokeh_figure
-
-    See Also
-    --------
-    ess : Calculate estimate of the effective sample size.
-
-    References
-    ----------
-    .. [1] Vehtari et al. (2021). Rank-normalization, folding, and
-        localization: An improved Rhat for assessing convergence of
-        MCMC. Bayesian analysis, 16(2):667-718.
-
-    Examples
-    --------
-    Plot local ESS.
-
-    .. plot::
-        :context: close-figs
-
-        >>> import arviz as az
-        >>> idata = az.load_arviz_data("centered_eight")
-        >>> coords = {"school": ["Choate", "Lawrenceville"]}
-        >>> az.plot_ess(
-        ...     idata, kind="local", var_names=["mu", "theta"], coords=coords
-        ... )
-
-    Plot ESS evolution as the number of samples increase. When the model is converging properly,
-    both lines in this plot should be roughly linear.
-
-    .. plot::
-        :context: close-figs
-
-        >>> az.plot_ess(
-        ...     idata, kind="evolution", var_names=["mu", "theta"], coords=coords
-        ... )
-
-    Customize local ESS plot to look like reference paper.
-
-    .. plot::
-        :context: close-figs
-
-        >>> az.plot_ess(
-        ...     idata, kind="local", var_names=["mu"], drawstyle="steps-mid", color="k",
-        ...     linestyle="-", marker=None, rug=True, rug_kwargs={"color": "r"}
-        ... )
-
-    Customize ESS evolution plot to look like reference paper.
-
-    .. plot::
-        :context: close-figs
-
-        >>> extra_kwargs = {"color": "lightsteelblue"}
-        >>> az.plot_ess(
-        ...     idata, kind="evolution", var_names=["mu"],
-        ...     color="royalblue", extra_kwargs=extra_kwargs
-        ... )
-
-    """
-    valid_kinds = ("local", "quantile", "evolution")
-    kind = kind.lower()
-    if kind not in valid_kinds:
-        raise ValueError(f"Invalid kind, kind must be one of {valid_kinds} not {kind}")
-
-    if coords is None:
-        coords = {}
-    if "chain" in coords or "draw" in coords:
-        raise ValueError("chain and draw are invalid coordinates for this kind of plot")
-    if labeller is None:
-        labeller = BaseLabeller()
-    extra_methods = False if kind == "evolution" else extra_methods
-
-    data = get_coords(convert_to_dataset(idata, group="posterior"), coords)
-    var_names = _var_names(var_names, data, filter_vars)
-    n_draws = data.sizes["draw"]
-    n_samples = n_draws * data.sizes["chain"]
-
-    ess_tail_dataset = None
-    mean_ess = None
-    sd_ess = None
-
-    if kind == "quantile":
-        probs = np.linspace(1 / n_points, 1 - 1 / n_points, n_points)
-        xdata = probs
-        ylabel = "{} for quantiles"
-        ess_dataset = xr.concat(
-            [
-                ess(data, var_names=var_names, relative=relative, method="quantile", prob=p)
-                for p in probs
-            ],
-            dim="ess_dim",
-        )
-    elif kind == "local":
-        probs = np.linspace(0, 1, n_points, endpoint=False)
-        xdata = probs
-        ylabel = "{} for small intervals"
-        ess_dataset = xr.concat(
-            [
-                ess(
-                    data,
-                    var_names=var_names,
-                    relative=relative,
-                    method="local",
-                    prob=[p, p + 1 / n_points],
-                )
-                for p in probs
-            ],
-            dim="ess_dim",
-        )
-    else:
-        first_draw = data.draw.values[0]
-        ylabel = "{}"
-        xdata = np.linspace(n_samples / n_points, n_samples, n_points)
-        draw_divisions = np.linspace(n_draws // n_points, n_draws, n_points, dtype=int)
-        ess_dataset = xr.concat(
-            [
-                ess(
-                    data.sel(draw=slice(first_draw + draw_div)),
-                    var_names=var_names,
-                    relative=relative,
-                    method="bulk",
-                )
-                for draw_div in draw_divisions
-            ],
-            dim="ess_dim",
-        )
-        ess_tail_dataset = xr.concat(
-            [
-                ess(
-                    data.sel(draw=slice(first_draw + draw_div)),
-                    var_names=var_names,
-                    relative=relative,
-                    method="tail",
-                )
-                for draw_div in draw_divisions
-            ],
-            dim="ess_dim",
-        )
-
-    plotters = filter_plotters_list(
-        list(xarray_var_iter(ess_dataset, var_names=var_names, skip_dims={"ess_dim"})), "plot_ess"
-    )
-    length_plotters = len(plotters)
-    rows, cols = default_grid(length_plotters, grid=grid)
-
-    if extra_methods:
-        mean_ess = ess(data, var_names=var_names, method="mean", relative=relative)
-        sd_ess = ess(data, var_names=var_names, method="sd", relative=relative)
-
-    essplot_kwargs = dict(
-        ax=ax,
-        plotters=plotters,
-        xdata=xdata,
-        ess_tail_dataset=ess_tail_dataset,
-        mean_ess=mean_ess,
-        sd_ess=sd_ess,
-        idata=idata,
-        data=data,
-        kind=kind,
-        extra_methods=extra_methods,
-        textsize=textsize,
-        rows=rows,
-        cols=cols,
-        figsize=figsize,
-        kwargs=kwargs,
-        extra_kwargs=extra_kwargs,
-        text_kwargs=text_kwargs,
-        n_samples=n_samples,
-        relative=relative,
-        min_ess=min_ess,
-        labeller=labeller,
-        ylabel=ylabel,
-        rug=rug,
-        rug_kind=rug_kind,
-        rug_kwargs=rug_kwargs,
-        hline_kwargs=hline_kwargs,
-        backend_kwargs=backend_kwargs,
-        show=show,
-    )
-
-    if backend is None:
-        backend = rcParams["plot.backend"]
-    backend = backend.lower()
-
-    # TODO: Add backend kwargs
-    plot = get_plotting_function("plot_ess", "essplot", backend)
-    ax = plot(**essplot_kwargs)
-    return ax

arviz/plots/forestplot.py

@@ -1,304 +0,0 @@
-"""Forest plot."""
-
-from ..data import convert_to_dataset
-from ..labels import BaseLabeller, NoModelLabeller
-from ..rcparams import rcParams
-from ..utils import _var_names, get_coords
-from .plot_utils import get_plotting_function
-
-
-def plot_forest(
-    data,
-    kind="forestplot",
-    model_names=None,
-    var_names=None,
-    filter_vars=None,
-    transform=None,
-    coords=None,
-    combined=False,
-    combine_dims=None,
-    hdi_prob=None,
-    rope=None,
-    quartiles=True,
-    ess=False,
-    r_hat=False,
-    colors="cycle",
-    textsize=None,
-    linewidth=None,
-    markersize=None,
-    legend=True,
-    labeller=None,
-    ridgeplot_alpha=None,
-    ridgeplot_overlap=2,
-    ridgeplot_kind="auto",
-    ridgeplot_truncate=True,
-    ridgeplot_quantiles=None,
-    figsize=None,
-    ax=None,
-    backend=None,
-    backend_config=None,
-    backend_kwargs=None,
-    show=None,
-):
-    r"""Forest plot to compare HDI intervals from a number of distributions.
-
-    Generate forest or ridge plots to compare distributions from a model or list of models.
-    Additionally, the function can display effective sample sizes (ess) and Rhats to visualize
-    convergence diagnostics alongside the distributions.
-
-    Parameters
-    ----------
-    data : InferenceData
-        Any object that can be converted to an :class:`arviz.InferenceData` object
-        Refer to documentation of :func:`arviz.convert_to_dataset` for details.
-    kind : {"forestplot", "ridgeplot"}, default "forestplot"
-        Specify the kind of plot:
-
-        * The ``kind="forestplot"`` generates credible intervals, where the central points are the
-          estimated posterior median, the thick lines are the central quartiles, and the thin lines
-          represent the :math:`100\times(hdi\_prob)\%` highest density intervals.
-        * The ``kind="ridgeplot"`` option generates density plots (kernel density estimate or
-          histograms) in the same graph. Ridge plots can be configured to have different overlap,
-          truncation bounds and quantile markers.
-
-    model_names : list of str, optional
-        List with names for the models in the list of data. Useful when plotting more that one
-        dataset.
-    var_names : list of str, optional
-        Variables to be plotted. Prefix the variables by ``~`` when you want to exclude
-        them from the plot. See :ref:`this section <common_var_names>` for usage examples.
-    combine_dims : set_like of str, optional
-        List of dimensions to reduce. Defaults to reducing only the "chain" and "draw" dimensions.
-        See :ref:`this section <common_combine_dims>` for usage examples.
-    filter_vars : {None, "like", "regex"}, default None
-        If `None` (default), interpret `var_names` as the real variables names. If "like",
-        interpret `var_names` as substrings of the real variables names. If "regex",
-        interpret `var_names` as regular expressions on the real variables names. See
-        :ref:`this section <common_filter_vars>` for usage examples.
-    transform : callable or dict, optional
-        Function to transform the data. Defaults to None, i.e., the identity function.
-    coords : dict, optional
-        Coordinates of ``var_names`` to be plotted. Passed to :meth:`xarray.Dataset.sel`.
-        See :ref:`this section <common_coords>` for usage examples.
-    combined : bool, default False
-        Flag for combining multiple chains into a single chain. If False, chains will
-        be plotted separately. See :ref:`this section <common_combine>` for usage examples.
-    hdi_prob : float, default 0.94
-        Plots highest posterior density interval for chosen percentage of density.
-        See :ref:`this section <common_ hdi_prob>` for usage examples.
-    rope : list, tuple or dictionary of {str : tuples or lists}, optional
-        A dictionary of tuples with the lower and upper values of the Region Of Practical
-        Equivalence. See :ref:`this section <common_rope>` for usage examples.
-    quartiles : bool, default True
-        Flag for plotting the interquartile range, in addition to the ``hdi_prob`` intervals.
-    r_hat : bool, default False
-        Flag for plotting Split R-hat statistics. Requires 2 or more chains.
-    ess : bool, default False
-        Flag for plotting the effective sample size.
-    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. Defaults to 'cycle'.
-    textsize : float, optional
-        Text size scaling factor for labels, titles and lines. If `None` it will be autoscaled based
-        on ``figsize``.
-    linewidth : int, optional
-        Line width throughout. If `None` it will be autoscaled based on ``figsize``.
-    markersize : int, optional
-        Markersize throughout. If `None` it will be autoscaled based on ``figsize``.
-    legend : bool, optional
-        Show a legend with the color encoded model information.
-        Defaults to True, if there are multiple models.
-    labeller : Labeller, optional
-        Class providing the method ``make_label_vert`` to generate the labels in the plot titles.
-        Read the :ref:`label_guide` for more details and usage examples.
-    ridgeplot_alpha: float, optional
-        Transparency for ridgeplot fill.  If ``ridgeplot_alpha=0``, border is colored by model,
-        otherwise a `black` outline is used.
-    ridgeplot_overlap : float, default 2
-        Overlap height for ridgeplots.
-    ridgeplot_kind : string, optional
-        By default ("auto") continuous variables are plotted using KDEs and discrete ones using
-        histograms. To override this use "hist" to plot histograms and "density" for KDEs.
-    ridgeplot_truncate : bool, default True
-        Whether to truncate densities according to the value of ``hdi_prob``.
-    ridgeplot_quantiles : list, optional
-        Quantiles in ascending order used to segment the KDE. Use [.25, .5, .75] for quartiles.
-    figsize : (float, float), optional
-        Figure size. If `None`, it will be defined automatically.
-    ax : axes, optional
-        :class:`matplotlib.axes.Axes` or :class:`bokeh.plotting.Figure`.
-    backend : {"matplotlib", "bokeh"}, default "matplotlib"
-        Select plotting backend.
-    backend_config : dict, optional
-        Currently specifies the bounds to use for bokeh axes. Defaults to value set in ``rcParams``.
-    backend_kwargs : dict, optional
-        These are kwargs specific to the backend being used, passed to
-        :func:`matplotlib.pyplot.subplots` or :class:`bokeh.plotting.figure`.
-        For additional documentation check the plotting method of the backend.
-    show : bool, optional
-        Call backend show function.
-
-    Returns
-    -------
-    1D ndarray of matplotlib_axes or bokeh_figures
-
-    See Also
-    --------
-    plot_posterior : Plot Posterior densities in the style of John K. Kruschke's book.
-    plot_density : Generate KDE plots for continuous variables and histograms for discrete ones.
-    summary : Create a data frame with summary statistics.
-
-    Examples
-    --------
-    Forestplot
-
-    .. plot::
-        :context: close-figs
-
-        >>> import arviz as az
-        >>> non_centered_data = az.load_arviz_data('non_centered_eight')
-        >>> axes = az.plot_forest(non_centered_data,
-        >>>                            kind='forestplot',
-        >>>                            var_names=["^the"],
-        >>>                            filter_vars="regex",
-        >>>                            combined=True,
-        >>>                            figsize=(9, 7))
-        >>> axes[0].set_title('Estimated theta for 8 schools model')
-
-    Forestplot with multiple datasets
-
-    .. plot::
-        :context: close-figs
-
-        >>> centered_data = az.load_arviz_data('centered_eight')
-        >>> axes = az.plot_forest([non_centered_data, centered_data],
-        >>>                            model_names = ["non centered eight", "centered eight"],
-        >>>                            kind='forestplot',
-        >>>                            var_names=["^the"],
-        >>>                            filter_vars="regex",
-        >>>                            combined=True,
-        >>>                            figsize=(9, 7))
-        >>> axes[0].set_title('Estimated theta for 8 schools models')
-
-    Ridgeplot
-
-    .. plot::
-        :context: close-figs
-
-        >>> 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')
-
-    Ridgeplot non-truncated and with quantiles
-
-    .. plot::
-        :context: close-figs
-
-        >>> axes = az.plot_forest(non_centered_data,
-        >>>                            kind='ridgeplot',
-        >>>                            var_names=['theta'],
-        >>>                            combined=True,
-        >>>                            ridgeplot_truncate=False,
-        >>>                            ridgeplot_quantiles=[.25, .5, .75],
-        >>>                            ridgeplot_overlap=0.7,
-        >>>                            colors='white',
-        >>>                            figsize=(9, 7))
-        >>> axes[0].set_title('Estimated theta for 8 schools model')
-    """
-    if not isinstance(data, (list, tuple)):
-        data = [data]
-    if len(data) == 1:
-        legend = False
-
-    if coords is None:
-        coords = {}
-
-    if combine_dims is None:
-        combine_dims = set()
-
-    if labeller is None:
-        labeller = NoModelLabeller() if legend else BaseLabeller()
-
-    datasets = [convert_to_dataset(datum) for datum in reversed(data)]
-    if transform is not None:
-        if callable(transform):
-            datasets = [transform(dataset) for dataset in datasets]
-        elif isinstance(transform, dict):
-            transformed_datasets = []
-            for dataset in datasets:
-                new_dataset = dataset.copy()
-                for var_name, func in transform.items():
-                    if var_name in new_dataset:
-                        new_dataset[var_name] = func(new_dataset[var_name])
-                transformed_datasets.append(new_dataset)
-            datasets = transformed_datasets
-        else:
-            raise ValueError("transform must be either a callable or a dict {var_name: callable}")
-    datasets = get_coords(
-        datasets, list(reversed(coords)) if isinstance(coords, (list, tuple)) else coords
-    )
-
-    var_names = _var_names(var_names, datasets, filter_vars)
-
-    ncols, width_ratios = 1, [3]
-
-    if ess:
-        ncols += 1
-        width_ratios.append(1)
-
-    if r_hat:
-        ncols += 1
-        width_ratios.append(1)
-
-    if hdi_prob is None:
-        hdi_prob = rcParams["stats.ci_prob"]
-    elif not 1 >= hdi_prob > 0:
-        raise ValueError("The value of hdi_prob should be in the interval (0, 1]")
-
-    plot_forest_kwargs = dict(
-        ax=ax,
-        datasets=datasets,
-        var_names=var_names,
-        model_names=model_names,
-        combined=combined,
-        combine_dims=combine_dims,
-        colors=colors,
-        figsize=figsize,
-        width_ratios=width_ratios,
-        linewidth=linewidth,
-        markersize=markersize,
-        kind=kind,
-        ncols=ncols,
-        hdi_prob=hdi_prob,
-        quartiles=quartiles,
-        rope=rope,
-        ridgeplot_overlap=ridgeplot_overlap,
-        ridgeplot_alpha=ridgeplot_alpha,
-        ridgeplot_kind=ridgeplot_kind,
-        ridgeplot_truncate=ridgeplot_truncate,
-        ridgeplot_quantiles=ridgeplot_quantiles,
-        textsize=textsize,
-        legend=legend,
-        labeller=labeller,
-        ess=ess,
-        r_hat=r_hat,
-        backend_kwargs=backend_kwargs,
-        backend_config=backend_config,
-        show=show,
-    )
-
-    if backend is None:
-        backend = rcParams["plot.backend"]
-    backend = backend.lower()
-
-    # TODO: Add backend kwargs
-    plot = get_plotting_function("plot_forest", "forestplot", backend)
-    axes = plot(**plot_forest_kwargs)
-    return axes