make dm-tree optional

arviz/data/base.py

@@ -9,9 +9,13 @@
 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 @@
         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 @@
     return xr.DataArray(ary, coords=coords, dims=dims)
 
 
-def pytree_to_dataset(
+def dict_to_dataset(
     data,
     *,
     attrs=None,
@@ -312,6 +319,8 @@
 ):
     """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 @@
     """
     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 @@
     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/tests/base_tests/test_data.py

@@ -44,6 +44,10 @@
     models,
 )
 
+# Check if dm-tree is installed
+dm_tree_installed = importlib.util.find_spec("tree") is not None  # pylint: disable=invalid-name
+skip_tests = (not dm_tree_installed) and ("ARVIZ_REQUIRE_ALL_DEPS" not in os.environ)
+
 
 @pytest.fixture(autouse=True)
 def no_remote_data(monkeypatch, tmpdir):
@@ -1076,6 +1080,7 @@ def test_dict_to_dataset():
     assert set(dataset.b.coords) == {"chain", "draw", "c"}
 
 
+@pytest.mark.skipif(skip_tests, reason="test requires dm-tree which is not installed")
 def test_nested_dict_to_dataset():
     datadict = {
         "top": {"a": np.random.randn(100), "b": np.random.randn(1, 100, 10)},

requirements-docs.txt

@@ -15,3 +15,4 @@ sphinx_design
 sphinx-codeautolink>=0.9.0
 jupyter-sphinx
 sphinxcontrib-youtube
+dm-tree>=0.1.8

requirements-optional.txt

@@ -6,3 +6,4 @@ ujson
 dask[distributed]
 zarr>=2.5.0,<3
 xarray-datatree
+dm-tree>=0.1.8

requirements.txt

@@ -4,7 +4,6 @@ numpy>=1.23.0
 scipy>=1.9.0
 packaging
 pandas>=1.5.0
-dm-tree>=0.1.8
 xarray>=2022.6.0
 h5netcdf>=1.0.2
 typing_extensions>=4.1.0