Make imports compatible with ArviZ 1.0

Closes #127

Author: michaelosthege

CITATION.cff

@@ -25,7 +25,7 @@ abstract: >-
   only MCMC draws but also sampler statistics, and are
   compatible with sparse data, or varying dimensionality.
   MCMC chains stored with McBackend can be queried directly,
-  or convert to the popular ArviZ InferenceData objects.
+  or convert to the popular ArviZ inference data objects.
 keywords:
   - mcmc
   - arviz

README.md

@@ -15,7 +15,7 @@ The `mcbackend` package consists of three parts:
 ### Part 1: A schema for MCMC run & chain metadata
 No matter which programming language your favorite PPL is written in, the [ProtocolBuffers](https://developers.google.com/protocol-buffers/) from McBackend can be used to generate code in languages like C++, C#, Python and many more to represent commonly used metadata about MCMC runs, chains and model variables.
 
-The definitions in [`protobufs/meta.proto`](./protobufs/meta.proto) are designed to maximize compatibility with [`ArviZ`](https://github.com/arviz-devs/arviz) objects, making it easy to transform MCMC draws stored according to the McBackend schema to `InferenceData` objects for plotting & analysis.
+The definitions in [`protobufs/meta.proto`](./protobufs/meta.proto) are designed to maximize compatibility with [`ArviZ`](https://github.com/arviz-devs/arviz) objects, making it easy to transform MCMC draws stored according to the McBackend schema to `xarray.DataTree` objects for plotting & analysis.
 
 ### Part 2: A storage backend interface
 The  `draws` and `stats` created by MCMC sampling algorithms at runtime need to be stored _somewhere_.
@@ -83,10 +83,10 @@ chain = run.get_chains()[0]
 chain.get_draws("my favorite variable")
 # >>> array([ ... ])
 
-# Convert everything to `InferenceData`
+# Convert everything to an inference data structure
 idata = run.to_inferencedata()
 print(idata)
-# >>> Inference data with groups:
+# >>> DataTree:
 # >>> 	> posterior
 # >>> 	> sample_stats
 # >>> 	> observed_data
@@ -113,7 +113,7 @@ Getting rid of `MultiTrace` was a [long-term goal](https://github.com/pymc-devs/
 First clone the repository and set up a development environment containing the protobuf compiler.
 
 ```bash
-mamba create -n mcb python=3.11 grpcio-tools protobuf -y
+mamba create -n mcb python=3.13 grpcio-tools protobuf -y
 activate mcb
 pip install -r requirements-dev.txt
 pip install --pre "betterproto[compiler]"

mcbackend/__init__.py

@@ -14,7 +14,7 @@
 except ModuleNotFoundError:
     pass
 
-__version__ = "0.5.3"
+__version__ = "0.5.4"
 __all__ = [
     "NumPyBackend",
     "NullBackend",

mcbackend/core.py

@@ -14,11 +14,16 @@
 from .utils import as_array_from_ragged
 
 try:
-    from arviz import InferenceData, from_dict
+    from arviz import from_dict
+
+    try:
+        from arviz import InferenceData as IData
+    except ImportError:
+        from xarray import DataTree as IData
 
     _HAS_ARVIZ = True
 except ModuleNotFoundError:
-    InferenceData = TypeVar("InferenceData")  # type: ignore
+    IData = TypeVar("IData")  # type: ignore
     _HAS_ARVIZ = False
 
 Shape = Sequence[int]
@@ -192,8 +197,8 @@ def constant_data(self) -> Dict[str, numpy.ndarray]:
     def observed_data(self) -> Dict[str, numpy.ndarray]:
         return {dv.name: ndarray_to_numpy(dv.value) for dv in self.meta.data if dv.is_observed}
 
-    def to_inferencedata(self, *, equalize_chain_lengths: bool = True, **kwargs) -> InferenceData:
-        """Creates an ArviZ ``InferenceData`` object from this run.
+    def to_inferencedata(self, *, equalize_chain_lengths: bool = True, **kwargs) -> IData:
+        """Creates an ArviZ inference data structure from this run.
 
         Parameters
         ----------
@@ -204,8 +209,9 @@ def to_inferencedata(self, *, equalize_chain_lengths: bool = True, **kwargs) ->
 
         Returns
         -------
-        idata : arviz.InferenceData
+        idata
             Samples and metadata of this inference run.
+            ``az.InferenceData`` (ArviZ <1) or ``xarray.DataTree`` (ArviZ >1).
         """
         if not _HAS_ARVIZ:
             raise ModuleNotFoundError("ArviZ is not installed.")
@@ -216,7 +222,7 @@ def to_inferencedata(self, *, equalize_chain_lengths: bool = True, **kwargs) ->
         nonrigid_vars = {var for var in variables if var.undefined_ndim or not is_rigid(var.shape)}
         if nonrigid_vars:
             raise NotImplementedError(
-                "Creating InferenceData from runs with non-rigid variables is not supported."
+                "Creating inference data from runs with non-rigid variables is not supported."
                 f" The non-rigid variables are: {nonrigid_vars}."
             )
 
@@ -226,11 +232,11 @@ def to_inferencedata(self, *, equalize_chain_lengths: bool = True, **kwargs) ->
             if not equalize_chain_lengths:
                 msg += (
                     "\nArviZ does not properly support uneven chain lengths (see ArviZ issue #2094)."
-                    "\nWe'll try to give you an InferenceData, but best case the chain & draw dimensions"
+                    "\nWe'll try to give you a DataTree, but best case the chain & draw dimensions"
                     " will be messed-up as {'chain': 1, 'draws': n_chains}."
-                    "\nYou won't be able to save this InferenceData to a file"
+                    "\nYou might not be able to save this DataTree to a file"
                     " and you should expect many ArviZ functions to choke on it."
-                    "\nSpecify `to_inferencedata(equalize_chain_lengths=True)` to get regular InferenceData."
+                    "\nSpecify `to_inferencedata(equalize_chain_lengths=True)` to get regular inference data."
                 )
             else:
                 msg += "\nTruncating to the length of the shortest chain."

mcbackend/test_utils.py

@@ -8,6 +8,7 @@
 import numpy
 import pandas
 import pytest
+import xarray
 
 import mcbackend
 from mcbackend import utils as mutils
@@ -299,7 +300,10 @@ def test__to_inferencedata(self, tstatname, caplog):
             chain.append(d, s)
 
         idata = run.to_inferencedata()
-        assert isinstance(idata, arviz.InferenceData)
+        try:
+            assert isinstance(idata, arviz.InferenceData)
+        except AttributeError:
+            assert isinstance(idata, xarray.DataTree)
         assert idata.warmup_posterior.dims["chain"] == 1
         assert idata.posterior.dims["chain"] == 1
         if tstatname == "nottune":

requirements-dev.txt

@@ -1,4 +1,4 @@
-arviz
+arviz>=1.0.0rc0
 clickhouse-driver
 flake8
 pre-commit