Use datatree from xarray and update example to latest einstats (#2458)

* Use datatree from xarray and update example to latest einstats

* update test

* black

Author: OriolAbril

arviz/data/inference_data.py

@@ -532,24 +532,27 @@ def to_netcdf(
         return filename
 
     def to_datatree(self):
-        """Convert InferenceData object to a :class:`~datatree.DataTree`."""
+        """Convert InferenceData object to a :class:`~xarray.DataTree`."""
         try:
-            from datatree import DataTree
-        except ModuleNotFoundError as err:
-            raise ModuleNotFoundError(
-                "datatree must be installed in order to use InferenceData.to_datatree"
+            from xarray import DataTree
+        except ImportError as err:
+            raise ImportError(
+                "xarray must be have DataTree in order to use InferenceData.to_datatree. "
+                "Update to xarray>=2024.11.0"
             ) from err
         return DataTree.from_dict({group: ds for group, ds in self.items()})
 
     @staticmethod
     def from_datatree(datatree):
-        """Create an InferenceData object from a :class:`~datatree.DataTree`.
+        """Create an InferenceData object from a :class:`~xarray.DataTree`.
 
         Parameters
         ----------
         datatree : DataTree
         """
-        return InferenceData(**{group: sub_dt.to_dataset() for group, sub_dt in datatree.items()})
+        return InferenceData(
+            **{group: child.to_dataset() for group, child in datatree.children.items()}
+        )
 
     def to_dict(self, groups=None, filter_groups=None):
         """Convert InferenceData to a dictionary following xarray naming conventions.
@@ -1531,9 +1534,8 @@ def add_groups(
             import xarray as xr
             from xarray_einstats.stats import XrDiscreteRV
             from scipy.stats import poisson
-            dist = XrDiscreteRV(poisson)
-            log_lik = xr.Dataset()
-            log_lik["home_points"] = dist.logpmf(obs["home_points"], np.exp(post["atts"]))
+            dist = XrDiscreteRV(poisson, np.exp(post["atts"]))
+            log_lik = dist.logpmf(obs["home_points"]).to_dataset(name="home_points")
             idata2.add_groups({"log_likelihood": log_lik})
             idata2
 

arviz/data/io_datatree.py

@@ -4,7 +4,7 @@
 
 
 def to_datatree(data):
-    """Convert InferenceData object to a :class:`~datatree.DataTree`.
+    """Convert InferenceData object to a :class:`~xarray.DataTree`.
 
     Parameters
     ----------
@@ -14,7 +14,7 @@ def to_datatree(data):
 
 
 def from_datatree(datatree):
-    """Create an InferenceData object from a :class:`~datatree.DataTree`.
+    """Create an InferenceData object from a :class:`~xarray.DataTree`.
 
     Parameters
     ----------

arviz/tests/base_tests/test_data.py

@@ -1501,10 +1501,6 @@ def test_json_converters(self, models):
         assert not os.path.exists(filepath)
 
 
-@pytest.mark.skipif(
-    not (importlib.util.find_spec("datatree") or "ARVIZ_REQUIRE_ALL_DEPS" in os.environ),
-    reason="test requires xarray-datatree library",
-)
 class TestDataTree:
     def test_datatree(self):
         idata = load_arviz_data("centered_eight")

doc/source/conf.py

@@ -115,6 +115,7 @@
 numpydoc_xref_aliases = {
     "DataArray": ":class:`~xarray.DataArray`",
     "Dataset": ":class:`~xarray.Dataset`",
+    "DataTree": ":class:`~xarray.DataTree`",
     "Labeller": ":ref:`Labeller <labeller_api>`",
     "ndarray": ":class:`~numpy.ndarray`",
     "InferenceData": ":class:`~arviz.InferenceData`",

requirements-optional.txt

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