Zarr 3 compatibility

[OriolAbril]

The reading of inferencedata objects from zarr uses API elements that are being modified/deprecated with zarr 3.0. Ref zarr-developers/zarr-python#1777. I will add a pin for now in #2356 but we should check the changes and decide how to support both (ideally) otherwise when to switch to 3.x

[ahartikainen]

Any ideas how much API will change?

[OriolAbril]

Haven't really checked, but it looks like some basic things like store/storage are changing

[kmiddleton]

I think the update to zarr is generating errors with arviz.to_zarr() now with pymc 5.20.1, arviz 0.20.0, and zarr 3.0.2.

Trying to save and InferenceData object to zarr gives the error:

    store = zarr.storage.DirectoryStore(path=store)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: module 'zarr.storage' has no attribute 'DirectoryStore'

I think it is the same issue as: zarr-developers/zarr-python#2699

Here is a minimal example:

import pymc as pm
import numpy as np
import arviz as az

np.random.seed(324777)
observed_data = np.random.normal(loc=5.0, scale=2.0, size=100)

with pm.Model() as model:
    mu = pm.Normal("mu", mu=0, sigma=10)
    sigma = pm.HalfNormal("sigma", sigma=10)
    likelihood = pm.Normal("obs", mu=mu, sigma=sigma, observed=observed_data)
    trace = pm.sample(1000)

trace.to_zarr("idata")

Versions of imported packages:

re: 2.2.1
logging: 0.5.1.2
ipaddress: 1.0
json: 2.0.9
platform: 1.0.8
zlib: 1.0
numpy.version: 1.26.4
numpy.core._multiarray_umath: 3.1
_ctypes: 1.1.0
ctypes: 1.1.0
numpy.core: 1.26.4
numpy.linalg._umath_linalg: 0.1.5
numpy: 1.26.4
scipy: 1.15.1
numpy._core._multiarray_umath: 3.1
scipy._lib.array_api_compat: 1.9.1
scipy._lib.array_api_compat.numpy: 1.26.4
scipy._lib.array_api_extra: 0.2.0
scipy.linalg._fblas: 2.0.2
scipy.linalg._flapack: 2.0.2
scipy._lib.decorator: 4.0.5
scipy.sparse.linalg._eigen.arpack._arpack: 2.0.2
scipy.sparse.linalg._propack._spropack: 2.0.2
scipy.sparse.linalg._propack._dpropack: 2.0.2
scipy.sparse.linalg._propack._cpropack: 2.0.2
scipy.sparse.linalg._propack._zpropack: 2.0.2
scipy.optimize._cobyla: 2.0.2
scipy.optimize._slsqp: 2.0.2
scipy._lib._uarray: 0.8.8.dev0+aa94c5a4.scipy
_decimal: 1.70
decimal: 1.70
scipy.integrate._vode: 2.0.2
scipy.integrate._dop: 2.0.2
scipy.integrate._lsoda: 2.0.2
scipy.interpolate._dfitpack: 2.0.2
scipy.stats._mvn: 2.0.2
six: 1.17.0
multipledispatch: 0.6.0
unification: 0.4.6
cons: 0.4.6
packaging: 24.2
etuples: 0.3.9
pytensor: 2.27.1
pytz: 2024.1
dateutil._version: 2.9.0.post0
dateutil: 2.9.0.post0
pyarrow._generated_version: 19.0.0
cloudpickle: 3.1.1
pyarrow: 19.0.0
_csv: 1.0
csv: 1.0
pandas._version_meson: 2.2.3
pandas: 2.2.3
PIL._version: 11.1.0
PIL: 11.1.0
PIL._deprecate: 11.1.0
PIL.Image: 11.1.0
pyparsing: 3.2.1
cycler: 0.12.1
kiwisolver._cext: 1.4.8
kiwisolver: 1.4.8
matplotlib: 3.10.0
xarray: 2025.1.2
arviz.data.base: 0.20.0
urllib.request: 3.12
yaml: 6.0.2
llvmlite: 0.44.0
colorama: 0.4.6
numba.core.types.scalars.np: 1.26.4
numba.cloudpickle: 3.0.0
numba.core.typing.builtins.np: 1.26.4
numba.cpython.builtins.np: 1.26.4
numba.misc.appdirs: 1.4.1
numba: 0.61.0
numba.cpython.mathimpl.np: 1.26.4
numba.cpython.mathimpl.llvmlite: 0.44.0
xarray_einstats: 0.8.0
cffi: 1.17.1
_cffi_backend: 1.17.1
pycparser.ply: 3.9
pycparser.ply.yacc: 3.10
pycparser.ply.lex: 3.10
pycparser: 2.22
numba.cpython.hashing.np: 1.26.4
numba.cpython.hashing.ctypes: 1.1.0
numba.cpython.numbers.np: 1.26.4
numba.np.arraymath.llvmlite: 0.44.0
numba.np.arraymath.np: 1.26.4
numba.np.random.distributions.np: 1.26.4
numba.np.random.random_methods.np: 1.26.4
arviz: 0.20.0
cachetools: 5.5.1
numcodecs.version: 0.15.1
wrapt: 1.17.2
deprecated: 1.2.18
numcodecs.blosc: 1.21.6
numcodecs.zstd: 1.5.6
numcodecs.lz4: 1.9.4
msgpack: 1.1.0
numcodecs: 0.15.1
zarr._version: 3.0.2
donfig: 0.8.1.post1
zarr: 3.0.2
ctypes.macholib: 1.0
threadpoolctl: 3.5.0
IPython.core.release: 8.32.0
traitlets._version: 5.14.3
traitlets: 5.14.3
socketserver: 0.4
argparse: 1.1
executing.version: 2.1.0
executing: 2.1.0
pure_eval.version: 0.2.3
pure_eval: 0.2.3
stack_data.version: 0.6.3
stack_data: 0.6.3
pygments: 2.19.1
pickleshare: 0.7.5
decorator: 5.1.1
wcwidth: 0.2.13
prompt_toolkit: 3.0.50
parso: 0.8.4
jedi: 0.19.2
IPython: 8.32.0
pymc: 5.20.1

[ahartikainen]

Yeah, the api is very different. We probably should restrict zarr<3.0.0 until we are ready to move to the new api.

Edit. We are restricting it