Support zarr v3

pyproject.toml

@@ -47,7 +47,7 @@ test = [
     "scipy",
     "pandas",
     "h5py",
-    "zarr>=2.16.0,<3.0.0"
+    "zarr>=3,<4"
 ]
 
 docs = [

src/probeinterface/__init__.py

@@ -1,7 +1,17 @@
 import importlib.metadata
+import importlib.util
+from packaging.version import parse
 
 __version__ = importlib.metadata.version("probeinterface")
 
+# If Zarr is installed, it must be >= 3.0.0
+ZARR_INSTALLED = importlib.util.find_spec("zarr") is not None
+if ZARR_INSTALLED:
+    import zarr
+
+    if parse(zarr.__version__) < parse("3.0.0"):
+        raise ImportError("zarr version must be >= 3.0.0")
+
 
 from .probe import Probe, select_axes
 from .probegroup import ProbeGroup

src/probeinterface/probe.py

@@ -1152,7 +1152,7 @@ def add_probe_to_zarr_group(self, group: "zarr.Group") -> None:
         # add fields and contact annotations
         for field_name, (dtype, offset) in probe_arr.dtype.fields.items():
             data = probe_arr[field_name]
-            group.create_dataset(name=field_name, data=data, dtype=dtype, chunks=False)
+            group.create_array(name=field_name, data=data, chunks=data.shape)
 
         # Annotations as a group (special attributes are stored as annotations)
         annotations_group = group.create_group("annotations")
@@ -1161,8 +1161,8 @@ def add_probe_to_zarr_group(self, group: "zarr.Group") -> None:
 
         # Add planar contour
         if self.probe_planar_contour is not None:
-            group.create_dataset(
-                name="probe_planar_contour", data=self.probe_planar_contour, dtype="float64", chunks=False
+            group.create_array(
+                name="probe_planar_contour", data=self.probe_planar_contour, chunks=self.probe_planar_contour.shape
             )
 
     def to_zarr(self, folder_path: str | Path) -> None:
@@ -1219,7 +1219,7 @@ def from_zarr_group(group: "zarr.Group") -> "Probe":
                 probe_arr_keys.append(key)
                 dtype.append((key, dset.dtype))
                 if num_contacts is None:
-                    num_contacts = len(dset)
+                    num_contacts = dset.shape[0]
 
         # Create a structured array from the datasets
         probe_arr = np.zeros(num_contacts, dtype=dtype)