Make from_dict more flexible, and add from_pytree

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## v0.x.x Unreleased
 
 ### New features
+- Support for `pytree`s and robust to nested dictionaries. (2291)
 
 ### Maintenance and fixes
 - Fix deprecations introduced in latest pandas and xarray versions, and prepare for numpy 2.0 ones ([2315](https://github.com/arviz-devs/arviz/pull/2315)))

arviz/data/__init__.py

@@ -7,7 +7,7 @@
 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,6 +38,7 @@
  "from_cmdstanpy",
  "from_datatree",
  "from_dict",
+ "from_pytree",
  "from_json",
  "from_pyro",
  "from_numpyro",

arviz/data/base.py

@@ -8,6 +8,7 @@
 from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar, Union
 
 import numpy as np
+import tree
 import xarray as xr
 
 try:
@@ -67,6 +68,48 @@ def wrapped(cls: RequiresArgTypeT) -> Optional[RequiresReturnTypeT]:
  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 +298,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,11 +309,34 @@ def dict_to_dataset(
  index_origin=None,
  skip_event_dims=None,
 ):
- """Convert a dictionary of numpy arrays to an xarray.Dataset.
+ """Convert a 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.
+
+ In case of nested pytrees, the variable name will be a tuple of individual names.
+
+ For example,
+
+ pytree_to_dataset({'top': {'second': 1.}, 'top2': 1.})
+
+ will have `var_names` `('top', 'second')` and `top2`.
+
+ Dimensions and co-ordinates can be defined as usual:
+
+ datadict = {
+ "top": {"a": np.random.randn(100), "b": np.random.randn(1, 100, 10)},
+ "d": np.random.randn(100),
+ }
+ dataset = convert_to_dataset(datadict,
+ coords={"c": np.arange(10)},
+ dims={("top", "b"): ["c"]})
+
+ Then `dataset.data_vars` will be `('top', 'a'), ('top', 'b'), 'd'`.
 
  Parameters
  ----------
- data : dict[str] -> ndarray
+ data : pytree
  Data to convert. Keys are variable names.
  attrs : dict
  Json serializable metadata to attach to the dataset, in addition to defaults.
@@ -302,6 +368,10 @@ def dict_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.
 
  data_vars = {
  key: numpy_to_data_array(
@@ -318,6 +388,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.
 

arviz/data/converters.py

@@ -1,5 +1,6 @@
 """High level conversion functions."""
 import numpy as np
+import tree
 import xarray as xr
 
 from .base import dict_to_dataset
@@ -105,6 +106,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 +121,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/io_dict.py

@@ -458,3 +458,6 @@ def from_dict(
  attrs=attrs,
  **kwargs,
  ).to_inference_data()
+
+
+from_pytree = from_dict

arviz/plots/backends/matplotlib/pairplot.py

@@ -333,7 +333,7 @@ def plot_pair(
  if reference_values:
  x_name = flat_var_names[i]
  y_name = flat_var_names[j + not_marginals]
- if x_name and y_name not in difference:
+ if (x_name not in difference) and (y_name not in difference):
  ax[j, i].plot(
  reference_values_copy[x_name],
  reference_values_copy[y_name],

arviz/tests/base_tests/test_data.py

@@ -1077,6 +1077,20 @@ def test_dict_to_dataset():
  assert set(dataset.b.coords) == {"chain", "draw", "c"}
 
 
+def test_nested_dict_to_dataset():
+ datadict = {
+ "top": {"a": np.random.randn(100), "b": np.random.randn(1, 100, 10)},
+ "d": np.random.randn(100),
+ }
+ dataset = convert_to_dataset(datadict, coords={"c": np.arange(10)}, dims={("top", "b"): ["c"]})
+ assert set(dataset.data_vars) == {("top", "a"), ("top", "b"), "d"}
+ assert set(dataset.coords) == {"chain", "draw", "c"}
+
+ assert set(dataset[("top", "a")].coords) == {"chain", "draw"}
+ assert set(dataset[("top", "b")].coords) == {"chain", "draw", "c"}
+ assert set(dataset.d.coords) == {"chain", "draw"}
+
+
 def test_dict_to_dataset_event_dims_error():
  datadict = {"a": np.random.randn(1, 100, 10)}
  coords = {"b": np.arange(10), "c": ["x", "y", "z"]}

requirements.txt

@@ -4,6 +4,7 @@ numpy>=1.22.0,<2.0
 scipy>=1.8.0
 packaging
 pandas>=1.4.0
+dm-tree>=0.1.8
 xarray>=0.21.0
 h5netcdf>=1.0.2
 typing_extensions>=4.1.0