Support zarr v3

.github/scripts/serialization/objects.py

@@ -17,10 +17,12 @@
     the full state (traces or spike trains, properties, annotations, probe) to disk.
     Targets the on-disk encoding axis: property and annotation preservation, and the
     probe representation. "binary" is recording-only, hence "numpy_folder" for sortings.
-"""
-from packaging.version import parse
-from spikeinterface import __version__ as si_version
 
+Additionally, the `check_extra_data` kwarg can be used to store additional information in a JSON file alongside the
+fixture, which is then passed to the check function. This is used for the SortingAnalyzer extension data keys,
+which are computed and stored in the fixture JSON to verify that they are reloaded correctly.
+
+"""
 
 # Filename suffix per format, relative to the fixtures dir (folder formats use a suffix
 # rather than an extension). Both the generator and the loader build a fixture path as
@@ -32,6 +34,7 @@
     "binary_parallel": "_binary_parallel",
     "numpy_folder": "_numpy_folder",
     "zarr": ".zarr",
+    "binary_folder": "_binary_folder",
     "zarr_parallel": "_parallel.zarr",
 }
 
@@ -57,7 +60,7 @@ def _build_noise_generator_recording():
     )
 
 
-def _check_noise_generator_recording(rec):
+def _check_noise_generator_recording(rec, check_extra_data=None):
     assert type(rec).__name__ == "NoiseGeneratorRecording", type(rec).__name__
     assert rec.get_num_channels() == 4
     assert rec.get_num_segments() == 2
@@ -70,7 +73,7 @@ def _build_mock_recording():
     return generate_recording(num_channels=4, durations=[DEFAULT_DURATION], sampling_frequency=30000.0, seed=0)
 
 
-def _check_mock_recording(rec):
+def _check_mock_recording(rec, check_extra_data=None):
     from spikeinterface.core import BaseRecording
 
     assert isinstance(rec, BaseRecording), type(rec).__name__
@@ -91,7 +94,7 @@ def _build_recording_with_properties():
     return rec
 
 
-def _check_recording_with_properties(rec):
+def _check_recording_with_properties(rec, check_extra_data=None):
     assert rec.get_num_channels() == 4
     assert list(rec.get_property("quality")) == ["good", "good", "bad", "good"]
     assert rec.get_annotation("experimenter") == "test"
@@ -105,7 +108,7 @@ def _build_recording_with_probe():
     rec = generate_recording(num_channels=8, durations=[DEFAULT_DURATION], sampling_frequency=30000.0, seed=0)
     probe = generate_linear_probe(num_elec=8)
     probe.set_device_channel_indices(np.arange(8))
-    # old API returns a new recording, new API >=0.105.0 returns None and sets the probe in place
+    # old API (in_place=False by default) returns a new recording; new API is always in place and returns None
     rec_with_probe = rec.set_probe(probe)
     return rec_with_probe if rec_with_probe is not None else rec
 
@@ -117,14 +120,14 @@ def _build_recording_with_timestamps():
     return rec
 
 
-def _check_recording_with_timestamps(rec):
+def _check_recording_with_timestamps(rec, check_extra_data=None):
     import numpy as np
     expected_times = np.arange(int(DEFAULT_DURATION * 30000)) / 30000.0 + 100
     times = rec.get_times(segment_index=0)
     assert np.allclose(times, expected_times)
 
 
-def _check_recording_with_probe(rec):
+def _check_recording_with_probe(rec, check_extra_data=None):
     import numpy as np
 
     assert rec.get_num_channels() == 8
@@ -149,12 +152,12 @@ def _build_recording_with_interleaved_probes():
     probegroup.add_probe(probe0)
     probegroup.add_probe(probe1)
     probegroup.set_global_device_channel_indices([0, 2, 4, 6, 1, 3, 5, 7])
-    # old API returns a new recording, new API >=0.105.0 returns None and sets the probe in place
+    # old API (in_place=False by default) returns a new recording; new API is always in place and returns None
     rec_with_probe = rec.set_probegroup(probegroup)
     return rec_with_probe if rec_with_probe is not None else rec
 
 
-def _check_recording_with_interleaved_probes(rec):
+def _check_recording_with_interleaved_probes(rec, check_extra_data=None):
     import numpy as np
 
     assert rec.get_num_channels() == 8
@@ -187,7 +190,7 @@ def _build_preprocessed_chain():
     return common_reference(scale(rec, gain=2.0))
 
 
-def _check_preprocessed_chain(rec):
+def _check_preprocessed_chain(rec, check_extra_data=None):
     # The outer wrapper and the recursive parent chain must both reload (the kwargs
     # embed the parent recording dict, so this exercises recursive deserialization).
     assert type(rec).__name__ == "CommonReferenceRecording", type(rec).__name__
@@ -201,7 +204,7 @@ def _build_sorting():
     return generate_sorting(num_units=5, sampling_frequency=30000.0, durations=[DEFAULT_DURATION])
 
 
-def _check_sorting(sorting):
+def _check_sorting(sorting, check_extra_data=None):
     assert sorting.get_num_units() == 5, sorting.get_num_units()
     assert sorting.get_num_segments() == 1
     spike_train = sorting.get_unit_spike_train(sorting.unit_ids[0], segment_index=0)
@@ -218,12 +221,41 @@ def _build_sorting_with_properties():
     return sorting
 
 
-def _check_sorting_with_properties(sorting):
+def _check_sorting_with_properties(sorting, check_extra_data=None):
     assert sorting.get_num_units() == 4
     assert list(sorting.get_property("quality")) == ["good", "good", "bad", "good"]
     assert sorting.get_annotation("experimenter") == "test"
 
 
+def _build_sorting_analyzer_with_extensions():
+    from spikeinterface.core import generate_ground_truth_recording, create_sorting_analyzer
+    recording, sorting = generate_ground_truth_recording(durations=[10, 5], num_channels=16, num_units=5, seed=0)
+    analyzer = create_sorting_analyzer(sorting, recording)
+    extensions = analyzer.get_computable_extensions()
+    analyzer.compute(extensions, n_jobs=-1)
+    # for SortingAnalyzer, we also save a JSON file with extension as
+    # keys and date entries as values to check that everything is reloaded correctly
+    check_extra_data = {}
+    for ext_name in analyzer.extensions:
+        ext = analyzer.get_extension(ext_name)
+        check_extra_data[ext_name] = list(ext.data.keys())
+    return analyzer, check_extra_data
+
+
+def _check_sorting_analyzer_with_extensions(analyzer, check_extra_data=None):
+    extensions = analyzer.get_saved_extension_names()
+    for extension in extensions:
+        extension = analyzer.get_extension(extension)  # just check it loads without error
+        ext_data = extension.get_data()
+        assert ext_data is not None, f"extension {extension} returned None data"
+    if check_extra_data is not None:
+        # Check that the extension data keys match what was saved in the fixture JSON.
+        for ext_name, expected_keys in check_extra_data.items():
+            ext = analyzer.get_extension(ext_name)
+            actual_keys = set(ext.data.keys())
+            assert actual_keys == set(expected_keys), f"extension {ext_name} keys mismatch: {actual_keys} != {set(expected_keys)}"
+
+
 OBJECTS = [
     {
         "id": "noise_generator_recording",
@@ -279,4 +311,10 @@ def _check_sorting_with_properties(sorting):
         "check": _check_sorting_with_properties,
         "formats": ["numpy_folder", "zarr"],
     },
+    {
+        "id": "sorting_analyzer_with_extensions",
+        "build": _build_sorting_analyzer_with_extensions,
+        "check": _check_sorting_analyzer_with_extensions,  # no check; the extensions are computed and stored, but the analyzer itself is not serialized
+        "formats": ["binary_folder", "zarr"],
+    },
 ]

.github/scripts/serialization/serialize_objects.py

@@ -12,35 +12,57 @@
 """
 
 import sys
+import json
 from pathlib import Path
 
-import spikeinterface
+from spikeinterface import __version__ as si_version
+from spikeinterface import SortingAnalyzer
+from spikeinterface.core.base import BaseExtractor
+
 
 sys.path.insert(0, str(Path(__file__).parent))
 from objects import OBJECTS, FIXTURE_SUFFIX  # noqa: E402
 
 out_dir = Path(sys.argv[1]) if len(sys.argv) > 1 else Path("serialization_fixtures")
 out_dir.mkdir(parents=True, exist_ok=True)  # do not rmtree an arbitrary path; overwrite in place
 
-print(f"Generating serialization fixtures with spikeinterface {spikeinterface.__version__}")
+print(f"Generating serialization fixtures with spikeinterface {si_version}")
 for entry in OBJECTS:
-    obj = entry["build"]()
+    objs = entry["build"]()
+    if isinstance(objs, tuple) and len(objs) == 2 and isinstance(objs[1], dict):
+        obj, check_extra_data = objs
+    else:
+        obj = objs
+        check_extra_data = None
     for fmt in entry["formats"]:
         dest = out_dir / f"{entry['id']}{FIXTURE_SUFFIX[fmt]}"
-        if fmt == "json":
-            obj.dump_to_json(dest)
-        elif fmt == "pickle":
-            obj.dump_to_pickle(dest)
-        elif fmt == "binary":
-            obj.save(folder=dest, format="binary", overwrite=True)
-        elif fmt == "binary_parallel":
-            obj.save(folder=dest, format="binary", overwrite=True, n_jobs=2)
-        elif fmt == "numpy_folder":
-            obj.save(folder=dest, format="numpy_folder", overwrite=True)
-        elif fmt == "zarr":
-            obj.save(folder=dest, format="zarr", overwrite=True)
-        elif fmt == "zarr_parallel":
-            obj.save(folder=dest, format="zarr", overwrite=True, n_jobs=2)
+        if isinstance(obj, BaseExtractor):
+            if fmt == "json":
+                obj.dump_to_json(dest)
+            elif fmt == "pickle":
+                obj.dump_to_pickle(dest)
+            elif fmt == "binary":
+                obj.save(folder=dest, format="binary", overwrite=True)
+            elif fmt == "binary_parallel":
+                obj.save(folder=dest, format="binary", overwrite=True, n_jobs=2)
+            elif fmt == "numpy_folder":
+                obj.save(folder=dest, format="numpy_folder", overwrite=True)
+            elif fmt == "zarr":
+                obj.save(folder=dest, format="zarr", overwrite=True)
+            elif fmt == "zarr_parallel":
+                obj.save(folder=dest, format="zarr", overwrite=True, n_jobs=2)
+        elif isinstance(obj, SortingAnalyzer):
+            if dest.is_dir():
+                import shutil
+                shutil.rmtree(dest)
+            obj.save_as(folder=dest, format=fmt)
+        else:
+            raise TypeError(f"{entry['id']}: build() returned unsupported object {type(obj).__name__}")
         print(f"  wrote {dest.name} ({fmt})")
+    if check_extra_data is not None:
+        json_dest = out_dir / f"{entry['id']}.json"
+        with open(json_dest, "w") as f:
+            json.dump(check_extra_data, f)
+        print(f"  wrote {json_dest.name} (extra data)")
 
 print(f"Fixtures written to: {out_dir.resolve()}")

.github/scripts/serialization/test_cross_version_compatibility.py

@@ -9,6 +9,7 @@
 """
 
 import os
+import json
 from pathlib import Path
 
 import pytest
@@ -27,4 +28,8 @@ def test_load_old_serialized_object(entry, fmt):
     fixture = FIXTURES_DIR / f"{entry['id']}{FIXTURE_SUFFIX[fmt]}"
     assert fixture.exists(), f"missing fixture {fixture}"
     obj = load(fixture)
-    entry["check"](obj)
+    check_extra_data = None
+    if (FIXTURES_DIR / f"{entry['id']}.json").exists():
+        with open(FIXTURES_DIR / f"{entry['id']}.json", "r") as f:
+            check_extra_data = json.load(f)
+    entry["check"](obj, check_extra_data=check_extra_data)

.github/workflows/core-test.yml

@@ -0,0 +1,48 @@
+name: Testing core
+
+on:
+  pull_request:
+    types: [synchronize, opened, reopened]
+    branches:
+      - main
+
+concurrency: # Cancel previous workflows on the same pull request
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  build-and-test:
+    name: Test on ${{ matrix.os }} OS
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: ["ubuntu-latest", "macos-latest", "windows-latest"]
+    steps:
+      - uses: actions/checkout@v6
+      - uses: actions/setup-python@v6
+        with:
+          python-version: "3.12"
+      - uses: astral-sh/setup-uv@v7
+        with:
+          python-version: "3.12"
+          enable-cache: true
+      - name: Install dependencies
+        run: |
+          git config --global user.email "CI@example.com"
+          git config --global user.name "CI Almighty"
+          uv pip install --system -e . --group test-core
+      - name: Test core with pytest
+        run: |
+          pytest -m "core" -vv -ra --durations=0 --durations-min=0.001 | tee report.txt; test $? -eq 0 || exit 1
+        shell: bash # Necessary for pipeline to work on windows
+      - name: Build test summary
+        run: |
+          uv pip install --system pandas
+          uv pip install --system tabulate
+          echo "# Timing profile of core tests in ${{matrix.os}}" >> $GITHUB_STEP_SUMMARY
+          # Outputs markdown summary to standard output
+          python ./.github/scripts/build_job_summary.py report.txt >> $GITHUB_STEP_SUMMARY
+          cat $GITHUB_STEP_SUMMARY
+          rm report.txt
+        shell: bash # Necessary for pipeline to work on windows

.github/workflows/cross_version_serialization.yml

@@ -63,7 +63,8 @@ jobs:
           candidates = [v for v in available if not v.is_prerelease and (v.major, v.minor) >= (floor.major, floor.minor)]
           minor_releases = [max(group) for _, group in groupby(candidates, key=lambda v: (v.major, v.minor))]
 
-          print(json.dumps([str(v) for v in minor_releases]))
+          # always round-trip against the branch itself as well
+          print(json.dumps([str(v) for v in minor_releases] + ['dev']))
           ")
 
           echo "list=$versions_to_test" >> "$GITHUB_OUTPUT"
@@ -94,12 +95,25 @@ jobs:
           enable-cache: false
 
       - name: Generate fixtures with spikeinterface ${{ matrix.si-version }}
+        env:
+          SI_VERSION: ${{ matrix.si-version }}
         run: |
+          # "dev" means the checked-out branch; anything else is a release from PyPI.
+          # [full] is needed because the SortingAnalyzer fixture computes every
+          # extension (scikit-learn, numba, ...).
+          if [ "$SI_VERSION" = "dev" ]; then
+            spec=".[full]"
+          else
+            spec="spikeinterface[full]==$SI_VERSION"
+          fi
           uv run --isolated --no-project --python 3.11 \
-            --with "spikeinterface[core]==${{ matrix.si-version }}" \
+            --with "$spec" \
+            --with "numpy<2.5" \
             python .github/scripts/serialization/serialize_objects.py "$SI_SERIALIZATION_FIXTURES_DIR"
 
       - name: Load fixtures with the branch
         run: |
-          uv pip install --system -e . --group test-core
+          # [full] here too: loading the analyzer extensions unpickles e.g. the
+          # scikit-learn PCA models written at generation time
+          uv pip install --system -e ".[full]" "numpy<2.5" --group test-core
           pytest .github/scripts/serialization/test_cross_version_compatibility.py -v

.github/workflows/deepinterpolation.yml

@@ -20,10 +20,10 @@ jobs:
       - uses: actions/checkout@v6
       - uses: actions/setup-python@v6
         with:
-          python-version: "3.10"
+          python-version: "3.12"
       - uses: astral-sh/setup-uv@v7
         with:
-          python-version: "3.10"
+          python-version: "3.12"
           enable-cache: false
       - name: Get changed files
         id: changed-files

conftest.py

@@ -47,16 +47,10 @@ def pytest_collection_modifyitems(config, items):
     This function marks (in the pytest sense) the tests according to their name and file_path location
     Marking them in turn allows the tests to be run by using the pytest -m marker_name option.
     """
-
-    from spikeinterface.core.core_tools import _is_zarr_write_supported
-
     rootdir = Path(config.rootdir)
     modules_location = rootdir / "src" / "spikeinterface"
     for item in items:
         # TODO: remove once writing to zarr is supported with zarr>=3
-        if item.get_closest_marker("requires_zarr_write") and not _is_zarr_write_supported():
-            item.add_marker(pytest.mark.skip(reason="Writing to zarr is not supported yet with zarr>=3"))
-
         if config.getoption("--mp-context") is not None and item.name == "test_global_job_kwargs":
             item.add_marker(pytest.mark.skip(reason="--mp-context changes the default job kwargs"))