arviz 0.18.0__py3-none-any.whl → 0.20.0__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.18.0"
+__version__ = "0.20.0"
 
 import logging
 import os
@@ -37,6 +37,7 @@ from .stats import *
 from .rcparams import rc_context, rcParams
 from .utils import Numba, Dask, interactive_backend
 from .wrappers import *
+from . import preview
 
 # add ArviZ's styles to matplotlib's styles
 _arviz_style_path = os.path.join(os.path.dirname(__file__), "plots", "styles")

arviz/data/base.py

@@ -9,9 +9,13 @@ from copy import deepcopy
 from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar, Union
 
 import numpy as np
-import tree
 import xarray as xr
 
+try:
+    import tree
+except ImportError:
+    tree = None
+
 try:
     import ujson as json
 except ImportError:
@@ -89,6 +93,9 @@ def _yield_flat_up_to(shallow_tree, input_tree, path=()):
         input_tree.
     """
     # pylint: disable=protected-access
+    if tree is None:
+        raise ImportError("Missing optional dependency 'dm-tree'. Use pip or conda to install it")
+
     if isinstance(shallow_tree, tree._TEXT_OR_BYTES) or not (
         isinstance(shallow_tree, tree.collections_abc.Mapping)
         or tree._is_namedtuple(shallow_tree)
@@ -299,7 +306,7 @@ def numpy_to_data_array(
     return xr.DataArray(ary, coords=coords, dims=dims)
 
 
-def pytree_to_dataset(
+def dict_to_dataset(
     data,
     *,
     attrs=None,
@@ -312,6 +319,8 @@ def pytree_to_dataset(
 ):
     """Convert a dictionary or pytree of numpy arrays to an xarray.Dataset.
 
+    ArviZ itself supports conversion of flat dictionaries.
+    Suport for pytrees requires ``dm-tree`` which is an optional dependency.
     See https://jax.readthedocs.io/en/latest/pytrees.html for what a pytree is, but
     this inclues at least dictionaries and tuple types.
 
@@ -386,10 +395,12 @@ def pytree_to_dataset(
     """
     if dims is None:
         dims = {}
-    try:
-        data = {k[0] if len(k) == 1 else k: v for k, v in _flatten_with_path(data)}
-    except TypeError:  # probably unsortable keys -- the function will still work if
-        pass  # it is an honest dictionary.
+
+    if tree is not None:
+        try:
+            data = {k[0] if len(k) == 1 else k: v for k, v in _flatten_with_path(data)}
+        except TypeError:  # probably unsortable keys -- the function will still work if
+            pass  # it is an honest dictionary.
 
     data_vars = {
         key: numpy_to_data_array(
@@ -406,7 +417,7 @@ def pytree_to_dataset(
     return xr.Dataset(data_vars=data_vars, attrs=make_attrs(attrs=attrs, library=library))
 
 
-dict_to_dataset = pytree_to_dataset
+pytree_to_dataset = dict_to_dataset
 
 
 def make_attrs(attrs=None, library=None):

arviz/data/converters.py

@@ -1,9 +1,13 @@
 """High level conversion functions."""
 
 import numpy as np
-import tree
 import xarray as xr
 
+try:
+    from tree import is_nested
+except ImportError:
+    is_nested = lambda obj: False
+
 from .base import dict_to_dataset
 from .inference_data import InferenceData
 from .io_beanmachine import from_beanmachine
@@ -107,7 +111,7 @@ def convert_to_inference_data(obj, *, group="posterior", coords=None, dims=None,
         dataset = obj.to_dataset()
     elif isinstance(obj, dict):
         dataset = dict_to_dataset(obj, coords=coords, dims=dims)
-    elif tree.is_nested(obj) and not isinstance(obj, (list, tuple)):
+    elif is_nested(obj) and not isinstance(obj, (list, tuple)):
         dataset = dict_to_dataset(obj, coords=coords, dims=dims)
     elif isinstance(obj, np.ndarray):
         dataset = dict_to_dataset({"x": obj}, coords=coords, dims=dims)
@@ -122,7 +126,7 @@ def convert_to_inference_data(obj, *, group="posterior", coords=None, dims=None,
             "xarray dataarray",
             "xarray dataset",
             "dict",
-            "pytree",
+            "pytree (if 'dm-tree' is installed)",
             "netcdf filename",
             "numpy array",
             "pystan fit",

arviz/data/inference_data.py

@@ -266,6 +266,14 @@ class InferenceData(Mapping[str, xr.Dataset]):
             raise KeyError(key)
         return getattr(self, key)
 
+    def __setitem__(self, key: str, value: xr.Dataset):
+        """Set item by key and update group list accordingly."""
+        if key.startswith(WARMUP_TAG):
+            self._groups_warmup.append(key)
+        else:
+            self._groups.append(key)
+        setattr(self, key, value)
+
     def groups(self) -> List[str]:
         """Return all groups present in InferenceData object."""
         return self._groups_all

arviz/data/io_cmdstan.py

@@ -738,6 +738,7 @@ def _process_configuration(comments):
         elif "=" in comment:
             match_int = re.search(r"^(\S+)\s*=\s*([-+]?[0-9]+)$", comment)
             match_float = re.search(r"^(\S+)\s*=\s*([-+]?[0-9]+\.[0-9]+)$", comment)
+            match_str_bool = re.search(r"^(\S+)\s*=\s*(true|false)$", comment)
             match_str = re.search(r"^(\S+)\s*=\s*(\S+)$", comment)
             match_empty = re.search(r"^(\S+)\s*=\s*$", comment)
             if match_int:
@@ -746,6 +747,9 @@ def _process_configuration(comments):
             elif match_float:
                 key, value = match_float.group(1), match_float.group(2)
                 results[key] = float(value)
+            elif match_str_bool:
+                key, value = match_str_bool.group(1), match_str_bool.group(2)
+                results[key] = int(value == "true")
             elif match_str:
                 key, value = match_str.group(1), match_str.group(2)
                 results[key] = value

arviz/data/io_numpyro.py

@@ -194,7 +194,7 @@ class NumPyroConverter:
             )
             for obs_name, log_like in log_likelihood_dict.items():
                 shape = (self.nchains, self.ndraws) + log_like.shape[1:]
-                data[obs_name] = np.reshape(log_like.copy(), shape)
+                data[obs_name] = np.reshape(np.asarray(log_like), shape)
         return dict_to_dataset(
             data,
             library=self.numpyro,

arviz/plots/backends/bokeh/ecdfplot.py

@@ -13,7 +13,6 @@ def plot_ecdf(
     x_bands,
     lower,
     higher,
-    confidence_bands,
     plot_kwargs,
     fill_kwargs,
     plot_outline_kwargs,
@@ -58,7 +57,7 @@ def plot_ecdf(
         plot_outline_kwargs.setdefault("color", to_hex("C0"))
         plot_outline_kwargs.setdefault("alpha", 0.2)
 
-    if confidence_bands:
+    if x_bands is not None:
         ax.step(x_coord, y_coord, **plot_kwargs)
 
         if fill_band:

arviz/plots/backends/bokeh/khatplot.py

@@ -21,6 +21,7 @@ def plot_khat(
     figsize,
     xdata,
     khats,
+    good_k,
     kwargs,
     threshold,
     coord_labels,
@@ -53,7 +54,11 @@ def plot_khat(
 
     if hlines_kwargs is None:
         hlines_kwargs = {}
-    hlines_kwargs.setdefault("hlines", [0, 0.5, 0.7, 1])
+
+    if good_k is None:
+        good_k = 0.7
+
+    hlines_kwargs.setdefault("hlines", [0, good_k, 1])
 
     cmap = None
     if isinstance(color, str):
@@ -75,7 +80,7 @@ def plot_khat(
             rgba_c = cmap(color)
 
     khats = khats if isinstance(khats, np.ndarray) else khats.values.flatten()
-    alphas = 0.5 + 0.2 * (khats > 0.5) + 0.3 * (khats > 1)
+    alphas = 0.5 + 0.2 * (khats > good_k) + 0.3 * (khats > 1)
 
     rgba_c = vectorized_to_hex(rgba_c)
 
@@ -130,7 +135,7 @@ def plot_khat(
     xmax = len(khats)
 
     if show_bins:
-        bin_edges = np.array([ymin, 0.5, 0.7, 1, ymax])
+        bin_edges = np.array([ymin, good_k, 1, ymax])
         bin_edges = bin_edges[(bin_edges >= ymin) & (bin_edges <= ymax)]
         hist, _, _ = histogram(khats, bin_edges)
         for idx, count in enumerate(hist):

arviz/plots/backends/bokeh/pairplot.py

@@ -174,12 +174,8 @@ def plot_pair(
     source = ColumnDataSource(data=source_dict)
 
     if divergences:
-        source_nondiv = CDSView(
-            source=source, filters=[GroupFilter(column_name=divergenve_name, group="0")]
-        )
-        source_div = CDSView(
-            source=source, filters=[GroupFilter(column_name=divergenve_name, group="1")]
-        )
+        source_nondiv = CDSView(filter=GroupFilter(column_name=divergenve_name, group="0"))
+        source_div = CDSView(filter=GroupFilter(column_name=divergenve_name, group="1"))
 
     def get_width_and_height(jointplot, rotate):
         """Compute subplots dimensions for two or more variables."""

arviz/plots/backends/matplotlib/ecdfplot.py

@@ -13,7 +13,6 @@ def plot_ecdf(
     x_bands,
     lower,
     higher,
-    confidence_bands,
     plot_kwargs,
     fill_kwargs,
     plot_outline_kwargs,
@@ -59,7 +58,7 @@ def plot_ecdf(
 
     ax.step(x_coord, y_coord, **plot_kwargs)
 
-    if confidence_bands:
+    if x_bands is not None:
         if fill_band:
             ax.fill_between(x_bands, lower, higher, **fill_kwargs)
         else:

arviz/plots/backends/matplotlib/khatplot.py

@@ -20,6 +20,7 @@ def plot_khat(
     figsize,
     xdata,
     khats,
+    good_k,
     kwargs,
     threshold,
     coord_labels,
@@ -61,8 +62,11 @@ def plot_khat(
     backend_kwargs.setdefault("figsize", figsize)
     backend_kwargs["squeeze"] = True
 
+    if good_k is None:
+        good_k = 0.7
+
     hlines_kwargs = matplotlib_kwarg_dealiaser(hlines_kwargs, "hlines")
-    hlines_kwargs.setdefault("hlines", [0, 0.5, 0.7, 1])
+    hlines_kwargs.setdefault("hlines", [0, good_k, 1])
     hlines_kwargs.setdefault("linestyle", [":", "-.", "--", "-"])
     hlines_kwargs.setdefault("alpha", 0.7)
     hlines_kwargs.setdefault("zorder", -1)
@@ -102,7 +106,7 @@ def plot_khat(
                 rgba_c = cmap(norm_fun(color))
 
         khats = khats if isinstance(khats, np.ndarray) else khats.values.flatten()
-        alphas = 0.5 + 0.2 * (khats > 0.5) + 0.3 * (khats > 1)
+        alphas = 0.5 + 0.2 * (khats > good_k) + 0.3 * (khats > 1)
         rgba_c[:, 3] = alphas
         rgba_c = vectorized_to_hex(rgba_c)
         kwargs["c"] = rgba_c
@@ -151,7 +155,7 @@ def plot_khat(
         )
 
     if show_bins:
-        bin_edges = np.array([ymin, 0.5, 0.7, 1, ymax])
+        bin_edges = np.array([ymin, good_k, 1, ymax])
         bin_edges = bin_edges[(bin_edges >= ymin) & (bin_edges <= ymax)]
         hist, _, _ = histogram(khats, bin_edges)
         for idx, count in enumerate(hist):

arviz/plots/backends/matplotlib/traceplot.py

@@ -440,7 +440,7 @@ def plot_trace(
                     [], [], label="combined", **dealiase_sel_kwargs(plot_kwargs, chain_prop, -1)
                 ),
             )
-        ax.figure.axes[0].legend(handles=handles, title="chain", loc="upper right")
+        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)

arviz/plots/bpvplot.py

@@ -80,7 +80,7 @@ def plot_bpv(
     hdi_prob : float, optional
         Probability for the highest density interval for the analytical reference distribution when
         ``kind=u_values``. Should be in the interval (0, 1]. Defaults to the
-        rcParam ``stats.hdi_prob``. See :ref:`this section <common_hdi_prob>` for usage examples.
+        rcParam ``stats.ci_prob``. See :ref:`this section <common_hdi_prob>` for usage examples.
     color : str, optional
         Matplotlib color
     grid : tuple, optional
@@ -202,7 +202,7 @@ def plot_bpv(
         raise TypeError("`reference` argument must be either `analytical`, `samples`, or `None`")
 
     if hdi_prob is None:
-        hdi_prob = rcParams["stats.hdi_prob"]
+        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]")
 

arviz/plots/compareplot.py

@@ -11,9 +11,9 @@ def plot_compare(
     comp_df,
     insample_dev=False,
     plot_standard_error=True,
-    plot_ic_diff=True,
+    plot_ic_diff=False,
     order_by_rank=True,
-    legend=True,
+    legend=False,
     title=True,
     figsize=None,
     textsize=None,
@@ -45,12 +45,12 @@ def plot_compare(
         penalization given by the effective number of parameters (p_loo or p_waic).
     plot_standard_error : bool, default True
         Plot the standard error of the ELPD.
-    plot_ic_diff : bool, default True
+    plot_ic_diff : bool, default False
         Plot standard error of the difference in ELPD between each model
         and the top-ranked model.
     order_by_rank : bool, default True
         If True ensure the best model is used as reference.
-    legend : bool, default True
+    legend : bool, default False
         Add legend to figure.
     figsize : (float, float), optional
         If `None`, size is (6, num of models) inches.

arviz/plots/densityplot.py

@@ -209,7 +209,7 @@ def plot_density(
         )
 
     if hdi_prob is None:
-        hdi_prob = rcParams["stats.hdi_prob"]
+        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]")
 

arviz/plots/dotplot.py

@@ -67,7 +67,7 @@ def plot_dot(
         The shape of the marker. Valid for matplotlib backend.
     hdi_prob : float, optional
         Valid only when point_interval is True. Plots HDI for chosen percentage of density.
-        Defaults to ``stats.hdi_prob`` rcParam. See :ref:`this section <common_hdi_prob>`
+        Defaults to ``stats.ci_prob`` rcParam. See :ref:`this section <common_hdi_prob>`
         for usage examples.
     rotated : bool, default False
         Whether to rotate the dot plot by 90 degrees.
@@ -151,7 +151,7 @@ def plot_dot(
     values.sort()
 
     if hdi_prob is None:
-        hdi_prob = rcParams["stats.hdi_prob"]
+        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]")