arviz 0.17.1__py3-none-any.whl → 0.18.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.17.1"
+__version__ = "0.18.0"
 
 import logging
 import os
@@ -315,7 +315,8 @@ _linear_grey_10_95_c0 = [
 
 def _mpl_cm(name, colorlist):
     cmap = LinearSegmentedColormap.from_list(name, colorlist, N=256)
-    mpl.colormaps.register(cmap, name="cet_" + name)
+    if "cet_" + name not in mpl.colormaps():
+        mpl.colormaps.register(cmap, name="cet_" + name)
 
 
 try:

arviz/data/__init__.py

@@ -1,5 +1,6 @@
 """Code for loading and manipulating data structures."""
-from .base import CoordSpec, DimSpec, dict_to_dataset, numpy_to_data_array
+
+from .base import CoordSpec, DimSpec, dict_to_dataset, numpy_to_data_array, pytree_to_dataset
 from .converters import convert_to_dataset, convert_to_inference_data
 from .datasets import clear_data_home, list_datasets, load_arviz_data
 from .inference_data import InferenceData, concat
@@ -7,7 +8,7 @@ from .io_beanmachine import from_beanmachine
 from .io_cmdstan import from_cmdstan
 from .io_cmdstanpy import from_cmdstanpy
 from .io_datatree import from_datatree, to_datatree
-from .io_dict import from_dict
+from .io_dict import from_dict, from_pytree
 from .io_emcee import from_emcee
 from .io_json import from_json, to_json
 from .io_netcdf import from_netcdf, to_netcdf
@@ -38,10 +39,12 @@ __all__ = [
     "from_cmdstanpy",
     "from_datatree",
     "from_dict",
+    "from_pytree",
     "from_json",
     "from_pyro",
     "from_numpyro",
     "from_netcdf",
+    "pytree_to_dataset",
     "to_datatree",
     "to_json",
     "to_netcdf",

arviz/data/base.py

@@ -1,4 +1,5 @@
 """Low level converters usually used by other functions."""
+
 import datetime
 import functools
 import importlib
@@ -8,6 +9,7 @@ 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:
@@ -67,6 +69,48 @@ class requires:  # pylint: disable=invalid-name
         return wrapped
 
 
+def _yield_flat_up_to(shallow_tree, input_tree, path=()):
+    """Yields (path, value) pairs of input_tree flattened up to shallow_tree.
+
+    Adapted from dm-tree (https://github.com/google-deepmind/tree) to allow
+    lists as leaves.
+
+    Args:
+        shallow_tree: Nested structure. Traverse no further than its leaf nodes.
+        input_tree: Nested structure. Return the paths and values from this tree.
+            Must have the same upper structure as shallow_tree.
+        path: Tuple. Optional argument, only used when recursing. The path from the
+            root of the original shallow_tree, down to the root of the shallow_tree
+            arg of this recursive call.
+
+    Yields:
+        Pairs of (path, value), where path the tuple path of a leaf node in
+        shallow_tree, and value is the value of the corresponding node in
+        input_tree.
+    """
+    # pylint: disable=protected-access
+    if isinstance(shallow_tree, tree._TEXT_OR_BYTES) or not (
+        isinstance(shallow_tree, tree.collections_abc.Mapping)
+        or tree._is_namedtuple(shallow_tree)
+        or tree._is_attrs(shallow_tree)
+    ):
+        yield (path, input_tree)
+    else:
+        input_tree = dict(tree._yield_sorted_items(input_tree))
+        for shallow_key, shallow_subtree in tree._yield_sorted_items(shallow_tree):
+            subpath = path + (shallow_key,)
+            input_subtree = input_tree[shallow_key]
+            for leaf_path, leaf_value in _yield_flat_up_to(
+                shallow_subtree, input_subtree, path=subpath
+            ):
+                yield (leaf_path, leaf_value)
+    # pylint: enable=protected-access
+
+
+def _flatten_with_path(structure):
+    return list(_yield_flat_up_to(structure, structure))
+
+
 def generate_dims_coords(
     shape,
     var_name,
@@ -255,7 +299,7 @@ def numpy_to_data_array(
     return xr.DataArray(ary, coords=coords, dims=dims)
 
 
-def dict_to_dataset(
+def pytree_to_dataset(
     data,
     *,
     attrs=None,
@@ -266,26 +310,29 @@ def dict_to_dataset(
     index_origin=None,
     skip_event_dims=None,
 ):
-    """Convert a dictionary of numpy arrays to an xarray.Dataset.
+    """Convert a dictionary or pytree of numpy arrays to an xarray.Dataset.
+
+    See https://jax.readthedocs.io/en/latest/pytrees.html for what a pytree is, but
+    this inclues at least dictionaries and tuple types.
 
     Parameters
     ----------
-    data : dict[str] -> ndarray
+    data : dict of {str : array_like or dict} or pytree
         Data to convert. Keys are variable names.
-    attrs : dict
+    attrs : dict, optional
         Json serializable metadata to attach to the dataset, in addition to defaults.
-    library : module
+    library : module, optional
         Library used for performing inference. Will be attached to the attrs metadata.
-    coords : dict[str] -> ndarray
+    coords : dict of {str : ndarray}, optional
         Coordinates for the dataset
-    dims : dict[str] -> list[str]
+    dims : dict of {str : list of str}, optional
         Dimensions of each variable. The keys are variable names, values are lists of
         coordinates.
     default_dims : list of str, optional
         Passed to :py:func:`numpy_to_data_array`
     index_origin : int, optional
         Passed to :py:func:`numpy_to_data_array`
-    skip_event_dims : bool
+    skip_event_dims : bool, optional
         If True, cut extra dims whenever present to match the shape of the data.
         Necessary for PPLs which have the same name in both observed data and log
         likelihood groups, to account for their different shapes when observations are
@@ -293,15 +340,56 @@ def dict_to_dataset(
 
     Returns
     -------
-    xr.Dataset
+    xarray.Dataset
+        In case of nested pytrees, the variable name will be a tuple of individual names.
+
+    Notes
+    -----
+    This function is available through two aliases: ``dict_to_dataset`` or ``pytree_to_dataset``.
 
     Examples
     --------
-    dict_to_dataset({'x': np.random.randn(4, 100), 'y': np.random.rand(4, 100)})
+    Convert a dictionary with two 2D variables to a Dataset.
+
+    .. ipython::
+
+        In [1]: import arviz as az
+           ...: import numpy as np
+           ...: az.dict_to_dataset({'x': np.random.randn(4, 100), 'y': np.random.rand(4, 100)})
+
+    Note that unlike the :class:`xarray.Dataset` constructor, ArviZ has added extra
+    information to the generated Dataset such as default dimension names for sampled
+    dimensions and some attributes.
+
+    The function is also general enough to work on pytrees such as nested dictionaries:
+
+    .. ipython::
+
+        In [1]: az.pytree_to_dataset({'top': {'second': 1.}, 'top2': 1.})
+
+    which has two variables (as many as leafs) named ``('top', 'second')`` and ``top2``.
+
+    Dimensions and co-ordinates can be defined as usual:
+
+    .. ipython::
+
+        In [1]: datadict = {
+           ...:     "top": {"a": np.random.randn(100), "b": np.random.randn(1, 100, 10)},
+           ...:     "d": np.random.randn(100),
+           ...: }
+           ...: az.dict_to_dataset(
+           ...:     datadict,
+           ...:     coords={"c": np.arange(10)},
+           ...:     dims={("top", "b"): ["c"]}
+           ...: )
 
     """
     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.
 
     data_vars = {
         key: numpy_to_data_array(
@@ -318,6 +406,9 @@ def dict_to_dataset(
     return xr.Dataset(data_vars=data_vars, attrs=make_attrs(attrs=attrs, library=library))
 
 
+dict_to_dataset = pytree_to_dataset
+
+
 def make_attrs(attrs=None, library=None):
     """Make standard attributes to attach to xarray datasets.
 
@@ -332,7 +423,7 @@ def make_attrs(attrs=None, library=None):
         attrs
     """
     default_attrs = {
-        "created_at": datetime.datetime.utcnow().isoformat(),
+        "created_at": datetime.datetime.now(datetime.timezone.utc).isoformat(),
         "arviz_version": __version__,
     }
     if library is not None:

arviz/data/converters.py

@@ -1,5 +1,7 @@
 """High level conversion functions."""
+
 import numpy as np
+import tree
 import xarray as xr
 
 from .base import dict_to_dataset
@@ -105,6 +107,8 @@ 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)):
+        dataset = dict_to_dataset(obj, coords=coords, dims=dims)
     elif isinstance(obj, np.ndarray):
         dataset = dict_to_dataset({"x": obj}, coords=coords, dims=dims)
     elif isinstance(obj, (list, tuple)) and isinstance(obj[0], str) and obj[0].endswith(".csv"):
@@ -118,6 +122,7 @@ def convert_to_inference_data(obj, *, group="posterior", coords=None, dims=None,
             "xarray dataarray",
             "xarray dataset",
             "dict",
+            "pytree",
             "netcdf filename",
             "numpy array",
             "pystan fit",

arviz/data/datasets.py

@@ -1,4 +1,5 @@
 """Base IO code for all datasets. Heavily influenced by scikit-learn's implementation."""
+
 import hashlib
 import itertools
 import json

arviz/data/example_data/data_remote.json

@@ -9,9 +9,16 @@
     {
         "name": "rugby",
         "filename": "rugby.nc",
-        "url": "http://ndownloader.figshare.com/files/16254359",
-        "checksum": "9eecd2c6317e45b0388dd97ae6326adecf94128b5a7d15a52c9fcfac0937e2a6",
-        "description": "The Six Nations Championship is a yearly rugby competition between Italy, Ireland, Scotland, England, France and Wales. Fifteen games are played each year, representing all combinations of the six teams.\n\nThis example uses and includes results from 2014 - 2017, comprising 60 total games. It models latent parameters for each team's attack and defense, as well as a parameter for home team advantage.\n\nSee https://docs.pymc.io/notebooks/rugby_analytics.html by Peader Coyle for more details and references."
+        "url": "http://figshare.com/ndownloader/files/44916469",
+        "checksum": "f4a5e699a8a4cc93f722eb97929dd7c4895c59a2183f05309f5082f3f81eb228",
+        "description": "The Six Nations Championship is a yearly rugby competition between Italy, Ireland, Scotland, England, France and Wales. Fifteen games are played each year, representing all combinations of the six teams.\n\nThis example uses and includes results from 2014 - 2017, comprising 60 total games. It models latent parameters for each team's attack and defense, as well as a global parameter for home team advantage.\n\nSee https://github.com/arviz-devs/arviz_example_data/blob/main/code/rugby/rugby.ipynb for the whole model specification."
+    },
+    {
+        "name": "rugby_field",
+        "filename": "rugby_field.nc",
+        "url": "http://figshare.com/ndownloader/files/44667112",
+        "checksum": "53a99da7ac40d82cd01bb0b089263b9633ee016f975700e941b4c6ea289a1fb0",
+        "description": "A variant of the 'rugby' example dataset. The Six Nations Championship is a yearly rugby competition between Italy, Ireland, Scotland, England, France and Wales. Fifteen games are played each year, representing all combinations of the six teams.\n\nThis example uses and includes results from 2014 - 2017, comprising 60 total games. It models latent parameters for each team's attack and defense, with each team having different values depending on them being home or away team.\n\nSee https://github.com/arviz-devs/arviz_example_data/blob/main/code/rugby_field/rugby_field.ipynb for the whole model specification."
     },
     {
         "name": "regression1d",

arviz/data/inference_data.py

@@ -9,7 +9,7 @@ from collections import OrderedDict, defaultdict
 from collections.abc import MutableMapping, Sequence
 from copy import copy as ccopy
 from copy import deepcopy
-from datetime import datetime
+import datetime
 from html import escape
 from typing import (
     TYPE_CHECKING,
@@ -394,8 +394,10 @@ class InferenceData(Mapping[str, xr.Dataset]):
             )
 
         try:
-            with h5netcdf.File(filename, mode="r") if engine == "h5netcdf" else nc.Dataset(
-                filename, mode="r"
+            with (
+                h5netcdf.File(filename, mode="r")
+                if engine == "h5netcdf"
+                else nc.Dataset(filename, mode="r")
             ) as file_handle:
                 if base_group == "/":
                     data = file_handle
@@ -744,11 +746,11 @@ class InferenceData(Mapping[str, xr.Dataset]):
         if len(dfs) > 1:
             for group, df in dfs.items():
                 df.columns = [
-                    col
-                    if col in ("draw", "chain")
-                    else (group, *col)
-                    if isinstance(col, tuple)
-                    else (group, col)
+                    (
+                        col
+                        if col in ("draw", "chain")
+                        else (group, *col) if isinstance(col, tuple) else (group, col)
+                    )
                     for col in df.columns
                 ]
             dfs, *dfs_tail = list(dfs.values())
@@ -1475,12 +1477,12 @@ class InferenceData(Mapping[str, xr.Dataset]):
         Examples
         --------
         Add a ``log_likelihood`` group to the "rugby" example InferenceData after loading.
-        It originally doesn't have the ``log_likelihood`` group:
 
         .. jupyter-execute::
 
             import arviz as az
             idata = az.load_arviz_data("rugby")
+            del idata.log_likelihood
             idata2 = idata.copy()
             post = idata.posterior
             obs = idata.observed_data
@@ -1609,13 +1611,13 @@ class InferenceData(Mapping[str, xr.Dataset]):
         .. jupyter-execute::
 
             import arviz as az
-            idata = az.load_arviz_data("rugby")
+            idata = az.load_arviz_data("radon")
 
         Second InferenceData:
 
         .. jupyter-execute::
 
-            other_idata = az.load_arviz_data("radon")
+            other_idata = az.load_arviz_data("rugby")
 
         Call the ``extend`` method:
 
@@ -1687,6 +1689,7 @@ class InferenceData(Mapping[str, xr.Dataset]):
     compute = _extend_xr_method(xr.Dataset.compute)
     persist = _extend_xr_method(xr.Dataset.persist)
     quantile = _extend_xr_method(xr.Dataset.quantile)
+    close = _extend_xr_method(xr.Dataset.close)
 
     # The following lines use methods on xr.Dataset that are dynamically defined and attached.
     # As a result mypy cannot see them, so we have to suppress the resulting mypy errors.
@@ -1918,8 +1921,7 @@ def concat(
     copy: bool = True,
     inplace: "Literal[True]",
     reset_dim: bool = True,
-) -> None:
-    ...
+) -> None: ...
 
 
 @overload
@@ -1929,8 +1931,7 @@ def concat(
     copy: bool = True,
     inplace: "Literal[False]",
     reset_dim: bool = True,
-) -> InferenceData:
-    ...
+) -> InferenceData: ...
 
 
 @overload
@@ -1941,8 +1942,7 @@ def concat(
     copy: bool = True,
     inplace: "Literal[False]",
     reset_dim: bool = True,
-) -> InferenceData:
-    ...
+) -> InferenceData: ...
 
 
 @overload
@@ -1953,8 +1953,7 @@ def concat(
     copy: bool = True,
     inplace: "Literal[True]",
     reset_dim: bool = True,
-) -> None:
-    ...
+) -> None: ...
 
 
 @overload
@@ -1965,8 +1964,7 @@ def concat(
     copy: bool = True,
     inplace: bool = False,
     reset_dim: bool = True,
-) -> Optional[InferenceData]:
-    ...
+) -> Optional[InferenceData]: ...
 
 
 # pylint: disable=protected-access, inconsistent-return-statements
@@ -2083,7 +2081,7 @@ def concat(*args, dim=None, copy=True, inplace=False, reset_dim=True):
             else:
                 return args[0]
 
-    current_time = str(datetime.now())
+    current_time = datetime.datetime.now(datetime.timezone.utc).isoformat()
     combined_attr = defaultdict(list)
     for idata in args:
         for key, val in idata.attrs.items():

arviz/data/io_cmdstan.py

@@ -732,9 +732,7 @@ def _process_configuration(comments):
             key = (
                 "warmup_time_seconds"
                 if "(Warm-up)" in comment
-                else "sampling_time_seconds"
-                if "(Sampling)" in comment
-                else "total_time_seconds"
+                else "sampling_time_seconds" if "(Sampling)" in comment else "total_time_seconds"
             )
             results[key] = float(value)
         elif "=" in comment:

arviz/data/io_datatree.py

@@ -1,4 +1,5 @@
 """Conversion between InferenceData and DataTree."""
+
 from .inference_data import InferenceData
 
 

arviz/data/io_dict.py

@@ -1,4 +1,5 @@
 """Dictionary specific conversion code."""
+
 import warnings
 from typing import Optional
 
@@ -59,9 +60,7 @@ class DictConverter:
         self.coords = (
             coords
             if pred_coords is None
-            else pred_coords
-            if coords is None
-            else {**coords, **pred_coords}
+            else pred_coords if coords is None else {**coords, **pred_coords}
         )
         self.index_origin = index_origin
         self.coords = coords
@@ -458,3 +457,6 @@ def from_dict(
         attrs=attrs,
         **kwargs,
     ).to_inference_data()
+
+
+from_pytree = from_dict

arviz/data/io_emcee.py

@@ -1,4 +1,5 @@
 """emcee-specific conversion code."""
+
 import warnings
 from collections import OrderedDict
 

arviz/data/io_numpyro.py

@@ -1,4 +1,5 @@
 """NumPyro-specific conversion code."""
+
 import logging
 from typing import Callable, Optional
 

arviz/data/io_pyjags.py

@@ -1,4 +1,5 @@
 """Convert PyJAGS sample dictionaries to ArviZ inference data objects."""
+
 import typing as tp
 from collections import OrderedDict
 from collections.abc import Iterable

arviz/data/io_pyro.py

@@ -1,4 +1,5 @@
 """Pyro-specific conversion code."""
+
 import logging
 from typing import Callable, Optional
 import warnings

arviz/data/utils.py

@@ -1,4 +1,5 @@
 """Data specific utilities."""
+
 import warnings
 import numpy as np
 

arviz/plots/__init__.py

@@ -1,4 +1,5 @@
 """Plotting functions."""
+
 from .autocorrplot import plot_autocorr
 from .bpvplot import plot_bpv
 from .bfplot import plot_bf

arviz/plots/autocorrplot.py

@@ -1,4 +1,5 @@
 """Autocorrelation plot of data."""
+
 from ..data import convert_to_dataset
 from ..labels import BaseLabeller
 from ..sel_utils import xarray_var_iter

arviz/plots/backends/bokeh/autocorrplot.py

@@ -1,4 +1,5 @@
 """Bokeh Autocorrplot."""
+
 import numpy as np
 from bokeh.models import DataRange1d, BoxAnnotation
 from bokeh.models.annotations import Title

arviz/plots/backends/bokeh/bpvplot.py

@@ -1,4 +1,5 @@
 """Bokeh Bayesian p-value Posterior predictive plot."""
+
 import numpy as np
 from bokeh.models import BoxAnnotation
 from bokeh.models.annotations import Title

arviz/plots/backends/bokeh/compareplot.py

@@ -1,4 +1,5 @@
 """Bokeh Compareplot."""
+
 from bokeh.models import Span
 from bokeh.models.annotations import Title, Legend
 

arviz/plots/backends/bokeh/densityplot.py

@@ -1,4 +1,5 @@
 """Bokeh Densityplot."""
+
 from collections import defaultdict
 from itertools import cycle
 

arviz/plots/backends/bokeh/distplot.py

@@ -1,4 +1,5 @@
 """Bokeh Distplot."""
+
 import matplotlib.pyplot as plt
 import numpy as np
 

arviz/plots/backends/bokeh/dotplot.py

@@ -1,4 +1,5 @@
 """Bokeh dotplot."""
+
 import math
 import warnings
 import numpy as np

arviz/plots/backends/bokeh/ecdfplot.py

@@ -1,4 +1,5 @@
 """Bokeh ecdfplot."""
+
 from matplotlib.colors import to_hex
 
 from ...plot_utils import _scale_fig_size

arviz/plots/backends/bokeh/elpdplot.py

@@ -1,4 +1,5 @@
 """Bokeh ELPDPlot."""
+
 import warnings
 
 import bokeh.plotting as bkp

arviz/plots/backends/bokeh/energyplot.py

@@ -1,4 +1,5 @@
 """Bokeh energyplot."""
+
 from itertools import cycle
 
 import numpy as np

arviz/plots/backends/bokeh/hdiplot.py

@@ -1,4 +1,5 @@
 """Bokeh hdiplot."""
+
 import numpy as np
 
 from ...plot_utils import _scale_fig_size, vectorized_to_hex

arviz/plots/backends/bokeh/kdeplot.py

@@ -6,7 +6,7 @@ from numbers import Integral
 import numpy as np
 from bokeh.models import ColumnDataSource
 from bokeh.models.glyphs import Scatter
-from matplotlib.cm import get_cmap
+from matplotlib import colormaps
 from matplotlib.colors import rgb2hex
 from matplotlib.pyplot import rcParams as mpl_rcParams
 
@@ -188,7 +188,7 @@ def plot_kde(
 
             cmap = contourf_kwargs.pop("cmap", "viridis")
             if isinstance(cmap, str):
-                cmap = get_cmap(cmap)
+                cmap = colormaps[cmap]
             if isinstance(cmap, Callable):
                 colors = [rgb2hex(item) for item in cmap(np.linspace(0, 1, len(levels_scaled) + 1))]
             else:
@@ -225,7 +225,7 @@ def plot_kde(
         else:
             cmap = pcolormesh_kwargs.pop("cmap", "viridis")
             if isinstance(cmap, str):
-                cmap = get_cmap(cmap)
+                cmap = colormaps[cmap]
             if isinstance(cmap, Callable):
                 colors = [rgb2hex(item) for item in cmap(np.linspace(0, 1, 256))]
             else:

arviz/plots/backends/bokeh/khatplot.py

@@ -1,4 +1,5 @@
 """Bokeh pareto shape plot."""
+
 from collections.abc import Iterable
 
 from matplotlib import cm

arviz/plots/backends/bokeh/lmplot.py

@@ -1,4 +1,5 @@
 """Bokeh linear regression plot."""
+
 import numpy as np
 from bokeh.models.annotations import Legend
 

arviz/plots/backends/bokeh/loopitplot.py

@@ -1,4 +1,5 @@
 """Bokeh loopitplot."""
+
 import numpy as np
 from bokeh.models import BoxAnnotation
 from matplotlib.colors import hsv_to_rgb, rgb_to_hsv, to_hex, to_rgb

arviz/plots/backends/bokeh/mcseplot.py

@@ -1,4 +1,5 @@
 """Bokeh mcseplot."""
+
 import numpy as np
 from bokeh.models import ColumnDataSource, Span
 from bokeh.models.glyphs import Scatter

arviz/plots/backends/bokeh/pairplot.py

@@ -1,4 +1,5 @@
 """Bokeh pairplot."""
+
 import warnings
 from copy import deepcopy
 from uuid import uuid4

arviz/plots/backends/bokeh/parallelplot.py

@@ -1,4 +1,5 @@
 """Bokeh Parallel coordinates plot."""
+
 import numpy as np
 from bokeh.models import DataRange1d
 from bokeh.models.tickers import FixedTicker

arviz/plots/backends/bokeh/posteriorplot.py

@@ -1,4 +1,5 @@
 """Bokeh Plot posterior densities."""
+
 from numbers import Number
 from typing import Optional