Support zarr v3

.github/workflows/full_tests.yml

@@ -17,14 +17,14 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.10", "3.14"]  # Lower and higher versions we support
+        python-version: ["3.11", "3.14"]  # Lower and higher versions we support
         os: [macos-latest, windows-latest, ubuntu-latest]
     steps:
     - uses: actions/checkout@v6
     - name: Set up Python
       uses: actions/setup-python@v6
       with:
-        python-version: "3.11"
+        python-version: ${{ matrix.python-version }}
     - name: Install package
       run: |
         python -m pip install --upgrade pip

.readthedocs.yaml

@@ -3,7 +3,7 @@ version: 2
 build:
   os: ubuntu-24.04
   tools:
-    python: "3.10"
+    python: "3.11"
   jobs:
     install:
       - pip install -U pip

pyproject.toml

@@ -9,7 +9,7 @@ authors = [
 
 description = "Python package to handle probe layout, geometry and wiring to device."
 readme = "README.md"
-requires-python = ">=3.10"
+requires-python = ">=3.11"
 classifiers = [
     "Programming Language :: Python :: 3",
     "License :: OSI Approved :: MIT License",
@@ -47,7 +47,7 @@ test = [
     "scipy",
     "pandas",
     "h5py",
-    "zarr>=2.16.0,<3.0.0",
+    "zarr>=3,<4"
 ]
 
 docs = [

src/probeinterface/probe.py

@@ -1240,12 +1240,15 @@ def add_probe_to_zarr_group(self, group: "zarr.Group") -> None:
         group : zarr.Group
             The target Zarr group where the probe's data will be stored.
         """
+        if not hasattr(group, "create_array"):
+            raise ImportError("Saving a probe to zarr requires zarr>=3")
+
         probe_arr = self.to_numpy(complete=True)
 
         # 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")
@@ -1254,8 +1257,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:
@@ -1312,7 +1315,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)

tests/test_probe.py

@@ -229,6 +229,21 @@ def test_save_to_zarr(tmp_path):
     assert probe == reloaded_probe, "Reloaded Probe object does not match the original"
 
 
+@pytest.mark.parametrize("zarr_format", [2, 3])
+def test_zarr_group_formats(tmp_path, zarr_format):
+    # A probe saved in a zarr format 2 group, as zarr-python 2 wrote it, is still readable
+    import zarr
+
+    probe = generate_dummy_probe()
+    folder_path = Path(tmp_path) / "probe.zarr"
+
+    group = zarr.open_group(folder_path, mode="w", zarr_format=zarr_format)
+    probe.add_probe_to_zarr_group(group)
+
+    reloaded_probe = Probe.from_zarr(folder_path=folder_path)
+    assert probe == reloaded_probe
+
+
 def test_position_uniqueness():
     """Test that the error message matches the full expected string for three duplicates using pytest's match regex."""
     import re