arviz 0.17.0__py3-none-any.whl → 0.17.1__py3-none-any.whl

arviz/__init__.py

@@ -1,6 +1,6 @@
 # pylint: disable=wildcard-import,invalid-name,wrong-import-position
 """ArviZ is a library for exploratory analysis of Bayesian models."""
-__version__ = "0.17.0"
+__version__ = "0.17.1"
 
 import logging
 import os

arviz/data/inference_data.py

@@ -56,6 +56,7 @@ SUPPORTED_GROUPS = [
     "posterior_predictive",
     "predictions",
     "log_likelihood",
+    "log_prior",
     "sample_stats",
     "prior",
     "prior_predictive",
@@ -63,6 +64,8 @@ SUPPORTED_GROUPS = [
     "observed_data",
     "constant_data",
     "predictions_constant_data",
+    "unconstrained_posterior",
+    "unconstrained_prior",
 ]
 
 WARMUP_TAG = "warmup_"
@@ -73,6 +76,7 @@ SUPPORTED_GROUPS_WARMUP = [
     f"{WARMUP_TAG}predictions",
     f"{WARMUP_TAG}sample_stats",
     f"{WARMUP_TAG}log_likelihood",
+    f"{WARMUP_TAG}log_prior",
 ]
 
 SUPPORTED_GROUPS_ALL = SUPPORTED_GROUPS + SUPPORTED_GROUPS_WARMUP
@@ -250,8 +254,7 @@ class InferenceData(Mapping[str, xr.Dataset]):
 
     def __iter__(self) -> Iterator[str]:
         """Iterate over groups in InferenceData object."""
-        for group in self._groups_all:
-            yield group
+        yield from self._groups_all
 
     def __contains__(self, key: object) -> bool:
         """Return True if the named item is present, and False otherwise."""
@@ -1490,7 +1493,7 @@ class InferenceData(Mapping[str, xr.Dataset]):
 
             import numpy as np
             rng = np.random.default_rng(73)
-            ary = rng.normal(size=(post.dims["chain"], post.dims["draw"], obs.dims["match"]))
+            ary = rng.normal(size=(post.sizes["chain"], post.sizes["draw"], obs.sizes["match"]))
             idata.add_groups(
                 log_likelihood={"home_points": ary},
                 dims={"home_points": ["match"]},

arviz/data/io_pystan.py

@@ -676,8 +676,7 @@ def get_draws(fit, variables=None, ignore=None, warmup=False, dtypes=None):
         for item in par_keys:
             _, shape = item.replace("]", "").split("[")
             shape_idx_min = min(int(shape_value) for shape_value in shape.split(","))
-            if shape_idx_min < shift:
-                shift = shape_idx_min
+            shift = min(shift, shape_idx_min)
         # If shift is higher than 1, this will probably mean that Stan
         # has implemented sparse structure (saves only non-zero parts),
         # but let's hope that dims are still corresponding to the full shape

arviz/plots/backends/bokeh/bpvplot.py

@@ -171,8 +171,13 @@ def plot_bpv(
                 ax_i.line(0, 0, legend_label=f"bpv={p_value:.2f}", alpha=0)
 
             if plot_mean:
-                ax_i.circle(
-                    obs_vals.mean(), 0, fill_color=color, line_color="black", size=markersize
+                ax_i.scatter(
+                    obs_vals.mean(),
+                    0,
+                    fill_color=color,
+                    line_color="black",
+                    size=markersize,
+                    marker="circle",
                 )
 
         _title = Title()

arviz/plots/backends/bokeh/compareplot.py

@@ -69,13 +69,14 @@ def plot_compare(
             err_ys.append((y, y))
 
         # plot them
-        dif_tri = ax.triangle(
+        dif_tri = ax.scatter(
             comp_df[information_criterion].iloc[1:],
             yticks_pos[1::2],
             line_color=plot_kwargs.get("color_dse", "grey"),
             fill_color=plot_kwargs.get("color_dse", "grey"),
             line_width=2,
             size=6,
+            marker="triangle",
         )
         dif_line = ax.multi_line(err_xs, err_ys, line_color=plot_kwargs.get("color_dse", "grey"))
 
@@ -85,13 +86,14 @@ def plot_compare(
         ax.yaxis.ticker = yticks_pos[::2]
         ax.yaxis.major_label_overrides = dict(zip(yticks_pos[::2], yticks_labels))
 
-    elpd_circ = ax.circle(
+    elpd_circ = ax.scatter(
         comp_df[information_criterion],
         yticks_pos[::2],
         line_color=plot_kwargs.get("color_ic", "black"),
         fill_color=None,
         line_width=2,
         size=6,
+        marker="circle",
     )
     elpd_label = [elpd_circ]
 
@@ -110,7 +112,7 @@ def plot_compare(
 
     labels.append(("ELPD", elpd_label))
 
-    scale = comp_df["scale"][0]
+    scale = comp_df["scale"].iloc[0]
 
     if insample_dev:
         p_ic = comp_df[f"p_{information_criterion.split('_')[1]}"]
@@ -120,13 +122,14 @@ def plot_compare(
             correction = -p_ic
         elif scale == "deviance":
             correction = -(2 * p_ic)
-        insample_circ = ax.circle(
+        insample_circ = ax.scatter(
             comp_df[information_criterion] + correction,
             yticks_pos[::2],
             line_color=plot_kwargs.get("color_insample_dev", "black"),
             fill_color=plot_kwargs.get("color_insample_dev", "black"),
             line_width=2,
             size=6,
+            marker="circle",
         )
         labels.append(("In-sample ELPD", [insample_circ]))
 

arviz/plots/backends/bokeh/forestplot.py

@@ -15,7 +15,6 @@ from ....rcparams import rcParams
 from ....stats import hdi
 from ....stats.density_utils import get_bins, histogram, kde
 from ....stats.diagnostics import _ess, _rhat
-from ....utils import conditional_jit
 from ...plot_utils import _scale_fig_size
 from .. import show_layout
 from . import backend_kwarg_defaults
@@ -277,7 +276,6 @@ class PlotHandler:
         """Collect labels and ticks from plotters."""
         val = self.plotters.values()
 
-        @conditional_jit(forceobj=True, nopython=False)
         def label_idxs():
             labels, idxs = [], []
             for plotter in val:
@@ -640,7 +638,7 @@ class VarHandler:
             grouped_data = [[(0, datum)] for datum in self.data]
             skip_dims = self.combine_dims.union({"chain"})
         else:
-            grouped_data = [datum.groupby("chain") for datum in self.data]
+            grouped_data = [datum.groupby("chain", squeeze=False) for datum in self.data]
             skip_dims = self.combine_dims
 
         label_dict = OrderedDict()
@@ -648,7 +646,7 @@ class VarHandler:
         for name, grouped_datum in zip(self.model_names, grouped_data):
             for _, sub_data in grouped_datum:
                 datum_iter = xarray_var_iter(
-                    sub_data,
+                    sub_data.squeeze(),
                     var_names=[self.var_name],
                     skip_dims=skip_dims,
                     reverse_selections=True,

arviz/plots/backends/matplotlib/compareplot.py

@@ -84,7 +84,7 @@ def plot_compare(
     else:
         ax.set_yticks(yticks_pos[::2])
 
-    scale = comp_df["scale"][0]
+    scale = comp_df["scale"].iloc[0]
 
     if insample_dev:
         p_ic = comp_df[f"p_{information_criterion.split('_')[1]}"]

arviz/plots/backends/matplotlib/forestplot.py

@@ -11,7 +11,6 @@ from ....stats import hdi
 from ....stats.density_utils import get_bins, histogram, kde
 from ....stats.diagnostics import _ess, _rhat
 from ....sel_utils import xarray_var_iter
-from ....utils import conditional_jit
 from ...plot_utils import _scale_fig_size
 from . import backend_kwarg_defaults, backend_show
 
@@ -236,7 +235,6 @@ class PlotHandler:
         """Collect labels and ticks from plotters."""
         val = self.plotters.values()
 
-        @conditional_jit(forceobj=True, nopython=False)
         def label_idxs():
             labels, idxs = [], []
             for plotter in val:
@@ -536,7 +534,7 @@ class VarHandler:
             grouped_data = [[(0, datum)] for datum in self.data]
             skip_dims = self.combine_dims.union({"chain"})
         else:
-            grouped_data = [datum.groupby("chain") for datum in self.data]
+            grouped_data = [datum.groupby("chain", squeeze=False) for datum in self.data]
             skip_dims = self.combine_dims
 
         label_dict = OrderedDict()
@@ -544,7 +542,7 @@ class VarHandler:
         for name, grouped_datum in zip(self.model_names, grouped_data):
             for _, sub_data in grouped_datum:
                 datum_iter = xarray_var_iter(
-                    sub_data,
+                    sub_data.squeeze(),
                     var_names=[self.var_name],
                     skip_dims=skip_dims,
                     reverse_selections=True,

arviz/plots/backends/matplotlib/traceplot.py

@@ -430,7 +430,7 @@ def plot_trace(
             Line2D(
                 [], [], label=chain_id, **dealiase_sel_kwargs(legend_kwargs, chain_prop, chain_id)
             )
-            for chain_id in range(data.dims["chain"])
+            for chain_id in range(data.sizes["chain"])
         ]
         if combined:
             handles.insert(

arviz/plots/bfplot.py

@@ -38,7 +38,7 @@ def plot_bf(
     algorithm presented in [1]_.
 
     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.
@@ -52,16 +52,16 @@ def plot_bf(
         Tuple of valid Matplotlib colors. First element for the prior, second for the posterior.
     figsize : (float, float), optional
         Figure size. If `None` it will be defined automatically.
-    textsize: float, optional
+    textsize : float, optional
         Text size scaling factor for labels, titles and lines. If `None` it will be auto
         scaled based on `figsize`.
-    plot_kwargs : dicts, optional
+    plot_kwargs : dict, optional
         Additional keywords passed to :func:`matplotlib.pyplot.plot`.
-    hist_kwargs : dicts, optional
+    hist_kwargs : dict, optional
         Additional keywords passed to :func:`arviz.plot_dist`. Only works for discrete variables.
     ax : axes, optional
         :class:`matplotlib.axes.Axes` or :class:`bokeh.plotting.Figure`.
-    backend :{"matplotlib", "bokeh"}, default "matplotlib"
+    backend : {"matplotlib", "bokeh"}, default "matplotlib"
         Select plotting backend.
     backend_kwargs : dict, optional
         These are kwargs specific to the backend being used, passed to
@@ -78,7 +78,7 @@ def plot_bf(
     References
     ----------
     .. [1] Heck, D., 2019. A caveat on the Savage-Dickey density ratio:
-    The case of computing Bayes factors for regression parameters.
+       The case of computing Bayes factors for regression parameters.
 
     Examples
     --------
@@ -92,6 +92,7 @@ def plot_bf(
         >>> idata = az.from_dict(posterior={"a":np.random.normal(1, 0.5, 5000)},
         ...     prior={"a":np.random.normal(0, 1, 5000)})
         >>> az.plot_bf(idata, var_name="a", ref_val=0)
+
     """
     posterior = extract(idata, var_names=var_name).values
 

arviz/plots/bpvplot.py

@@ -230,11 +230,11 @@ def plot_bpv(
 
     if flatten_pp is None:
         if flatten is None:
-            flatten_pp = list(predictive_dataset.dims.keys())
+            flatten_pp = list(predictive_dataset.dims)
         else:
             flatten_pp = flatten
     if flatten is None:
-        flatten = list(observed.dims.keys())
+        flatten = list(observed.dims)
 
     if coords is None:
         coords = {}

arviz/plots/compareplot.py

@@ -90,10 +90,10 @@ def plot_compare(
     References
     ----------
     .. [1] Vehtari et al. (2016). Practical Bayesian model evaluation using leave-one-out
-    cross-validation and WAIC https://arxiv.org/abs/1507.04544
+       cross-validation and WAIC https://arxiv.org/abs/1507.04544
 
     .. [2] McElreath R. (2022). Statistical Rethinking A Bayesian Course with Examples in
-    R and Stan, Second edition, CRC Press.
+       R and Stan, Second edition, CRC Press.
 
     Examples
     --------

arviz/plots/ecdfplot.py

@@ -1,8 +1,9 @@
 """Plot ecdf or ecdf-difference plot with confidence bands."""
 import numpy as np
-from scipy.stats import uniform, binom
+from scipy.stats import uniform
 
 from ..rcparams import rcParams
+from ..stats.ecdf_utils import compute_ecdf, ecdf_confidence_band, _get_ecdf_points
 from .plot_utils import get_plotting_function
 
 
@@ -26,7 +27,7 @@ def plot_ecdf(
     show=None,
     backend=None,
     backend_kwargs=None,
-    **kwargs
+    **kwargs,
 ):
     r"""Plot ECDF or ECDF-Difference Plot with Confidence bands.
 
@@ -48,6 +49,7 @@ def plot_ecdf(
         Values to compare to the original sample.
     cdf : callable, optional
         Cumulative distribution function of the distribution to compare the original sample.
+        The function must take as input a numpy array of draws from the distribution.
     difference : bool, default False
         If True then plot ECDF-difference plot otherwise ECDF plot.
     pit : bool, default False
@@ -180,75 +182,47 @@ def plot_ecdf(
     values = np.ravel(values)
     values.sort()
 
-    ## This block computes gamma and uses it to get the upper and lower confidence bands
-    ## Here we check if we want confidence bands or not
-    if confidence_bands:
-        ## If plotting PIT then we find the PIT values of sample.
-        ## Basically here we generate the evaluation points(x) and find the PIT values.
-        ## z is the evaluation point for our uniform distribution in compute_gamma()
-        if pit:
-            x = np.linspace(1 / npoints, 1, npoints)
-            z = x
-            ## Finding PIT for our sample
-            probs = cdf(values) if cdf else compute_ecdf(values2, values) / len(values2)
-        else:
-            ## If not PIT use sample for plots and for evaluation points(x) use equally spaced
-            ## points between minimum and maximum of sample
-            ## For z we have used cdf(x)
-            x = np.linspace(values[0], values[-1], npoints)
-            z = cdf(x) if cdf else compute_ecdf(values2, x)
-            probs = values
-
-        n = len(values)  # number of samples
-        ## Computing gamma
-        gamma = fpr if pointwise else compute_gamma(n, z, npoints, num_trials, fpr)
-        ## Using gamma to get the confidence intervals
-        lower, higher = get_lims(gamma, n, z)
-
-        ## This block is for whether to plot ECDF or ECDF-difference
-        if not difference:
-            ## We store the coordinates of our ecdf in x_coord, y_coord
-            x_coord, y_coord = get_ecdf_points(x, probs, difference)
+    if pit:
+        eval_points = np.linspace(1 / npoints, 1, npoints)
+        if cdf:
+            sample = cdf(values)
         else:
-            ## Here we subtract the ecdf value as here we are plotting the ECDF-difference
-            x_coord, y_coord = get_ecdf_points(x, probs, difference)
-            for i, x_i in enumerate(x):
-                y_coord[i] = y_coord[i] - (
-                    x_i if pit else cdf(x_i) if cdf else compute_ecdf(values2, x_i)
-                )
-
-            ## Similarly we subtract from the upper and lower bounds
-            if pit:
-                lower = lower - x
-                higher = higher - x
-            else:
-                lower = lower - (cdf(x) if cdf else compute_ecdf(values2, x))
-                higher = higher - (cdf(x) if cdf else compute_ecdf(values2, x))
-
+            sample = compute_ecdf(values2, values) / len(values2)
+        cdf_at_eval_points = eval_points
+        rvs = uniform(0, 1).rvs
     else:
-        if pit:
-            x = np.linspace(1 / npoints, 1, npoints)
-            probs = cdf(values)
+        eval_points = np.linspace(values[0], values[-1], npoints)
+        sample = values
+        if confidence_bands or difference:
+            if cdf:
+                cdf_at_eval_points = cdf(eval_points)
+            else:
+                cdf_at_eval_points = compute_ecdf(values2, eval_points)
         else:
-            x = np.linspace(values[0], values[-1], npoints)
-            probs = values
+            cdf_at_eval_points = np.zeros_like(eval_points)
+        rvs = None
+
+    x_coord, y_coord = _get_ecdf_points(sample, eval_points, difference)
 
+    if difference:
+        y_coord -= cdf_at_eval_points
+
+    if confidence_bands:
+        ndraws = len(values)
+        band_kwargs = {"prob": 1 - fpr, "num_trials": num_trials, "rvs": rvs, "random_state": None}
+        band_kwargs["method"] = "pointwise" if pointwise else "simulated"
+        lower, higher = ecdf_confidence_band(ndraws, eval_points, cdf_at_eval_points, **band_kwargs)
+
+        if difference:
+            lower -= cdf_at_eval_points
+            higher -= cdf_at_eval_points
+    else:
         lower, higher = None, None
-        ## This block is for whether to plot ECDF or ECDF-difference
-        if not difference:
-            x_coord, y_coord = get_ecdf_points(x, probs, difference)
-        else:
-            ## Here we subtract the ecdf value as here we are plotting the ECDF-difference
-            x_coord, y_coord = get_ecdf_points(x, probs, difference)
-            for i, x_i in enumerate(x):
-                y_coord[i] = y_coord[i] - (
-                    x_i if pit else cdf(x_i) if cdf else compute_ecdf(values2, x_i)
-                )
 
     ecdf_plot_args = dict(
         x_coord=x_coord,
         y_coord=y_coord,
-        x_bands=x,
+        x_bands=eval_points,
         lower=lower,
         higher=higher,
         confidence_bands=confidence_bands,
@@ -260,7 +234,7 @@ def plot_ecdf(
         ax=ax,
         show=show,
         backend_kwargs=backend_kwargs,
-        **kwargs
+        **kwargs,
     )
 
     if backend is None:
@@ -271,52 +245,3 @@ def plot_ecdf(
     ax = plot(**ecdf_plot_args)
 
     return ax
-
-
-def compute_ecdf(sample, z):
-    """Compute ECDF.
-
-    This function computes the ecdf value at the evaluation point
-        or a sorted set of evaluation points.
-    """
-    return np.searchsorted(sample, z, side="right") / len(sample)
-
-
-def get_ecdf_points(x, probs, difference):
-    """Compute the coordinates for the ecdf points using compute_ecdf."""
-    y = compute_ecdf(probs, x)
-
-    if not difference:
-        x = np.insert(x, 0, x[0])
-        y = np.insert(y, 0, 0)
-    return x, y
-
-
-def compute_gamma(n, z, npoints=None, num_trials=1000, fpr=0.05):
-    """Compute gamma for confidence interval calculation.
-
-    This function simulates an adjusted value of gamma to account for multiplicity
-    when forming an 1-fpr level confidence envelope for the ECDF of a sample.
-    """
-    if npoints is None:
-        npoints = n
-    gamma = []
-    for _ in range(num_trials):
-        unif_samples = uniform.rvs(0, 1, n)
-        unif_samples = np.sort(unif_samples)
-        gamma_m = 1000
-        ## Can compute ecdf for all the z together or one at a time.
-        f_z = compute_ecdf(unif_samples, z)
-        f_z = compute_ecdf(unif_samples, z)
-        gamma_m = 2 * min(
-            np.amin(binom.cdf(n * f_z, n, z)), np.amin(1 - binom.cdf(n * f_z - 1, n, z))
-        )
-        gamma.append(gamma_m)
-    return np.quantile(gamma, fpr)
-
-
-def get_lims(gamma, n, z):
-    """Compute the simultaneous 1 - fpr level confidence bands."""
-    lower = binom.ppf(gamma / 2, n, z)
-    upper = binom.ppf(1 - gamma / 2, n, z)
-    return lower / n, upper / n

arviz/plots/elpdplot.py

@@ -98,7 +98,7 @@ def plot_elpd(
     References
     ----------
     .. [1] Vehtari et al. (2016). Practical Bayesian model evaluation using leave-one-out
-    cross-validation and WAIC https://arxiv.org/abs/1507.04544
+       cross-validation and WAIC https://arxiv.org/abs/1507.04544
 
     Examples
     --------

arviz/plots/essplot.py

@@ -202,8 +202,8 @@ def plot_ess(
 
     data = get_coords(convert_to_dataset(idata, group="posterior"), coords)
     var_names = _var_names(var_names, data, filter_vars)
-    n_draws = data.dims["draw"]
-    n_samples = n_draws * data.dims["chain"]
+    n_draws = data.sizes["draw"]
+    n_samples = n_draws * data.sizes["chain"]
 
     ess_tail_dataset = None
     mean_ess = None

arviz/plots/pairplot.py

@@ -229,7 +229,7 @@ def plot_pair(
             )
 
     if gridsize == "auto":
-        gridsize = int(dataset.dims["draw"] ** 0.35)
+        gridsize = int(dataset.sizes["draw"] ** 0.35)
 
     numvars = len(flat_var_names)
 

arviz/plots/ppcplot.py

@@ -19,7 +19,7 @@ def plot_ppc(
     kind="kde",
     alpha=None,
     mean=True,
-    observed=True,
+    observed=None,
     observed_rug=False,
     color=None,
     colors=None,
@@ -60,8 +60,9 @@ def plot_ppc(
         Defaults to 0.2 for ``kind = kde`` and cumulative, for scatter defaults to 0.7.
     mean : bool, default True
         Whether or not to plot the mean posterior/prior predictive distribution.
-    observed : bool, default True
-        Whether or not to plot the observed data.
+    observed : bool, optional
+        Whether or not to plot the observed data. Defaults to True for ``group = posterior``
+        and False for ``group = prior``.
     observed_rug : bool, default False
         Whether or not to plot a rug plot for the observed data. Only valid if `observed` is
         `True` and for kind `kde` or `cumulative`.
@@ -253,8 +254,12 @@ def plot_ppc(
 
     if group == "posterior":
         predictive_dataset = data.posterior_predictive
+        if observed is None:
+            observed = True
     elif group == "prior":
         predictive_dataset = data.prior_predictive
+        if observed is None:
+            observed = False
 
     if var_names is None:
         var_names = list(observed_data.data_vars)
@@ -264,11 +269,11 @@ def plot_ppc(
 
     if flatten_pp is None:
         if flatten is None:
-            flatten_pp = list(predictive_dataset.dims.keys())
+            flatten_pp = list(predictive_dataset.dims)
         else:
             flatten_pp = flatten
     if flatten is None:
-        flatten = list(observed_data.dims.keys())
+        flatten = list(observed_data.dims)
 
     if coords is None:
         coords = {}

arviz/stats/density_utils.py

@@ -231,8 +231,8 @@ def _fixed_point(t, N, k_sq, a_sq):
        Z. I. Botev, J. F. Grotowski, and D. P. Kroese.
        Ann. Statist. 38 (2010), no. 5, 2916--2957.
     """
-    k_sq = np.asfarray(k_sq, dtype=np.float64)
-    a_sq = np.asfarray(a_sq, dtype=np.float64)
+    k_sq = np.asarray(k_sq, dtype=np.float64)
+    a_sq = np.asarray(a_sq, dtype=np.float64)
 
     l = 7
     f = np.sum(np.power(k_sq, l) * a_sq * np.exp(-k_sq * np.pi**2 * t))

arviz/stats/diagnostics.py

@@ -457,10 +457,10 @@ def ks_summary(pareto_tail_indices):
     """
     _numba_flag = Numba.numba_flag
     if _numba_flag:
-        bins = np.asarray([-np.Inf, 0.5, 0.7, 1, np.Inf])
+        bins = np.asarray([-np.inf, 0.5, 0.7, 1, np.inf])
         kcounts, *_ = _histogram(pareto_tail_indices, bins)
     else:
-        kcounts, *_ = _histogram(pareto_tail_indices, bins=[-np.Inf, 0.5, 0.7, 1, np.Inf])
+        kcounts, *_ = _histogram(pareto_tail_indices, bins=[-np.inf, 0.5, 0.7, 1, np.inf])
     kprop = kcounts / len(pareto_tail_indices) * 100
     df_k = pd.DataFrame(
         dict(_=["(good)", "(ok)", "(bad)", "(very bad)"], Count=kcounts, Pct=kprop)