Use datatree from xarray (#35)

* update pyproject to use github arviz-base and datatree from xarray

* update leftover datatree.DataTree references

* use xarray.DataTree in tests

docs/source/conf.py

@@ -86,14 +86,13 @@
 numpydoc_xref_aliases = {
     "DataArray": ":class:`xarray.DataArray`",
     "Dataset": ":class:`xarray.Dataset`",
-    "DataTree": ":class:`datatree.DataTree`",
+    "DataTree": ":class:`xarray.DataTree`",
     **{f"{singular}s": f":any:`{singular}s <{singular}>`" for singular in singulars},
 }
 
 intersphinx_mapping = {
     "arviz_org": ("https://www.arviz.org/en/latest/", None),
     "dask": ("https://docs.dask.org/en/latest/", None),
-    "datatree": ("https://xarray-datatree.readthedocs.io/en/latest/", None),
     "numpy": ("https://numpy.org/doc/stable/", None),
     "python": ("https://docs.python.org/3/", None),
     "xarray": ("https://docs.xarray.dev/en/stable/", None),

pyproject.toml

@@ -40,7 +40,7 @@ funding = "https://opencollective.com/arviz"
 
 [project.optional-dependencies]
 xarray = [
-  "arviz-base @ git+https://github.com/arviz-devs/arviz-base",
+  "arviz-base @ git+https://github.com/arviz-devs/arviz-base@xarray_datatree",
   "xarray-einstats",
   "numba",
 ]

src/arviz_stats/accessors.py

@@ -6,7 +6,6 @@
 import numpy as np
 import xarray as xr
 from arviz_base.utils import _var_names
-from datatree import DataTree, register_datatree_accessor
 from xarray_einstats.numba import ecdf
 
 from arviz_stats.utils import get_function
@@ -38,7 +37,7 @@ def update_kwargs_with_dims(da, kwargs):
 def check_var_name_subset(obj, var_name):
     if isinstance(obj, xr.Dataset):
         return obj[var_name]
-    if isinstance(obj, DataTree):
+    if isinstance(obj, xr.DataTree):
         return obj.ds[var_name]
     return obj
 
@@ -228,7 +227,7 @@ def power_scale_sense(self, dims=None, **kwargs):
         return self._apply("power_scale_sense", dims=dims, **kwargs)
 
 
-@register_datatree_accessor("azstats")
+@xr.register_datatree_accessor("azstats")
 class AzStatsDtAccessor(_BaseAccessor):
     """ArviZ stats accessor class for DataTrees."""
 
@@ -253,7 +252,7 @@ def _apply(self, func_name, group, **kwargs):
         if isinstance(group, Hashable):
             group = [group]
             hashable_group = True
-        out_dt = DataTree.from_dict(
+        out_dt = xr.DataTree.from_dict(
             {
                 group_i: apply_function_to_dataset(
                     get_function(func_name),

src/arviz_stats/numba/diagnostics.py

@@ -4,7 +4,6 @@
 import numpy as np
 import scipy
 import xarray as xr
-from datatree import DataTree
 from scipy.fftpack import next_fast_len
 from xarray_einstats import stats
 from xarray_einstats.einops import raw_rearrange
@@ -118,7 +117,7 @@ def rhat(ds, group="posterior", method="rank", **kwargs):
     if method not in func_map:
         raise ValueError("method not recognized")
     rhat_func = func_map[method]
-    if isinstance(ds, DataTree):
+    if isinstance(ds, xr.DataTree):
         ds = ds[group]
         return ds.map(rhat_func, **kwargs)
     if isinstance(ds, xr.Dataset):
@@ -283,7 +282,7 @@ def ess(ds, group="posterior", method="bulk", **kwargs):
     if method not in func_map:
         raise ValueError("method not recognized")
     ess_func = func_map[method]
-    if isinstance(ds, DataTree):
+    if isinstance(ds, xr.DataTree):
         ds = ds[group]
         return ds.map(ess_func, **kwargs)
     if isinstance(ds, xr.Dataset):

tests/test_accessors.py

@@ -9,7 +9,7 @@
 import numpy as np
 import pytest
 from arviz_base import from_dict
-from datatree import DataTree
+from xarray import DataTree
 
 
 @pytest.fixture(scope="module")