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

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
@@ -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):