From adf8c9eeddfe3302595764eba2da038c44f2c126 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Thu, 11 Dec 2025 13:14:02 +0100 Subject: [PATCH 01/59] wip: support zarr v3 --- pyproject.toml | 3 +- src/spikeinterface/core/sortinganalyzer.py | 94 ++++++++++--------- src/spikeinterface/core/template.py | 12 ++- .../tests/test_analyzer_extension_core.py | 2 +- .../core/tests/test_baserecording.py | 4 +- .../core/tests/test_sortinganalyzer.py | 45 ++++----- .../core/tests/test_zarrextractors.py | 35 ++++--- src/spikeinterface/core/zarrextractors.py | 91 ++++++++++++------ 8 files changed, 169 insertions(+), 117 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 8c3a3cf3b1..a753911fa7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,12 +24,11 @@ dependencies = [ "numpy>=2.0.0;python_version>='3.13'", "threadpoolctl>=3.0.0", "tqdm", - "zarr>=2.18,<3", + "zarr>=3,<4", "neo>=0.14.3", "probeinterface>=0.3.1", "packaging", "pydantic", - "numcodecs<0.16.0", # For supporting zarr < 3 ] [build-system] diff --git a/src/spikeinterface/core/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py index 1870c24e7a..31cb317565 100644 --- a/src/spikeinterface/core/sortinganalyzer.py +++ b/src/spikeinterface/core/sortinganalyzer.py @@ -1,5 +1,5 @@ from __future__ import annotations -from typing import Literal, Optional, Any +from typing import Literal, Optional, Any, Iterable from pathlib import Path from itertools import chain @@ -621,6 +621,7 @@ def _get_zarr_root(self, mode="r+"): assert mode in ("r+", "a", "r"), "mode must be 'r+', 'a' or 'r'" storage_options = self._backend_options.get("storage_options", {}) + zarr_root = super_zarr_open(self.folder, mode=mode, storage_options=storage_options) return zarr_root @@ -644,7 +645,12 @@ def create_zarr(cls, folder, sorting, recording, sparsity, return_in_uV, rec_att storage_options = backend_options.get("storage_options", {}) saving_options = backend_options.get("saving_options", {}) - zarr_root = zarr.open(folder, mode="w", storage_options=storage_options) + if not is_path_remote(str(folder)): + storage_options_kwargs = {} + else: + storage_options_kwargs = storage_options + + zarr_root = zarr.open(folder, mode="w", **storage_options_kwargs) info = dict(version=spikeinterface.__version__, dev_mode=spikeinterface.DEV_MODE, object="SortingAnalyzer") zarr_root.attrs["spikeinterface_info"] = check_json(info) @@ -657,13 +663,8 @@ def create_zarr(cls, folder, sorting, recording, sparsity, return_in_uV, rec_att if recording is not None: rec_dict = recording.to_dict(relative_to=relative_to, recursive=True) if recording.check_serializability("json"): - # zarr_root.create_dataset("recording", data=rec_dict, object_codec=numcodecs.JSON()) - zarr_rec = np.array([check_json(rec_dict)], dtype=object) - zarr_root.create_dataset("recording", data=zarr_rec, object_codec=numcodecs.JSON()) - elif recording.check_serializability("pickle"): - # zarr_root.create_dataset("recording", data=rec_dict, object_codec=numcodecs.Pickle()) - zarr_rec = np.array([rec_dict], dtype=object) - zarr_root.create_dataset("recording", data=zarr_rec, object_codec=numcodecs.Pickle()) + # In zarr v3, store JSON-serializable data in attributes instead of using object_codec + zarr_root.attrs["recording"] = check_json(rec_dict) else: warnings.warn("The Recording is not serializable! The recording link will be lost for future load") else: @@ -673,11 +674,8 @@ def create_zarr(cls, folder, sorting, recording, sparsity, return_in_uV, rec_att # sorting provenance sort_dict = sorting.to_dict(relative_to=relative_to, recursive=True) if sorting.check_serializability("json"): - zarr_sort = np.array([check_json(sort_dict)], dtype=object) - zarr_root.create_dataset("sorting_provenance", data=zarr_sort, object_codec=numcodecs.JSON()) - elif sorting.check_serializability("pickle"): - zarr_sort = np.array([sort_dict], dtype=object) - zarr_root.create_dataset("sorting_provenance", data=zarr_sort, object_codec=numcodecs.Pickle()) + # In zarr v3, store JSON-serializable data in attributes instead of using object_codec + zarr_root.attrs["sorting_provenance"] = check_json(sort_dict) else: warnings.warn( "The sorting provenance is not serializable! The sorting provenance link will be lost for future load" @@ -698,12 +696,13 @@ def create_zarr(cls, folder, sorting, recording, sparsity, return_in_uV, rec_att recording_info.attrs["probegroup"] = check_json(probegroup.to_dict()) if sparsity is not None: - zarr_root.create_dataset("sparsity_mask", data=sparsity.mask, **saving_options) + zarr_root.create_array("sparsity_mask", data=sparsity.mask, **saving_options) add_sorting_to_zarr_group(sorting, zarr_root.create_group("sorting"), **saving_options) recording_info = zarr_root.create_group("extensions") + # consolidate metadata for zarr v3 zarr.consolidate_metadata(zarr_root.store) return cls.load_from_zarr(folder, recording=recording, backend_options=backend_options) @@ -715,6 +714,10 @@ def load_from_zarr(cls, folder, recording=None, backend_options=None): backend_options = {} if backend_options is None else backend_options storage_options = backend_options.get("storage_options", {}) + if not is_path_remote(str(folder)): + storage_options_kwargs = {} + else: + storage_options_kwargs = storage_options zarr_root = super_zarr_open(str(folder), mode="r", storage_options=storage_options) @@ -723,7 +726,7 @@ def load_from_zarr(cls, folder, recording=None, backend_options=None): # v0.101.0 did not have a consolidate metadata step after computing extensions. # Here we try to consolidate the metadata and throw a warning if it fails. try: - zarr_root_a = zarr.open(str(folder), mode="a", storage_options=storage_options) + zarr_root_a = zarr.open(str(folder), mode="a", **storage_options_kwargs) zarr.consolidate_metadata(zarr_root_a.store) except Exception as e: warnings.warn( @@ -741,9 +744,9 @@ def load_from_zarr(cls, folder, recording=None, backend_options=None): # load recording if possible if recording is None: - rec_field = zarr_root.get("recording") - if rec_field is not None: - rec_dict = rec_field[0] + # In zarr v3, recording is stored in attributes + rec_dict = zarr_root.attrs.get("recording", None) + if rec_dict is not None: try: recording = load(rec_dict, base_folder=folder) except: @@ -859,7 +862,7 @@ def set_sorting_property( if key in zarr_root["sorting"]["properties"]: zarr_root["sorting"]["properties"][key][:] = prop_values else: - zarr_root["sorting"]["properties"].create_dataset(name=key, data=prop_values, compressor=None) + zarr_root["sorting"]["properties"].create_array(name=key, data=prop_values, compressors=None) # IMPORTANT: we need to re-consolidate the zarr store! zarr.consolidate_metadata(zarr_root.store) @@ -1531,12 +1534,13 @@ def get_sorting_provenance(self): elif self.format == "zarr": zarr_root = self._get_zarr_root(mode="r") sorting_provenance = None - if "sorting_provenance" in zarr_root.keys(): + # In zarr v3, sorting_provenance is stored in attributes + sort_dict = zarr_root.attrs.get("sorting_provenance", None) + if sort_dict is not None: # try-except here is because it's not required to be able # to load the sorting provenance, as the user might have deleted # the original sorting folder try: - sort_dict = zarr_root["sorting_provenance"][0] sorting_provenance = load(sort_dict, base_folder=self.folder) except: pass @@ -2479,8 +2483,9 @@ def load_data(self): extension_group = self._get_zarr_extension_group(mode="r") for ext_data_name in extension_group.keys(): ext_data_ = extension_group[ext_data_name] - if "dict" in ext_data_.attrs: - ext_data = ext_data_[0] + # In zarr v3, check if it's a group with dict_data attribute + if "dict_data" in ext_data_.attrs: + ext_data = ext_data_.attrs["dict_data"] elif "dataframe" in ext_data_.attrs: import pandas as pd @@ -2565,9 +2570,10 @@ def run(self, save=True, **kwargs): if self.format == "zarr": import zarr - zarr.consolidate_metadata(self.sorting_analyzer._get_zarr_root().store) + zarr.consolidate_metadata(self.sorting_analyzer._get_zarr_root(mode="r+").store) def save(self): + self._reset_extension_folder() self._save_params() self._save_importing_provenance() self._save_run_info() @@ -2576,7 +2582,7 @@ def save(self): if self.format == "zarr": import zarr - zarr.consolidate_metadata(self.sorting_analyzer._get_zarr_root().store) + zarr.consolidate_metadata(self.sorting_analyzer._get_zarr_root(mode="r+").store) def _save_data(self): if self.format == "memory": @@ -2623,40 +2629,44 @@ def _save_data(self): extension_group = self._get_zarr_extension_group(mode="r+") # if compression is not externally given, we use the default - if "compressor" not in saving_options: - saving_options["compressor"] = get_default_zarr_compressor() + if "compressors" not in saving_options and "compressor" not in saving_options: + saving_options["compressors"] = get_default_zarr_compressor() + if "compressor" in saving_options: + saving_options["compressors"] = [saving_options["compressor"]] + del saving_options["compressor"] for ext_data_name, ext_data in self.data.items(): if ext_data_name in extension_group: del extension_group[ext_data_name] if isinstance(ext_data, dict): - extension_group.create_dataset( - name=ext_data_name, data=np.array([ext_data], dtype=object), object_codec=numcodecs.JSON() - ) + # In zarr v3, store dict in a subgroup with attributes + dict_group = extension_group.create_group(ext_data_name) + dict_group.attrs["dict_data"] = check_json(ext_data) elif isinstance(ext_data, np.ndarray): - extension_group.create_dataset(name=ext_data_name, data=ext_data, **saving_options) + extension_group.create_array(name=ext_data_name, data=ext_data, **saving_options) elif HAS_PANDAS and isinstance(ext_data, pd.DataFrame): df_group = extension_group.create_group(ext_data_name) # first we save the index indices = ext_data.index.to_numpy() if indices.dtype.kind == "O": indices = indices.astype(str) - df_group.create_dataset(name="index", data=indices) + df_group.create_array(name="index", data=indices) for col in ext_data.columns: col_data = ext_data[col].to_numpy() if col_data.dtype.kind == "O": col_data = col_data.astype(str) - df_group.create_dataset(name=col, data=col_data) + df_group.create_array(name=col, data=col_data) df_group.attrs["dataframe"] = True else: # any object - try: - extension_group.create_dataset( - name=ext_data_name, data=np.array([ext_data], dtype=object), object_codec=numcodecs.Pickle() - ) - except: - raise Exception(f"Could not save {ext_data_name} as extension data") - extension_group[ext_data_name].attrs["object"] = True + # try: + # extension_group.create_array( + # name=ext_data_name, data=np.array([ext_data], dtype=object), object_codec=numcodecs.Pickle() + # ) + # except: + # raise Exception(f"Could not save {ext_data_name} as extension data") + # extension_group[ext_data_name].attrs["object"] = True + warnings.warn(f"Data type of {ext_data_name} not supported for zarr saving, skipping.") def _reset_extension_folder(self): """ @@ -2734,8 +2744,6 @@ def set_params(self, save=True, **params): def _save_params(self): params_to_save = self.params.copy() - self._reset_extension_folder() - # TODO make sparsity local Result specific # if "sparsity" in params_to_save and params_to_save["sparsity"] is not None: # assert isinstance( diff --git a/src/spikeinterface/core/template.py b/src/spikeinterface/core/template.py index 91d25bece6..50ae6cbfdf 100644 --- a/src/spikeinterface/core/template.py +++ b/src/spikeinterface/core/template.py @@ -321,17 +321,19 @@ def add_templates_to_zarr_group(self, zarr_group: "zarr.Group") -> None: """ # Saves one chunk per unit - arrays_chunk = (1, None, None) - zarr_group.create_dataset("templates_array", data=self.templates_array, chunks=arrays_chunk) - zarr_group.create_dataset("channel_ids", data=self.channel_ids) - zarr_group.create_dataset("unit_ids", data=self.unit_ids) + # In zarr v3, chunks must be a full tuple with actual dimensions + num_units, num_samples, num_channels = self.templates_array.shape + arrays_chunk = (1, num_samples, num_channels) + zarr_group.create_array("templates_array", data=self.templates_array, chunks=arrays_chunk) + zarr_group.create_array("channel_ids", data=self.channel_ids) + zarr_group.create_array("unit_ids", data=self.unit_ids) zarr_group.attrs["sampling_frequency"] = self.sampling_frequency zarr_group.attrs["nbefore"] = self.nbefore zarr_group.attrs["is_in_uV"] = self.is_in_uV if self.sparsity_mask is not None: - zarr_group.create_dataset("sparsity_mask", data=self.sparsity_mask) + zarr_group.create_array("sparsity_mask", data=self.sparsity_mask) if self.probe is not None: probe_group = zarr_group.create_group("probe") diff --git a/src/spikeinterface/core/tests/test_analyzer_extension_core.py b/src/spikeinterface/core/tests/test_analyzer_extension_core.py index 6f5bef3c6c..574cc89e10 100644 --- a/src/spikeinterface/core/tests/test_analyzer_extension_core.py +++ b/src/spikeinterface/core/tests/test_analyzer_extension_core.py @@ -91,7 +91,7 @@ def test_ComputeRandomSpikes(format, sparse, create_cache_folder): print("Checking results") _check_result_extension(sorting_analyzer, "random_spikes", cache_folder) - print("Delering extension") + print("Deleting extension") sorting_analyzer.delete_extension("random_spikes") print("Re-computing random spikes") diff --git a/src/spikeinterface/core/tests/test_baserecording.py b/src/spikeinterface/core/tests/test_baserecording.py index 9de800b33d..186767c026 100644 --- a/src/spikeinterface/core/tests/test_baserecording.py +++ b/src/spikeinterface/core/tests/test_baserecording.py @@ -324,7 +324,7 @@ def test_BaseRecording(create_cache_folder): # test save to zarr compressor = get_default_zarr_compressor() - rec_zarr = rec2.save(format="zarr", folder=cache_folder / "recording", compressor=compressor) + rec_zarr = rec2.save(format="zarr", folder=cache_folder / "recording", compressors=compressor) rec_zarr_loaded = load(cache_folder / "recording.zarr") # annotations is False because Zarr adds compression ratios check_recordings_equal(rec2, rec_zarr, return_in_uV=False, check_annotations=False, check_properties=True) @@ -336,7 +336,7 @@ def test_BaseRecording(create_cache_folder): assert rec2.get_annotation(annotation_name) == rec_zarr_loaded.get_annotation(annotation_name) rec_zarr2 = rec2.save( - format="zarr", folder=cache_folder / "recording_channel_chunk", compressor=compressor, channel_chunk_size=2 + format="zarr", folder=cache_folder / "recording_channel_chunk", compressors=compressor, channel_chunk_size=2 ) rec_zarr2_loaded = load(cache_folder / "recording_channel_chunk.zarr") diff --git a/src/spikeinterface/core/tests/test_sortinganalyzer.py b/src/spikeinterface/core/tests/test_sortinganalyzer.py index ab0b071df4..3a1c73d746 100644 --- a/src/spikeinterface/core/tests/test_sortinganalyzer.py +++ b/src/spikeinterface/core/tests/test_sortinganalyzer.py @@ -17,6 +17,7 @@ AnalyzerExtension, _sort_extensions_by_dependency, ) +from spikeinterface.core.zarrextractors import check_compressors_match import numpy as np @@ -38,6 +39,8 @@ def get_dataset(): integer_unit_ids = [int(id) for id in sorting.get_unit_ids()] recording = recording.rename_channels(new_channel_ids=integer_channel_ids) + # make sure the recording is serializable + recording = recording.save() sorting = sorting.rename_units(new_unit_ids=integer_unit_ids) return recording, sorting @@ -133,13 +136,12 @@ def test_SortingAnalyzer_zarr(tmp_path, dataset): _check_sorting_analyzers(sorting_analyzer, sorting, cache_folder=tmp_path) # check that compression is applied - assert ( - sorting_analyzer._get_zarr_root()["extensions"]["random_spikes"]["random_spikes_indices"].compressor.codec_id - == default_compressor.codec_id + check_compressors_match( + default_compressor, + sorting_analyzer._get_zarr_root()["extensions"]["random_spikes"]["random_spikes_indices"].compressors[0], ) - assert ( - sorting_analyzer._get_zarr_root()["extensions"]["templates"]["average"].compressor.codec_id - == default_compressor.codec_id + check_compressors_match( + default_compressor, sorting_analyzer._get_zarr_root()["extensions"]["templates"]["average"].compressors[0] ) # test select_units see https://github.com/SpikeInterface/spikeinterface/issues/3041 @@ -160,35 +162,34 @@ def test_SortingAnalyzer_zarr(tmp_path, dataset): sparsity=None, return_in_uV=False, overwrite=True, - backend_options={"saving_options": {"compressor": None}}, + backend_options={"saving_options": {"compressors": None}}, ) print(sorting_analyzer_no_compression._backend_options) sorting_analyzer_no_compression.compute(["random_spikes", "templates"]) assert ( - sorting_analyzer_no_compression._get_zarr_root()["extensions"]["random_spikes"][ - "random_spikes_indices" - ].compressor - is None + len( + sorting_analyzer_no_compression._get_zarr_root()["extensions"]["random_spikes"][ + "random_spikes_indices" + ].compressors + ) + == 0 ) - assert sorting_analyzer_no_compression._get_zarr_root()["extensions"]["templates"]["average"].compressor is None + assert len(sorting_analyzer_no_compression._get_zarr_root()["extensions"]["templates"]["average"].compressors) == 0 # test a different compressor - from numcodecs import LZMA + from zarr.codecs.numcodecs import LZMA lzma_compressor = LZMA() folder = tmp_path / "test_SortingAnalyzer_zarr_lzma.zarr" sorting_analyzer_lzma = sorting_analyzer_no_compression.save_as( - format="zarr", folder=folder, backend_options={"saving_options": {"compressor": lzma_compressor}} + format="zarr", folder=folder, backend_options={"saving_options": {"compressors": lzma_compressor}} ) - assert ( - sorting_analyzer_lzma._get_zarr_root()["extensions"]["random_spikes"][ - "random_spikes_indices" - ].compressor.codec_id - == LZMA.codec_id + check_compressors_match( + lzma_compressor, + sorting_analyzer_lzma._get_zarr_root()["extensions"]["random_spikes"]["random_spikes_indices"].compressors[0], ) - assert ( - sorting_analyzer_lzma._get_zarr_root()["extensions"]["templates"]["average"].compressor.codec_id - == LZMA.codec_id + check_compressors_match( + lzma_compressor, sorting_analyzer_lzma._get_zarr_root()["extensions"]["templates"]["average"].compressors[0] ) # test set_sorting_property diff --git a/src/spikeinterface/core/tests/test_zarrextractors.py b/src/spikeinterface/core/tests/test_zarrextractors.py index cc0c60721e..a52d456594 100644 --- a/src/spikeinterface/core/tests/test_zarrextractors.py +++ b/src/spikeinterface/core/tests/test_zarrextractors.py @@ -10,50 +10,55 @@ generate_sorting, load, ) -from spikeinterface.core.zarrextractors import add_sorting_to_zarr_group, get_default_zarr_compressor +from spikeinterface.core.zarrextractors import ( + add_sorting_to_zarr_group, + get_default_zarr_compressor, + check_compressors_match, +) def test_zarr_compression_options(tmp_path): - from numcodecs import Blosc, Delta, FixedScaleOffset + from zarr.codecs.numcodecs import Delta, FixedScaleOffset + from zarr.codecs import BloscCodec, BloscShuffle recording = generate_recording(durations=[2]) recording.set_times(recording.get_times() + 100) # store in root standard normal way # default compressor - defaut_compressor = get_default_zarr_compressor() + default_compressor = get_default_zarr_compressor() # other compressor - other_compressor1 = Blosc(cname="zlib", clevel=3, shuffle=Blosc.NOSHUFFLE) - other_compressor2 = Blosc(cname="blosclz", clevel=8, shuffle=Blosc.AUTOSHUFFLE) + other_compressor1 = BloscCodec(cname="zlib", clevel=3, shuffle=BloscShuffle.noshuffle) + other_compressor2 = BloscCodec(cname="blosclz", clevel=8, shuffle=BloscShuffle.shuffle) # timestamps compressors / filters default_filters = None - other_filters1 = [FixedScaleOffset(scale=5, offset=2, dtype=recording.get_dtype())] + other_filters1 = [FixedScaleOffset(scale=5, offset=2, dtype=recording.get_dtype().str)] other_filters2 = [Delta(dtype="float64")] # default ZarrRecordingExtractor.write_recording(recording, tmp_path / "rec_default.zarr") rec_default = ZarrRecordingExtractor(tmp_path / "rec_default.zarr") - assert rec_default._root["traces_seg0"].compressor == defaut_compressor - assert rec_default._root["traces_seg0"].filters == default_filters - assert rec_default._root["times_seg0"].compressor == defaut_compressor - assert rec_default._root["times_seg0"].filters == default_filters + check_compressors_match(rec_default._root["traces_seg0"].compressors[0], default_compressor) + check_compressors_match(rec_default._root["times_seg0"].compressors[0], default_compressor) + check_compressors_match(rec_default._root["traces_seg0"].filters, default_filters) + check_compressors_match(rec_default._root["times_seg0"].filters, default_filters) # now with other compressor ZarrRecordingExtractor.write_recording( recording, tmp_path / "rec_other.zarr", - compressor=defaut_compressor, + compressors=default_compressor, filters=default_filters, compressor_by_dataset={"traces": other_compressor1, "times": other_compressor2}, filters_by_dataset={"traces": other_filters1, "times": other_filters2}, ) rec_other = ZarrRecordingExtractor(tmp_path / "rec_other.zarr") - assert rec_other._root["traces_seg0"].compressor == other_compressor1 - assert rec_other._root["traces_seg0"].filters == other_filters1 - assert rec_other._root["times_seg0"].compressor == other_compressor2 - assert rec_other._root["times_seg0"].filters == other_filters2 + check_compressors_match(rec_other._root["traces_seg0"].compressors[0], other_compressor1) + check_compressors_match(rec_other._root["traces_seg0"].filters, other_filters1) + check_compressors_match(rec_other._root["times_seg0"].compressors[0], other_compressor2) + check_compressors_match(rec_other._root["times_seg0"].filters, other_filters2) def test_ZarrSortingExtractor(tmp_path): diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py index 162d67a458..6b20c6bbaa 100644 --- a/src/spikeinterface/core/zarrextractors.py +++ b/src/spikeinterface/core/zarrextractors.py @@ -48,11 +48,13 @@ def super_zarr_open(folder_path: str | Path, mode: str = "r", storage_options: d import zarr # if mode is append or read/write, we try to open the folder with zarr.open - # since zarr.open_consolidated does not support creating new groups/datasets + # In zarr v3, we use use_consolidated parameter instead of open_consolidated if mode in ("a", "r+"): open_funcs = (zarr.open,) + use_consolidated_options = (False,) else: - open_funcs = (zarr.open_consolidated, zarr.open) + open_funcs = (zarr.open,) + use_consolidated_options = (True, False) # if storage_options is None, we try to open the folder with and without anonymous access # if storage_options is not None, we try to open the folder with the given storage options @@ -64,12 +66,14 @@ def super_zarr_open(folder_path: str | Path, mode: str = "r", storage_options: d root = None exception = None if is_path_remote(str(folder_path)): - for open_func in open_funcs: + for use_consolidated in use_consolidated_options: if root is not None: break for storage_options in storage_options_to_test: try: - root = open_func(str(folder_path), mode=mode, storage_options=storage_options) + root = zarr.open( + str(folder_path), mode=mode, storage_options=storage_options, use_consolidated=use_consolidated + ) break except Exception as e: exception = e @@ -77,9 +81,9 @@ def super_zarr_open(folder_path: str | Path, mode: str = "r", storage_options: d else: if not Path(folder_path).is_dir(): raise ValueError(f"Folder {folder_path} does not exist") - for open_func in open_funcs: + for use_consolidated in use_consolidated_options: try: - root = open_func(str(folder_path), mode=mode, storage_options=storage_options) + root = zarr.open(str(folder_path), mode=mode, use_consolidated=use_consolidated) break except Exception as e: exception = e @@ -91,6 +95,34 @@ def super_zarr_open(folder_path: str | Path, mode: str = "r", storage_options: d return root +def check_compressors_match(comp1, comp2, skip_typesize=True): + """ + Check if two compressor objects match. + + Parameters + ---------- + comp1 : zarr.Codec | Tuple[zarr.Codec] + The first compressor object to compare. + comp2 : zarr.Codec | Tuple[zarr.Codec] + The second compressor object to compare. + skip_typesize : bool, optional + Whether to skip the typesize check, default: True + """ + if not isinstance(comp1, (list, tuple)): + assert not isinstance(comp2, list) + comp1 = [comp1] + comp2 = [comp2] + for i in range(len(comp1)): + comp1_dict = comp1[i].to_dict() + comp2_dict = comp2[i].to_dict() + if skip_typesize: + if "typesize" in comp1_dict["configuration"]: + comp1_dict["configuration"].pop("typesize", None) + if "typesize" in comp2_dict["configuration"]: + comp2_dict["configuration"].pop("typesize", None) + assert comp1_dict == comp2_dict + + class ZarrRecordingExtractor(BaseRecording): """ RecordingExtractor for a zarr format @@ -289,7 +321,7 @@ def __init__(self, folder_path: Path | str, storage_options: dict | None = None, BaseSorting.__init__(self, sampling_frequency, unit_ids) - spikes = np.zeros(len(spikes_group["sample_index"]), dtype=minimum_spike_dtype) + spikes = np.zeros(spikes_group["sample_index"].shape[0], dtype=minimum_spike_dtype) spikes["sample_index"] = spikes_group["sample_index"][:] spikes["unit_index"] = spikes_group["unit_index"][:] for i, (start, end) in enumerate(segment_slices_list): @@ -392,9 +424,9 @@ def get_default_zarr_compressor(clevel: int = 5): Blosc.compressor The compressor object that can be used with the save to zarr function """ - from numcodecs import Blosc + from zarr.codecs import BloscCodec, BloscShuffle - return Blosc(cname="zstd", clevel=clevel, shuffle=Blosc.BITSHUFFLE) + return BloscCodec(cname="zstd", clevel=clevel, shuffle=BloscShuffle.bitshuffle) def add_properties_and_annotations(zarr_group: zarr.hierarchy.Group, recording_or_sorting: BaseRecording | BaseSorting): @@ -405,7 +437,7 @@ def add_properties_and_annotations(zarr_group: zarr.hierarchy.Group, recording_o if values.dtype.kind == "O": warnings.warn(f"Property {key} not saved because it is a python Object type") continue - prop_group.create_dataset(name=key, data=values, compressor=None) + prop_group.create_array(name=key, data=values, compressors=None) # save annotations zarr_group.attrs["annotations"] = check_json(recording_or_sorting._annotations) @@ -424,12 +456,12 @@ def add_sorting_to_zarr_group(sorting: BaseSorting, zarr_group: zarr.hierarchy.G kwargs : dict Other arguments passed to the zarr compressor """ - from numcodecs import Delta + from zarr.codecs.numcodecs import Delta num_segments = sorting.get_num_segments() zarr_group.attrs["sampling_frequency"] = float(sorting.sampling_frequency) zarr_group.attrs["num_segments"] = int(num_segments) - zarr_group.create_dataset(name="unit_ids", data=sorting.unit_ids, compressor=None) + zarr_group.create_array(name="unit_ids", data=sorting.unit_ids, compressors=None) compressor = kwargs.get("compressor", get_default_zarr_compressor()) @@ -438,18 +470,21 @@ def add_sorting_to_zarr_group(sorting: BaseSorting, zarr_group: zarr.hierarchy.G spikes = sorting.to_spike_vector() for field in spikes.dtype.fields: if field != "segment_index": - spikes_group.create_dataset( + dtype = spikes[field].dtype + spikes_data = spikes[field] + spikes_group.create_array( name=field, - data=spikes[field], - compressor=compressor, - filters=[Delta(dtype=spikes[field].dtype)], + data=spikes_data, + compressors=compressor, + filters=[Delta(dtype=spikes[field].dtype.str)], ) else: segment_slices = [] for segment_index in range(num_segments): i0, i1 = np.searchsorted(spikes["segment_index"], [segment_index, segment_index + 1]) segment_slices.append([i0, i1]) - spikes_group.create_dataset(name="segment_slices", data=segment_slices, compressor=None) + segment_slices = np.array(segment_slices, dtype="int64") + spikes_group.create_array(name="segment_slices", data=segment_slices, compressors=None) add_properties_and_annotations(zarr_group, sorting) @@ -468,7 +503,7 @@ def add_recording_to_zarr_group( # save data (done the subclass) zarr_group.attrs["sampling_frequency"] = float(recording.get_sampling_frequency()) zarr_group.attrs["num_segments"] = int(recording.get_num_segments()) - zarr_group.create_dataset(name="channel_ids", data=recording.get_channel_ids(), compressor=None) + zarr_group.create_array(name="channel_ids", data=recording.get_channel_ids(), compressors=None) dataset_paths = [f"traces_seg{i}" for i in range(recording.get_num_segments())] dtype = recording.get_dtype() if dtype is None else dtype @@ -484,7 +519,7 @@ def add_recording_to_zarr_group( recording=recording, zarr_group=zarr_group, dataset_paths=dataset_paths, - compressor=compressor_traces, + compressors=compressor_traces, filters=filters_traces, dtype=dtype, channel_chunk_size=channel_chunk_size, @@ -507,17 +542,17 @@ def add_recording_to_zarr_group( filters_times = filters_by_dataset.get("times", global_filters) if time_vector is not None: - _ = zarr_group.create_dataset( + _ = zarr_group.create_array( name=f"times_seg{segment_index}", data=time_vector, filters=filters_times, - compressor=compressor_times, + compressors=compressor_times, ) elif d["t_start"] is not None: t_starts[segment_index] = d["t_start"] if np.any(~np.isnan(t_starts)): - zarr_group.create_dataset(name="t_starts", data=t_starts, compressor=None) + zarr_group.create_array(name="t_starts", data=t_starts, compressors=None) add_properties_and_annotations(zarr_group, recording) @@ -528,7 +563,7 @@ def add_traces_to_zarr( dataset_paths, channel_chunk_size=None, dtype=None, - compressor=None, + compressors=None, filters=None, verbose=False, **job_kwargs, @@ -548,7 +583,7 @@ def add_traces_to_zarr( Channels per chunk dtype : dtype, default: None Type of the saved data - compressor : zarr compressor or None, default: None + compressors : zarr compressor or None, default: None Zarr compressor filters : list, default: None List of zarr filters @@ -581,13 +616,15 @@ def add_traces_to_zarr( num_channels = recording.get_num_channels() dset_name = dataset_paths[segment_index] shape = (num_frames, num_channels) - dset = zarr_group.create_dataset( + # In zarr v3, chunks must be a tuple of integers (no None allowed) + chunks = (chunk_size, channel_chunk_size if channel_chunk_size is not None else num_channels) + dset = zarr_group.create_array( name=dset_name, shape=shape, - chunks=(chunk_size, channel_chunk_size), + chunks=chunks, dtype=dtype, filters=filters, - compressor=compressor, + compressors=compressors, ) zarr_datasets.append(dset) # synchronizer=zarr.ThreadSynchronizer()) From 4dbb7b38fab00795aebb3644ba4bb325052607cf Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Fri, 12 Dec 2025 16:49:23 +0100 Subject: [PATCH 02/59] wip --- src/spikeinterface/core/zarrextractors.py | 23 ++++++++++++++--------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py index 6b20c6bbaa..0ec28d544a 100644 --- a/src/spikeinterface/core/zarrextractors.py +++ b/src/spikeinterface/core/zarrextractors.py @@ -14,6 +14,9 @@ from .core_tools import is_path_remote +zarr.config.set({"default_zarr_version": 3}) + + def super_zarr_open(folder_path: str | Path, mode: str = "r", storage_options: dict | None = None): """ Open a zarr folder with super powers. @@ -463,7 +466,9 @@ def add_sorting_to_zarr_group(sorting: BaseSorting, zarr_group: zarr.hierarchy.G zarr_group.attrs["num_segments"] = int(num_segments) zarr_group.create_array(name="unit_ids", data=sorting.unit_ids, compressors=None) - compressor = kwargs.get("compressor", get_default_zarr_compressor()) + compressor = kwargs.get("compressors") or kwargs.get("compressor") + if compressor is None: + compressor = get_default_zarr_compressor() # save sub fields spikes_group = zarr_group.create_group(name="spikes") @@ -508,7 +513,9 @@ def add_recording_to_zarr_group( dtype = recording.get_dtype() if dtype is None else dtype channel_chunk_size = zarr_kwargs.get("channel_chunk_size", None) - global_compressor = zarr_kwargs.pop("compressor", get_default_zarr_compressor()) + global_compressor = kwargs.get("compressors") or kwargs.get("compressor") + if global_compressor is None: + global_compressor = get_default_zarr_compressor() compressor_by_dataset = zarr_kwargs.pop("compressor_by_dataset", {}) global_filters = zarr_kwargs.pop("filters", None) filters_by_dataset = zarr_kwargs.pop("filters_by_dataset", {}) @@ -609,6 +616,9 @@ def add_traces_to_zarr( job_kwargs = fix_job_kwargs(job_kwargs) chunk_size = ensure_chunk_size(recording, **job_kwargs) + if not isinstance(compressors, (list, tuple)): + compressors = [compressors] + # create zarr datasets files zarr_datasets = [] for segment_index in range(recording.get_num_segments()): @@ -618,13 +628,8 @@ def add_traces_to_zarr( shape = (num_frames, num_channels) # In zarr v3, chunks must be a tuple of integers (no None allowed) chunks = (chunk_size, channel_chunk_size if channel_chunk_size is not None else num_channels) - dset = zarr_group.create_array( - name=dset_name, - shape=shape, - chunks=chunks, - dtype=dtype, - filters=filters, - compressors=compressors, + dset = zarr_group.create( + name=dset_name, shape=shape, chunks=chunks, dtype=dtype, filters=filters, codecs=compressors, zarr_format=3 ) zarr_datasets.append(dset) # synchronizer=zarr.ThreadSynchronizer()) From 4bc45daa956d1fb2f7cdb5f6b27147cc708ad2ff Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Mon, 23 Mar 2026 12:43:42 +0100 Subject: [PATCH 03/59] fix: zarr.Group --- src/spikeinterface/core/zarrextractors.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py index 43a7f55f3d..4e467ebdbd 100644 --- a/src/spikeinterface/core/zarrextractors.py +++ b/src/spikeinterface/core/zarrextractors.py @@ -38,7 +38,7 @@ def super_zarr_open(folder_path: str | Path, mode: str = "r", storage_options: d Returns ------- - root: zarr.hierarchy.Group + root: zarr.Group The zarr root group object Raises @@ -496,7 +496,7 @@ def build_codec_pipeline(filters=None, compressors=None): return codecs if codecs else None -def add_properties_and_annotations(zarr_group: zarr.hierarchy.Group, recording_or_sorting: BaseRecording | BaseSorting): +def add_properties_and_annotations(zarr_group: zarr.Group, recording_or_sorting: BaseRecording | BaseSorting): # save properties prop_group = zarr_group.create_group("properties") for key in recording_or_sorting.get_property_keys(): @@ -510,7 +510,7 @@ def add_properties_and_annotations(zarr_group: zarr.hierarchy.Group, recording_o zarr_group.attrs["annotations"] = check_json(recording_or_sorting._annotations) -def add_sorting_to_zarr_group(sorting: BaseSorting, zarr_group: zarr.hierarchy.Group, **kwargs): +def add_sorting_to_zarr_group(sorting: BaseSorting, zarr_group: zarr.Group, **kwargs): """ Add a sorting extractor to a zarr group. @@ -518,7 +518,7 @@ def add_sorting_to_zarr_group(sorting: BaseSorting, zarr_group: zarr.hierarchy.G ---------- sorting : BaseSorting The sorting extractor object to be added to the zarr group - zarr_group : zarr.hierarchy.Group + zarr_group : zarr.Group The zarr group kwargs : dict Other arguments passed to the zarr compressor @@ -556,9 +556,7 @@ def add_sorting_to_zarr_group(sorting: BaseSorting, zarr_group: zarr.hierarchy.G # Recording -def add_recording_to_zarr_group( - recording: BaseRecording, zarr_group: zarr.hierarchy.Group, verbose=False, dtype=None, **kwargs -): +def add_recording_to_zarr_group(recording: BaseRecording, zarr_group: zarr.Group, verbose=False, dtype=None, **kwargs): zarr_kwargs, job_kwargs = split_job_kwargs(kwargs) if recording.check_if_json_serializable(): From e9a567f29a9574d8af8f2a95db2903ec62c0c427 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Mon, 23 Mar 2026 13:12:22 +0100 Subject: [PATCH 04/59] wip: fix v3 --- src/spikeinterface/core/zarrextractors.py | 88 +++++++++++++++-------- 1 file changed, 58 insertions(+), 30 deletions(-) diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py index 4e467ebdbd..ba3d2f1318 100644 --- a/src/spikeinterface/core/zarrextractors.py +++ b/src/spikeinterface/core/zarrextractors.py @@ -163,7 +163,8 @@ def __init__( assert sampling_frequency is not None, "'sampling_frequency' attiribute not found!" assert num_segments is not None, "'num_segments' attiribute not found!" - channel_ids = np.array(channel_ids) + # zarr returns vlen-utf8 as StringDType (numpy 2.0); convert via list to classic unicode array. + channel_ids = np.array(channel_ids.tolist()) dtype = self._root["traces_seg0"].dtype @@ -201,7 +202,7 @@ def __init__( if load_compression_ratio: nbytes_segment = self._root[trace_name].nbytes - nbytes_stored_segment = self._root[trace_name].nbytes_stored + nbytes_stored_segment = self._root[trace_name].nbytes_stored() if nbytes_stored_segment > 0: cr_by_segment[segment_index] = nbytes_segment / nbytes_stored_segment else: @@ -220,7 +221,11 @@ def __init__( if "properties" in self._root: prop_group = self._root["properties"] for key in prop_group.keys(): - values = self._root["properties"][key] + values = self._root["properties"][key][:] + # zarr returns vlen-utf8 as StringDType (numpy 2.0); convert via list to classic unicode array. + if hasattr(values.dtype, "na_object") or values.dtype.kind == "O": + if values.size > 0 and isinstance(values.tolist()[0], str): + values = np.array(values.tolist()) self.set_property(key, values) # load annotations @@ -338,7 +343,11 @@ def __init__(self, folder_path: Path | str, storage_options: dict | None = None, if "properties" in self._root: prop_group = self._root["properties"] for key in prop_group.keys(): - values = self._root["properties"][key] + values = self._root["properties"][key][:] + # zarr returns vlen-utf8 as StringDType (numpy 2.0); convert via list to classic unicode array. + if hasattr(values.dtype, "na_object") or values.dtype.kind == "O": + if values.size > 0 and isinstance(values.tolist()[0], str): + values = np.array(values.tolist()) self.set_property(key, values) # load annotations @@ -434,12 +443,12 @@ def get_default_zarr_compressor(clevel: int = 5): def build_codec_pipeline(filters=None, compressors=None): """ - Build a zarr v3 codecs list from filters and compressors. + Build zarr v3 codec kwargs from filters and compressors. - Assembles a valid zarr v3 codec pipeline in the required order: - 1. ArrayArrayCodec (filters, e.g. Delta) - 2. ArrayBytesCodec (serializer, e.g. WavPack, BytesCodec) - 3. BytesBytesCodec (compressors, e.g. BloscCodec, ZstdCodec) + Classifies codecs into the three slots accepted by ``zarr.Group.create_array()``: + 1. ``filters`` — ArrayArrayCodec (e.g. Delta) + 2. ``serializer`` — ArrayBytesCodec (e.g. WavPack, BytesCodec) + 3. ``compressors``— BytesBytesCodec (e.g. BloscCodec, ZstdCodec) This allows callers to pass an ArrayBytesCodec (e.g. WavPack) as a compressor and have it placed in the correct serializer slot automatically. @@ -454,10 +463,10 @@ def build_codec_pipeline(filters=None, compressors=None): Returns ------- - list of codecs or None - Full codec pipeline suitable for the ``codecs=`` parameter of - ``zarr.create()``. Returns None when both inputs are empty/None, - letting zarr use its defaults. + dict + Keyword arguments to unpack into ``zarr.Group.create_array()``. + Only keys with explicit values are included; omitted keys let zarr + use its defaults. Raises ------ @@ -492,8 +501,18 @@ def build_codec_pipeline(filters=None, compressors=None): if len(serializers) > 1: raise ValueError("Only one ArrayBytesCodec (serializer) is allowed in the codec pipeline.") - codecs = filters + serializers + byte_compressors - return codecs if codecs else None + codec_kwargs = {} + codec_kwargs["filters"] = filters + codec_kwargs["serializer"] = serializers[0] + codec_kwargs["compressors"] = byte_compressors + return codec_kwargs + + +def _has_string_fields(dtype: np.dtype) -> bool: + """Return True if dtype is or contains fixed-length unicode (U) sub-fields.""" + if dtype.names: + return any(_has_string_fields(dtype.fields[name][0]) for name in dtype.names) + return dtype.kind == "U" def add_properties_and_annotations(zarr_group: zarr.Group, recording_or_sorting: BaseRecording | BaseSorting): @@ -504,7 +523,20 @@ def add_properties_and_annotations(zarr_group: zarr.Group, recording_or_sorting: if values.dtype.kind == "O": warnings.warn(f"Property {key} not saved because it is a python Object type") continue - prop_group.create_array(name=key, data=values, compressors=None) + if values.dtype.names and _has_string_fields(values.dtype): + # Structured arrays with unicode sub-fields have no stable zarr v3 spec; skip them. + # Probe geometry (contact_vector) is already persisted via zarr_group.attrs["probe"]. + warnings.warn( + f"Property '{key}' not saved because it is a structured array with unicode fields, " + "which do not have a stable zarr V3 specification." + ) + continue + # Use variable-length UTF-8 (stable zarr v3 spec) for unicode arrays. + if values.dtype.kind == "U": + arr = prop_group.create_array(name=key, shape=values.shape, dtype=str, compressors=None) + arr[:] = values + else: + prop_group.create_array(name=key, data=values, compressors=None) # save annotations zarr_group.attrs["annotations"] = check_json(recording_or_sorting._annotations) @@ -541,9 +573,8 @@ def add_sorting_to_zarr_group(sorting: BaseSorting, zarr_group: zarr.Group, **kw if field != "segment_index": dtype = spikes[field].dtype spikes_data = spikes[field] - codecs = build_codec_pipeline(filters=[Delta(dtype=spikes[field].dtype.str)], compressors=compressor) - arr = spikes_group.create(name=field, shape=spikes_data.shape, dtype=spikes_data.dtype, codecs=codecs) - arr[:] = spikes_data + codec_kwargs = build_codec_pipeline(filters=[Delta(dtype=spikes[field].dtype.str)], compressors=compressor) + spikes_group.create_array(name=field, data=spikes_data, **codec_kwargs) else: segment_slices = [] for segment_index in range(num_segments): @@ -567,7 +598,10 @@ def add_recording_to_zarr_group(recording: BaseRecording, zarr_group: zarr.Group # save data (done the subclass) zarr_group.attrs["sampling_frequency"] = float(recording.get_sampling_frequency()) zarr_group.attrs["num_segments"] = int(recording.get_num_segments()) - zarr_group.create_array(name="channel_ids", data=recording.get_channel_ids(), compressors=None) + # Use variable-length UTF-8 (stable zarr v3 spec) instead of fixed-length unicode. + channel_ids = recording.get_channel_ids() + arr = zarr_group.create_array(name="channel_ids", shape=channel_ids.shape, dtype=str, compressors=None) + arr[:] = channel_ids dataset_paths = [f"traces_seg{i}" for i in range(recording.get_num_segments())] dtype = recording.get_dtype() if dtype is None else dtype @@ -608,14 +642,8 @@ def add_recording_to_zarr_group(recording: BaseRecording, zarr_group: zarr.Group filters_times = filters_by_dataset.get("times", global_filters) if time_vector is not None: - codecs = build_codec_pipeline(filters=filters_times, compressors=compressor_times) - arr = zarr_group.create( - name=f"times_seg{segment_index}", - shape=time_vector.shape, - dtype=time_vector.dtype, - codecs=codecs, - ) - arr[:] = time_vector + codec_kwargs = build_codec_pipeline(filters=filters_times, compressors=compressor_times) + zarr_group.create_array(name=f"times_seg{segment_index}", data=time_vector, **codec_kwargs) elif d["t_start"] is not None: t_starts[segment_index] = d["t_start"] @@ -677,7 +705,7 @@ def add_traces_to_zarr( job_kwargs = fix_job_kwargs(job_kwargs) chunk_size = ensure_chunk_size(recording, **job_kwargs) - codecs = build_codec_pipeline(filters=filters, compressors=compressors) + codec_kwargs = build_codec_pipeline(filters=filters, compressors=compressors) # create zarr datasets files zarr_datasets = [] @@ -688,7 +716,7 @@ def add_traces_to_zarr( shape = (num_frames, num_channels) # In zarr v3, chunks must be a tuple of integers (no None allowed) chunks = (chunk_size, channel_chunk_size if channel_chunk_size is not None else num_channels) - dset = zarr_group.create(name=dset_name, shape=shape, chunks=chunks, dtype=dtype, codecs=codecs, zarr_format=3) + dset = zarr_group.create_array(name=dset_name, shape=shape, chunks=chunks, dtype=dtype, **codec_kwargs) zarr_datasets.append(dset) # synchronizer=zarr.ThreadSynchronizer()) From 39a75883f71dadeaa294ceb2db8a59a94be33b71 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Wed, 25 Mar 2026 10:47:08 +0100 Subject: [PATCH 05/59] Fix v2/v3 test, add zarr_class_info, pin to dev probeinterface for testing --- .github/scripts/generate_zarr_v2_fixtures.py | 37 +++++++++---------- pyproject.toml | 16 ++++++-- src/spikeinterface/core/sortinganalyzer.py | 8 +++- .../core/tests/test_zarr_backwards_compat.py | 16 +++----- src/spikeinterface/core/zarrextractors.py | 26 +++++++++---- 5 files changed, 60 insertions(+), 43 deletions(-) diff --git a/.github/scripts/generate_zarr_v2_fixtures.py b/.github/scripts/generate_zarr_v2_fixtures.py index 90d0515da8..e555873458 100644 --- a/.github/scripts/generate_zarr_v2_fixtures.py +++ b/.github/scripts/generate_zarr_v2_fixtures.py @@ -12,64 +12,61 @@ - expected_values.json : key values used to verify correct loading """ import argparse +import shutil import json from pathlib import Path import numpy as np +import zarr +import spikeinterface as si -def main(output_dir: Path) -> None: - import spikeinterface - - print(f"spikeinterface version : {spikeinterface.__version__}") - - import zarr +def main(output_dir: Path) -> None: + print(f"spikeinterface version : {si.__version__}") print(f"zarr version : {zarr.__version__}") - from spikeinterface.core import generate_recording, generate_sorting - from spikeinterface.core import ZarrRecordingExtractor, ZarrSortingExtractor - from spikeinterface.core import create_sorting_analyzer, load_sorting_analyzer output_dir.mkdir(parents=True, exist_ok=True) - recording = generate_recording(num_channels=4, num_segments=2, seed=0) - sorting = generate_sorting(num_units=3, num_segments=2, seed=0) - + recording, sorting = si.generate_ground_truth_recording(durations=[10, 5],num_channels=32, num_units=10, seed=0) + # save to binary to make them JSON serializable for later expected values extraction + recording = recording.save(folder=output_dir / "recording_binary", overwrite=True) + sorting = sorting.save(folder=output_dir / "sorting_binary", overwrite=True) # --- save recording --- recording_path = output_dir / "recording.zarr" - ZarrRecordingExtractor.write_recording(recording, recording_path) + recording_zarr = recording.save(format="zarr", folder=recording_path, overwrite=True) print(f"Saved recording -> {recording_path}") # --- save sorting --- sorting_path = output_dir / "sorting.zarr" - ZarrSortingExtractor.write_sorting(sorting, sorting_path) + sorting_zarr = sorting.save(format="zarr", folder=sorting_path, overwrite=True) print(f"Saved sorting -> {sorting_path}") # --- save SortingAnalyzer --- # Reload the recording from zarr so it is a serializable ZarrRecordingExtractor, # which the analyzer can store as provenance. - recording_zarr = ZarrRecordingExtractor(recording_path) analyzer_path = output_dir / "analyzer.zarr" - analyzer = create_sorting_analyzer( - sorting, recording_zarr, format="zarr", folder=analyzer_path, sparse=False, sparsity=None + if analyzer_path.is_dir(): + shutil.rmtree(analyzer_path) + analyzer = si.create_sorting_analyzer( + sorting_zarr, recording_zarr, format="zarr", folder=analyzer_path, overwrite=True ) analyzer.compute(["random_spikes", "templates"]) print(f"Saved analyzer -> {analyzer_path}") # Reload to verify templates are accessible before writing expected values - analyzer = load_sorting_analyzer(analyzer_path) templates_array = analyzer.get_extension("templates").get_data() # --- capture expected values for later assertion --- expected = { - "spikeinterface_version": spikeinterface.__version__, + "spikeinterface_version": si.__version__, "zarr_version": zarr.__version__, "recording": { "num_channels": int(recording.get_num_channels()), "num_segments": int(recording.get_num_segments()), "sampling_frequency": float(recording.get_sampling_frequency()), - "num_frames_per_segment": [int(recording.get_num_frames(seg)) for seg in range(recording.get_num_segments())], + "num_samples_per_segment": [int(recording.get_num_samples(seg)) for seg in range(recording.get_num_segments())], "channel_ids": recording.get_channel_ids().tolist(), "dtype": str(recording.get_dtype()), # first 10 frames of segment 0 for all channels diff --git a/pyproject.toml b/pyproject.toml index 7978862df4..0fa6d12fd8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -126,7 +126,9 @@ test_core = [ # for github test : probeinterface and neo from master # for release we need pypi, so this need to be commented - "probeinterface @ git+https://github.com/SpikeInterface/probeinterface.git", + # FOR TESTING: use probeinterface zarrv3 branch + "probeinterface @ git+https://github.com/alejoe91/probeinterface.git@zarrv3", + # "probeinterface @ git+https://github.com/SpikeInterface/probeinterface.git", "neo @ git+https://github.com/NeuralEnsemble/python-neo.git", # for slurm jobs, @@ -138,7 +140,9 @@ test_extractors = [ "pooch>=1.8.2", "datalad>=1.0.2", # Commenting out for release - "probeinterface @ git+https://github.com/SpikeInterface/probeinterface.git", + # FOR TESTING: use probeinterface zarrv3 branch + "probeinterface @ git+https://github.com/alejoe91/probeinterface.git@zarrv3", + # "probeinterface @ git+https://github.com/SpikeInterface/probeinterface.git", "neo @ git+https://github.com/NeuralEnsemble/python-neo.git", ] @@ -189,7 +193,9 @@ test = [ # for github test : probeinterface and neo from master # for release we need pypi, so this need to be commented - "probeinterface @ git+https://github.com/SpikeInterface/probeinterface.git", + # FOR TESTING: use probeinterface zarrv3 branch + "probeinterface @ git+https://github.com/alejoe91/probeinterface.git@zarrv3", + # "probeinterface @ git+https://github.com/SpikeInterface/probeinterface.git", "neo @ git+https://github.com/NeuralEnsemble/python-neo.git", # for slurm jobs @@ -218,7 +224,9 @@ docs = [ "huggingface_hub", # For automated curation # for release we need pypi, so this needs to be commented - "probeinterface @ git+https://github.com/SpikeInterface/probeinterface.git", # We always build from the latest version + # FOR TESTING: use probeinterface zarrv3 branch + "probeinterface @ git+https://github.com/alejoe91/probeinterface.git@zarrv3", + # "probeinterface @ git+https://github.com/SpikeInterface/probeinterface.git", "neo @ git+https://github.com/NeuralEnsemble/python-neo.git", # We always build from the latest version ] diff --git a/src/spikeinterface/core/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py index 6f47790b18..6e26508c81 100644 --- a/src/spikeinterface/core/sortinganalyzer.py +++ b/src/spikeinterface/core/sortinganalyzer.py @@ -1887,8 +1887,14 @@ def get_saved_extension_names(self): elif self.format == "zarr": zarr_root = self._get_zarr_root(mode="r") - if "extensions" in zarr_root.keys(): + # Avoid iterating zarr_root.keys() because legacy v2 stores may contain + # object-dtype arrays (e.g. "recording", "sorting_provenance") that zarr v3 + # cannot parse, causing ValueError on enumeration. + try: extension_group = zarr_root["extensions"] + except KeyError: + extension_group = None + if extension_group is not None: for extension_name in extension_group.keys(): if "params" in extension_group[extension_name].attrs.keys(): saved_extension_names.append(extension_name) diff --git a/src/spikeinterface/core/tests/test_zarr_backwards_compat.py b/src/spikeinterface/core/tests/test_zarr_backwards_compat.py index d4f49d6fd8..49d99ce8eb 100644 --- a/src/spikeinterface/core/tests/test_zarr_backwards_compat.py +++ b/src/spikeinterface/core/tests/test_zarr_backwards_compat.py @@ -18,6 +18,8 @@ import numpy as np import pytest +import spikeinterface as si + FIXTURES_PATH = os.environ.get("ZARR_V2_FIXTURES_PATH") pytestmark = pytest.mark.skipif( @@ -38,9 +40,7 @@ def expected(fixtures_dir: Path) -> dict: def test_load_recording(fixtures_dir, expected): - from spikeinterface.core import read_zarr_recording - - recording = read_zarr_recording(fixtures_dir / "recording.zarr") + recording = si.load(fixtures_dir / "recording.zarr") exp = expected["recording"] assert recording.get_num_channels() == exp["num_channels"] @@ -49,7 +49,7 @@ def test_load_recording(fixtures_dir, expected): assert str(recording.get_dtype()) == exp["dtype"] for seg in range(recording.get_num_segments()): - assert recording.get_num_frames(seg) == exp["num_frames_per_segment"][seg] + assert recording.get_num_samples(seg) == exp["num_samples_per_segment"][seg] assert list(recording.get_channel_ids()) == exp["channel_ids"] @@ -58,9 +58,7 @@ def test_load_recording(fixtures_dir, expected): def test_load_sorting(fixtures_dir, expected): - from spikeinterface.core import read_zarr_sorting - - sorting = read_zarr_sorting(fixtures_dir / "sorting.zarr") + sorting = si.load(fixtures_dir / "sorting.zarr") exp = expected["sorting"] assert sorting.get_num_segments() == exp["num_segments"] @@ -73,9 +71,7 @@ def test_load_sorting(fixtures_dir, expected): def test_load_sorting_analyzer(fixtures_dir, expected): - from spikeinterface.core import load_sorting_analyzer - - analyzer = load_sorting_analyzer(fixtures_dir / "analyzer.zarr") + analyzer = si.load(fixtures_dir / "analyzer.zarr") exp = expected["analyzer"] assert analyzer.get_num_units() == exp["num_units"] diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py index ba3d2f1318..4c3058cde0 100644 --- a/src/spikeinterface/core/zarrextractors.py +++ b/src/spikeinterface/core/zarrextractors.py @@ -8,7 +8,7 @@ from .base import minimum_spike_dtype from .baserecording import BaseRecording, BaseRecordingSegment from .basesorting import BaseSorting, SpikeVectorSortingSegment -from .core_tools import define_function_from_class, check_json +from .core_tools import define_function_from_class, check_json, retrieve_importing_provenance from .job_tools import split_job_kwargs from .core_tools import is_path_remote @@ -251,6 +251,7 @@ def write_recording( recording: BaseRecording, folder_path: str | Path, storage_options: dict | None = None, **kwargs ): zarr_root = zarr.open(str(folder_path), mode="w", storage_options=storage_options) + zarr_root.attrs["zarr_class_info"] = retrieve_importing_provenance(ZarrRecordingExtractor) add_recording_to_zarr_group(recording, zarr_root, **kwargs) @@ -363,6 +364,7 @@ def write_sorting(sorting: BaseSorting, folder_path: str | Path, storage_options Write a sorting extractor to zarr format. """ zarr_root = zarr.open(str(folder_path), mode="w", storage_options=storage_options) + zarr_root.attrs["zarr_class_info"] = retrieve_importing_provenance(ZarrSortingExtractor) add_sorting_to_zarr_group(sorting, zarr_root, **kwargs) @@ -388,15 +390,23 @@ def read_zarr( extractor : ZarrExtractor The loaded extractor """ - # TODO @alessio : we should have something more explicit in our zarr format to tell which object it is. - # for the futur SortingAnalyzer we will have this 2 fields!!! root = super_zarr_open(folder_path, mode="r", storage_options=storage_options) - if "channel_ids" in root.keys(): - return read_zarr_recording(folder_path, storage_options=storage_options) - elif "unit_ids" in root.keys(): - return read_zarr_sorting(folder_path, storage_options=storage_options) + zarr_class_info = root.attrs.get("zarr_class_info", None) + if zarr_class_info is not None: + class_name = zarr_class_info["class"] + extractor_class = _get_class_from_string(class_name) + return extractor_class(folder_path, storage_options=storage_options) else: - raise ValueError("Cannot find 'channel_ids' or 'unit_ids' in zarr root. Not a valid SpikeInterface zarr format") + # For v<0.105.0 and old zarr files, revert to old way of loading based on the presence of "channel_ids" + # or "unit_ids" in the root + if "channel_ids" in root.keys(): + return read_zarr_recording(folder_path, storage_options=storage_options) + elif "unit_ids" in root.keys(): + return read_zarr_sorting(folder_path, storage_options=storage_options) + else: + raise ValueError( + "Cannot find 'channel_ids' or 'unit_ids' in zarr root. Not a valid SpikeInterface zarr format" + ) ### UTILITY FUNCTIONS ### From 7be375b0fb80bf8948a40ae739016e677e6aa50a Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Wed, 25 Mar 2026 12:21:56 +0100 Subject: [PATCH 06/59] Python>=3.11 and disable deepinterpolation action --- .github/workflows/all-tests.yml | 2 +- .github/workflows/deepinterpolation.yml | 6 ++---- .github/workflows/test_containers_docker.yml | 2 +- .github/workflows/test_containers_singularity.yml | 2 +- pyproject.toml | 2 +- 5 files changed, 6 insertions(+), 8 deletions(-) diff --git a/.github/workflows/all-tests.yml b/.github/workflows/all-tests.yml index 0d242b759a..41c3f81054 100644 --- a/.github/workflows/all-tests.yml +++ b/.github/workflows/all-tests.yml @@ -24,7 +24,7 @@ jobs: strategy: fail-fast: false matrix: - python-version: ["3.10", "3.13"] # Lower and higher versions we support + python-version: ["3.11", "3.13"] # Lower and higher versions we support os: [macos-latest, windows-latest, ubuntu-latest] steps: - uses: actions/checkout@v4 diff --git a/.github/workflows/deepinterpolation.yml b/.github/workflows/deepinterpolation.yml index be003da742..2e7b8d03eb 100644 --- a/.github/workflows/deepinterpolation.yml +++ b/.github/workflows/deepinterpolation.yml @@ -1,10 +1,8 @@ name: Testing deepinterpolation +# Manual only — deepinterpolation requires Python 3.10, incompatible with 3.11+ required by Zarr 3.0.0+ on: - pull_request: - types: [synchronize, opened, reopened] - branches: - - main + workflow_dispatch: concurrency: # Cancel previous workflows on the same pull request group: ${{ github.workflow }}-${{ github.ref }} diff --git a/.github/workflows/test_containers_docker.yml b/.github/workflows/test_containers_docker.yml index 211db5f775..73a194efb3 100644 --- a/.github/workflows/test_containers_docker.yml +++ b/.github/workflows/test_containers_docker.yml @@ -15,7 +15,7 @@ jobs: - uses: actions/checkout@v4 - uses: actions/setup-python@v5 with: - python-version: '3.10' + python-version: '3.11' - name: Python version run: python --version diff --git a/.github/workflows/test_containers_singularity.yml b/.github/workflows/test_containers_singularity.yml index 00941215b1..0554a0060c 100644 --- a/.github/workflows/test_containers_singularity.yml +++ b/.github/workflows/test_containers_singularity.yml @@ -16,7 +16,7 @@ jobs: - uses: actions/checkout@v4 - uses: actions/setup-python@v5 with: - python-version: '3.10' + python-version: '3.11' - uses: eWaterCycle/setup-singularity@v7 with: singularity-version: 3.8.7 diff --git a/pyproject.toml b/pyproject.toml index ebe6f3ba1a..2edee0194c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ authors = [ ] description = "Python toolkit for analysis, visualization, and comparison of spike sorting output" readme = "README.md" -requires-python = ">=3.10" +requires-python = ">=3.11" classifiers = [ "Programming Language :: Python :: 3 :: Only", "License :: OSI Approved :: MIT License", From 38987afec555dcfc020c69ab92be3440ffe5933b Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Wed, 25 Mar 2026 14:45:09 +0100 Subject: [PATCH 07/59] Add sharding option --- src/spikeinterface/core/base.py | 42 +++---- src/spikeinterface/core/baserecording.py | 47 +++++++- src/spikeinterface/core/testing_tools.py | 2 +- .../core/tests/test_sortinganalyzer.py | 2 +- .../core/tests/test_zarrextractors.py | 43 ++++++- src/spikeinterface/core/zarr_tools.py | 26 ++++ src/spikeinterface/core/zarrextractors.py | 114 ++++++++++-------- .../preprocessing/tests/test_scaling.py | 2 +- 8 files changed, 197 insertions(+), 81 deletions(-) create mode 100644 src/spikeinterface/core/zarr_tools.py diff --git a/src/spikeinterface/core/base.py b/src/spikeinterface/core/base.py index 9dc270d38d..c16cfe80e3 100644 --- a/src/spikeinterface/core/base.py +++ b/src/spikeinterface/core/base.py @@ -873,6 +873,8 @@ def save(self, **kwargs) -> BaseExtractor: * dump_ext: "json" or "pkl", default "json" (if format is "folder") * verbose: if True output is verbose * **save_kwargs: additional kwargs format-dependent and job kwargs for recording + (check `save_to_memory()`, `save_to_folder()`, `save_to_zarr()` for more details on format-dependent + kwargs) {} Returns @@ -892,13 +894,27 @@ def save(self, **kwargs) -> BaseExtractor: save.__doc__ = save.__doc__.format(_shared_job_kwargs_doc) def save_to_memory(self, sharedmem=True, **save_kwargs) -> BaseExtractor: + """ + Save the object to memory. + + Parameters + ---------- + sharedmem : bool, default: True + If True, the object is saved to shared memory, allowing it to be accessed by multiple processes without + copying. If False, the object is saved to regular memory, which may involve copying when accessed by + multiple processes. + + Returns + ------- + BaseExtractor + A saved copy of the extractor in memory. + """ save_kwargs.pop("format", None) cached = self._save(format="memory", sharedmem=sharedmem, **save_kwargs) self.copy_metadata(cached) return cached - # TODO rename to saveto_binary_folder def save_to_folder( self, name: str | None = None, @@ -944,8 +960,7 @@ def save_to_folder( If True, an existing folder at the specified path will be deleted before saving. verbose : bool, default: True If True, print information about the cache folder being used. - **save_kwargs - Additional keyword arguments to be passed to the underlying save method. + {} Returns ------- @@ -1010,7 +1025,6 @@ def save_to_zarr( folder=None, overwrite=False, storage_options=None, - channel_chunk_size=None, verbose=True, **save_kwargs, ): @@ -1030,26 +1044,9 @@ def save_to_zarr( storage_options: dict or None, default: None Storage options for zarr `store`. E.g., if "s3://" or "gcs://" they can provide authentication methods, etc. For cloud storage locations, this should not be None (in case of default values, use an empty dict) - channel_chunk_size: int or None, default: None - Channels per chunk (only for BaseRecording) - compressor: numcodecs.Codec or None, default: None - Global compressor. If None, Blosc-zstd, level 5, with bit shuffle is used - filters: list[numcodecs.Codec] or None, default: None - Global filters for zarr (global) - compressor_by_dataset: dict or None, default: None - Optional compressor per dataset: - - traces - - times - If None, the global compressor is used - filters_by_dataset: dict or None, default: None - Optional filters per dataset: - - traces - - times - If None, the global filters are used verbose: bool, default: True If True, the output is verbose - auto_cast_uint: bool, default: True - If True, unsigned integers are cast to signed integers to avoid issues with zarr (only for BaseRecording) + {} Returns ------- @@ -1085,7 +1082,6 @@ def save_to_zarr( assert not zarr_path.exists(), f"Path {zarr_path} already exists, choose another name" save_kwargs["zarr_path"] = zarr_path save_kwargs["storage_options"] = storage_options - save_kwargs["channel_chunk_size"] = channel_chunk_size cached = self._save(format="zarr", verbose=verbose, **save_kwargs) cached = read_zarr(zarr_path) diff --git a/src/spikeinterface/core/baserecording.py b/src/spikeinterface/core/baserecording.py index 75bd47597b..a2ccb62937 100644 --- a/src/spikeinterface/core/baserecording.py +++ b/src/spikeinterface/core/baserecording.py @@ -5,10 +5,10 @@ import numpy as np from probeinterface import read_probeinterface, write_probeinterface -from .base import BaseSegment +from .base import BaseSegment, BaseExtractor from .baserecordingsnippets import BaseRecordingSnippets from .core_tools import convert_bytes_to_str, convert_seconds_to_str -from .job_tools import split_job_kwargs +from .job_tools import split_job_kwargs, _shared_job_kwargs_doc from .recording_tools import write_binary_recording @@ -39,6 +39,41 @@ class BaseRecording(BaseRecordingSnippets): "noise_level_rms_scaled", ] + _save_to_folder_docs_params = """dtype: np.dtype | None, default: None + The dtype to use for saving the binary file. If None, the dtype of the recording is used. +""" + _shared_job_kwargs_doc + + _save_to_zarr_docs_params = """ +channel_chunk_size: int | None, default: None + Chunk size for the channel dimension. If None, no chunking is done on the channel dimension. +chunks: tuple | None, default: None + Chunks for the traces dataset. If None, no chunking is done. Note that sharding requires chunking to be specified + and that chunk dimensions need to be larger than shard dimensions (if shards is not None). + If `chunks` is not None, it needs to be a tuple of length 2 with the chunk size for the time and channel + dimensions respectively and `channel_chunk_size` should not be specified. +shard_factor: int | None, default: None + If specified, the shard size will be set to chunk_size * shard_factor in the first dimension (time), + and to be the at most the total number of channels in the second dimension. Note that `shard_factor` cannot + be specified together with `shards`. +shards: tuple | None, default: None + Number of shard size. If None, no sharding is done. Note that shards dimensions need to be larger than + chunk dimensions (if chunks is not None) and that sharding is only done on the first dimension. +compressors: list[numcodecs.Codec] | None, default: None + Global compressor. If None, Blosc-zstd, level 5, with bit shuffle is used +filters: list[numcodecs.Codec] | None, default: None + Global filters for zarr (global) +compressors_by_dataset: dict | None, default: None + Optional compressor per dataset: + - traces + - times + If None, the global compressor is used +filters_by_dataset: dict | None, default: None + Optional filters per dataset: + - traces + - times + If None, the global filters are used +""" + _shared_job_kwargs_doc + def __init__(self, sampling_frequency: float, channel_ids: list, dtype): BaseRecordingSnippets.__init__( self, channel_ids=channel_ids, sampling_frequency=sampling_frequency, dtype=dtype @@ -592,8 +627,8 @@ def _save(self, format="binary", verbose: bool = False, **save_kwargs): if format == "binary": folder = kwargs["folder"] - file_paths = [folder / f"traces_cached_seg{i}.raw" for i in range(self.get_num_segments())] dtype = kwargs.get("dtype", None) or self.get_dtype() + file_paths = [folder / f"traces_cached_seg{i}.raw" for i in range(self.get_num_segments())] t_starts = self._get_t_starts() write_binary_recording(self, file_paths=file_paths, dtype=dtype, verbose=verbose, **job_kwargs) @@ -904,6 +939,12 @@ def astype(self, dtype, round: bool | None = None): return astype(self, dtype=dtype, round=round) +BaseRecording.save_to_folder.__doc__ = BaseExtractor.save_to_folder.__doc__.format( + BaseRecording._save_to_folder_docs_params +) +BaseRecording.save_to_zarr.__doc__ = BaseExtractor.save_to_zarr.__doc__.format(BaseRecording._save_to_zarr_docs_params) + + class BaseRecordingSegment(BaseSegment): """ Abstract class representing a multichannel timeseries, or block of raw ephys traces diff --git a/src/spikeinterface/core/testing_tools.py b/src/spikeinterface/core/testing_tools.py index 899aa3852f..0169d5e50e 100644 --- a/src/spikeinterface/core/testing_tools.py +++ b/src/spikeinterface/core/testing_tools.py @@ -1,7 +1,7 @@ import warnings warnings.warn( - "The 'testing_tools' submodule is deprecated. " "Use spikeinterface.core.generate instead", + "The 'testing_tools' submodule is deprecated. Use spikeinterface.core.testing instead", DeprecationWarning, stacklevel=2, ) diff --git a/src/spikeinterface/core/tests/test_sortinganalyzer.py b/src/spikeinterface/core/tests/test_sortinganalyzer.py index c280fda672..7f4820d2e8 100644 --- a/src/spikeinterface/core/tests/test_sortinganalyzer.py +++ b/src/spikeinterface/core/tests/test_sortinganalyzer.py @@ -17,7 +17,7 @@ AnalyzerExtension, _sort_extensions_by_dependency, ) -from spikeinterface.core.zarrextractors import check_compressors_match +from spikeinterface.core.zarr_utils import check_compressors_match from spikeinterface.core.analyzer_extension_core import BaseSpikeVectorExtension # to test basespikevectorextension with node pipeline diff --git a/src/spikeinterface/core/tests/test_zarrextractors.py b/src/spikeinterface/core/tests/test_zarrextractors.py index a52d456594..0eedb645c8 100644 --- a/src/spikeinterface/core/tests/test_zarrextractors.py +++ b/src/spikeinterface/core/tests/test_zarrextractors.py @@ -10,10 +10,11 @@ generate_sorting, load, ) +from spikeinterface.core.testing import check_recordings_equal +from spikeinterface.core.zarr_tools import check_compressors_match from spikeinterface.core.zarrextractors import ( add_sorting_to_zarr_group, get_default_zarr_compressor, - check_compressors_match, ) @@ -80,6 +81,46 @@ def test_ZarrSortingExtractor(tmp_path): sorting = load(sorting.to_dict()) +def test_sharding_options(tmp_path): + recording = generate_recording(durations=[10], num_channels=20) + folder = tmp_path / "zarr_sharding" + + # explicitly specify chunks and shards + ZarrRecordingExtractor.write_recording(recording, folder, chunks=(1000, 5), shards=(5000, 10), n_jobs=2) + recording_zarr = ZarrRecordingExtractor(folder) + assert recording_zarr._root["traces_seg0"].chunks == (1000, 5) + assert recording_zarr._root["traces_seg0"].shards == (5000, 10) + check_recordings_equal(recording, recording_zarr) + + # specify shard_factor and chunk_size + folder = tmp_path / "zarr_sharding_factor" + ZarrRecordingExtractor.write_recording( + recording, folder, chunk_size=1000, channel_chunk_size=2, shard_factor=5, n_jobs=2 + ) + recording_zarr = ZarrRecordingExtractor(folder) + assert recording_zarr._root["traces_seg0"].chunks == (1000, 2) + assert recording_zarr._root["traces_seg0"].shards == (5000, 10) + check_recordings_equal(recording, recording_zarr) + + # raise error if both shards and shard_factor are provided + with pytest.raises(ValueError): + ZarrRecordingExtractor.write_recording( + recording, folder, chunk_size=1000, channel_chunk_size=2, shard_factor=5, shards=(5000, 10), n_jobs=2 + ) + + # raise error if shards is smaller than chunks + with pytest.raises(AssertionError): + ZarrRecordingExtractor.write_recording( + recording, folder, chunk_size=1000, channel_chunk_size=2, shards=(500, 10), n_jobs=2 + ) + + # raise error if shards is not a multiple of chunks + with pytest.raises(AssertionError): + ZarrRecordingExtractor.write_recording( + recording, folder, chunk_size=1000, channel_chunk_size=2, shards=(5500, 10), n_jobs=2 + ) + + if __name__ == "__main__": tmp_path = Path("tmp") test_zarr_compression_options(tmp_path) diff --git a/src/spikeinterface/core/zarr_tools.py b/src/spikeinterface/core/zarr_tools.py new file mode 100644 index 0000000000..d97630535e --- /dev/null +++ b/src/spikeinterface/core/zarr_tools.py @@ -0,0 +1,26 @@ +def check_compressors_match(comp1, comp2, skip_typesize=True): + """ + Check if two compressor objects match. + + Parameters + ---------- + comp1 : zarr.Codec | Tuple[zarr.Codec] + The first compressor object to compare. + comp2 : zarr.Codec | Tuple[zarr.Codec] + The second compressor object to compare. + skip_typesize : bool, optional + Whether to skip the typesize check, default: True + """ + if not isinstance(comp1, (list, tuple)): + assert not isinstance(comp2, list) + comp1 = [comp1] + comp2 = [comp2] + for i in range(len(comp1)): + comp1_dict = comp1[i].to_dict() + comp2_dict = comp2[i].to_dict() + if skip_typesize: + if "typesize" in comp1_dict["configuration"]: + comp1_dict["configuration"].pop("typesize", None) + if "typesize" in comp2_dict["configuration"]: + comp2_dict["configuration"].pop("typesize", None) + assert comp1_dict == comp2_dict, f"Compressor {i} does not match: {comp1_dict} != {comp2_dict}" diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py index 4c3058cde0..a7c6b8fc7d 100644 --- a/src/spikeinterface/core/zarrextractors.py +++ b/src/spikeinterface/core/zarrextractors.py @@ -5,12 +5,11 @@ from probeinterface import ProbeGroup -from .base import minimum_spike_dtype +from .base import minimum_spike_dtype, _get_class_from_string from .baserecording import BaseRecording, BaseRecordingSegment from .basesorting import BaseSorting, SpikeVectorSortingSegment -from .core_tools import define_function_from_class, check_json, retrieve_importing_provenance -from .job_tools import split_job_kwargs -from .core_tools import is_path_remote +from .core_tools import define_function_from_class, check_json, is_path_remote, retrieve_importing_provenance +from .job_tools import split_job_kwargs, fix_job_kwargs, ensure_chunk_size, ChunkRecordingExecutor zarr.config.set({"default_zarr_version": 3}) @@ -96,34 +95,6 @@ def super_zarr_open(folder_path: str | Path, mode: str = "r", storage_options: d return root -def check_compressors_match(comp1, comp2, skip_typesize=True): - """ - Check if two compressor objects match. - - Parameters - ---------- - comp1 : zarr.Codec | Tuple[zarr.Codec] - The first compressor object to compare. - comp2 : zarr.Codec | Tuple[zarr.Codec] - The second compressor object to compare. - skip_typesize : bool, optional - Whether to skip the typesize check, default: True - """ - if not isinstance(comp1, (list, tuple)): - assert not isinstance(comp2, list) - comp1 = [comp1] - comp2 = [comp2] - for i in range(len(comp1)): - comp1_dict = comp1[i].to_dict() - comp2_dict = comp2[i].to_dict() - if skip_typesize: - if "typesize" in comp1_dict["configuration"]: - comp1_dict["configuration"].pop("typesize", None) - if "typesize" in comp2_dict["configuration"]: - comp2_dict["configuration"].pop("typesize", None) - assert comp1_dict == comp2_dict - - class ZarrRecordingExtractor(BaseRecording): """ RecordingExtractor for a zarr format @@ -513,7 +484,7 @@ def build_codec_pipeline(filters=None, compressors=None): codec_kwargs = {} codec_kwargs["filters"] = filters - codec_kwargs["serializer"] = serializers[0] + codec_kwargs["serializer"] = serializers[0] if len(serializers) == 1 else "auto" codec_kwargs["compressors"] = byte_compressors return codec_kwargs @@ -576,14 +547,22 @@ def add_sorting_to_zarr_group(sorting: BaseSorting, zarr_group: zarr.Group, **kw if compressor is None: compressor = get_default_zarr_compressor() - # save sub fields + # Save sub fields of spikes as separate arrays to allow for more efficient compression and to + # avoid issues with structured arrays with unicode fields in zarr v3. + # The "segment_index" field is saved as "segment_slices" which contains the start and end indices of spikes for + # each segment, to avoid having a large array of segment indices when there are many spikes. spikes_group = zarr_group.create_group(name="spikes") spikes = sorting.to_spike_vector() for field in spikes.dtype.fields: if field != "segment_index": dtype = spikes[field].dtype spikes_data = spikes[field] - codec_kwargs = build_codec_pipeline(filters=[Delta(dtype=spikes[field].dtype.str)], compressors=compressor) + if field == "sample_index": + # Delta filter is very effective for spike times (sample_index) + filters = [Delta(dtype=spikes[field].dtype.str)] + else: + filters = None + codec_kwargs = build_codec_pipeline(filters=filters, compressors=compressor) spikes_group.create_array(name=field, data=spikes_data, **codec_kwargs) else: segment_slices = [] @@ -599,6 +578,7 @@ def add_sorting_to_zarr_group(sorting: BaseSorting, zarr_group: zarr.Group, **kw # Recording def add_recording_to_zarr_group(recording: BaseRecording, zarr_group: zarr.Group, verbose=False, dtype=None, **kwargs): zarr_kwargs, job_kwargs = split_job_kwargs(kwargs) + job_kwargs = fix_job_kwargs(job_kwargs) if recording.check_if_json_serializable(): zarr_group.attrs["provenance"] = check_json(recording.to_dict(recursive=True)) @@ -614,17 +594,53 @@ def add_recording_to_zarr_group(recording: BaseRecording, zarr_group: zarr.Group arr[:] = channel_ids dataset_paths = [f"traces_seg{i}" for i in range(recording.get_num_segments())] + num_channels = recording.get_num_channels() dtype = recording.get_dtype() if dtype is None else dtype - channel_chunk_size = zarr_kwargs.get("channel_chunk_size", None) + + # Compressors and filters global_compressor = kwargs.get("compressors") or kwargs.get("compressor") if global_compressor is None: global_compressor = get_default_zarr_compressor() compressor_by_dataset = zarr_kwargs.pop("compressor_by_dataset", {}) global_filters = zarr_kwargs.pop("filters", None) filters_by_dataset = zarr_kwargs.pop("filters_by_dataset", {}) - compressor_traces = compressor_by_dataset.get("traces", global_compressor) filters_traces = filters_by_dataset.get("traces", global_filters) + + # Chunking and sharding + chunks = zarr_kwargs.get("chunks", None) + channel_chunk_size = zarr_kwargs.get("channel_chunk_size", None) + shards = zarr_kwargs.get("shards", None) + shard_factor = zarr_kwargs.get("shard_factor", None) + if shards is not None and shard_factor is not None: + raise ValueError("Cannot specify both 'shards' and 'shard_factor' in zarr_kwargs") + if chunks is not None and channel_chunk_size is not None: + raise ValueError("Cannot specify both 'chunks' and 'channel_chunk_size' in zarr_kwargs") + + # If not specified by chunk, we set the chunk size in the first dimension (time) to be the chunk size that we use + # for the job executor, and the chunk size in the second dimension (channels) to be either the provided + # channel_chunk_size or the total number of channels (no chunking in channels). + if chunks is not None: + job_kwargs["chunk_size"] = chunks[0] + else: + chunk_size = ensure_chunk_size(recording, **job_kwargs) + chunks = (chunk_size, channel_chunk_size if channel_chunk_size is not None else num_channels) + + if shards is not None: + assert len(shards) == len(chunks), "Shards and chunks must have the same number of dimensions" + for dim in range(len(chunks)): + assert ( + shards[dim] >= chunks[dim] and shards[dim] % chunks[dim] == 0 + ), "Shard size must be a multiple of chunk size" + # When sharding is used, chunk_size in job_kwargs is used to determine the number of samples per chunk to + # write in each job. Each process will write all chunks in a shard. + job_kwargs["chunk_size"] = shards[0] + elif shard_factor is not None: + # If shard_factor is provided, we set the shard size to be chunk_size * shard_factor in the first dimension (time), + # and to be the at most the total number of channels in the second dimension. + shards = (chunks[0] * shard_factor, min(chunks[1] * shard_factor, num_channels)) + job_kwargs["chunk_size"] = shards[0] + add_traces_to_zarr( recording=recording, zarr_group=zarr_group, @@ -632,7 +648,8 @@ def add_recording_to_zarr_group(recording: BaseRecording, zarr_group: zarr.Group compressors=compressor_traces, filters=filters_traces, dtype=dtype, - channel_chunk_size=channel_chunk_size, + chunks=chunks, + shards=shards, verbose=verbose, **job_kwargs, ) @@ -667,7 +684,8 @@ def add_traces_to_zarr( recording, zarr_group, dataset_paths, - channel_chunk_size=None, + chunks=None, + shards=None, dtype=None, compressors=None, filters=None, @@ -685,8 +703,10 @@ def add_traces_to_zarr( The zarr group to add traces to dataset_paths : list List of paths to traces datasets in the zarr group - channel_chunk_size : int or None, default: None (chunking in time only) + chunks : tuple or None, default: None (chunking in time only) Channels per chunk + shards : tuple or None, default: None + If not None, a tuple of (time, num_chunks_per_shard) to dtype : dtype, default: None Type of the saved data compressors : zarr compressor or None, default: None @@ -697,12 +717,6 @@ def add_traces_to_zarr( If True, output is verbose (when chunks are used) {} """ - from .job_tools import ( - ensure_chunk_size, - fix_job_kwargs, - ChunkRecordingExecutor, - ) - assert dataset_paths is not None, "Provide 'file_path'" if not isinstance(dataset_paths, list): @@ -712,9 +726,6 @@ def add_traces_to_zarr( if dtype is None: dtype = recording.get_dtype() - job_kwargs = fix_job_kwargs(job_kwargs) - chunk_size = ensure_chunk_size(recording, **job_kwargs) - codec_kwargs = build_codec_pipeline(filters=filters, compressors=compressors) # create zarr datasets files @@ -725,8 +736,9 @@ def add_traces_to_zarr( dset_name = dataset_paths[segment_index] shape = (num_frames, num_channels) # In zarr v3, chunks must be a tuple of integers (no None allowed) - chunks = (chunk_size, channel_chunk_size if channel_chunk_size is not None else num_channels) - dset = zarr_group.create_array(name=dset_name, shape=shape, chunks=chunks, dtype=dtype, **codec_kwargs) + dset = zarr_group.create_array( + name=dset_name, shape=shape, chunks=chunks, shards=shards, dtype=dtype, **codec_kwargs + ) zarr_datasets.append(dset) # synchronizer=zarr.ThreadSynchronizer()) diff --git a/src/spikeinterface/preprocessing/tests/test_scaling.py b/src/spikeinterface/preprocessing/tests/test_scaling.py index a19d116b16..cc06d88960 100644 --- a/src/spikeinterface/preprocessing/tests/test_scaling.py +++ b/src/spikeinterface/preprocessing/tests/test_scaling.py @@ -1,6 +1,6 @@ import pytest import numpy as np -from spikeinterface.core.testing_tools import generate_recording +from spikeinterface.core.testing import generate_recording from spikeinterface.preprocessing.preprocessing_classes import scale_to_uV, CenterRecording, scale_to_physical_units From f1737863e694b327d1b2925d8d5f06ff2411bb67 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Wed, 25 Mar 2026 14:53:42 +0100 Subject: [PATCH 08/59] wrong imports --- src/spikeinterface/core/tests/test_sortinganalyzer.py | 2 +- src/spikeinterface/preprocessing/tests/test_scaling.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/spikeinterface/core/tests/test_sortinganalyzer.py b/src/spikeinterface/core/tests/test_sortinganalyzer.py index 7f4820d2e8..2912d4f5a1 100644 --- a/src/spikeinterface/core/tests/test_sortinganalyzer.py +++ b/src/spikeinterface/core/tests/test_sortinganalyzer.py @@ -17,7 +17,7 @@ AnalyzerExtension, _sort_extensions_by_dependency, ) -from spikeinterface.core.zarr_utils import check_compressors_match +from spikeinterface.core.zarr_tools import check_compressors_match from spikeinterface.core.analyzer_extension_core import BaseSpikeVectorExtension # to test basespikevectorextension with node pipeline diff --git a/src/spikeinterface/preprocessing/tests/test_scaling.py b/src/spikeinterface/preprocessing/tests/test_scaling.py index cc06d88960..27f1de8542 100644 --- a/src/spikeinterface/preprocessing/tests/test_scaling.py +++ b/src/spikeinterface/preprocessing/tests/test_scaling.py @@ -1,6 +1,6 @@ import pytest import numpy as np -from spikeinterface.core.testing import generate_recording +from spikeinterface.core.generate import generate_recording from spikeinterface.preprocessing.preprocessing_classes import scale_to_uV, CenterRecording, scale_to_physical_units From 42bb52157c5f675036b8634a30f97fb0c1399205 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Thu, 26 Mar 2026 09:25:53 +0100 Subject: [PATCH 09/59] fix: channel_ids dtype --- src/spikeinterface/core/zarrextractors.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py index a7c6b8fc7d..b50ac66021 100644 --- a/src/spikeinterface/core/zarrextractors.py +++ b/src/spikeinterface/core/zarrextractors.py @@ -590,8 +590,7 @@ def add_recording_to_zarr_group(recording: BaseRecording, zarr_group: zarr.Group zarr_group.attrs["num_segments"] = int(recording.get_num_segments()) # Use variable-length UTF-8 (stable zarr v3 spec) instead of fixed-length unicode. channel_ids = recording.get_channel_ids() - arr = zarr_group.create_array(name="channel_ids", shape=channel_ids.shape, dtype=str, compressors=None) - arr[:] = channel_ids + arr = zarr_group.create_array(name="channel_ids", data=channel_ids, compressors=None) dataset_paths = [f"traces_seg{i}" for i in range(recording.get_num_segments())] num_channels = recording.get_num_channels() From 0de855bc4a22f683531acecb963f96cf5da4448b Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Thu, 26 Mar 2026 17:40:16 +0100 Subject: [PATCH 10/59] Fix NWB-zarr tests --- src/spikeinterface/extractors/nwbextractors.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/src/spikeinterface/extractors/nwbextractors.py b/src/spikeinterface/extractors/nwbextractors.py index b89999d088..c5c7639209 100644 --- a/src/spikeinterface/extractors/nwbextractors.py +++ b/src/spikeinterface/extractors/nwbextractors.py @@ -307,8 +307,8 @@ def _get_backend_from_local_file(file_path: str | Path) -> str: try: import zarr - with zarr.open(file_path, "r") as f: - backend = "zarr" + _ = zarr.open(file_path, mode="r") + backend = "zarr" except: raise RuntimeError(f"{file_path} is not a valid Zarr folder!") else: @@ -333,7 +333,8 @@ def _find_neurodata_type_from_backend(group, path="", result=None, neurodata_typ if result is None: result = [] - for neurodata_name, value in group.items(): + for neurodata_name in group.keys(): + value = group[neurodata_name] # Check if it's a group and if it has the neurodata_type if isinstance(value, group_class): current_path = f"{path}/{neurodata_name}" if path else neurodata_name @@ -1409,7 +1410,8 @@ def _find_timeseries_from_backend(group, path="", result=None, backend="hdf5"): if result is None: result = [] - for name, value in group.items(): + for name in group.keys(): + value = group[name] if isinstance(value, group_class): current_path = f"{path}/{name}" if path else name if value.attrs.get("neurodata_type") == "TimeSeries": From 79a104d40234cc420cc1cd03412fbace5d2effaf Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Tue, 31 Mar 2026 09:41:22 +0200 Subject: [PATCH 11/59] update s3fs --- pyproject.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 5e71f930e1..c2a58ef527 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -86,7 +86,7 @@ streaming_extractors = [ "hdmf-zarr @ git+https://github.com/hdmf-dev/hdmf-zarr.git@zarr-v3-migration", # "hdmf-zarr>=0.11.0", "remfile", - "s3fs" + "s3fs>=2025.7.0" ] @@ -168,7 +168,7 @@ test = [ "ibllib>=3.4.1;python_version>='3.10'", # streaming templates - "s3fs", + "s3fs>=2025.7.0", # exporters "pynapple", From 008b5e2048c9216eaea93c7d43c0ae4f8c886654 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Tue, 31 Mar 2026 12:56:06 +0200 Subject: [PATCH 12/59] Fix nwb reading --- src/spikeinterface/extractors/nwbextractors.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/spikeinterface/extractors/nwbextractors.py b/src/spikeinterface/extractors/nwbextractors.py index c5c7639209..a4d640ad16 100644 --- a/src/spikeinterface/extractors/nwbextractors.py +++ b/src/spikeinterface/extractors/nwbextractors.py @@ -1254,7 +1254,7 @@ def _fetch_sorting_segment_info_backend( spike_times_index_data = units_table["spike_times_index"] if "unit_name" in units_table: - unit_ids = units_table["unit_name"] + unit_ids = np.asarray(units_table["unit_name"][:].tolist()) else: unit_ids = units_table["id"] From 87841e488fb7c5222d08f51a0537b9de5d43d4f0 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Tue, 31 Mar 2026 14:40:28 +0200 Subject: [PATCH 13/59] Fix streaming extractors tests --- src/spikeinterface/core/zarrextractors.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py index ebe1241d3d..f4374af5b3 100644 --- a/src/spikeinterface/core/zarrextractors.py +++ b/src/spikeinterface/core/zarrextractors.py @@ -66,14 +66,15 @@ def super_zarr_open(folder_path: str | Path, mode: str = "r", storage_options: d root = None exception = None if is_path_remote(str(folder_path)): + from zarr.storage import FsspecStore + for use_consolidated in use_consolidated_options: if root is not None: break for storage_options in storage_options_to_test: try: - root = zarr.open( - str(folder_path), mode=mode, storage_options=storage_options, use_consolidated=use_consolidated - ) + store = FsspecStore.from_url(str(folder_path), storage_options=storage_options) + root = zarr.open(store, mode=mode, use_consolidated=use_consolidated) break except Exception as e: exception = e From 53d931503087a6bd0ab114e2738b6ed8e4e1912a Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Tue, 9 Jun 2026 15:38:12 +0200 Subject: [PATCH 14/59] fix: imports --- src/spikeinterface/core/baserecording.py | 3 ++- src/spikeinterface/core/zarrextractors.py | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/src/spikeinterface/core/baserecording.py b/src/spikeinterface/core/baserecording.py index 8aa19aac36..28ae7437ef 100644 --- a/src/spikeinterface/core/baserecording.py +++ b/src/spikeinterface/core/baserecording.py @@ -5,10 +5,11 @@ import numpy as np from probeinterface import read_probeinterface, write_probeinterface +from .base import BaseExtractor from .time_series import TimeSeriesSegment, TimeSeries from .baserecordingsnippets import BaseRecordingSnippets from .core_tools import convert_bytes_to_str, convert_seconds_to_str -from .job_tools import split_job_kwargs +from .job_tools import split_job_kwargs, _shared_job_kwargs_doc class BaseRecording(BaseRecordingSnippets, TimeSeries): diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py index 2c27094814..601606e4f9 100644 --- a/src/spikeinterface/core/zarrextractors.py +++ b/src/spikeinterface/core/zarrextractors.py @@ -9,7 +9,7 @@ from .baserecording import BaseRecording, BaseRecordingSegment from .basesorting import BaseSorting, SpikeVectorSortingSegment from .core_tools import define_function_from_class, check_json, is_path_remote, retrieve_importing_provenance -from .job_tools import split_job_kwargs, fix_job_kwargs, ensure_chunk_size, ChunkRecordingExecutor +from .job_tools import split_job_kwargs, fix_job_kwargs, ensure_chunk_size, TimeSeriesChunkExecutor zarr.config.set({"default_zarr_version": 3}) From 14f272eb6879174cb53ed61ff78c05eb0905def3 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Tue, 9 Jun 2026 16:44:24 +0200 Subject: [PATCH 15/59] fix: update test dependencies --- pyproject.toml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 9f2c5257bd..3b9b94c95d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -250,7 +250,9 @@ test-common = [ "pytest-cov", "pytest-mock", "psutil", - "probeinterface @ git+https://github.com/SpikeInterface/probeinterface.git", + # FOR TESTING: use probeinterface zarrv3 branch + "probeinterface @ git+https://github.com/alejoe91/probeinterface.git@zarrv3", + # "probeinterface @ git+https://github.com/SpikeInterface/probeinterface.git", "neo @ git+https://github.com/NeuralEnsemble/python-neo.git", ] From 4e35afb382085105a5650751a6a2db3404036b67 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Wed, 24 Jun 2026 12:23:49 +0200 Subject: [PATCH 16/59] fix: pipeline from analyzer --- src/spikeinterface/preprocessing/pipeline.py | 8 +++----- src/spikeinterface/preprocessing/tests/test_pipeline.py | 2 ++ 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/spikeinterface/preprocessing/pipeline.py b/src/spikeinterface/preprocessing/pipeline.py index 4d90007964..1512d033d0 100644 --- a/src/spikeinterface/preprocessing/pipeline.py +++ b/src/spikeinterface/preprocessing/pipeline.py @@ -222,11 +222,9 @@ def get_preprocessing_dict_from_analyzer(analyzer_folder, format="auto", backend storage_options = backend_options.get("storage_options", {}) zarr_root = super_zarr_open(str(analyzer_folder), mode="r", storage_options=storage_options) - rec_field = zarr_root.get("recording") - if rec_field is not None: - recording_dict = rec_field[0] - else: - recording_dict = {} + recording_dict = zarr_root.attrs.get("recording") + if recording_dict is None: + raise ValueError(f"Cannot find `recording` attribute in {analyzer_folder}.") preprocessing_dict = _make_pipeline_dict_from_recording_dict(recording_dict) diff --git a/src/spikeinterface/preprocessing/tests/test_pipeline.py b/src/spikeinterface/preprocessing/tests/test_pipeline.py index 6a96d4a66c..d3d3b9370e 100644 --- a/src/spikeinterface/preprocessing/tests/test_pipeline.py +++ b/src/spikeinterface/preprocessing/tests/test_pipeline.py @@ -195,6 +195,8 @@ def test_loading_from_analyzer(create_cache_folder): cache_folder = create_cache_folder recording, sorting = generate_ground_truth_recording() + # Make it JSON-serializable by saving it to a folder and reloading it + recording = recording.save(folder=cache_folder / "recording") preprocessing_dict = {"common_reference": {}, "highpass_filter": {"freq_min": 301.0}} pp_recording = apply_preprocessing_pipeline(recording, preprocessing_dict) From 3aee605056fd2b77d3aad28aba6b40a7419d2ff0 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Wed, 24 Jun 2026 12:27:19 +0200 Subject: [PATCH 17/59] fix: zarr backward compatibility tests --- .github/workflows/test_zarr_compat.yml | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/.github/workflows/test_zarr_compat.yml b/.github/workflows/test_zarr_compat.yml index 27be8d633d..7cf5739a1d 100644 --- a/.github/workflows/test_zarr_compat.yml +++ b/.github/workflows/test_zarr_compat.yml @@ -28,15 +28,19 @@ jobs: uses: actions/setup-python@v5 with: python-version: "3.11" + - uses: astral-sh/setup-uv@v7 + with: + python-version: ${{ matrix.python-version }} + enable-cache: true - name: Install SI 0.104.0 with zarr v2 - run: pip install "spikeinterface==0.104.0" "zarr<3" + run: uv pip install --system "spikeinterface==0.104.0" "zarr<3" - name: Generate zarr v2 fixtures run: python .github/scripts/generate_zarr_v2_fixtures.py --output /tmp/zarr_v2_fixtures - name: Install current SI with zarr v3 - run: pip install -e ".[test_core]" + run: uv pip install --system -e . --group test-core - name: Check zarr version is v3 run: python -c "import zarr; v = zarr.__version__; print(f'zarr {v}'); assert int(v.split('.')[0]) >= 3" From b62c0bca009272941dca9d9266b2db6661cf2c72 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Wed, 24 Jun 2026 12:37:08 +0200 Subject: [PATCH 18/59] docs: fix build --- pyproject.toml | 2 +- readthedocs.yml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 444d05deed..145c22e06d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -226,7 +226,7 @@ docs = [ "MEArec", # Use as an example "pandas<3", # in the modules gallery comparison tutorial "hdbscan>=0.8.33", # For sorters spykingcircus2 + tridesclous - "numba", # For many postprocessing functions + "numba>=0.59", # For many postprocessing functions "networkx", "seaborn", "skops", # For automated curation diff --git a/readthedocs.yml b/readthedocs.yml index 486ace7d53..c63f9e5053 100644 --- a/readthedocs.yml +++ b/readthedocs.yml @@ -4,7 +4,7 @@ version: 2 build: os: "ubuntu-24.04" tools: - python: "3.10" + python: "3.13" commands: - asdf plugin add uv - asdf install uv latest From 054ecf9be365690e04c36034e8cfe81eb9b6351e Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Wed, 15 Jul 2026 13:27:26 +0200 Subject: [PATCH 19/59] fix: sorting analzyer zarr tests --- src/spikeinterface/core/tests/test_sortinganalyzer.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/spikeinterface/core/tests/test_sortinganalyzer.py b/src/spikeinterface/core/tests/test_sortinganalyzer.py index 6e63f01c67..db1554afa2 100644 --- a/src/spikeinterface/core/tests/test_sortinganalyzer.py +++ b/src/spikeinterface/core/tests/test_sortinganalyzer.py @@ -138,6 +138,9 @@ def test_SortingAnalyzer_binary_folder(tmp_path, dataset): def test_SortingAnalyzer_zarr(tmp_path, dataset): recording, sorting = dataset + # make recording JSON serializable + recording = recording.save(folder=tmp_path / "recording_for_zarr", overwrite=True) + folder = tmp_path / "test_SortingAnalyzer_zarr.zarr" default_compressor = get_default_zarr_compressor() @@ -177,7 +180,6 @@ def test_SortingAnalyzer_zarr(tmp_path, dataset): overwrite=True, backend_options={"saving_options": {"compressors": None}}, ) - print(sorting_analyzer_no_compression._backend_options) sorting_analyzer_no_compression.compute(["random_spikes", "templates"]) assert ( len( From 9bd8f94db41a6318572dc2c8eb4f00273473496b Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Wed, 15 Jul 2026 14:41:11 +0200 Subject: [PATCH 20/59] fix: try/except when iterating zarr children --- .../extractors/nwbextractors.py | 33 +++++++++++++++++-- 1 file changed, 31 insertions(+), 2 deletions(-) diff --git a/src/spikeinterface/extractors/nwbextractors.py b/src/spikeinterface/extractors/nwbextractors.py index eb009f0c36..b96b125fde 100644 --- a/src/spikeinterface/extractors/nwbextractors.py +++ b/src/spikeinterface/extractors/nwbextractors.py @@ -313,6 +313,27 @@ def _get_backend_from_local_file(file_path: str | Path) -> str: return backend +def _zarr_group_child_names(group): + """ + Return the names of the immediate children of a zarr group without parsing their metadata. + + zarr-python 3.x eagerly reads and validates every child's metadata when iterating + ``group.keys()``. Some arrays written by hdmf-zarr (e.g. variable-length string columns + with an integer ``fill_value``) cannot be parsed by zarr-python 3.x and make the whole + iteration fail. Listing the store directly avoids touching the children's metadata. + """ + if hasattr(group, "store_path"): # zarr v3 + from zarr.core.sync import sync + + async def _collect(): + return [key async for key in group.store.list_dir(group.path)] + + # Filter out this group's own metadata files (".zgroup", ".zattrs", "zarr.json", ...). + return [name for name in sync(_collect()) if not name.startswith(".") and name != "zarr.json"] + else: # zarr v2 + return list(group.keys()) + + def _find_neurodata_type_from_backend(group, path="", result=None, neurodata_type="ElectricalSeries", backend="hdf5"): """ Recursively searches for groups with the specified neurodata_type hdf5 or zarr object, @@ -322,16 +343,24 @@ def _find_neurodata_type_from_backend(group, path="", result=None, neurodata_typ import h5py group_class = h5py.Group + child_names = list(group.keys()) else: import zarr group_class = zarr.Group + child_names = _zarr_group_child_names(group) if result is None: result = [] - for neurodata_name in group.keys(): - value = group[neurodata_name] + for neurodata_name in child_names: + try: + value = group[neurodata_name] + except Exception: + # Skip children whose metadata cannot be parsed (e.g. hdmf-zarr arrays with a + # fill_value that zarr-python 3.x rejects). These are never groups, so skipping + # them is safe when searching for a neurodata_type. + continue # Check if it's a group and if it has the neurodata_type if isinstance(value, group_class): current_path = f"{path}/{neurodata_name}" if path else neurodata_name From 253241a03f6dd39a06cab40c6b8ccdc0d4632d2b Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Thu, 16 Jul 2026 10:00:12 +0200 Subject: [PATCH 21/59] ci: test python 3.14 --- .github/workflows/all-tests.yml | 2 +- src/spikeinterface/extractors/nwbextractors.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/.github/workflows/all-tests.yml b/.github/workflows/all-tests.yml index b8ccf8910b..b20a5a3b8b 100644 --- a/.github/workflows/all-tests.yml +++ b/.github/workflows/all-tests.yml @@ -24,7 +24,7 @@ jobs: strategy: fail-fast: false matrix: - python-version: ["3.11", "3.13"] # 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 diff --git a/src/spikeinterface/extractors/nwbextractors.py b/src/spikeinterface/extractors/nwbextractors.py index b96b125fde..26df52dad9 100644 --- a/src/spikeinterface/extractors/nwbextractors.py +++ b/src/spikeinterface/extractors/nwbextractors.py @@ -329,7 +329,9 @@ async def _collect(): return [key async for key in group.store.list_dir(group.path)] # Filter out this group's own metadata files (".zgroup", ".zattrs", "zarr.json", ...). - return [name for name in sync(_collect()) if not name.startswith(".") and name != "zarr.json"] + # list_dir does not guarantee an order, so sort for deterministic traversal (matches h5py). + names = [name for name in sync(_collect()) if not name.startswith(".") and name != "zarr.json"] + return sorted(names) else: # zarr v2 return list(group.keys()) From b5bdfa2a00040b18053519e9d37ac93ad6a7f6f8 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Thu, 16 Jul 2026 10:19:16 +0200 Subject: [PATCH 22/59] ci: py3.13, one thing at a time --- .github/workflows/all-tests.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/all-tests.yml b/.github/workflows/all-tests.yml index b20a5a3b8b..b8ccf8910b 100644 --- a/.github/workflows/all-tests.yml +++ b/.github/workflows/all-tests.yml @@ -24,7 +24,7 @@ jobs: strategy: fail-fast: false matrix: - python-version: ["3.11", "3.14"] # Lower and higher versions we support + python-version: ["3.11", "3.13"] # Lower and higher versions we support os: [macos-latest, windows-latest, ubuntu-latest] steps: - uses: actions/checkout@v6 From a8bd62792dae246eb82c24fcce53f998264cf208 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Thu, 30 Jul 2026 12:35:52 +0200 Subject: [PATCH 23/59] fix: storage_options --- src/spikeinterface/core/sortinganalyzer.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/spikeinterface/core/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py index ee735a7994..d8ea875a77 100644 --- a/src/spikeinterface/core/sortinganalyzer.py +++ b/src/spikeinterface/core/sortinganalyzer.py @@ -1039,10 +1039,12 @@ def create_zarr( saving_options = backend_options.get("saving_options", {}) if not is_remote: - storage_options = {} + storage_options_kwargs = {} + else: + storage_options_kwargs = {"storage_options": storage_options} # Create zarr root group (and subgroups) - zarr_root = zarr.open(folder, mode="w", storage_options=storage_options) + zarr_root = zarr.open(folder, mode="w", **storage_options_kwargs) sorting_group = zarr_root.create_group("sorting") # for sorting output recording_info_group = zarr_root.create_group("recording_info") # rec_attributes and probe group zarr_root.create_group("extensions") # used later From 78c10e2949499892ed863085e0f8dc204f0b989c Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Thu, 30 Jul 2026 14:20:12 +0200 Subject: [PATCH 24/59] build: update to python 3.11 in all actions --- .github/workflows/caches_cron_job.yml | 4 ++-- .github/workflows/core-test.yml | 4 ++-- .github/workflows/cross_version_serialization.yml | 8 ++++---- .github/workflows/deepinterpolation.yml | 6 +++--- .github/workflows/test_imports.yml | 4 ++-- 5 files changed, 13 insertions(+), 13 deletions(-) diff --git a/.github/workflows/caches_cron_job.yml b/.github/workflows/caches_cron_job.yml index b0587186d9..92d87eac63 100644 --- a/.github/workflows/caches_cron_job.yml +++ b/.github/workflows/caches_cron_job.yml @@ -19,10 +19,10 @@ jobs: steps: - uses: actions/setup-python@v6 with: - python-version: "3.10" + python-version: "3.11" - uses: astral-sh/setup-uv@v7 with: - python-version: "3.10" + python-version: "3.11" enable-cache: true ignore-nothing-to-cache: true # some runs do not require building a cache (ie macOS) this says this is okay - name: Create the directory to store the data diff --git a/.github/workflows/core-test.yml b/.github/workflows/core-test.yml index e64c49851e..500be7d7fb 100644 --- a/.github/workflows/core-test.yml +++ b/.github/workflows/core-test.yml @@ -22,10 +22,10 @@ jobs: - uses: actions/checkout@v6 - uses: actions/setup-python@v6 with: - python-version: "3.10" + python-version: "3.11" - uses: astral-sh/setup-uv@v7 with: - python-version: "3.10" + python-version: "3.11" enable-cache: true - name: Install dependencies run: | diff --git a/.github/workflows/cross_version_serialization.yml b/.github/workflows/cross_version_serialization.yml index 4712a3976d..1578a1bc05 100644 --- a/.github/workflows/cross_version_serialization.yml +++ b/.github/workflows/cross_version_serialization.yml @@ -43,7 +43,7 @@ jobs: steps: - uses: actions/setup-python@v6 with: - python-version: "3.10" + python-version: "3.11" - id: set env: GITHUB_EVENT_NAME: ${{ github.event_name }} @@ -93,17 +93,17 @@ jobs: - name: Set up Python uses: actions/setup-python@v6 with: - python-version: "3.10" + python-version: "3.11" - name: Set up uv uses: astral-sh/setup-uv@v7 with: - python-version: "3.10" + python-version: "3.11" enable-cache: false - name: Generate fixtures with spikeinterface ${{ matrix.si-version }} run: | - uv run --isolated --no-project --python 3.10 \ + uv run --isolated --no-project --python 3.11 \ --with "spikeinterface[core]==${{ matrix.si-version }}" \ python .github/scripts/serialization/serialize_objects.py "$SI_SERIALIZATION_FIXTURES_DIR" diff --git a/.github/workflows/deepinterpolation.yml b/.github/workflows/deepinterpolation.yml index 9801c63c56..a67000bce3 100644 --- a/.github/workflows/deepinterpolation.yml +++ b/.github/workflows/deepinterpolation.yml @@ -1,6 +1,6 @@ name: Testing deepinterpolation -# Manual only — deepinterpolation requires Python 3.10, incompatible with 3.11+ required by Zarr 3.0.0+ +# Manual only — deepinterpolation requires Python 3.11, incompatible with 3.11+ required by Zarr 3.0.0+ on: workflow_dispatch: @@ -20,10 +20,10 @@ jobs: - uses: actions/checkout@v6 - uses: actions/setup-python@v6 with: - python-version: "3.10" + python-version: "3.11" - uses: astral-sh/setup-uv@v7 with: - python-version: "3.10" + python-version: "3.11" enable-cache: false - name: Get changed files id: changed-files diff --git a/.github/workflows/test_imports.yml b/.github/workflows/test_imports.yml index d3d5b5dbb1..6af1e1a253 100644 --- a/.github/workflows/test_imports.yml +++ b/.github/workflows/test_imports.yml @@ -22,10 +22,10 @@ jobs: - uses: actions/checkout@v6 - uses: actions/setup-python@v6 with: - python-version: "3.10" + python-version: "3.11" - uses: astral-sh/setup-uv@v7 with: - python-version: "3.10" + python-version: "3.11" enable-cache: true - name: Install Spikeinterface with only core dependencies run: | From b16209542b788de7b2a634512f278f0cce8e94e0 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Thu, 30 Jul 2026 14:24:02 +0200 Subject: [PATCH 25/59] fix: convert channel_ids to StringDtype before saving to zarr --- src/spikeinterface/core/zarrextractors.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py index 025efa6ea8..8e84dd8c20 100644 --- a/src/spikeinterface/core/zarrextractors.py +++ b/src/spikeinterface/core/zarrextractors.py @@ -746,7 +746,9 @@ def add_recording_to_zarr_group(recording: BaseRecording, zarr_group: zarr.Group zarr_group.attrs["sampling_frequency"] = float(recording.get_sampling_frequency()) zarr_group.attrs["num_segments"] = int(recording.get_num_segments()) # Use variable-length UTF-8 (stable zarr v3 spec) instead of fixed-length unicode. - channel_ids = recording.get_channel_ids() + channel_ids = recording.channel_ids + if channel_ids.dtype.kind in ("U", "S"): + channel_ids = channel_ids.astype("T") arr = zarr_group.create_array(name="channel_ids", data=channel_ids, compressors=None) dataset_paths = [f"traces_seg{i}" for i in range(recording.get_num_segments())] From c82ea59fc92006cef666d48137eb092c39b0417e Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Fri, 31 Jul 2026 11:47:58 +0200 Subject: [PATCH 26/59] fix: load v2 array/objects compatibility and add SortingAnalyzer with exts to serialization check --- .github/scripts/generate_zarr_v2_fixtures.py | 101 ----- .github/scripts/serialization/objects.py | 57 ++- .../serialization/serialize_objects.py | 47 ++- .../test_cross_version_compatibility.py | 7 +- .github/workflows/test_zarr_compat.yml | 51 --- src/spikeinterface/core/loading.py | 10 +- src/spikeinterface/core/sortinganalyzer.py | 52 ++- .../core/tests/test_zarr_backwards_compat.py | 84 ---- src/spikeinterface/core/zarr_tools.py | 378 ++++++++++++++++++ src/spikeinterface/core/zarrextractors.py | 10 +- .../metrics/spiketrain/spiketrain_metrics.py | 2 +- 11 files changed, 513 insertions(+), 286 deletions(-) delete mode 100644 .github/scripts/generate_zarr_v2_fixtures.py delete mode 100644 .github/workflows/test_zarr_compat.yml delete mode 100644 src/spikeinterface/core/tests/test_zarr_backwards_compat.py diff --git a/.github/scripts/generate_zarr_v2_fixtures.py b/.github/scripts/generate_zarr_v2_fixtures.py deleted file mode 100644 index e555873458..0000000000 --- a/.github/scripts/generate_zarr_v2_fixtures.py +++ /dev/null @@ -1,101 +0,0 @@ -#!/usr/bin/env python3 -""" -Generate zarr v2 fixtures for backward compatibility tests. - -Run this script with an old spikeinterface version and zarr<3, e.g.: - pip install "spikeinterface==0.104.0" "zarr<3" - python generate_zarr_v2_fixtures.py --output /tmp/zarr_v2_fixtures - -The script saves: - - recording.zarr : a small ZarrRecordingExtractor - - sorting.zarr : a small ZarrSortingExtractor - - expected_values.json : key values used to verify correct loading -""" -import argparse -import shutil -import json -from pathlib import Path - -import numpy as np -import zarr - -import spikeinterface as si - - -def main(output_dir: Path) -> None: - print(f"spikeinterface version : {si.__version__}") - print(f"zarr version : {zarr.__version__}") - - - output_dir.mkdir(parents=True, exist_ok=True) - - recording, sorting = si.generate_ground_truth_recording(durations=[10, 5],num_channels=32, num_units=10, seed=0) - # save to binary to make them JSON serializable for later expected values extraction - recording = recording.save(folder=output_dir / "recording_binary", overwrite=True) - sorting = sorting.save(folder=output_dir / "sorting_binary", overwrite=True) - # --- save recording --- - recording_path = output_dir / "recording.zarr" - recording_zarr = recording.save(format="zarr", folder=recording_path, overwrite=True) - print(f"Saved recording -> {recording_path}") - - # --- save sorting --- - sorting_path = output_dir / "sorting.zarr" - sorting_zarr = sorting.save(format="zarr", folder=sorting_path, overwrite=True) - print(f"Saved sorting -> {sorting_path}") - - # --- save SortingAnalyzer --- - # Reload the recording from zarr so it is a serializable ZarrRecordingExtractor, - # which the analyzer can store as provenance. - analyzer_path = output_dir / "analyzer.zarr" - if analyzer_path.is_dir(): - shutil.rmtree(analyzer_path) - analyzer = si.create_sorting_analyzer( - sorting_zarr, recording_zarr, format="zarr", folder=analyzer_path, overwrite=True - ) - analyzer.compute(["random_spikes", "templates"]) - print(f"Saved analyzer -> {analyzer_path}") - - # Reload to verify templates are accessible before writing expected values - templates_array = analyzer.get_extension("templates").get_data() - - # --- capture expected values for later assertion --- - expected = { - "spikeinterface_version": si.__version__, - "zarr_version": zarr.__version__, - "recording": { - "num_channels": int(recording.get_num_channels()), - "num_segments": int(recording.get_num_segments()), - "sampling_frequency": float(recording.get_sampling_frequency()), - "num_samples_per_segment": [int(recording.get_num_samples(seg)) for seg in range(recording.get_num_segments())], - "channel_ids": recording.get_channel_ids().tolist(), - "dtype": str(recording.get_dtype()), - # first 10 frames of segment 0 for all channels - "traces_seg0_first10": recording.get_traces(start_frame=0, end_frame=10, segment_index=0).tolist(), - }, - "sorting": { - "num_segments": int(sorting.get_num_segments()), - "sampling_frequency": float(sorting.get_sampling_frequency()), - "unit_ids": sorting.get_unit_ids().tolist(), - "spike_trains_seg0": { - str(uid): sorting.get_unit_spike_train(unit_id=uid, segment_index=0).tolist() - for uid in sorting.unit_ids - }, - }, - "analyzer": { - "num_units": int(analyzer.get_num_units()), - "num_channels": int(analyzer.get_num_channels()), - "templates_shape": list(templates_array.shape), - }, - } - - expected_path = output_dir / "expected_values.json" - with open(expected_path, "w") as f: - json.dump(expected, f, indent=2) - print(f"Saved expected -> {expected_path}") - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Generate zarr v2 fixtures for backward compatibility tests") - parser.add_argument("--output", type=Path, required=True, help="Directory to write fixtures into") - args = parser.parse_args() - main(args.output) diff --git a/.github/scripts/serialization/objects.py b/.github/scripts/serialization/objects.py index ce3d215260..19d6448194 100644 --- a/.github/scripts/serialization/objects.py +++ b/.github/scripts/serialization/objects.py @@ -17,6 +17,11 @@ 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. + +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. + """ from packaging.version import parse from spikeinterface import __version__ as si_version @@ -31,6 +36,7 @@ "binary": "_binary", "numpy_folder": "_numpy_folder", "zarr": ".zarr", + "binary_folder": "_binary_folder", } @@ -49,7 +55,7 @@ def _build_noise_generator_recording(): return NoiseGeneratorRecording(num_channels=4, sampling_frequency=30000.0, durations=[1.0, 1.5], seed=0) -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 @@ -62,7 +68,7 @@ def _build_mock_recording(): return generate_recording(num_channels=4, durations=[1.0], 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__ @@ -83,7 +89,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" @@ -104,7 +110,7 @@ def _build_recording_with_probe(): return rec_with_probe -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 @@ -137,7 +143,7 @@ def _build_recording_with_interleaved_probes(): return rec_with_probe -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 @@ -170,7 +176,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__ @@ -184,7 +190,7 @@ def _build_sorting(): return generate_sorting(num_units=5, sampling_frequency=30000.0, durations=[1.0]) -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) @@ -201,12 +207,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", @@ -256,4 +291,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"], + }, ] diff --git a/.github/scripts/serialization/serialize_objects.py b/.github/scripts/serialization/serialize_objects.py index 44a68ca26e..823a7170fc 100644 --- a/.github/scripts/serialization/serialize_objects.py +++ b/.github/scripts/serialization/serialize_objects.py @@ -12,9 +12,13 @@ """ 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 @@ -22,21 +26,38 @@ 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 == "numpy_folder": - obj.save(folder=dest, format="numpy_folder", overwrite=True) - elif fmt == "zarr": - obj.save(folder=dest, format="zarr", overwrite=True) + 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 == "numpy_folder": + obj.save(folder=dest, format="numpy_folder", overwrite=True) + elif fmt == "zarr": + obj.save(folder=dest, format="zarr", overwrite=True) + elif isinstance(obj, SortingAnalyzer): + if dest.is_dir(): + import shutil + shutil.rmtree(dest) + obj.save_as(folder=dest, format=fmt) + 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()}") diff --git a/.github/scripts/serialization/test_cross_version_compatibility.py b/.github/scripts/serialization/test_cross_version_compatibility.py index cd485f7a1e..bdb7d019ba 100644 --- a/.github/scripts/serialization/test_cross_version_compatibility.py +++ b/.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) diff --git a/.github/workflows/test_zarr_compat.yml b/.github/workflows/test_zarr_compat.yml deleted file mode 100644 index 7cf5739a1d..0000000000 --- a/.github/workflows/test_zarr_compat.yml +++ /dev/null @@ -1,51 +0,0 @@ -name: Test zarr backwards compatibility - -on: - workflow_dispatch: - pull_request: - types: [synchronize, opened, reopened] - branches: - - main - paths: - - "src/spikeinterface/core/zarrextractors.py" - - "src/spikeinterface/core/zarrrecordingextractor.py" - - "src/spikeinterface/core/tests/test_zarr_backwards_compat.py" - - ".github/workflows/test_zarr_compat.yml" - - ".github/scripts/generate_zarr_v2_fixtures.py" - -concurrency: - group: ${{ github.workflow }}-${{ github.ref }} - cancel-in-progress: true - -jobs: - test-zarr-compat: - name: zarr v2 -> v3 backwards compatibility - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v4 - - - name: Set up Python - uses: actions/setup-python@v5 - with: - python-version: "3.11" - - uses: astral-sh/setup-uv@v7 - with: - python-version: ${{ matrix.python-version }} - enable-cache: true - - - name: Install SI 0.104.0 with zarr v2 - run: uv pip install --system "spikeinterface==0.104.0" "zarr<3" - - - name: Generate zarr v2 fixtures - run: python .github/scripts/generate_zarr_v2_fixtures.py --output /tmp/zarr_v2_fixtures - - - name: Install current SI with zarr v3 - run: uv pip install --system -e . --group test-core - - - name: Check zarr version is v3 - run: python -c "import zarr; v = zarr.__version__; print(f'zarr {v}'); assert int(v.split('.')[0]) >= 3" - - - name: Run backward compatibility tests - env: - ZARR_V2_FIXTURES_PATH: /tmp/zarr_v2_fixtures - run: pytest src/spikeinterface/core/tests/test_zarr_backwards_compat.py -v diff --git a/src/spikeinterface/core/loading.py b/src/spikeinterface/core/loading.py index 93ac7c2a56..1792c399d6 100644 --- a/src/spikeinterface/core/loading.py +++ b/src/spikeinterface/core/loading.py @@ -270,11 +270,15 @@ def _guess_object_from_zarr(zarr_folder): return _guess_object_from_dict(spikeinterface_info) # here it is the old fashion and a bit ambiguous - if "templates_array" in zarr_root.keys(): + # get_zarr_group_keys is used to be compatible with both zarr v2 and zarr v3 groups + from .zarr_tools import get_zarr_group_keys + + root_keys = get_zarr_group_keys(zarr_root) + if "templates_array" in root_keys: return "Templates" - elif "channel_ids" in zarr_root.keys() and "unit_ids" not in zarr_root.keys(): + elif "channel_ids" in root_keys and "unit_ids" not in root_keys: return "Recording" - elif "unit_ids" in zarr_root.keys() and "channel_ids" not in zarr_root.keys(): + elif "unit_ids" in root_keys and "channel_ids" not in root_keys: return "Sorting" diff --git a/src/spikeinterface/core/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py index d8ea875a77..6f1a631ec7 100644 --- a/src/spikeinterface/core/sortinganalyzer.py +++ b/src/spikeinterface/core/sortinganalyzer.py @@ -39,6 +39,13 @@ from .sparsity import ChannelSparsity, estimate_sparsity from .sortingfolder import NumpyFolderSorting from .zarrextractors import get_default_zarr_compressor, ZarrSortingExtractor, super_zarr_open +from .zarr_tools import ( + iterate_zarr_group, + get_zarr_attr_or_legacy_object, + is_sklearn_estimator, + save_sklearn_model_to_zarr_group, + load_sklearn_model_from_zarr_group, +) from .node_pipeline import run_node_pipeline # Typing hints @@ -1168,8 +1175,8 @@ def load_from_zarr( # Load recording (if available) if recording is None: - # In zarr v3, recording is stored in attributes - rec_dict = zarr_root.attrs.get("recording", None) + # In zarr v3, recording is stored in attributes (in zarr v2 it was an object array) + rec_dict = get_zarr_attr_or_legacy_object(zarr_root, "recording") if rec_dict is not None: try: recording = load(rec_dict, base_folder=folder) @@ -2136,8 +2143,8 @@ def get_sorting_provenance(self): elif self.format == "zarr": zarr_root = self._get_zarr_root(mode="r") sorting_provenance = None - # In zarr v3, sorting_provenance is stored in attributes - sort_dict = zarr_root.attrs.get("sorting_provenance", None) + # In zarr v3, sorting_provenance is stored in attributes (in zarr v2 it was an object array) + sort_dict = get_zarr_attr_or_legacy_object(zarr_root, "sorting_provenance") if sort_dict is not None: # try-except here is because it's not required to be able # to load the sorting provenance, as the user might have deleted @@ -2516,8 +2523,8 @@ def get_saved_extension_names(self): except KeyError: extension_group = None if extension_group is not None: - for extension_name in extension_group.keys(): - if "params" in extension_group[extension_name].attrs.keys(): + for extension_name, extension in iterate_zarr_group(extension_group): + if "params" in extension.attrs.keys(): saved_extension_names.append(extension_name) else: @@ -3207,21 +3214,26 @@ def load_data(self, lazy=False): self.set_data(ext_data_name, ext_data) elif self.format == "zarr": extension_group = self._get_zarr_extension_group(mode="r") - for ext_data_name in extension_group.keys(): - ext_data_ = extension_group[ext_data_name] + # iterate_zarr_group is used to be compatible with both zarr v2 (saved by + # spikeinterface < 0.105) and zarr v3 groups + for ext_data_name, ext_data_ in iterate_zarr_group(extension_group): # In zarr v3, check if it's a group with dict_data attribute if "dict_data" in ext_data_.attrs: ext_data = ext_data_.attrs["dict_data"] + elif "sklearn_model" in ext_data_.attrs: + ext_data = load_sklearn_model_from_zarr_group(ext_data_) elif "dataframe" in ext_data_.attrs: import pandas as pd index = ext_data_["index"] ext_data = pd.DataFrame(index=index) - for col in ext_data_.keys(): + for col, col_data in iterate_zarr_group(ext_data_): if col != "index": - ext_data.loc[:, col] = ext_data_[col][:] + ext_data.loc[:, col] = col_data[:] ext_data = ext_data.convert_dtypes() - elif "object" in ext_data_.attrs: + elif "object" in ext_data_.attrs or "dict" in ext_data_.attrs: + # "dict" is the zarr v2 (spikeinterface < 0.105) flag for dict/list data, + # which was saved as a length-1 object array ext_data = ext_data_[0] else: ext_data = ext_data_ if lazy else np.array(ext_data_[:]) @@ -3402,15 +3414,17 @@ def _save_data(self): col_data = col_data.astype(str) df_group.create_array(name=col, data=col_data) df_group.attrs["dataframe"] = True + elif is_sklearn_estimator(ext_data): + # sklearn models (e.g. the PCA models) are saved as arrays + attributes, + # so that no pickle is needed (in zarr v2 they were pickled object arrays) + try: + save_sklearn_model_to_zarr_group(extension_group, ext_data_name, ext_data, **saving_options) + except Exception as e: + warnings.warn(f"Could not save the {ext_data_name} model to zarr, skipping: {e}") + if ext_data_name in extension_group: + del extension_group[ext_data_name] else: - # any object - # try: - # extension_group.create_array( - # name=ext_data_name, data=np.array([ext_data], dtype=object), object_codec=numcodecs.Pickle() - # ) - # except: - # raise Exception(f"Could not save {ext_data_name} as extension data") - # extension_group[ext_data_name].attrs["object"] = True + # any other object warnings.warn(f"Data type of {ext_data_name} not supported for zarr saving, skipping.") def _reset_extension_folder(self): diff --git a/src/spikeinterface/core/tests/test_zarr_backwards_compat.py b/src/spikeinterface/core/tests/test_zarr_backwards_compat.py deleted file mode 100644 index 49d99ce8eb..0000000000 --- a/src/spikeinterface/core/tests/test_zarr_backwards_compat.py +++ /dev/null @@ -1,84 +0,0 @@ -""" -Tests for zarr format backward compatibility. - -Loads zarr v2 fixtures (generated with spikeinterface==0.104.0 and zarr<3) using -the current spikeinterface version, which uses zarr>=3. - -The fixtures directory is passed via the ZARR_V2_FIXTURES_PATH environment variable. -These tests are skipped when that variable is not set (i.e. in normal CI runs). - -To run locally: - ZARR_V2_FIXTURES_PATH=/tmp/zarr_v2_fixtures pytest test_zarr_backwards_compat.py -v -""" - -import json -import os -from pathlib import Path - -import numpy as np -import pytest - -import spikeinterface as si - -FIXTURES_PATH = os.environ.get("ZARR_V2_FIXTURES_PATH") - -pytestmark = pytest.mark.skipif( - FIXTURES_PATH is None, - reason="ZARR_V2_FIXTURES_PATH environment variable not set", -) - - -@pytest.fixture(scope="module") -def fixtures_dir() -> Path: - return Path(FIXTURES_PATH) - - -@pytest.fixture(scope="module") -def expected(fixtures_dir: Path) -> dict: - with open(fixtures_dir / "expected_values.json") as f: - return json.load(f) - - -def test_load_recording(fixtures_dir, expected): - recording = si.load(fixtures_dir / "recording.zarr") - exp = expected["recording"] - - assert recording.get_num_channels() == exp["num_channels"] - assert recording.get_num_segments() == exp["num_segments"] - assert recording.get_sampling_frequency() == exp["sampling_frequency"] - assert str(recording.get_dtype()) == exp["dtype"] - - for seg in range(recording.get_num_segments()): - assert recording.get_num_samples(seg) == exp["num_samples_per_segment"][seg] - - assert list(recording.get_channel_ids()) == exp["channel_ids"] - - traces = recording.get_traces(start_frame=0, end_frame=10, segment_index=0) - np.testing.assert_array_equal(traces, np.array(exp["traces_seg0_first10"])) - - -def test_load_sorting(fixtures_dir, expected): - sorting = si.load(fixtures_dir / "sorting.zarr") - exp = expected["sorting"] - - assert sorting.get_num_segments() == exp["num_segments"] - assert sorting.get_sampling_frequency() == exp["sampling_frequency"] - assert list(sorting.get_unit_ids()) == exp["unit_ids"] - - for uid in sorting.unit_ids: - spike_train = sorting.get_unit_spike_train(unit_id=uid, segment_index=0) - np.testing.assert_array_equal(spike_train, np.array(exp["spike_trains_seg0"][str(uid)])) - - -def test_load_sorting_analyzer(fixtures_dir, expected): - analyzer = si.load(fixtures_dir / "analyzer.zarr") - exp = expected["analyzer"] - - assert analyzer.get_num_units() == exp["num_units"] - assert analyzer.get_num_channels() == exp["num_channels"] - - templates_ext = analyzer.get_extension("templates") - assert templates_ext is not None, "templates extension not found in analyzer" - - templates = templates_ext.get_data() - assert list(templates.shape) == exp["templates_shape"] diff --git a/src/spikeinterface/core/zarr_tools.py b/src/spikeinterface/core/zarr_tools.py index d97630535e..fc3ef8e061 100644 --- a/src/spikeinterface/core/zarr_tools.py +++ b/src/spikeinterface/core/zarr_tools.py @@ -1,3 +1,12 @@ +import warnings + +import numpy as np +import zarr + +# metadata keys that are not members of a group (zarr v2 and v3) +_ZARR_METADATA_KEYS = (".zarray", ".zattrs", ".zgroup", ".zmetadata", "zarr.json") + + def check_compressors_match(comp1, comp2, skip_typesize=True): """ Check if two compressor objects match. @@ -24,3 +33,372 @@ def check_compressors_match(comp1, comp2, skip_typesize=True): if "typesize" in comp2_dict["configuration"]: comp2_dict["configuration"].pop("typesize", None) assert comp1_dict == comp2_dict, f"Compressor {i} does not match: {comp1_dict} != {comp2_dict}" + + +class LegacyZarrObjectArray: + """ + Read-only stand-in for a zarr v2 object-dtype array that zarr-python >= 3 cannot open. + + spikeinterface < 0.105 saved python objects (dicts, lists, provenance dicts, ...) as + length-1 object-dtype arrays using the `numcodecs` JSON or Pickle object codecs. + zarr-python >= 3 dropped support for these dtype/codec combinations, so such arrays + are decoded "manually" with `numcodecs` and exposed through this minimal wrapper, + which mimics the small part of the `zarr.Array` API used to read them + (`attrs`, `__getitem__`, `shape`, `dtype`, `len`). + + Parameters + ---------- + values : np.ndarray + The decoded object-dtype array. + attrs : dict + The zarr attributes of the array (content of `.zattrs`). + """ + + def __init__(self, values: np.ndarray, attrs: dict): + self._values = values + self._attrs = dict(attrs) + + @property + def attrs(self) -> dict: + return self._attrs + + @property + def shape(self): + return self._values.shape + + @property + def dtype(self): + return self._values.dtype + + def __len__(self): + return len(self._values) + + def __getitem__(self, key): + return self._values[key] + + def __array__(self, dtype=None, copy=None): + if dtype is None: + return self._values + return self._values.astype(dtype) + + def __repr__(self): + return f"LegacyZarrObjectArray(shape={self.shape}, attrs={list(self._attrs.keys())})" + + +def _sync(coroutine): + """Run a zarr async coroutine synchronously (zarr-python >= 3).""" + from zarr.core.sync import sync + + return sync(coroutine) + + +def _store_get_json(store, key: str): + """Read and json-decode a metadata key from a zarr store. Return None if missing.""" + import json + + from zarr.core.buffer import default_buffer_prototype + + buffer = _sync(store.get(key, prototype=default_buffer_prototype())) + if buffer is None: + return None + return json.loads(buffer.to_bytes().decode()) + + +def _store_get_bytes(store, key: str): + """Read raw bytes from a zarr store. Return None if the key is missing.""" + from zarr.core.buffer import default_buffer_prototype + + buffer = _sync(store.get(key, prototype=default_buffer_prototype())) + if buffer is None: + return None + return buffer.to_bytes() + + +def read_legacy_zarr_object_array(store, path: str) -> LegacyZarrObjectArray: + """ + Read a zarr v2 object-dtype array written with a `numcodecs` object codec + (JSON, Pickle, MsgPack), which zarr-python >= 3 cannot open. + + Only single-chunk arrays are supported, which is what spikeinterface < 0.105 wrote. + + Parameters + ---------- + store : zarr.abc.store.Store + The store containing the array. + path : str + The path of the array inside the store. + + Returns + ------- + legacy_array : LegacyZarrObjectArray + The decoded array. + """ + import numcodecs + + path = path.strip("/") + zarray = _store_get_json(store, f"{path}/.zarray") + if zarray is None: + raise KeyError(f"No zarr v2 array metadata found at {path}") + zattrs = _store_get_json(store, f"{path}/.zattrs") or {} + + shape = tuple(zarray["shape"]) + chunks = tuple(zarray["chunks"]) + if any(chunk_size < dim for chunk_size, dim in zip(chunks, shape)): + raise NotImplementedError( + f"Legacy zarr v2 object array at {path} has more than one chunk, which is not supported" + ) + + separator = zarray.get("dimension_separator", ".") + chunk_key = separator.join(["0"] * max(len(shape), 1)) + chunk_bytes = _store_get_bytes(store, f"{path}/{chunk_key}") + if chunk_bytes is None: + # array was never written to: return an empty object array + return LegacyZarrObjectArray(np.empty(shape, dtype=object), zattrs) + + if zarray.get("compressor", None) is not None: + chunk_bytes = numcodecs.get_codec(zarray["compressor"]).decode(chunk_bytes) + values = chunk_bytes + # filters are applied on encode, so they are decoded in reverse order. + # for object arrays the object codec is the (only) filter + for filter_config in reversed(zarray.get("filters", None) or []): + values = numcodecs.get_codec(filter_config).decode(values) + values = np.asarray(values, dtype=object).reshape(shape) + + return LegacyZarrObjectArray(values, zattrs) + + +def get_zarr_attr_or_legacy_object(zarr_group, name: str): + """ + Get a value saved either in the attributes of a zarr group (spikeinterface >= 0.105) + or as a legacy zarr v2 length-1 object array with the same name + (spikeinterface < 0.105, e.g. "recording" and "sorting_provenance"). + + Parameters + ---------- + zarr_group : zarr.Group + The zarr group to read from. + name : str + The name of the attribute / legacy array. + + Returns + ------- + value : Any | None + The value, or None if it is not found (or cannot be decoded). + """ + value = zarr_group.attrs.get(name, None) + if value is not None: + return value + try: + legacy_array = read_legacy_zarr_object_array(zarr_group.store, f"{zarr_group.path}/{name}") + except Exception: + return None + if len(legacy_array) == 0: + return None + return legacy_array[0] + + +def _list_member_names(zarr_group) -> list[str]: + """List the member names of a zarr group without opening them.""" + consolidated = getattr(zarr_group.metadata, "consolidated_metadata", None) + if consolidated is not None: + return list(consolidated.metadata.keys()) + + store = zarr_group.store + if not store.supports_listing: + raise ValueError( + f"The store associated to this group ({type(store).__name__}) does not support listing, " + "so its members cannot be listed without consolidated metadata." + ) + + async def _list_dir(): + return [key async for key in store.list_dir(zarr_group.path)] + + keys = _sync(_list_dir()) + # skip zarr metadata documents and hidden files (e.g. AppleDouble "._*" files) + return sorted(key for key in keys if key not in _ZARR_METADATA_KEYS and not key.startswith(".")) + + +def iterate_zarr_group(zarr_group, skip_unreadable: bool = True): + """ + Iterate over the members (arrays and sub-groups) of a zarr group, transparently + handling zarr v2 and zarr v3 groups. + + Unlike `zarr.Group.keys()` / `zarr.Group.members()`, a single member that + zarr-python >= 3 cannot open does not make the whole iteration fail. This happens for + zarr v2 data saved by spikeinterface < 0.105, where python objects were stored as + object-dtype arrays with the `numcodecs` JSON/Pickle object codecs. Such arrays are + decoded with `numcodecs` and returned as a `LegacyZarrObjectArray`. + + Parameters + ---------- + zarr_group : zarr.Group + The zarr group to iterate over. + skip_unreadable : bool, default: True + If True, members that cannot be opened nor decoded are skipped with a warning. + If False, an error is raised instead. + + Yields + ------ + name : str + The name of the member. + member : zarr.Group | zarr.Array | LegacyZarrObjectArray + The member itself. + """ + for name in _list_member_names(zarr_group): + # note: members are retrieved outside of the yield statement, so that exceptions + # raised by the consumer of this generator are not caught here + member = None + open_error = None + try: + member = zarr_group[name] + except KeyError: + # the key is an object in the store (e.g. a stray file), not a zarr node + continue + except Exception as e: + open_error = e + + if open_error is not None: + # the member exists but zarr-python cannot open it: try the legacy object array path + try: + member = read_legacy_zarr_object_array(zarr_group.store, f"{zarr_group.path}/{name}") + except Exception as legacy_error: + if not skip_unreadable: + raise ValueError( + f"Cannot read member '{name}' of zarr group '{zarr_group.path}': " + f"{open_error}\nLegacy object array fallback failed with: {legacy_error}" + ) from open_error + warnings.warn( + f"Skipping member '{name}' of zarr group '{zarr_group.path}' because it cannot be read " + f"with zarr v{zarr.__version__}: {open_error} ({legacy_error})" + ) + continue + + yield name, member + + +def get_zarr_group_keys(zarr_group) -> list[str]: + """ + Return the names of the members of a zarr group, for both zarr v2 and zarr v3 groups. + + Contrary to `zarr.Group.keys()`, the members are not opened, so this also works for + groups containing legacy zarr v2 arrays that zarr-python >= 3 cannot open + (see `iterate_zarr_group`). + + Parameters + ---------- + zarr_group : zarr.Group + The zarr group to list. + + Returns + ------- + keys : list[str] + The names of the members of the group. + """ + return _list_member_names(zarr_group) + + +def is_sklearn_estimator(obj) -> bool: + """ + Check whether an object looks like a fitted scikit-learn estimator + (e.g. the PCA models of the "principal_components" extension). + """ + return ( + callable(getattr(obj, "get_params", None)) + and callable(getattr(obj, "set_params", None)) + and hasattr(obj, "__dict__") + ) + + +def save_sklearn_model_to_zarr_group(parent_group, name: str, model, **saving_options) -> None: + """ + Save a scikit-learn estimator in a zarr sub-group without using pickle. + + The state of the estimator (`vars(model)`) is split in two: numpy arrays are saved as + zarr arrays and the remaining (json-serializable) values are saved in the + "sklearn_model" attribute of the group, together with the class of the estimator. + See `load_sklearn_model_from_zarr_group` for the reverse operation. + + Parameters + ---------- + parent_group : zarr.Group + The zarr group in which the sub-group is created. + name : str + The name of the sub-group. + model : sklearn estimator + The estimator to save. + **saving_options : dict + Options passed to `zarr.Group.create_array` for the array parts of the state. + + Raises + ------ + ValueError + If part of the state of the estimator can neither be saved as a zarr array nor + serialized to json (e.g. a nested estimator or an arbitrary python object). + """ + import json + + state = {} + arrays = {} + for key, value in vars(model).items(): + if isinstance(value, np.ndarray): + if value.dtype.kind == "O": + raise ValueError(f"Cannot save object-dtype array '{key}' of {type(model).__name__} to zarr") + arrays[key] = value + else: + if isinstance(value, np.generic): + value = value.item() + try: + json.dumps(value) + except TypeError: + raise ValueError( + f"Cannot save attribute '{key}' of {type(model).__name__} to zarr: " + f"{type(value).__name__} is not json-serializable" + ) + state[key] = value + + model_group = parent_group.create_group(name) + for key, value in arrays.items(): + model_group.create_array(name=key, data=value, **saving_options) + + model_info = { + "class": f"{type(model).__module__}.{type(model).__qualname__}", + "state": state, + "array_keys": sorted(arrays.keys()), + } + try: + import sklearn + + model_info["sklearn_version"] = sklearn.__version__ + except ImportError: + pass + model_group.attrs["sklearn_model"] = model_info + + +def load_sklearn_model_from_zarr_group(model_group): + """ + Rebuild a scikit-learn estimator saved by `save_sklearn_model_to_zarr_group`. + + Parameters + ---------- + model_group : zarr.Group + The zarr group containing the estimator state. + + Returns + ------- + model : sklearn estimator + The estimator, with the same state it had when saved. + """ + import importlib + + model_info = model_group.attrs["sklearn_model"] + module_name, _, class_name = model_info["class"].rpartition(".") + model_class = getattr(importlib.import_module(module_name), class_name) + + # the full state is restored, so the constructor is bypassed on purpose + model = model_class.__new__(model_class) + for key, value in model_info["state"].items(): + setattr(model, key, value) + for key in model_info["array_keys"]: + setattr(model, key, np.asarray(model_group[key][...])) + + return model diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py index 8e84dd8c20..ae9ff48d2d 100644 --- a/src/spikeinterface/core/zarrextractors.py +++ b/src/spikeinterface/core/zarrextractors.py @@ -11,6 +11,7 @@ from .basesorting import BaseSorting, SpikeVectorSortingSegment from .core_tools import define_function_from_class, check_json, is_path_remote, retrieve_importing_provenance from .job_tools import split_job_kwargs, fix_job_kwargs, ensure_chunk_size +from .zarr_tools import iterate_zarr_group zarr.config.set({"default_zarr_version": 3}) @@ -214,16 +215,15 @@ def __init__( # load properties if "properties" in self._root: prop_group = self._root["properties"] - for key in prop_group.keys(): + for key, values in iterate_zarr_group(prop_group): # Skip contact_vector property since it is not used anymore to represent probegroup if key == "contact_vector": continue - values = self._root["properties"][key][:] + values = values[:] # zarr returns vlen-utf8 as StringDType (numpy 2.0); convert via list to classic unicode array. if hasattr(values.dtype, "na_object") or values.dtype.kind == "O": if values.size > 0 and isinstance(values.tolist()[0], str): values = np.array(values.tolist()) - values = self._root["properties"][key] self.set_property(key, values) # load annotations @@ -467,8 +467,8 @@ def __init__( # load properties if "properties" in self._root: prop_group = self._root["properties"] - for key in prop_group.keys(): - values = self._root["properties"][key][:] + for key, values in iterate_zarr_group(prop_group): + values = values[:] # zarr returns vlen-utf8 as StringDType (numpy 2.0); convert via list to classic unicode array. if hasattr(values.dtype, "na_object") or values.dtype.kind == "O": if values.size > 0 and isinstance(values.tolist()[0], str): diff --git a/src/spikeinterface/metrics/spiketrain/spiketrain_metrics.py b/src/spikeinterface/metrics/spiketrain/spiketrain_metrics.py index 79ee3e1b5e..375cd42b1c 100644 --- a/src/spikeinterface/metrics/spiketrain/spiketrain_metrics.py +++ b/src/spikeinterface/metrics/spiketrain/spiketrain_metrics.py @@ -38,7 +38,7 @@ class ComputeSpikeTrainMetrics(BaseMetricExtension): extension_name = "spiketrain_metrics" depend_on = [] - need_backward_compatibility_on_load = True + need_backward_compatibility_on_load = False metric_list = spiketrain_metrics From 64a145b81597e798a4e806f226704cc85ca04a89 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Fri, 31 Jul 2026 11:54:23 +0200 Subject: [PATCH 27/59] fix: cross serialization workflow includes dev --- .../workflows/cross_version_serialization.yml | 23 +++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/.github/workflows/cross_version_serialization.yml b/.github/workflows/cross_version_serialization.yml index 1578a1bc05..cdc228313f 100644 --- a/.github/workflows/cross_version_serialization.yml +++ b/.github/workflows/cross_version_serialization.yml @@ -8,7 +8,9 @@ name: Cross-version serialization # run in an isolated uv environment (nothing committed). The version list is computed # from PyPI (latest patch of each minor at or above MIN_VERSION_TO_TEST): on a pull # request only the latest released minor is tested; the full matrix runs weekly and on -# manual dispatch. +# manual dispatch. The pseudo-version "dev" is always included: it generates the fixtures +# with the branch under test, so a round-trip break (write and read with the same code) +# is caught independently of any release. on: pull_request: @@ -71,7 +73,8 @@ jobs: if os.environ.get('GITHUB_EVENT_NAME') == 'pull_request': minor_releases = minor_releases[-1:] - 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" @@ -102,12 +105,24 @@ 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" \ 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]" --group test-core pytest .github/scripts/serialization/test_cross_version_compatibility.py -v From fd2ea1eb5f5e9dc31ef161cd94b7bc3b647027e3 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Fri, 31 Jul 2026 12:43:41 +0200 Subject: [PATCH 28/59] ci: always cross serializaion tests on all versions --- .github/workflows/cross_version_serialization.yml | 7 ------- 1 file changed, 7 deletions(-) diff --git a/.github/workflows/cross_version_serialization.yml b/.github/workflows/cross_version_serialization.yml index cdc228313f..32ab92cd3c 100644 --- a/.github/workflows/cross_version_serialization.yml +++ b/.github/workflows/cross_version_serialization.yml @@ -25,8 +25,6 @@ on: - "src/spikeinterface/core/**" - ".github/scripts/serialization/**" - ".github/workflows/cross_version_serialization.yml" - schedule: - - cron: "0 4 * * 0" # weekly, Sunday 04:00 UTC workflow_dispatch: concurrency: @@ -48,7 +46,6 @@ jobs: python-version: "3.11" - id: set env: - GITHUB_EVENT_NAME: ${{ github.event_name }} # Support floor: oldest minor to test. Below this, installs fail on the CI # Python (numcodecs dependency rot in 0.100/0.101). Raise when 0.102 rots. MIN_VERSION_TO_TEST: "0.102" # We will change this when making breaking changes (hopefully not too often) @@ -69,10 +66,6 @@ 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))] - # on a pull request, test only the latest minor - if os.environ.get('GITHUB_EVENT_NAME') == 'pull_request': - minor_releases = minor_releases[-1:] - # always round-trip against the branch itself as well print(json.dumps([str(v) for v in minor_releases] + ['dev'])) ") From c3a4236314e169a37560b2b88b860dce008319e3 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Mon, 28 Sep 2026 12:28:53 +0200 Subject: [PATCH 29/59] test: add test on sorting chunks and shards --- .../core/tests/test_basesorting.py | 103 +++++++++++------- src/spikeinterface/core/zarrextractors.py | 30 +++-- 2 files changed, 86 insertions(+), 47 deletions(-) diff --git a/src/spikeinterface/core/tests/test_basesorting.py b/src/spikeinterface/core/tests/test_basesorting.py index 25ba1a3c2d..51f20433db 100644 --- a/src/spikeinterface/core/tests/test_basesorting.py +++ b/src/spikeinterface/core/tests/test_basesorting.py @@ -27,10 +27,35 @@ from spikeinterface.core.testing import check_sorted_arrays_equal, check_sortings_equal +def _make_sorting_with_shuffled_ties(num_units, num_segments, seed=42): + """Build a NumpySorting whose cotemporal spikes are in arbitrary unit_index order. + + A spike vector is only guaranteed to be segment-blocked and sample_index-ascending within each + segment; the unit_index order among spikes sharing a sample_index is unspecified (see #4606). + Building via `NumpySorting.from_unit_dict` happens to produce unit-ascending ties, so it + can't test the shuffled tie case. + """ + rng = np.random.default_rng(seed) + num_spikes = 2_000 + + # A sample range far smaller than num_spikes, so cotemporal spikes are abundant -- including + # repeats of the same (segment, sample, unit), the tie that np.lexsort itself cannot break. + spikes = np.empty(num_spikes, dtype=minimum_spike_dtype) + spikes["sample_index"] = rng.integers(0, 200, size=num_spikes) + spikes["unit_index"] = rng.integers(0, num_units, size=num_spikes) + spikes["segment_index"] = rng.integers(0, num_segments, size=num_spikes) + + # Order by segment then sample, breaking ties randomly rather than by unit_index. + spikes = spikes[np.lexsort((rng.random(num_spikes), spikes["sample_index"], spikes["segment_index"]))] + + sorting = NumpySorting(spikes, 30_000.0, np.arange(num_units)) + assert sorting.get_num_segments() == num_segments + return sorting + + def test_BaseSorting(create_cache_folder): - cache_folder = create_cache_folder num_seg = 2 - file_path = cache_folder / "test_BaseSorting.npz" + file_path = create_cache_folder / "test_BaseSorting.npz" file_path.parent.mkdir(exist_ok=True) create_sorting_npz(num_seg, file_path) @@ -57,28 +82,28 @@ def test_BaseSorting(create_cache_folder): check_sortings_equal(sorting, sorting3, check_annotations=True, check_properties=True) # dump/load json - sorting.dump_to_json(cache_folder / "test_BaseSorting.json") - sorting2 = load(cache_folder / "test_BaseSorting.json") - sorting3 = load(cache_folder / "test_BaseSorting.json") + sorting.dump_to_json(create_cache_folder / "test_BaseSorting.json") + sorting2 = load(create_cache_folder / "test_BaseSorting.json") + sorting3 = load(create_cache_folder / "test_BaseSorting.json") check_sortings_equal(sorting, sorting2, check_annotations=True, check_properties=False) check_sortings_equal(sorting, sorting3, check_annotations=True, check_properties=False) # dump/load pickle - sorting.dump_to_pickle(cache_folder / "test_BaseSorting.pkl") - sorting2 = load(cache_folder / "test_BaseSorting.pkl") - sorting3 = load(cache_folder / "test_BaseSorting.pkl") + sorting.dump_to_pickle(create_cache_folder / "test_BaseSorting.pkl") + sorting2 = load(create_cache_folder / "test_BaseSorting.pkl") + sorting3 = load(create_cache_folder / "test_BaseSorting.pkl") check_sortings_equal(sorting, sorting2, check_annotations=True, check_properties=True) check_sortings_equal(sorting, sorting3, check_annotations=True, check_properties=True) # cache old format : npz_folder - folder = cache_folder / "simple_sorting_npz_folder" + folder = create_cache_folder / "simple_sorting_npz_folder" sorting.set_property("test", np.ones(len(sorting.unit_ids))) sorting.save(folder=folder, format="npz_folder") sorting2 = load(folder) assert isinstance(sorting2, NpzFolderSorting) # cache new format : binary - folder = cache_folder / "simple_sorting_binary" + folder = create_cache_folder / "simple_sorting_binary" sorting.set_property("test", np.ones(len(sorting.unit_ids))) sorting.save(folder=folder, format="numpy_folder") sorting2 = load(folder) @@ -145,8 +170,8 @@ def test_BaseSorting(create_cache_folder): # test save to zarr # compressor = get_default_zarr_compressor() - sorting_zarr = sorting.save(format="zarr", folder=cache_folder / "sorting.zarr") - sorting_zarr_loaded = load(cache_folder / "sorting.zarr") + sorting_zarr = sorting.save(format="zarr", folder=create_cache_folder / "sorting.zarr") + sorting_zarr_loaded = load(create_cache_folder / "sorting.zarr") # annotations is False because Zarr adds compression ratios check_sortings_equal(sorting, sorting_zarr, check_annotations=False, check_properties=True) check_sortings_equal(sorting_zarr, sorting_zarr_loaded, check_annotations=False, check_properties=True) @@ -155,30 +180,34 @@ def test_BaseSorting(create_cache_folder): assert sorting.get_annotation(annotation_name) == sorting_zarr_loaded.get_annotation(annotation_name) -def _make_sorting_with_shuffled_ties(num_units, num_segments, seed=42): - """Build a NumpySorting whose cotemporal spikes are in arbitrary unit_index order. - - A spike vector is only guaranteed to be segment-blocked and sample_index-ascending within each - segment; the unit_index order among spikes sharing a sample_index is unspecified (see #4606). - Building via `NumpySorting.from_unit_dict` happens to produce unit-ascending ties, so it - can't test the shuffled tie case. - """ - rng = np.random.default_rng(seed) - num_spikes = 2_000 - - # A sample range far smaller than num_spikes, so cotemporal spikes are abundant -- including - # repeats of the same (segment, sample, unit), the tie that np.lexsort itself cannot break. - spikes = np.empty(num_spikes, dtype=minimum_spike_dtype) - spikes["sample_index"] = rng.integers(0, 200, size=num_spikes) - spikes["unit_index"] = rng.integers(0, num_units, size=num_spikes) - spikes["segment_index"] = rng.integers(0, num_segments, size=num_spikes) - - # Order by segment then sample, breaking ties randomly rather than by unit_index. - spikes = spikes[np.lexsort((rng.random(num_spikes), spikes["sample_index"], spikes["segment_index"]))] +def test_zarr_save_with_sharding(create_cache_folder): + sorting = generate_sorting(durations=[600, 600]) + # 1 MB target chunk size for Zarr + zarr_target_bytes = 1 * 1024 * 1024 + shard_factor = 3 + sorting_zarr = sorting.save( + format="zarr", + folder=create_cache_folder / "sorting_sharded.zarr", + target_chunk_size_bytes=zarr_target_bytes, + shard_factor=shard_factor, + ) + sorting_zarr_loaded = load(create_cache_folder / "sorting_sharded.zarr") + check_sortings_equal(sorting, sorting_zarr, check_annotations=False, check_properties=True) + check_sortings_equal(sorting_zarr, sorting_zarr_loaded, check_annotations=False, check_properties=True) - sorting = NumpySorting(spikes, 30_000.0, np.arange(num_units)) - assert sorting.get_num_segments() == num_segments - return sorting + # check that chunks and shards are correctly set + spikes_group = sorting_zarr._root["spikes"] + for field in spikes_group: + if field != "segment_slices": + array = spikes_group[field] + assert array.chunks is not None + assert array.shards is not None + assert array.shards[0] == array.chunks[0] * shard_factor + assert array.chunks[0] * array.dtype.itemsize <= zarr_target_bytes + print(f"Field: {field}, Chunks: {array.chunks}, Shards: {array.shards}") + print( + f"Field: {field}, Chunk size in bytes: {array.chunks[0] * array.dtype.itemsize} - Target: {zarr_target_bytes}" + ) @pytest.mark.parametrize("use_numba", [True, False], ids=["numba", "numpy"]) @@ -415,9 +444,9 @@ def test_select_periods(): from pathlib import Path with tempfile.TemporaryDirectory() as tmpdirname: - cache_folder = Path(tmpdirname) + create_cache_folder = Path(tmpdirname) - test_BaseSorting(cache_folder) + test_BaseSorting(create_cache_folder) test_npy_sorting() test_empty_sorting() test_select_periods() diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py index 82512cc93b..445e110491 100644 --- a/src/spikeinterface/core/zarrextractors.py +++ b/src/spikeinterface/core/zarrextractors.py @@ -841,19 +841,19 @@ def add_sorting_to_zarr_group(sorting: BaseSorting, zarr_group: zarr.Group, **kw chunks = None shards = None target_chunk_size_bytes = kwargs.get("target_chunk_size_bytes") - if target_chunk_size_bytes is not None: - spike_num_bytes = spikes.dtype.itemsize - target_chunk_size = target_chunk_size_bytes // spike_num_bytes - chunks = (target_chunk_size,) - shard_factor = kwargs.get("shard_factor") - if shard_factor is not None: - shards = (target_chunk_size * shard_factor,) - else: - chunks = (len(spikes),) # We need https://github.com/zarr-developers/zarr-python/pull/3996 released before being able to # access the structured array lazily. Until then, we always materialise it. # For now, let's keep the old by field implementation + # if target_chunk_size_bytes is not None: + # spike_num_bytes = spikes.dtype.itemsize + # target_chunk_size = target_chunk_size_bytes // spike_num_bytes + # chunks = (target_chunk_size,) + # shard_factor = kwargs.get("shard_factor") + # if shard_factor is not None: + # shards = (target_chunk_size * shard_factor,) + # else: + # chunks = (len(spikes),) # spikes = zarr_group.create_array( # name="spikes", # data=spikes, @@ -861,11 +861,21 @@ def add_sorting_to_zarr_group(sorting: BaseSorting, zarr_group: zarr.Group, **kw # shards=shards, # **codec_kwargs # ) - # save sub fields + + # Save sub fields: in this case chunks and shards are set per field spikes_group = zarr_group.create_group(name="spikes") spikes = sorting.to_spike_vector() for field in spikes.dtype.fields: if field != "segment_index": + if target_chunk_size_bytes is not None: + spike_num_bytes = spikes[field].dtype.itemsize + target_chunk_size = target_chunk_size_bytes // spike_num_bytes + chunks = (target_chunk_size,) + shard_factor = kwargs.get("shard_factor") + if shard_factor is not None: + shards = (target_chunk_size * shard_factor,) + else: + chunks = (len(spikes),) codec_kwargs = build_codec_pipeline(filters=Delta(dtype=spikes[field].dtype.str), compressors=compressor) spikes_group.create_array(name=field, data=spikes[field], chunks=chunks, shards=shards, **codec_kwargs) From ab77d6b2178d5e77937a7e49a8e1e5a62795132f Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Mon, 28 Sep 2026 12:51:40 +0200 Subject: [PATCH 30/59] fix: tuple generation --- pyproject.toml | 8 ++------ src/spikeinterface/core/zarr_tools.py | 2 +- 2 files changed, 3 insertions(+), 7 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index b32f26a3c9..c02738fa90 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,9 +67,7 @@ changelog = "https://spikeinterface.readthedocs.io/en/latest/whatisnew.html" extractors = [ "MEArec>=1.8", "pynwb>=2.6.0", - # FOR TESTING - "hdmf-zarr @ git+https://github.com/hdmf-dev/hdmf-zarr.git@zarr-v3-migration", - # "hdmf-zarr>=0.11.0", + "hdmf-zarr>=0.14.0", "pyedflib>=0.1.30", "lxml", # lxml for neuroscope "scipy", @@ -85,9 +83,7 @@ streaming_extractors = [ "fsspec", "aiohttp", "requests", - # FOR TESTING - "hdmf-zarr @ git+https://github.com/hdmf-dev/hdmf-zarr.git@zarr-v3-migration", - # "hdmf-zarr>=0.11.0", + "hdmf-zarr>=0.14.0", "remfile", "s3fs>=2025.7.0" ] diff --git a/src/spikeinterface/core/zarr_tools.py b/src/spikeinterface/core/zarr_tools.py index b1b3c6e7e2..38d43de7ad 100644 --- a/src/spikeinterface/core/zarr_tools.py +++ b/src/spikeinterface/core/zarr_tools.py @@ -54,7 +54,7 @@ def adjust_chunks_shards_and_job_kwargs( else: chunk_size = ensure_chunk_size(time_series, **job_kwargs) chunks = ( - (chunk_size,) + extra_chunks if extra_chunks is not None else (chunk_size, time_series.get_shape(0)[1:]) + (chunk_size,) + extra_chunks if extra_chunks is not None else (chunk_size,) + time_series.get_shape(0)[1:] ) if shards is not None: From f1a9d552b280d641670358bfb6f55b0bce9ca9cd Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Mon, 28 Sep 2026 13:03:33 +0200 Subject: [PATCH 31/59] build: pythonn >=3.12 --- .github/workflows/all-tests.yml | 2 +- .github/workflows/caches_cron_job.yml | 4 ++-- .github/workflows/core-test.yml | 4 ++-- .github/workflows/cross_version_serialization.yml | 8 ++++---- .github/workflows/deepinterpolation.yml | 6 +++--- .github/workflows/test_containers_docker.yml | 4 ++-- .github/workflows/test_containers_singularity.yml | 4 ++-- .github/workflows/test_imports.yml | 4 ++-- pyproject.toml | 2 +- 9 files changed, 19 insertions(+), 19 deletions(-) diff --git a/.github/workflows/all-tests.yml b/.github/workflows/all-tests.yml index c18ad13113..295be91f5a 100644 --- a/.github/workflows/all-tests.yml +++ b/.github/workflows/all-tests.yml @@ -25,7 +25,7 @@ jobs: strategy: fail-fast: false matrix: - python-version: ["3.11", "3.13"] # Lower and higher versions we support + python-version: ["3.12", "3.13"] # Lower and higher versions we support os: [macos-latest, windows-latest, ubuntu-latest] steps: - uses: actions/checkout@v6 diff --git a/.github/workflows/caches_cron_job.yml b/.github/workflows/caches_cron_job.yml index 74a02eee3f..38d0f5bcbe 100644 --- a/.github/workflows/caches_cron_job.yml +++ b/.github/workflows/caches_cron_job.yml @@ -15,10 +15,10 @@ jobs: steps: - uses: actions/setup-python@v6 with: - python-version: "3.11" + python-version: "3.12" - uses: astral-sh/setup-uv@v7 with: - python-version: "3.11" + python-version: "3.12" enable-cache: true ignore-nothing-to-cache: true # some runs do not require building a cache (ie macOS) this says this is okay - name: Create the directory to store the data diff --git a/.github/workflows/core-test.yml b/.github/workflows/core-test.yml index 500be7d7fb..b65606588b 100644 --- a/.github/workflows/core-test.yml +++ b/.github/workflows/core-test.yml @@ -22,10 +22,10 @@ jobs: - uses: actions/checkout@v6 - uses: actions/setup-python@v6 with: - python-version: "3.11" + python-version: "3.12" - uses: astral-sh/setup-uv@v7 with: - python-version: "3.11" + python-version: "3.12" enable-cache: true - name: Install dependencies run: | diff --git a/.github/workflows/cross_version_serialization.yml b/.github/workflows/cross_version_serialization.yml index 1422c26f69..21843416f5 100644 --- a/.github/workflows/cross_version_serialization.yml +++ b/.github/workflows/cross_version_serialization.yml @@ -40,7 +40,7 @@ jobs: steps: - uses: actions/setup-python@v6 with: - python-version: "3.11" + python-version: "3.12" - id: set env: # Support floor: oldest minor to test. Below this, installs fail on the CI @@ -86,12 +86,12 @@ jobs: - name: Set up Python uses: actions/setup-python@v6 with: - python-version: "3.11" + python-version: "3.12" - name: Set up uv uses: astral-sh/setup-uv@v7 with: - python-version: "3.11" + python-version: "3.12" enable-cache: false - name: Generate fixtures with spikeinterface ${{ matrix.si-version }} @@ -106,7 +106,7 @@ jobs: else spec="spikeinterface[full]==$SI_VERSION" fi - uv run --isolated --no-project --python 3.11 \ + uv run --isolated --no-project --python 3.12 \ --with "$spec" \ python .github/scripts/serialization/serialize_objects.py "$SI_SERIALIZATION_FIXTURES_DIR" diff --git a/.github/workflows/deepinterpolation.yml b/.github/workflows/deepinterpolation.yml index 3c2cd70b35..c5a3af31b9 100644 --- a/.github/workflows/deepinterpolation.yml +++ b/.github/workflows/deepinterpolation.yml @@ -1,6 +1,6 @@ name: Testing deepinterpolation -# Manual only — deepinterpolation requires Python 3.11, incompatible with 3.11+ required by Zarr 3.0.0+ +# Manual only — deepinterpolation requires Python 3.12, incompatible with 3.12+ required by Zarr 3.0.0+ on: workflow_dispatch: @@ -20,10 +20,10 @@ jobs: - uses: actions/checkout@v6 - uses: actions/setup-python@v6 with: - python-version: "3.11" + python-version: "3.12" - uses: astral-sh/setup-uv@v7 with: - python-version: "3.11" + python-version: "3.12" enable-cache: false - name: Get changed files id: changed-files diff --git a/.github/workflows/test_containers_docker.yml b/.github/workflows/test_containers_docker.yml index 7b85a8e2dd..1b794d4eb7 100644 --- a/.github/workflows/test_containers_docker.yml +++ b/.github/workflows/test_containers_docker.yml @@ -14,10 +14,10 @@ jobs: - uses: actions/checkout@v6 - uses: actions/setup-python@v6 with: - python-version: '3.11' + python-version: '3.12' - uses: astral-sh/setup-uv@v7 with: - python-version: '3.11' + python-version: '3.12' enable-cache: true - name: Python version run: python --version diff --git a/.github/workflows/test_containers_singularity.yml b/.github/workflows/test_containers_singularity.yml index 1b85040628..184c50aeff 100644 --- a/.github/workflows/test_containers_singularity.yml +++ b/.github/workflows/test_containers_singularity.yml @@ -15,10 +15,10 @@ jobs: - uses: actions/checkout@v6 - uses: actions/setup-python@v6 with: - python-version: '3.11' + python-version: '3.12' - uses: astral-sh/setup-uv@v7 with: - python-version: '3.11' + python-version: '3.12' enable-cache: true - uses: eWaterCycle/setup-singularity@v7 with: diff --git a/.github/workflows/test_imports.yml b/.github/workflows/test_imports.yml index 6af1e1a253..d2bd1fadcf 100644 --- a/.github/workflows/test_imports.yml +++ b/.github/workflows/test_imports.yml @@ -22,10 +22,10 @@ jobs: - uses: actions/checkout@v6 - uses: actions/setup-python@v6 with: - python-version: "3.11" + python-version: "3.12" - uses: astral-sh/setup-uv@v7 with: - python-version: "3.11" + python-version: "3.12" enable-cache: true - name: Install Spikeinterface with only core dependencies run: | diff --git a/pyproject.toml b/pyproject.toml index c02738fa90..fa5e711587 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ authors = [ ] description = "Python toolkit for analysis, visualization, and comparison of spike sorting output" readme = "README.md" -requires-python = ">=3.11" +requires-python = ">=3.12" classifiers = [ "Programming Language :: Python :: 3 :: Only", "License :: OSI Approved :: MIT License", From 1483a02c2597e7c310d24bfbb69a1d2dbd8459a2 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Mon, 28 Sep 2026 13:11:02 +0200 Subject: [PATCH 32/59] fix: cross-serialization tests --- .../serialization/serialize_objects.py | 34 +++++++++++-------- 1 file changed, 20 insertions(+), 14 deletions(-) diff --git a/.github/scripts/serialization/serialize_objects.py b/.github/scripts/serialization/serialize_objects.py index 74b2773e01..b22e012210 100644 --- a/.github/scripts/serialization/serialize_objects.py +++ b/.github/scripts/serialization/serialize_objects.py @@ -36,20 +36,26 @@ 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) print(f" wrote {dest.name} ({fmt})") if check_extra_data is not None: json_dest = out_dir / f"{entry['id']}.json" From 416a8f2747fb5e46519f5ce9df13e67fdedd88e4 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Mon, 28 Sep 2026 14:58:42 +0200 Subject: [PATCH 33/59] fix: use _LINKS for zarr links in pynwb --- src/spikeinterface/extractors/nwbextractors.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/spikeinterface/extractors/nwbextractors.py b/src/spikeinterface/extractors/nwbextractors.py index b9b15fac02..9f860d4c13 100644 --- a/src/spikeinterface/extractors/nwbextractors.py +++ b/src/spikeinterface/extractors/nwbextractors.py @@ -379,14 +379,16 @@ def _find_neurodata_type_from_backend(group, path="", result=None, neurodata_typ def _retrieve_electrodes_indices_from_electrical_series_backend(open_file, electrical_series, backend="hdf5"): """ Retrieves the indices of the electrodes from the electrical series. - For the Zarr backend, the electrodes are stored in the electrical_series.attrs["zarr_link"]. + For the Zarr backend, the electrodes are stored in the electrical_series.attrs["_LINKS"] + or legacy electrical_series.attrs["zarr_link"]. + See https://github.com/hdmf-dev/hdmf-zarr/pull/336 """ if "electrodes" not in electrical_series: if backend == "zarr": import zarr # links must be resolved - zarr_links = electrical_series.attrs["zarr_link"] + zarr_links = electrical_series.attrs.get("_LINKS", electrical_series.attrs.get("zarr_link", [])) electrodes_path = None for zarr_link in zarr_links: if zarr_link["name"] == "electrodes": From e83581eee270433711cbeec29895da2f15508409 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Tue, 29 Sep 2026 17:04:30 +0200 Subject: [PATCH 34/59] buiold: pin numpy<2.5 in cross-serialization tests --- .github/workflows/cross_version_serialization.yml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/cross_version_serialization.yml b/.github/workflows/cross_version_serialization.yml index 21843416f5..b7d7b88c61 100644 --- a/.github/workflows/cross_version_serialization.yml +++ b/.github/workflows/cross_version_serialization.yml @@ -108,11 +108,12 @@ jobs: fi uv run --isolated --no-project --python 3.12 \ --with "$spec" \ + --with "numpy<2.5" \ python .github/scripts/serialization/serialize_objects.py "$SI_SERIALIZATION_FIXTURES_DIR" - name: Load fixtures with the branch run: | # [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]" --group test-core + uv pip install --system -e ".[full]" "numpy<2.5" --group test-core pytest .github/scripts/serialization/test_cross_version_compatibility.py -v From fa1cdbf6273f04bba980fb34e4e778963e7f02d9 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Tue, 29 Sep 2026 17:19:44 +0200 Subject: [PATCH 35/59] fix: missing check_extra_data in serialize tests --- .github/scripts/serialization/objects.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/scripts/serialization/objects.py b/.github/scripts/serialization/objects.py index bb0f6b86bb..fb9e0a4175 100644 --- a/.github/scripts/serialization/objects.py +++ b/.github/scripts/serialization/objects.py @@ -125,14 +125,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 From 2fe08da55a1cb3657e4fa35088ed56a0b9208e18 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Tue, 29 Sep 2026 17:25:35 +0200 Subject: [PATCH 36/59] fix: cross-serialization tests --- .github/scripts/serialization/objects.py | 19 ++++++------------- .../serialization/serialize_objects.py | 2 ++ 2 files changed, 8 insertions(+), 13 deletions(-) diff --git a/.github/scripts/serialization/objects.py b/.github/scripts/serialization/objects.py index fb9e0a4175..d65856a498 100644 --- a/.github/scripts/serialization/objects.py +++ b/.github/scripts/serialization/objects.py @@ -23,9 +23,6 @@ which are computed and stored in the fixture JSON to verify that they are reloaded correctly. """ -from packaging.version import parse -from spikeinterface import __version__ as si_version - # 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 @@ -111,11 +108,9 @@ 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)) - if parse(si_version) <= parse("0.105.0"): - rec_with_probe = rec.set_probe(probe, in_place=False) # old API returns a new recording; portable across versions - else: - rec_with_probe = rec.set_probe(probe) # new API returns a new recording; portable across versions - return rec_with_probe + # 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 def _build_recording_with_timestamps(): @@ -158,11 +153,9 @@ def _build_recording_with_interleaved_probes(): probegroup.add_probe(probe1) probegroup.set_global_device_channel_indices([0, 2, 4, 6, 1, 3, 5, 7]) # Interleave the two probes' channels: channel i alternates between probe0 and probe1. - if parse(si_version) <= parse("0.105.0"): - rec_with_probe = rec.set_probegroup(probegroup, in_place=False) # old API returns a new recording; portable across versions - else: - rec_with_probe = rec.set_probegroup(probegroup) # new API returns a new recording; portable across versions - return rec_with_probe + # 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, check_extra_data=None): diff --git a/.github/scripts/serialization/serialize_objects.py b/.github/scripts/serialization/serialize_objects.py index b22e012210..9e9ebe9bfe 100644 --- a/.github/scripts/serialization/serialize_objects.py +++ b/.github/scripts/serialization/serialize_objects.py @@ -56,6 +56,8 @@ 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" From 9b4abc9bab4fc7bf8d0f0445775ee85022355439 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Tue, 29 Sep 2026 17:54:31 +0200 Subject: [PATCH 37/59] fix: postprocessing trsts --- .../postprocessing/tests/test_multi_extensions.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/src/spikeinterface/postprocessing/tests/test_multi_extensions.py b/src/spikeinterface/postprocessing/tests/test_multi_extensions.py index 1bb3d383a2..3e3566f9fa 100644 --- a/src/spikeinterface/postprocessing/tests/test_multi_extensions.py +++ b/src/spikeinterface/postprocessing/tests/test_multi_extensions.py @@ -69,7 +69,7 @@ } -def get_dataset_to_merge(): +def get_dataset_to_merge(folder): # generate a dataset with some split units to minimize merge errors recording, sorting = generate_ground_truth_recording( durations=[30.0], @@ -81,6 +81,7 @@ def get_dataset_to_merge(): generate_unit_locations_kwargs=dict(margin_um=10.0, minimum_z=2.0, maximum_z=15.0, minimum_distance=20), seed=2205, ) + recording = recording.save(folder / "recording_for_merge") # since templates are going to be averaged and this might be a problem for amplitude scaling # we select the 3 units with the largest templates to split @@ -96,7 +97,7 @@ def get_dataset_to_merge(): return recording, sorting_with_splits, split_unit_ids -def get_dataset_to_split(): +def get_dataset_to_split(folder): # generate a dataset and return large unit to split to minimize split errors recording, sorting = generate_ground_truth_recording( durations=[30.0], @@ -107,6 +108,7 @@ def get_dataset_to_split(): noise_kwargs=dict(noise_levels=5.0, strategy="tile_pregenerated"), seed=2205, ) + recording = recording.save(folder / "recording_for_split") # since templates are going to be averaged and this might be a problem for amplitude scaling # we select the 3 units with the largest templates to split @@ -120,13 +122,13 @@ def get_dataset_to_split(): @pytest.fixture(scope="module") -def dataset_to_merge(): - return get_dataset_to_merge() +def dataset_to_merge(create_cache_folder): + return get_dataset_to_merge(create_cache_folder) @pytest.fixture(scope="module") -def dataset_to_split(): - return get_dataset_to_split() +def dataset_to_split(create_cache_folder): + return get_dataset_to_split(create_cache_folder) @pytest.mark.parametrize("lazy", [False, True]) From 734368fce83c0994ff78d6b36ccb2d562b7b1267 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Tue, 29 Sep 2026 18:29:44 +0200 Subject: [PATCH 38/59] oups --- .../postprocessing/tests/test_multi_extensions.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/spikeinterface/postprocessing/tests/test_multi_extensions.py b/src/spikeinterface/postprocessing/tests/test_multi_extensions.py index 3e3566f9fa..7cf052913f 100644 --- a/src/spikeinterface/postprocessing/tests/test_multi_extensions.py +++ b/src/spikeinterface/postprocessing/tests/test_multi_extensions.py @@ -81,7 +81,7 @@ def get_dataset_to_merge(folder): generate_unit_locations_kwargs=dict(margin_um=10.0, minimum_z=2.0, maximum_z=15.0, minimum_distance=20), seed=2205, ) - recording = recording.save(folder / "recording_for_merge") + recording = recording.save(folder=folder / "recording_for_merge") # since templates are going to be averaged and this might be a problem for amplitude scaling # we select the 3 units with the largest templates to split @@ -108,7 +108,7 @@ def get_dataset_to_split(folder): noise_kwargs=dict(noise_levels=5.0, strategy="tile_pregenerated"), seed=2205, ) - recording = recording.save(folder / "recording_for_split") + recording = recording.save(folder=folder / "recording_for_split") # since templates are going to be averaged and this might be a problem for amplitude scaling # we select the 3 units with the largest templates to split From f5311b16336a830e41a2636dccb56679580a1fdd Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Wed, 30 Sep 2026 11:13:57 +0200 Subject: [PATCH 39/59] fix: build_codec_pipeline in _write_time_series_to_zarr --- src/spikeinterface/core/time_series_tools.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/src/spikeinterface/core/time_series_tools.py b/src/spikeinterface/core/time_series_tools.py index deb8634623..5d40733e2f 100644 --- a/src/spikeinterface/core/time_series_tools.py +++ b/src/spikeinterface/core/time_series_tools.py @@ -362,6 +362,12 @@ def _write_time_series_to_zarr( job_kwargs = fix_job_kwargs(job_kwargs) + # place an ArrayBytesCodec passed as a compressor (e.g. WavPack) in the serializer slot + from .zarrextractors import build_codec_pipeline + + codec_kwargs_data = build_codec_pipeline(filters=filters_data, compressors=compressor_data) + codec_kwargs_times = build_codec_pipeline(filters=filters_times, compressors=compressor_times) + # create zarr datasets files zarr_datasets = [] zarr_timestamps_datasets = [] @@ -376,8 +382,7 @@ def _write_time_series_to_zarr( chunks=chunks, shards=shards, dtype=dtype, - filters=filters_data, - compressors=compressor_data, + **codec_kwargs_data, ) zarr_datasets.append(dset) if dataset_timestamps_paths[segment_index] is not None: @@ -391,8 +396,7 @@ def _write_time_series_to_zarr( chunks=chunks_times, shards=shards_times, dtype="float64", - filters=filters_times, - compressors=compressor_times, + **codec_kwargs_times, ) ) else: From 00514996dc50c9f97946df0e21590fce4ba0ec00 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Wed, 30 Sep 2026 11:15:52 +0200 Subject: [PATCH 40/59] fix: save params for zarr --- src/spikeinterface/core/baserecording.py | 54 ++++++------------------ 1 file changed, 13 insertions(+), 41 deletions(-) diff --git a/src/spikeinterface/core/baserecording.py b/src/spikeinterface/core/baserecording.py index 002616dac7..f9df170686 100644 --- a/src/spikeinterface/core/baserecording.py +++ b/src/spikeinterface/core/baserecording.py @@ -38,41 +38,6 @@ class BaseRecording(BaseRecordingSnippets, TimeSeries): "noise_level_rms_scaled", ] - _save_to_folder_docs_params = """dtype: np.dtype | None, default: None - The dtype to use for saving the binary file. If None, the dtype of the recording is used. -""" + _shared_job_kwargs_doc - - _save_to_zarr_docs_params = """ -channel_chunk_size: int | None, default: None - Chunk size for the channel dimension. If None, no chunking is done on the channel dimension. -chunks: tuple | None, default: None - Chunks for the traces dataset. If None, no chunking is done. Note that sharding requires chunking to be specified - and that chunk dimensions need to be larger than shard dimensions (if shards is not None). - If `chunks` is not None, it needs to be a tuple of length 2 with the chunk size for the time and channel - dimensions respectively and `channel_chunk_size` should not be specified. -shard_factor: int | None, default: None - If specified, the shard size will be set to chunk_size * shard_factor in the first dimension (time), - and to be the at most the total number of channels in the second dimension. Note that `shard_factor` cannot - be specified together with `shards`. -shards: tuple | None, default: None - Number of shard size. If None, no sharding is done. Note that shards dimensions need to be larger than - chunk dimensions (if chunks is not None) and that sharding is only done on the first dimension. -compressors: list[numcodecs.Codec] | None, default: None - Global compressor. If None, Blosc-zstd, level 5, with bit shuffle is used -filters: list[numcodecs.Codec] | None, default: None - Global filters for zarr (global) -compressors_by_dataset: dict | None, default: None - Optional compressor per dataset: - - traces - - times - If None, the global compressor is used -filters_by_dataset: dict | None, default: None - Optional filters per dataset: - - traces - - times - If None, the global filters are used -""" + _shared_job_kwargs_doc - def __init__(self, sampling_frequency: float, channel_ids: list, dtype): BaseRecordingSnippets.__init__( self, channel_ids=channel_ids, sampling_frequency=sampling_frequency, dtype=dtype @@ -382,6 +347,18 @@ def save(self, format="binary", verbose: bool = False, **save_kwargs): For cloud storage locations, this should not be None (in case of default values, use an empty dict) - channel_chunk_size: int or None, default: None Channels per chunk (only for BaseRecording) + - chunks: tuple | None, default: None + Chunks for the traces dataset. If None, no chunking is done. Note that sharding requires chunking to be specified + and that chunk dimensions need to be larger than shard dimensions (if shards is not None). + If `chunks` is not None, it needs to be a tuple of length 2 with the chunk size for the time and channel + dimensions respectively and `channel_chunk_size` should not be specified. + - shard_factor: int | tuple[int] | None, default: None + If integer, the shard size will be set to chunk_size * shard_factor in the first dimension (time). + If tuple, the shard_factor to be applied to each dimension. Note that `shard_factor` cannot + be specified together with `shards`. + -shards: tuple | None, default: None + Number of shard size. If None, no sharding is done. Note that shards dimensions need to be larger than + chunk dimensions (if chunks is not None) and that sharding is only done on the first dimension. - compressor: numcodecs.Codec or None, default: None Global compressor. If None, Blosc-zstd, level 5, with bit shuffle is used - filters: list[numcodecs.Codec] or None, default: None @@ -400,6 +377,7 @@ def save(self, format="binary", verbose: bool = False, **save_kwargs): - times If None, the global filters are used + * "memory" format: - sharedmem : bool, default: True If True, the recording is saved in shared memory. If False, it is saved as @@ -690,12 +668,6 @@ def astype(self, dtype, round: bool | None = None): return astype(self, dtype=dtype, round=round) -BaseRecording.save_to_folder.__doc__ = BaseExtractor.save_to_folder.__doc__.format( - BaseRecording._save_to_folder_docs_params -) -BaseRecording.save_to_zarr.__doc__ = BaseExtractor.save_to_zarr.__doc__.format(BaseRecording._save_to_zarr_docs_params) - - class BaseRecordingSegment(TimeSeriesSegment): """ Abstract class representing a multichannel timeseries, or block of raw ephys traces From b5a7bb113ea45b25d46b30de1b4b5ee104216e88 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Thu, 8 Oct 2026 17:45:01 +0200 Subject: [PATCH 41/59] fix: remove all requires_zarr_write --- pyproject.toml | 1 - .../core/tests/test_analyzer_extension_core.py | 17 ++++------------- src/spikeinterface/core/tests/test_loading.py | 9 +++------ .../core/tests/test_node_pipeline.py | 2 -- .../core/tests/test_sortinganalyzer.py | 13 ++++--------- .../core/tests/test_template_class.py | 1 - .../core/tests/test_time_handling.py | 4 ++-- .../core/tests/test_waveform_tools.py | 2 -- .../core/tests/test_zarrextractors.py | 2 -- .../tests/test_multi_extensions.py | 8 ++------ .../sorters/internal/tests/test_lupin.py | 2 +- .../sorters/internal/tests/test_tridesclous2.py | 2 +- .../sorters/tests/test_launcher.py | 2 -- .../sorters/tests/test_runsorter.py | 1 - 14 files changed, 17 insertions(+), 49 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index ae4138da05..7b3a1fa385 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -366,7 +366,6 @@ markers = [ "widgets", "sortingcomponents", "streaming_extractors: extractors that require streaming such as ross and fsspec", - "requires_zarr_write: tests that write to zarr, skipped with zarr>=3 where writing is not supported yet", ] filterwarnings =[ 'ignore:.*distutils Version classes are deprecated.*:DeprecationWarning', diff --git a/src/spikeinterface/core/tests/test_analyzer_extension_core.py b/src/spikeinterface/core/tests/test_analyzer_extension_core.py index ab8d3104e0..6423ce11ce 100644 --- a/src/spikeinterface/core/tests/test_analyzer_extension_core.py +++ b/src/spikeinterface/core/tests/test_analyzer_extension_core.py @@ -74,9 +74,7 @@ def _check_result_extension(sorting_analyzer, extension_name, cache_folder): # print(k, arr.shape) -@pytest.mark.parametrize( - "format", ["memory", "binary_folder", pytest.param("zarr", marks=pytest.mark.requires_zarr_write)] -) +@pytest.mark.parametrize("format", ["memory", "binary_folder", "zarr"]) @pytest.mark.parametrize( "sparse", [ @@ -106,9 +104,7 @@ def test_ComputeRandomSpikes(format, sparse, create_cache_folder): _check_result_extension(sorting_analyzer, "random_spikes", cache_folder) -@pytest.mark.parametrize( - "format", ["memory", "binary_folder", pytest.param("zarr", marks=pytest.mark.requires_zarr_write)] -) +@pytest.mark.parametrize("format", ["memory", "binary_folder", "zarr"]) @pytest.mark.parametrize("sparse", [True, False]) def test_ComputeWaveforms(format, sparse, create_cache_folder): cache_folder = create_cache_folder @@ -121,7 +117,6 @@ def test_ComputeWaveforms(format, sparse, create_cache_folder): _check_result_extension(sorting_analyzer, "waveforms", cache_folder) -@pytest.mark.requires_zarr_write @pytest.mark.parametrize("sparse", [False, True]) def test_ComputeWaveforms_consistent_across_formats(create_cache_folder, sparse): # Computing waveforms (and templates) on memory / binary_folder / zarr analyzers with the same @@ -177,9 +172,7 @@ def test_ComputeWaveforms_consistent_across_formats(create_cache_folder, sparse) assert np.array_equal(zarr_waveforms, wfs_mem) -@pytest.mark.parametrize( - "format", ["memory", "binary_folder", pytest.param("zarr", marks=pytest.mark.requires_zarr_write)] -) +@pytest.mark.parametrize("format", ["memory", "binary_folder", "zarr"]) @pytest.mark.parametrize("sparse", [True, False]) def test_ComputeTemplates(format, sparse, create_cache_folder): cache_folder = create_cache_folder @@ -268,9 +261,7 @@ def test_ComputeTemplates(format, sparse, create_cache_folder): _check_result_extension(sorting_analyzer, "templates", cache_folder) -@pytest.mark.parametrize( - "format", ["memory", "binary_folder", pytest.param("zarr", marks=pytest.mark.requires_zarr_write)] -) +@pytest.mark.parametrize("format", ["memory", "binary_folder", "zarr"]) @pytest.mark.parametrize("sparse", [True, False]) def test_ComputeNoiseLevels(format, sparse, create_cache_folder): cache_folder = create_cache_folder diff --git a/src/spikeinterface/core/tests/test_loading.py b/src/spikeinterface/core/tests/test_loading.py index 364c5bd910..8d57a3412f 100644 --- a/src/spikeinterface/core/tests/test_loading.py +++ b/src/spikeinterface/core/tests/test_loading.py @@ -90,7 +90,7 @@ def generate_motion_object(): return motion -@pytest.mark.parametrize("output_format", ["binary", pytest.param("zarr", marks=pytest.mark.requires_zarr_write)]) +@pytest.mark.parametrize("output_format", ["binary", "zarr"]) def test_load_binary_recording(generate_recording_sorting, tmp_path, output_format): rec, _ = generate_recording_sorting _ = rec.save(folder=tmp_path / "test_recording", format=output_format, overwrite=True) @@ -103,7 +103,7 @@ def test_load_binary_recording(generate_recording_sorting, tmp_path, output_form check_recordings_equal(rec, rec_loaded) -@pytest.mark.parametrize("output_format", ["numpy_folder", pytest.param("zarr", marks=pytest.mark.requires_zarr_write)]) +@pytest.mark.parametrize("output_format", ["numpy_folder", "zarr"]) def test_load_binary_sorting(generate_recording_sorting, tmp_path, output_format): _, sort = generate_recording_sorting _ = sort.save(folder=tmp_path / "test_sorting", format=output_format, overwrite=True) @@ -141,9 +141,7 @@ def test_load_ext_extractors(generate_recording_sorting, tmp_path, extension): check_sortings_equal(sort, sort_loaded, check_properties=False) -@pytest.mark.parametrize( - "output_format", ["binary_folder", pytest.param("zarr", marks=pytest.mark.requires_zarr_write)] -) +@pytest.mark.parametrize("output_format", ["binary_folder", "zarr"]) def test_load_sorting_analyzer(generate_sorting_analyzer, tmp_path, output_format): analyzer = generate_sorting_analyzer _ = analyzer.save_as(folder=tmp_path / "analyzer", format=output_format) @@ -160,7 +158,6 @@ def test_load_sorting_analyzer(generate_sorting_analyzer, tmp_path, output_forma assert ext in analyzer_loaded.extensions -@pytest.mark.requires_zarr_write def test_load_templates(tmp_path, generate_templates_object): templates = generate_templates_object templates_dict = templates.to_dict() diff --git a/src/spikeinterface/core/tests/test_node_pipeline.py b/src/spikeinterface/core/tests/test_node_pipeline.py index 7cfb9d47a5..a721fc3367 100644 --- a/src/spikeinterface/core/tests/test_node_pipeline.py +++ b/src/spikeinterface/core/tests/test_node_pipeline.py @@ -271,7 +271,6 @@ def test_run_node_pipeline(cache_folder_creation): unpickled_node = pickle.loads(pickled_node) -@pytest.mark.requires_zarr_write def test_gather_to_zarr_chunking(tmp_path): # the zarr chunk size along the first axis must be picked from a byte target (not from the # size of the first gathered buffer), so it stays sensible for billions of spikes and never @@ -320,7 +319,6 @@ def test_gather_to_zarr_chunking(tmp_path): assert np.array_equal(waveforms[:], waveforms2[:]) -@pytest.mark.requires_zarr_write def test_gather_to_zarr_chunk_bytes_per_name(tmp_path): # `zarr_target_chunk_bytes` can also be a dict to use a different byte target per array recording, sorting = generate_ground_truth_recording(num_channels=8, num_units=5, durations=[20.0], seed=7) diff --git a/src/spikeinterface/core/tests/test_sortinganalyzer.py b/src/spikeinterface/core/tests/test_sortinganalyzer.py index 9dc34903f5..1de350651c 100644 --- a/src/spikeinterface/core/tests/test_sortinganalyzer.py +++ b/src/spikeinterface/core/tests/test_sortinganalyzer.py @@ -141,7 +141,6 @@ def test_SortingAnalyzer_binary_folder(tmp_path, dataset): assert "number" in sorting_analyzer_reloded.sorting.get_property_keys() -@pytest.mark.requires_zarr_write def test_SortingAnalyzer_zarr(tmp_path, dataset): recording, sorting = dataset recording = recording.save(folder=tmp_path / "recording_zarr") @@ -370,7 +369,7 @@ def test_SortingAnalyzer_interleaved_probegroup(dataset): assert np.array_equal(recording.get_channel_locations(), sorting_analyzer.get_channel_locations()) -@pytest.mark.parametrize("format", ["binary_folder", pytest.param("zarr", marks=pytest.mark.requires_zarr_write)]) +@pytest.mark.parametrize("format", ["binary_folder", "zarr"]) def test_load_in_lazy_mode(tmp_path, dataset, format): recording, sorting = dataset @@ -862,9 +861,7 @@ def _compute_reference_pipeline_data(dataset): return analyzer.get_extension("dummy_pipeline").get_data() -@pytest.mark.parametrize( - "format", ["memory", "binary_folder", pytest.param("zarr", marks=pytest.mark.requires_zarr_write)] -) +@pytest.mark.parametrize("format", ["memory", "binary_folder", "zarr"]) @pytest.mark.parametrize("lazy", [True, False]) def test_compute_pipeline_extension_gather_to_disk_lazy(tmp_path, dataset, format, lazy): """ @@ -929,7 +926,7 @@ def test_compute_pipeline_extension_gather_to_disk_lazy(tmp_path, dataset, forma assert np.array_equal(load_sorting_analyzer(folder).get_extension("dummy_pipeline").get_data(), amp_ref) -@pytest.mark.parametrize("format", ["binary_folder", pytest.param("zarr", marks=pytest.mark.requires_zarr_write)]) +@pytest.mark.parametrize("format", ["binary_folder", "zarr"]) def test_compute_pipeline_extension_save_false(tmp_path, dataset, format): """ With save=False on a disk-backed analyzer, node-pipeline extensions are computed in memory @@ -951,9 +948,7 @@ def test_compute_pipeline_extension_save_false(tmp_path, dataset, format): assert not analyzer_reloaded.has_extension("dummy_pipeline") -@pytest.mark.parametrize( - "format", ["memory", "binary_folder", pytest.param("zarr", marks=pytest.mark.requires_zarr_write)] -) +@pytest.mark.parametrize("format", ["memory", "binary_folder", "zarr"]) @pytest.mark.parametrize("lazy", [True, False]) def test_compute_one_pipeline_extension_gather_to_disk(tmp_path, dataset, format, lazy): """ diff --git a/src/spikeinterface/core/tests/test_template_class.py b/src/spikeinterface/core/tests/test_template_class.py index 20f88817cb..953527c8d5 100644 --- a/src/spikeinterface/core/tests/test_template_class.py +++ b/src/spikeinterface/core/tests/test_template_class.py @@ -115,7 +115,6 @@ def test_initialization_fail_with_dense_templates(): template = generate_test_template(template_type="sparse_with_dense_templates") -@pytest.mark.requires_zarr_write @pytest.mark.parametrize("is_in_uV", [True, False]) @pytest.mark.parametrize("template_type", ["dense", "sparse"]) def test_save_and_load_zarr(template_type, is_in_uV, tmp_path): diff --git a/src/spikeinterface/core/tests/test_time_handling.py b/src/spikeinterface/core/tests/test_time_handling.py index 4a08a1c7c8..2f3cc1a6bf 100644 --- a/src/spikeinterface/core/tests/test_time_handling.py +++ b/src/spikeinterface/core/tests/test_time_handling.py @@ -114,7 +114,7 @@ def test_has_time_vector(self, time_vector_recording): assert raw_recording.has_time_vector(segment_idx) is False assert times_recording.has_time_vector(segment_idx) is True - @pytest.mark.parametrize("mode", ["binary", pytest.param("zarr", marks=pytest.mark.requires_zarr_write)]) + @pytest.mark.parametrize("mode", ["binary", "zarr"]) @pytest.mark.parametrize("fixture_name", ["time_vector_recording", "t_start_recording"]) def test_times_propagated_to_save_folder(self, request, fixture_name, mode, tmp_path): """ @@ -375,7 +375,7 @@ def test_save_and_load_time_shift(self, request, fixture_name, tmp_path): times_recording.get_times(segment_index=idx), loaded_recording.get_times(segment_index=idx) ) - @pytest.mark.parametrize("save_format", ["binary", pytest.param("zarr", marks=pytest.mark.requires_zarr_write)]) + @pytest.mark.parametrize("save_format", ["binary", "zarr"]) def test_shift_times_after_load(self, request, save_format, tmp_path): """ Shift times on a recording loaded from disk as a read-only np.memmap diff --git a/src/spikeinterface/core/tests/test_waveform_tools.py b/src/spikeinterface/core/tests/test_waveform_tools.py index c2750774b6..6623c92046 100644 --- a/src/spikeinterface/core/tests/test_waveform_tools.py +++ b/src/spikeinterface/core/tests/test_waveform_tools.py @@ -157,7 +157,6 @@ def test_waveform_tools(create_cache_folder): _check_all_wf_equal(list_wfs_sparse) -@pytest.mark.requires_zarr_write @pytest.mark.parametrize("sparse", [False, True]) def test_extract_waveforms_to_single_buffer_zarr(tmp_path, sparse): # the "zarr" mode writes waveforms directly to a zarr dataset. Workers return their block and @@ -211,7 +210,6 @@ def test_extract_waveforms_to_single_buffer_zarr(tmp_path, sparse): assert np.array_equal(reference, reloaded[:]) -@pytest.mark.requires_zarr_write def test_waveforms_at_segment_borders(tmp_path): # spikes near the segment borders are partially filled: samples outside of the segment are 0 from spikeinterface.core import NumpySorting diff --git a/src/spikeinterface/core/tests/test_zarrextractors.py b/src/spikeinterface/core/tests/test_zarrextractors.py index fbe2cccd93..11e1458ad3 100644 --- a/src/spikeinterface/core/tests/test_zarrextractors.py +++ b/src/spikeinterface/core/tests/test_zarrextractors.py @@ -18,7 +18,6 @@ ) -@pytest.mark.requires_zarr_write def test_zarr_compression_options(tmp_path): from zarr.codecs.numcodecs import Delta, FixedScaleOffset from zarr.codecs import BloscCodec, BloscShuffle @@ -63,7 +62,6 @@ def test_zarr_compression_options(tmp_path): check_compressors_match(rec_other._root["times_seg0"].filters, other_filters2) -@pytest.mark.requires_zarr_write def test_ZarrSortingExtractor(tmp_path): np_sorting = generate_sorting() diff --git a/src/spikeinterface/postprocessing/tests/test_multi_extensions.py b/src/spikeinterface/postprocessing/tests/test_multi_extensions.py index 6ae72e374d..7b53f24f29 100644 --- a/src/spikeinterface/postprocessing/tests/test_multi_extensions.py +++ b/src/spikeinterface/postprocessing/tests/test_multi_extensions.py @@ -133,9 +133,7 @@ def dataset_to_split(create_cache_folder): @pytest.mark.parametrize("lazy", [False, True]) @pytest.mark.parametrize("sparse", [False, True]) -@pytest.mark.parametrize( - "format", ["memory", "binary_folder", pytest.param("zarr", marks=pytest.mark.requires_zarr_write)] -) +@pytest.mark.parametrize("format", ["memory", "binary_folder", "zarr"]) def test_SortingAnalyzer_merge_all_extensions(dataset_to_merge, lazy, sparse, format, tmp_path): if format == "memory" and lazy: pytest.skip("lazy has no effect for format='memory' (nothing on disk to load lazily)") @@ -267,9 +265,7 @@ def test_SortingAnalyzer_merge_all_extensions(dataset_to_merge, lazy, sparse, fo @pytest.mark.parametrize("lazy", [False, True]) @pytest.mark.parametrize("sparse", [False, True]) -@pytest.mark.parametrize( - "format", ["memory", "binary_folder", pytest.param("zarr", marks=pytest.mark.requires_zarr_write)] -) +@pytest.mark.parametrize("format", ["memory", "binary_folder", "zarr"]) def test_SortingAnalyzer_split_all_extensions(dataset_to_split, lazy, sparse, format, tmp_path): if format == "memory" and lazy: pytest.skip("lazy has no effect for format='memory' (nothing on disk to load lazily)") diff --git a/src/spikeinterface/sorters/internal/tests/test_lupin.py b/src/spikeinterface/sorters/internal/tests/test_lupin.py index c4eb788a6d..f838f9f4eb 100644 --- a/src/spikeinterface/sorters/internal/tests/test_lupin.py +++ b/src/spikeinterface/sorters/internal/tests/test_lupin.py @@ -13,7 +13,7 @@ class LupinSorterCommonTestSuite(SorterCommonTestSuite, unittest.TestCase): SorterClass = LupinSorter # TODO: remove once writing to zarr is supported with zarr>=3 (save_array writes the templates to zarr) - @pytest.mark.requires_zarr_write + def test_with_run(self): super().test_with_run() diff --git a/src/spikeinterface/sorters/internal/tests/test_tridesclous2.py b/src/spikeinterface/sorters/internal/tests/test_tridesclous2.py index cb6ed09ce4..6431532b87 100644 --- a/src/spikeinterface/sorters/internal/tests/test_tridesclous2.py +++ b/src/spikeinterface/sorters/internal/tests/test_tridesclous2.py @@ -13,7 +13,7 @@ class Tridesclous2SorterCommonTestSuite(SorterCommonTestSuite, unittest.TestCase SorterClass = Tridesclous2Sorter # TODO: remove once writing to zarr is supported with zarr>=3 (save_array writes the templates to zarr) - @pytest.mark.requires_zarr_write + def test_with_run(self): super().test_with_run() diff --git a/src/spikeinterface/sorters/tests/test_launcher.py b/src/spikeinterface/sorters/tests/test_launcher.py index 35db70e677..6619431866 100644 --- a/src/spikeinterface/sorters/tests/test_launcher.py +++ b/src/spikeinterface/sorters/tests/test_launcher.py @@ -52,7 +52,6 @@ def job_list(create_cache_folder): return get_job_list(folder) -@pytest.mark.requires_zarr_write def test_run_sorter_jobs_loop(job_list): sortings = run_sorter_jobs(job_list, engine="loop", return_output=True) print(sortings) @@ -202,7 +201,6 @@ def test_run_sorter_jobs_slurm_kwargs(mocker, tmp_path, job_list): assert str(tmp_script_folder) in mock_subprocess_run.call_args_list[-1].args[0][5] -@pytest.mark.requires_zarr_write def test_run_sorter_by_property(create_cache_folder): cache_folder = create_cache_folder working_folder1 = cache_folder / "test_run_sorter_by_property_1" diff --git a/src/spikeinterface/sorters/tests/test_runsorter.py b/src/spikeinterface/sorters/tests/test_runsorter.py index 1316430e55..c07cd94c9b 100644 --- a/src/spikeinterface/sorters/tests/test_runsorter.py +++ b/src/spikeinterface/sorters/tests/test_runsorter.py @@ -24,7 +24,6 @@ def generate_recording(create_cache_folder): return _generate_recording(create_cache_folder) -@pytest.mark.requires_zarr_write @pytest.mark.xfail( platform.system() == "Windows" and parse(platform.python_version()) > parse("3.12"), reason="3rd parth threadpoolctl issue: OSError('GetModuleFileNameEx failed')", From 6d9babe8d5d2f10a300239cb3b0f518cb37b6820 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Thu, 8 Oct 2026 17:49:56 +0200 Subject: [PATCH 42/59] fix: create_array --- src/spikeinterface/core/zarrextractors.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py index 7bea9bacc5..ef95ee673e 100644 --- a/src/spikeinterface/core/zarrextractors.py +++ b/src/spikeinterface/core/zarrextractors.py @@ -992,7 +992,7 @@ def add_sorting_to_zarr_group( # first sample_index of every zarr chunk: lets a lazy reader search sample_index # one chunk at a time (see ZarrSampleIndexSearch) instead of materialising it chunk_length = spikes_group["sample_index"].chunks[0] - spikes_group.create_dataset( + spikes_group.create_array( name="sample_index_chunk_firsts", data=np.asarray(spikes["sample_index"][::chunk_length], dtype="int64"), compressor=None, From ebcb617614b58642a0f3ecf44578bc2f5aa8561e Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Thu, 8 Oct 2026 19:08:51 +0200 Subject: [PATCH 43/59] test: fix create_array in new tests --- src/spikeinterface/core/tests/test_zarrextractors.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/spikeinterface/core/tests/test_zarrextractors.py b/src/spikeinterface/core/tests/test_zarrextractors.py index d34c630243..4f17cb28e4 100644 --- a/src/spikeinterface/core/tests/test_zarrextractors.py +++ b/src/spikeinterface/core/tests/test_zarrextractors.py @@ -155,8 +155,8 @@ def test_ZarrSortingExtractor_lazy_search(tmp_path): spikes_group = zarr.open(folder, mode="a")["spikes"] sample_index = spikes_group["sample_index"][:] del spikes_group["sample_index"], spikes_group["sample_index_chunk_firsts"] - spikes_group.create_dataset("sample_index", data=sample_index, chunks=(97,)) - spikes_group.create_dataset("sample_index_chunk_firsts", data=sample_index[::97], compressor=None) + spikes_group.create_array("sample_index", data=sample_index, chunks=(97,)) + spikes_group.create_array("sample_index_chunk_firsts", data=sample_index[::97], compressor=None) assert spikes_group["sample_index"].nchunks > 3 in_ram = ZarrSortingExtractor(folder) From 58d29cff870172abfd1d01fa7893bee07a950753 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Fri, 9 Oct 2026 10:17:32 +0200 Subject: [PATCH 44/59] fix: basesorting test --- src/spikeinterface/core/tests/test_basesorting.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/spikeinterface/core/tests/test_basesorting.py b/src/spikeinterface/core/tests/test_basesorting.py index 51f20433db..09c82a9ecb 100644 --- a/src/spikeinterface/core/tests/test_basesorting.py +++ b/src/spikeinterface/core/tests/test_basesorting.py @@ -198,7 +198,7 @@ def test_zarr_save_with_sharding(create_cache_folder): # check that chunks and shards are correctly set spikes_group = sorting_zarr._root["spikes"] for field in spikes_group: - if field != "segment_slices": + if field not in ("segment_slices", "sample_index_chunk_firsts"): array = spikes_group[field] assert array.chunks is not None assert array.shards is not None From 7ab4050d003e74346c9dee15624ee4a548ac6e7c Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Fri, 9 Oct 2026 10:28:01 +0200 Subject: [PATCH 45/59] test: remove requires_zarr_write after merge --- src/spikeinterface/core/tests/test_zarrextractors.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/spikeinterface/core/tests/test_zarrextractors.py b/src/spikeinterface/core/tests/test_zarrextractors.py index e851d6750d..4f17cb28e4 100644 --- a/src/spikeinterface/core/tests/test_zarrextractors.py +++ b/src/spikeinterface/core/tests/test_zarrextractors.py @@ -146,7 +146,6 @@ def test_ZarrSampleIndexSearch(tmp_path): np.testing.assert_array_equal(search.searchsorted([5], 10, 10), [0]) -@pytest.mark.requires_zarr_write def test_ZarrSortingExtractor_lazy_search(tmp_path): sorting = generate_sorting(num_units=10, durations=[5.0, 3.0, 4.0], firing_rates=40.0, seed=0) folder = tmp_path / "sorting.zarr" From 7d00017a647b6dc3fde8f7821c620671cba85692 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Fri, 9 Oct 2026 10:33:13 +0200 Subject: [PATCH 46/59] test: remove requires_zarr_write after merge 2 --- src/spikeinterface/core/tests/test_unitsselectionsorting.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/spikeinterface/core/tests/test_unitsselectionsorting.py b/src/spikeinterface/core/tests/test_unitsselectionsorting.py index 309563df62..c240d25d88 100644 --- a/src/spikeinterface/core/tests/test_unitsselectionsorting.py +++ b/src/spikeinterface/core/tests/test_unitsselectionsorting.py @@ -248,7 +248,6 @@ def test_non_identity_selection_does_not_share(unit_ids): assert len(parent._cached_lexsorted_spike_vector) == 0 -@pytest.mark.requires_zarr_write def test_identity_selection_keeps_lazy_zarr_vector(tmp_path): """A lazy parent spike vector should stay lazy through an identity selection.""" from spikeinterface.core.zarrextractors import ZarrSortingExtractor, ZarrSpikeVector From 727395d193b71ed1d5803b76cea6bd12b3bc4e55 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Fri, 9 Oct 2026 11:14:00 +0200 Subject: [PATCH 47/59] doc: update compression section --- doc/modules/core.rst | 20 ++++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/doc/modules/core.rst b/doc/modules/core.rst index 54960e8c4c..d8e7e42ed7 100644 --- a/doc/modules/core.rst +++ b/doc/modules/core.rst @@ -644,8 +644,8 @@ used when creating a :py:class:`~spikeinterface.core.SortingAnalyzer` which caus .. _save_load: -Saving, loading, and compression --------------------------------- +Saving and loading +------------------ The Base SpikeInterface objects (:py:class:`~spikeinterface.core.BaseRecording`, :py:class:`~spikeinterface.core.BaseSorting`, and @@ -678,6 +678,9 @@ This saving/loading features enables us to store SpikeInterface objects efficien # save sorting to NPZ sorting_saved = sorting.save(folder="sorting") +Compression +----------- + **NOTE:** the Zarr format by default applies data compression with :code:`Blosc.Zstandard` codec with BIT shuffling. Any other Zarr-compatible compressors and filters can be applied using the :code:`compressor` and :code:`filters` arguments. For example, in this case we apply `LZMA `_ @@ -695,6 +698,19 @@ and use a `Delta `_ filte compressors=compressor, filters=filters, **job_kwargs) +Zarr stores each dataset as a grid of **chunks**: the unit of compression and of reading, so a chunk is always +decompressed as a whole. When saving a recording, the chunk size along time is set by the job parameters +(e.g. :code:`chunk_duration="1s"`). Zarr v3 also supports **shards**, which group many chunks in a single file/object. +Chunks are still compressed and read independently, but sharding drastically reduces the number of files, which is +useful on cluster file systems and cloud storage. A shard must be a multiple of the chunk size, which is set with +:code:`shard_factor` (shard size = :code:`shard_factor` x chunk size). For example, 1 s chunks and 20 s shards: + +.. code-block:: python + + job_kwargs = dict(n_jobs=8, chunk_duration="1s") + recording_sharded = recording.save(folder="recording", format="zarr", + shard_factor=20, **job_kwargs) + Parallel processing and job_kwargs ---------------------------------- From 66e57a9c6b7809fec110e05eb4e928cc09513c33 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Mon, 19 Oct 2026 10:53:07 +0200 Subject: [PATCH 48/59] remove core tests --- .github/workflows/core-test.yml | 48 --------------------------------- 1 file changed, 48 deletions(-) delete mode 100644 .github/workflows/core-test.yml diff --git a/.github/workflows/core-test.yml b/.github/workflows/core-test.yml deleted file mode 100644 index b65606588b..0000000000 --- a/.github/workflows/core-test.yml +++ /dev/null @@ -1,48 +0,0 @@ -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 From 361f56747f42ec0e23708d5b0e77a05cc0cad480 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Mon, 19 Oct 2026 10:54:16 +0200 Subject: [PATCH 49/59] fix: remove unused default_zarr_version --- src/spikeinterface/core/zarrextractors.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py index ef95ee673e..b92fea1a24 100644 --- a/src/spikeinterface/core/zarrextractors.py +++ b/src/spikeinterface/core/zarrextractors.py @@ -21,8 +21,6 @@ from .zarr_tools import iterate_zarr_group, adjust_chunks_shards_and_job_kwargs from .time_series_tools import _write_time_series_to_zarr -zarr.config.set({"default_zarr_version": 3}) - def super_zarr_open(folder_path: str | Path, mode: str = "r", storage_options: dict | None = None): """ From 7eb166ff6c23c0f95ab750e9f9537101e1ebb80f Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Mon, 19 Oct 2026 11:18:03 +0200 Subject: [PATCH 50/59] test: skip streaming zarr if python<3.12 and save docstring --- src/spikeinterface/core/baserecording.py | 4 ++-- .../extractors/tests/test_nwbextractors_streaming.py | 4 ++++ 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/src/spikeinterface/core/baserecording.py b/src/spikeinterface/core/baserecording.py index c2769f4846..750a219d77 100644 --- a/src/spikeinterface/core/baserecording.py +++ b/src/spikeinterface/core/baserecording.py @@ -357,10 +357,10 @@ def save(self, format="binary", verbose: bool = False, **save_kwargs): If integer, the shard size will be set to chunk_size * shard_factor in the first dimension (time). If tuple, the shard_factor to be applied to each dimension. Note that `shard_factor` cannot be specified together with `shards`. - -shards: tuple | None, default: None + - shards: tuple | None, default: None Number of shard size. If None, no sharding is done. Note that shards dimensions need to be larger than chunk dimensions (if chunks is not None) and that sharding is only done on the first dimension. - - compressor: numcodecs.Codec or None, default: None + - compressors: zarr.Codec or None, default: None Global compressor. If None, Blosc-zstd, level 5, with bit shuffle is used - filters: list[numcodecs.Codec] or None, default: None Global filters for zarr (global) diff --git a/src/spikeinterface/extractors/tests/test_nwbextractors_streaming.py b/src/spikeinterface/extractors/tests/test_nwbextractors_streaming.py index 161040cbaf..9cdf2f4499 100644 --- a/src/spikeinterface/extractors/tests/test_nwbextractors_streaming.py +++ b/src/spikeinterface/extractors/tests/test_nwbextractors_streaming.py @@ -1,3 +1,4 @@ +import sys from pathlib import Path import pickle @@ -9,6 +10,8 @@ from spikeinterface.core.testing import check_recordings_equal, check_sortings_equal from spikeinterface.extractors.extractor_classes import NwbRecordingExtractor, NwbSortingExtractor +py_version = sys.version_info + @pytest.mark.streaming_extractors @pytest.mark.parametrize("cache", [True, False]) # Test with and without cache @@ -185,6 +188,7 @@ def test_sorting_s3_nwb_remfile(tmp_path): @pytest.mark.streaming_extractors +@pytest.mark.skipif(py_version < (3, 12), reason="requires Python >= 3.12") def test_sorting_s3_nwb_zarr(tmp_path): file_path = ( "s3://aind-open-data/ecephys_625749_2022-08-03_15-15-06_nwb_2023-05-16_16-34-55/" From ac73333579e9b473d1b5e9f92d77d7be1b45059a Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Tue, 20 Oct 2026 17:51:41 +0200 Subject: [PATCH 51/59] fix: remove is_sklearn_estimator --- src/spikeinterface/core/sortinganalyzer.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/spikeinterface/core/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py index 8e9ffc7cd1..ad4c256277 100644 --- a/src/spikeinterface/core/sortinganalyzer.py +++ b/src/spikeinterface/core/sortinganalyzer.py @@ -43,7 +43,6 @@ from .zarr_tools import ( iterate_zarr_group, get_zarr_attr_or_legacy_object, - is_sklearn_estimator, write_json_array, write_pickle_array, read_encoded_array, @@ -3603,9 +3602,8 @@ def _save_data(self): col_data = col_data.astype(str) df_group.create_array(name=col, data=col_data) df_group.attrs["dataframe"] = True - elif is_sklearn_estimator(ext_data): - # sklearn models (e.g. the PCA models) are saved as arrays + attributes, - # so that no pickle is needed (in zarr v2 they were pickled object arrays) + else: + # any object try: write_pickle_array( extension_group, ext_data_name, ext_data, compressor=saving_options["compressor"] From bd3875bcf934ee7eed7a47eb6d3bdedb00773fc8 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Tue, 20 Oct 2026 18:54:22 +0200 Subject: [PATCH 52/59] fix: compressors arg --- src/spikeinterface/core/sortinganalyzer.py | 1 + src/spikeinterface/core/waveform_tools.py | 4 ++-- src/spikeinterface/core/zarr_tools.py | 12 ++++++------ 3 files changed, 9 insertions(+), 8 deletions(-) diff --git a/src/spikeinterface/core/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py index ad4c256277..a22871e440 100644 --- a/src/spikeinterface/core/sortinganalyzer.py +++ b/src/spikeinterface/core/sortinganalyzer.py @@ -3609,6 +3609,7 @@ def _save_data(self): extension_group, ext_data_name, ext_data, compressor=saving_options["compressor"] ) except: + # TODO: fix raise Exception(f"Could not save {ext_data_name} as extension data") def _reset_extension_folder(self): diff --git a/src/spikeinterface/core/waveform_tools.py b/src/spikeinterface/core/waveform_tools.py index bcbdaf7e14..a557c52ba6 100644 --- a/src/spikeinterface/core/waveform_tools.py +++ b/src/spikeinterface/core/waveform_tools.py @@ -725,12 +725,12 @@ def _create_empty_zarr_waveforms(file_path, shape, dtype, zarr_target_chunk_byte row_nbytes = int(np.prod(shape[1:])) * dtype.itemsize chunk0 = max(1, zarr_target_chunk_bytes // max(1, row_nbytes)) zarr_root = zarr.open(str(store_path), mode="a") - return zarr_root.create_dataset( + return zarr_root.create_array( name=dataset_path, shape=shape, chunks=(chunk0,) + shape[1:], dtype=dtype, - compressor=get_default_zarr_compressor(), + compressors=get_default_zarr_compressor(), overwrite=True, ) diff --git a/src/spikeinterface/core/zarr_tools.py b/src/spikeinterface/core/zarr_tools.py index 34f68a265e..4866043c66 100644 --- a/src/spikeinterface/core/zarr_tools.py +++ b/src/spikeinterface/core/zarr_tools.py @@ -371,7 +371,7 @@ def get_zarr_group_keys(zarr_group) -> list[str]: return _list_member_names(zarr_group) -def write_json_array(group, name: str, data, compressor=None): +def write_json_array(group, name: str, data, compressors=None): """ Save json-serializable data as the utf-8 bytes of its json text in a uint8 array. @@ -389,7 +389,7 @@ def write_json_array(group, name: str, data, compressor=None): The name of the array. data : dict | list The json-serializable data to save. - compressor : numcodecs.Codec | None, default: None + compressors : numcodecs.Codec | None, default: None The compressor of the array. Returns @@ -398,12 +398,12 @@ def write_json_array(group, name: str, data, compressor=None): The created array. """ payload = np.frombuffer(json.dumps(data).encode("utf-8"), dtype="uint8") - array = group.create_dataset(name=name, data=payload, compressor=compressor) + array = group.create_array(name=name, data=payload, compressors=compressors) array.attrs["encoding"] = "json" return array -def write_pickle_array(group, name: str, data, compressor=None): +def write_pickle_array(group, name: str, data, compressors=None): """ Save a python object as its pickle bytes in a uint8 array. @@ -419,7 +419,7 @@ def write_pickle_array(group, name: str, data, compressor=None): The name of the array. data : Any The picklable object to save. - compressor : numcodecs.Codec | None, default: None + compressors : numcodecs.Codec | None, default: None The compressor of the array. Returns @@ -428,7 +428,7 @@ def write_pickle_array(group, name: str, data, compressor=None): The created array. """ payload = np.frombuffer(pickle.dumps(data), dtype="uint8") - array = group.create_dataset(name=name, data=payload, compressor=compressor) + array = group.create_array(name=name, data=payload, compressors=compressors) array.attrs["encoding"] = "pickle" return array From 5f87e5724739c684aa2419e76fc7615fe96fbc7f Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Wed, 21 Oct 2026 11:07:16 +0200 Subject: [PATCH 53/59] fix: serialization tests and sorting provenance --- src/spikeinterface/core/sortinganalyzer.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/src/spikeinterface/core/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py index a22871e440..cb64f1c9f6 100644 --- a/src/spikeinterface/core/sortinganalyzer.py +++ b/src/spikeinterface/core/sortinganalyzer.py @@ -2183,8 +2183,8 @@ def get_sorting_provenance(self): zarr_root = self._get_zarr_root(mode="r") sorting_provenance = None # In zarr v3, sorting_provenance is stored in attributes (in zarr v2 it was an object array) - sort_dict = get_zarr_attr_or_legacy_object(zarr_root, "provenance") - if sort_dict is not None: + sort_dict = zarr_root["sorting"].attrs.get("provenance") + if sort_dict is None: # Legacy sort_dict = get_zarr_attr_or_legacy_object(zarr_root, "sorting_provenance") if sort_dict is not None: @@ -3388,8 +3388,7 @@ def load_data(self, lazy=False): self.set_data(ext_data_name, ext_data) elif self.format == "zarr": extension_group = self._get_zarr_extension_group(mode="r") - for ext_data_name in extension_group.keys(): - ext_data_ = extension_group[ext_data_name] + for ext_data_name, ext_data_ in iterate_zarr_group(extension_group): if "encoding" in ext_data_.attrs: if lazy and ext_data_.attrs["encoding"] == "pickle": continue @@ -3582,7 +3581,7 @@ def _save_data(self): del extension_group[ext_data_name] if isinstance(ext_data, (dict, list)): write_json_array( - extension_group, ext_data_name, check_json(ext_data), compressor=saving_options["compressor"] + extension_group, ext_data_name, check_json(ext_data), compressors=saving_options["compressors"] ) elif isinstance(ext_data, np.ndarray): # only save the array if the dataset does not already exist, since it is created directly @@ -3606,10 +3605,9 @@ def _save_data(self): # any object try: write_pickle_array( - extension_group, ext_data_name, ext_data, compressor=saving_options["compressor"] + extension_group, ext_data_name, ext_data, compressors=saving_options["compressors"] ) except: - # TODO: fix raise Exception(f"Could not save {ext_data_name} as extension data") def _reset_extension_folder(self): From 26f72ae19290cc36c6d5338087b581a9c1bd5970 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Wed, 21 Oct 2026 11:11:39 +0200 Subject: [PATCH 54/59] test: add test on provenance save in analyzer --- .../core/tests/test_sortinganalyzer.py | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/src/spikeinterface/core/tests/test_sortinganalyzer.py b/src/spikeinterface/core/tests/test_sortinganalyzer.py index 9e080e720b..e38648cae9 100644 --- a/src/spikeinterface/core/tests/test_sortinganalyzer.py +++ b/src/spikeinterface/core/tests/test_sortinganalyzer.py @@ -95,6 +95,26 @@ def test_SortingAnalyzer_memory(tmp_path, dataset): _check_sorting_analyzers(sorting_analyzer, sorting2, cache_folder=tmp_path) +@pytest.mark.parametrize("format", ["binary_folder", "zarr"]) +def test_provenance_saved(tmp_path, dataset, format): + recording, sorting = dataset + recording = recording.save(folder=tmp_path / "recording_binary") + sorting = sorting.save(folder=tmp_path / "sorting_npz") + folder = tmp_path / f"analyzer_{format}" if format == "binary_folder" else tmp_path / "analyzer.zarr" + create_sorting_analyzer(sorting, recording, format=format, folder=folder, sparse=False) + + loaded = load_sorting_analyzer(folder) + assert loaded.has_recording() + assert np.array_equal(loaded.recording.channel_ids, recording.channel_ids) + assert np.array_equal(loaded.recording.get_traces(end_frame=100), recording.get_traces(end_frame=100)) + + sorting_provenance = loaded.get_sorting_provenance() + assert sorting_provenance is not None + assert np.array_equal(sorting_provenance.unit_ids, sorting.unit_ids) + for unit_id in sorting.unit_ids: + assert np.array_equal(sorting_provenance.get_unit_spike_train(unit_id), sorting.get_unit_spike_train(unit_id)) + + def test_SortingAnalyzer_binary_folder(tmp_path, dataset): recording, sorting = dataset recording = recording.save(folder=tmp_path / "recording_binary") From 658d5d1ec1a12987458ba15fada347879aeec7bf Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Wed, 21 Oct 2026 11:41:08 +0200 Subject: [PATCH 55/59] fix: serialization tests --- src/spikeinterface/core/sortinganalyzer.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/spikeinterface/core/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py index cb64f1c9f6..f85fdb7bbe 100644 --- a/src/spikeinterface/core/sortinganalyzer.py +++ b/src/spikeinterface/core/sortinganalyzer.py @@ -1195,9 +1195,13 @@ def load_from_zarr( # Load recording (if available) if recording is None: - rec_field = zarr_root.get("recording") - if rec_field is not None: - rec_dict = read_encoded_array(rec_field) + try: + rec_field = zarr_root.get("recording") + rec_dict = read_encoded_array(rec_field) if rec_field is not None else None + except Exception: + # legacy zarr v2 object array, which zarr v3 cannot open natively + rec_dict = get_zarr_attr_or_legacy_object(zarr_root, "recording") + if rec_dict is not None: try: recording = load(rec_dict, base_folder=folder) except: From 45a18ebdb35b14feaac8dc30c72887b230817178 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Thu, 22 Oct 2026 14:50:35 +0200 Subject: [PATCH 56/59] fis: serialization tests and skip streaming zarr if hdmf-zarr<0.14.1 --- src/spikeinterface/core/zarr_tools.py | 7 ++++++- .../extractors/tests/test_nwbextractors_streaming.py | 11 +++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/src/spikeinterface/core/zarr_tools.py b/src/spikeinterface/core/zarr_tools.py index 4866043c66..6a65f3a168 100644 --- a/src/spikeinterface/core/zarr_tools.py +++ b/src/spikeinterface/core/zarr_tools.py @@ -235,7 +235,12 @@ def read_legacy_zarr_object_array(store, path: str) -> LegacyZarrObjectArray: # filters are applied on encode, so they are decoded in reverse order. # for object arrays the object codec is the (only) filter for filter_config in reversed(zarray.get("filters", None) or []): - values = numcodecs.get_codec(filter_config).decode(values) + if filter_config.get("id") == "pickle": + # numcodecs>=0.16 disables Pickle.decode by default; we already trust pickles in the analyzer folder + values = pickle.loads(bytes(values)) + else: + values = numcodecs.get_codec(filter_config).decode(values) + values = np.asarray(values, dtype=object).reshape(shape) return LegacyZarrObjectArray(values, zattrs) diff --git a/src/spikeinterface/extractors/tests/test_nwbextractors_streaming.py b/src/spikeinterface/extractors/tests/test_nwbextractors_streaming.py index 9cdf2f4499..7d43253acf 100644 --- a/src/spikeinterface/extractors/tests/test_nwbextractors_streaming.py +++ b/src/spikeinterface/extractors/tests/test_nwbextractors_streaming.py @@ -1,6 +1,9 @@ import sys from pathlib import Path import pickle +from importlib.metadata import version, PackageNotFoundError + +from packaging.version import parse import pytest import numpy as np @@ -12,6 +15,11 @@ py_version = sys.version_info +try: + hdmf_zarr_version = parse(version("hdmf_zarr")) +except PackageNotFoundError: + hdmf_zarr_version = None + @pytest.mark.streaming_extractors @pytest.mark.parametrize("cache", [True, False]) # Test with and without cache @@ -189,6 +197,9 @@ def test_sorting_s3_nwb_remfile(tmp_path): @pytest.mark.streaming_extractors @pytest.mark.skipif(py_version < (3, 12), reason="requires Python >= 3.12") +@pytest.mark.skipif( + hdmf_zarr_version is None or hdmf_zarr_version < parse("0.14.1"), reason="requires hdmf_zarr >= 0.14.1" +) def test_sorting_s3_nwb_zarr(tmp_path): file_path = ( "s3://aind-open-data/ecephys_625749_2022-08-03_15-15-06_nwb_2023-05-16_16-34-55/" From db295ad2e2a65c21351065346bf3fed864b0f0a4 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Fri, 23 Oct 2026 11:30:34 +0200 Subject: [PATCH 57/59] fix: codecs docstrings, clean up spikes group, add back default chunking for spikes (10MB default) --- doc/modules/core.rst | 18 ++- src/spikeinterface/core/baserecording.py | 7 +- src/spikeinterface/core/basesorting.py | 13 +- src/spikeinterface/core/node_pipeline.py | 30 ++-- .../core/tests/test_sortinganalyzer.py | 2 +- .../core/tests/test_zarrextractors.py | 3 +- src/spikeinterface/core/time_series_tools.py | 8 +- src/spikeinterface/core/zarr_tools.py | 93 +++++++++-- src/spikeinterface/core/zarrextractors.py | 150 ++++-------------- 9 files changed, 163 insertions(+), 161 deletions(-) diff --git a/doc/modules/core.rst b/doc/modules/core.rst index d8e7e42ed7..93ca94df56 100644 --- a/doc/modules/core.rst +++ b/doc/modules/core.rst @@ -683,20 +683,24 @@ Compression **NOTE:** the Zarr format by default applies data compression with :code:`Blosc.Zstandard` codec with BIT shuffling. Any other Zarr-compatible compressors and filters can be applied using the :code:`compressor` and :code:`filters` -arguments. For example, in this case we apply `LZMA `_ -and use a `Delta `_ filter: +arguments. For example, in this case we apply ``LZMA `_ +and use a `Delta `_ filter: .. code-block:: python - from numcodecs import LZMA, Delta + from zarr.codecs import LZMA, Delta - compressor = LZMA() + compressors = [LZMA()] filters = [Delta(dtype="int16")] - recording_custom_comp = recording.save(folder="recording", format="zarr", - compressors=compressor, filters=filters, - **job_kwargs) + recording_custom_comp = recording.save( + folder="recording", + format="zarr", + compressor=compressors, + filters=filters, + **job_kwargs + ) Zarr stores each dataset as a grid of **chunks**: the unit of compression and of reading, so a chunk is always decompressed as a whole. When saving a recording, the chunk size along time is set by the job parameters diff --git a/src/spikeinterface/core/baserecording.py b/src/spikeinterface/core/baserecording.py index 750a219d77..ed7c37dae0 100644 --- a/src/spikeinterface/core/baserecording.py +++ b/src/spikeinterface/core/baserecording.py @@ -360,9 +360,10 @@ def save(self, format="binary", verbose: bool = False, **save_kwargs): - shards: tuple | None, default: None Number of shard size. If None, no sharding is done. Note that shards dimensions need to be larger than chunk dimensions (if chunks is not None) and that sharding is only done on the first dimension. - - compressors: zarr.Codec or None, default: None - Global compressor. If None, Blosc-zstd, level 5, with bit shuffle is used - - filters: list[numcodecs.Codec] or None, default: None + - compressors: BytesBytesCodec | ArrayBytesCodec | list of them, or None, default: None + Global compressors. Only one ArrayBytesCodec (serializer) is supported. + If None, Blosc-zstd, level 5, with bit shuffle is used. + - filters: ArrayArrayCodec | list[ArrayArrayCodec] or None, default: None Global filters for zarr (global) - compressor_by_dataset: dict or None, default: None Optional compressor per dataset: diff --git a/src/spikeinterface/core/basesorting.py b/src/spikeinterface/core/basesorting.py index fdd15e08e6..84257592ef 100644 --- a/src/spikeinterface/core/basesorting.py +++ b/src/spikeinterface/core/basesorting.py @@ -590,8 +590,17 @@ def save(self, format="numpy_folder", **save_kwargs): - storage_options: dict or None, default: None Storage options for zarr `store`. E.g., if "s3://" or "gcs://" they can provide authentication methods, etc. For cloud storage locations, this should not be None (in case of default values, use an empty dict) - - compressor: numcodecs.Codec or None, default: None - Global compressor. If None, Blosc-zstd, level 5, with bit shuffle is used + - target_chunk_size_bytes: int or None, default: 10 * 1024 * 1024 (10 MB) + The target chunk size for each spike field in the zarr arrays. + If explicitly set to None, no chunking is used. + - compressors: zarr compressors or str or None, default: "default" + The compressor used for each spike field in the zarr arrays. + If "default", the SpikeInterface default zarr compressor is used (Blosc-zstd, level 5, bitshuffle). + - filters: list of zarr filters or None, default: None + The list of zarr filters to apply to each spike field in the zarr arrays. + If None, a Delta filter is applied. + - shard_factor: int or None, default: None + If given, a shard will be shard_factor * chunk size. * "memory" format: - sharedmem : bool, default: True If True, the recording is saved in shared memory. If False, it is saved as diff --git a/src/spikeinterface/core/node_pipeline.py b/src/spikeinterface/core/node_pipeline.py index afa786e8e8..ea5370cfb1 100644 --- a/src/spikeinterface/core/node_pipeline.py +++ b/src/spikeinterface/core/node_pipeline.py @@ -14,6 +14,15 @@ from spikeinterface.core import get_channel_distances from spikeinterface.core.core_tools import ms_to_samples, samples_to_ms +compressors_type = ( + zarr.abc.codec.BytesBytesCodec + | zarr.abc.codec.ArrayBytesCodec + | list[zarr.abc.codec.BytesBytesCodec | zarr.abc.codec.ArrayBytesCodec] + | str + | None +) +filters_type = zarr.abc.codec.ArrayArrayCodec | list[zarr.abc.codec.ArrayArrayCodec] | None + class PipelineNode: @@ -1076,10 +1085,13 @@ class GatherToZarr: Note that in both modes the `dest` must contain a ".zarr" suffix. names : list of str | None Names of the outputs. Can be None when `dest` is a list of explicit destinations. - compressor : numcodecs codec | "default" | None, default: "default" - The compressor used for every array. If "default", the SpikeInterface default + compressors : zarr.abc.codec.BytesBytesCodec | zarr.abc.codec.ArrayBytesCodec | list of them | str, or None, default: "default" + The compressors used for every array. If "default", the SpikeInterface default zarr compressor is used (Blosc-zstd, level 5, bitshuffle). If None, no compression. Ignored for destinations that are already a ``zarr.Array`` (they keep their own compressor). + filters : zarr.abc.codec.ArrayArrayCodec | list[zarr.abc.codec.ArrayArrayCodec] or None, default: None + The filters used for every array. If None, no filters are applied. + Ignored for destinations that are already a ``zarr.Array`` (they keep their own filters). zarr_target_chunk_bytes : int | dict[str, int], default: 10485760 (10 MiB) Target (uncompressed) size in bytes of one zarr chunk, used to compute the number of rows per chunk. An int applies the same target to every array, a dict sets it per array @@ -1091,16 +1103,16 @@ def __init__( self, dest: str | Path | list[str | Path | zarr.Array] | None = None, names: list[str] | None = None, - compressor: "numcodecs.abc.Codec | str | None" = "default", + compressors: compressors_type = "default", + filters: filters_type = None, zarr_target_chunk_bytes: int | dict[str, int] = 10 * 1024 * 1024, ): - import zarr - from spikeinterface.core.zarrextractors import get_default_zarr_compressor + from spikeinterface.core.zarr_tools import build_codec_pipeline - if compressor == "default": - compressor = get_default_zarr_compressor() - self.compressor = compressor + if compressors == "default": + compressors = get_default_zarr_compressor() + self.codec_kwargs = build_codec_pipeline(filters=filters, compressors=compressors) if isinstance(dest, (list, tuple)): # append to explicit destinations (paths inside a store or pre-created zarr.Arrays) @@ -1187,8 +1199,8 @@ def __call__(self, res): shape=(0,) + trailing_shape, chunks=(chunk0,) + trailing_shape, dtype=buf.dtype, - compressors=self.compressor, overwrite=True, + **self.codec_kwargs, ) self.arrays[i_name].append(buf, axis=0) diff --git a/src/spikeinterface/core/tests/test_sortinganalyzer.py b/src/spikeinterface/core/tests/test_sortinganalyzer.py index e38648cae9..34acba91d4 100644 --- a/src/spikeinterface/core/tests/test_sortinganalyzer.py +++ b/src/spikeinterface/core/tests/test_sortinganalyzer.py @@ -220,7 +220,7 @@ def test_SortingAnalyzer_zarr(tmp_path, dataset): assert len(sorting_analyzer_no_compression._get_zarr_root()["extensions"]["templates"]["average"].compressors) == 0 # test a different compressor - from zarr.codecs.numcodecs import LZMA + from zarr.codecs import LZMA lzma_compressor = LZMA() folder = tmp_path / "test_SortingAnalyzer_zarr_lzma.zarr" diff --git a/src/spikeinterface/core/tests/test_zarrextractors.py b/src/spikeinterface/core/tests/test_zarrextractors.py index 4f17cb28e4..bd6bbd45d2 100644 --- a/src/spikeinterface/core/tests/test_zarrextractors.py +++ b/src/spikeinterface/core/tests/test_zarrextractors.py @@ -21,8 +21,7 @@ def test_zarr_compression_options(tmp_path): - from zarr.codecs.numcodecs import Delta, FixedScaleOffset - from zarr.codecs import BloscCodec, BloscShuffle + from zarr.codecs import BloscCodec, BloscShuffle, Delta, FixedScaleOffset recording = generate_recording(durations=[2]) recording.set_times(recording.get_times() + 100) diff --git a/src/spikeinterface/core/time_series_tools.py b/src/spikeinterface/core/time_series_tools.py index 5d40733e2f..1ebf0374df 100644 --- a/src/spikeinterface/core/time_series_tools.py +++ b/src/spikeinterface/core/time_series_tools.py @@ -2,7 +2,6 @@ import warnings import numpy as np -import numpy.typing from .core_tools import add_suffix, make_shared_array from .job_tools import ( @@ -340,11 +339,13 @@ def _write_time_series_to_zarr( {} """ from .job_tools import ( - ensure_chunk_size, fix_job_kwargs, TimeSeriesChunkExecutor, ) + # place an ArrayBytesCodec passed as a compressor (e.g. WavPack) in the serializer slot + from .zarr_tools import build_codec_pipeline + assert dataset_paths is not None, "Provide 'dataset_paths' to save data in zarr format" if dataset_timestamps_paths is not None: assert ( @@ -362,9 +363,6 @@ def _write_time_series_to_zarr( job_kwargs = fix_job_kwargs(job_kwargs) - # place an ArrayBytesCodec passed as a compressor (e.g. WavPack) in the serializer slot - from .zarrextractors import build_codec_pipeline - codec_kwargs_data = build_codec_pipeline(filters=filters_data, compressors=compressor_data) codec_kwargs_times = build_codec_pipeline(filters=filters_times, compressors=compressor_times) diff --git a/src/spikeinterface/core/zarr_tools.py b/src/spikeinterface/core/zarr_tools.py index 6a65f3a168..41ff57a358 100644 --- a/src/spikeinterface/core/zarr_tools.py +++ b/src/spikeinterface/core/zarr_tools.py @@ -87,9 +87,9 @@ def check_compressors_match(comp1, comp2, skip_typesize=True): Parameters ---------- - comp1 : zarr.Codec | Tuple[zarr.Codec] + comp1 : zarr.abc.codec.Codec | Tuple[zarr.abc.codec.Codec] The first compressor object to compare. - comp2 : zarr.Codec | Tuple[zarr.Codec] + comp2 : zarr.abc.codec.Codec | Tuple[zarr.abc.codec.Codec] The second compressor object to compare. skip_typesize : bool, optional Whether to skip the typesize check, default: True @@ -376,7 +376,7 @@ def get_zarr_group_keys(zarr_group) -> list[str]: return _list_member_names(zarr_group) -def write_json_array(group, name: str, data, compressors=None): +def write_json_array(group, name: str, data, compressors=None, filters=None): """ Save json-serializable data as the utf-8 bytes of its json text in a uint8 array. @@ -394,8 +394,10 @@ def write_json_array(group, name: str, data, compressors=None): The name of the array. data : dict | list The json-serializable data to save. - compressors : numcodecs.Codec | None, default: None - The compressor of the array. + compressors: BytesBytesCodec | ArrayBytesCodec | list of them, or None, default: None + Compressors to use for the array. + filters: ArrayArrayCodec | list[ArrayArrayCodec] or None, default: None + Filters to use for the array. Returns ------- @@ -403,12 +405,13 @@ def write_json_array(group, name: str, data, compressors=None): The created array. """ payload = np.frombuffer(json.dumps(data).encode("utf-8"), dtype="uint8") - array = group.create_array(name=name, data=payload, compressors=compressors) + codec_kwargs = build_codec_pipeline(filters=filters, compressors=compressors) + array = group.create_array(name=name, data=payload, **codec_kwargs) array.attrs["encoding"] = "json" return array -def write_pickle_array(group, name: str, data, compressors=None): +def write_pickle_array(group, name: str, data, compressors=None, filters=None): """ Save a python object as its pickle bytes in a uint8 array. @@ -424,8 +427,10 @@ def write_pickle_array(group, name: str, data, compressors=None): The name of the array. data : Any The picklable object to save. - compressors : numcodecs.Codec | None, default: None - The compressor of the array. + compressors: BytesBytesCodec | ArrayBytesCodec | list of them, or None, default: None + Compressors to use for the array. + filters: ArrayArrayCodec | list[ArrayArrayCodec] or None, default: None + Filters to use for the array. Returns ------- @@ -433,7 +438,8 @@ def write_pickle_array(group, name: str, data, compressors=None): The created array. """ payload = np.frombuffer(pickle.dumps(data), dtype="uint8") - array = group.create_array(name=name, data=payload, compressors=compressors) + codec_kwargs = build_codec_pipeline(filters=filters, compressors=compressors) + array = group.create_array(name=name, data=payload, **codec_kwargs) array.attrs["encoding"] = "pickle" return array @@ -464,3 +470,70 @@ def read_encoded_array(array): raise ValueError( f"Unknown encoding {encoding!r} for {array.name}. It was probably saved by a newer version of spikeinterface." ) + + +def build_codec_pipeline(filters=None, compressors=None): + """ + Build zarr v3 codec kwargs from filters and compressors. + + Classifies codecs into the three slots accepted by ``zarr.Group.create_array()``: + 1. ``filters`` — ArrayArrayCodec (e.g. Delta) + 2. ``serializer`` — ArrayBytesCodec (e.g. WavPack, BytesCodec) + 3. ``compressors``— BytesBytesCodec (e.g. BloscCodec, ZstdCodec) + + This allows callers to pass an ArrayBytesCodec (e.g. WavPack) as a + compressor and have it placed in the correct serializer slot automatically. + + Parameters + ---------- + filters : ArrayArrayCodec or list of ArrayArrayCodec or None + Codec(s) applied before serialization. + compressors : codec or list of codecs or None + Can be a mix of ArrayBytesCodec (serializer) and BytesBytesCodec + (byte-level compressors). At most one ArrayBytesCodec is allowed. + + Returns + ------- + dict + Keyword arguments to unpack into ``zarr.Group.create_array()``. + Only keys with explicit values are included; omitted keys let zarr + use its defaults. + + Raises + ------ + ValueError + If filters contain non-ArrayArrayCodec instances, if more than one + ArrayBytesCodec is provided, or if an unrecognised codec type is passed. + """ + from zarr.abc.codec import ArrayArrayCodec, ArrayBytesCodec, BytesBytesCodec + + if filters is None: + filters = [] + if not isinstance(filters, (list, tuple)): + filters = [filters] + + if compressors is None: + compressors = [] + if not isinstance(compressors, (list, tuple)): + compressors = [compressors] + + for f in filters: + if not isinstance(f, ArrayArrayCodec): + raise ValueError(f"All filters must be ArrayArrayCodec instances, got {type(f)}") + + serializers = [c for c in compressors if isinstance(c, ArrayBytesCodec)] + byte_compressors = [c for c in compressors if isinstance(c, BytesBytesCodec)] + invalid = [c for c in compressors if not isinstance(c, (ArrayBytesCodec, BytesBytesCodec))] + + if invalid: + raise ValueError( + f"Compressors must be ArrayBytesCodec or BytesBytesCodec instances, got {[type(c) for c in invalid]}" + ) + if len(serializers) > 1: + raise ValueError("Only one ArrayBytesCodec (serializer) is allowed in the codec pipeline.") + + codec_kwargs = {} + codec_kwargs["filters"] = filters + codec_kwargs["serializer"] = serializers[0] if len(serializers) == 1 else "auto" + codec_kwargs["compressors"] = byte_compressors + return codec_kwargs diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py index 4894ae9792..73c89379d9 100644 --- a/src/spikeinterface/core/zarrextractors.py +++ b/src/spikeinterface/core/zarrextractors.py @@ -18,7 +18,7 @@ is_path_remote, ) from .job_tools import split_job_kwargs, fix_job_kwargs -from .zarr_tools import iterate_zarr_group, adjust_chunks_shards_and_job_kwargs +from .zarr_tools import iterate_zarr_group, adjust_chunks_shards_and_job_kwargs, build_codec_pipeline from .time_series_tools import _write_time_series_to_zarr @@ -582,35 +582,20 @@ def __init__( unit_ids = np.array(unit_ids) assert "spikes" in self._root.keys(), "'spikes' dataset not found!" - spikes_item = self._root["spikes"] + spikes_group = self._root["spikes"] BaseSorting.__init__(self, sampling_frequency, unit_ids) - if isinstance(spikes_item, zarr.Group): - # Legacy format: individual field arrays + a "segment_slices" sub-array inside the group. - spikes_group = spikes_item - segment_slices_list = np.asarray(spikes_group["segment_slices"][:], dtype="int64") + segment_slices_list = np.asarray(spikes_group["segment_slices"][:], dtype="int64") - if lazy_spike_vector: - spikes = ZarrSpikeVector(spikes_group, segment_slices_list) - else: - spikes = np.zeros(spikes_group["sample_index"].shape[0], dtype=minimum_spike_dtype) - spikes["sample_index"] = spikes_group["sample_index"][:] - spikes["unit_index"] = spikes_group["unit_index"][:] - for i, (start, end) in enumerate(segment_slices_list): - spikes["segment_index"][start:end] = i + if lazy_spike_vector: + spikes = ZarrSpikeVector(spikes_group, segment_slices_list) else: - # New format: a single structured zarr array; segment_slices stored as an array attribute. - # We need https://github.com/zarr-developers/zarr-python/pull/3996 released before being able to - # access the structured array lazily. Until then, we always materialise it. - spikes = spikes_item - segment_slices_list = np.asarray(spikes.attrs["segment_slices"], dtype="int64") - - # we do not need to lexsort at init (very high cost) because spikes are already sorted by frame before saving. - # In version 0.104.X this was fully lexsorted, but we don't need it anymore because it's only important in the - # context of SpikeVectorBased extensions in the SortingAnalyzer, which stores its own copy of the Sorting - # object. This makes the extension data and the spike vector always matching their order. - # spikes = spikes[np.lexsort((spikes["unit_index"], spikes["sample_index"], spikes["segment_index"]))] + spikes = np.zeros(spikes_group["sample_index"].shape[0], dtype=minimum_spike_dtype) + spikes["sample_index"] = spikes_group["sample_index"][:] + spikes["unit_index"] = spikes_group["unit_index"][:] + for i, (start, end) in enumerate(segment_slices_list): + spikes["segment_index"][start:end] = i self._lazy_spike_vector = lazy_spike_vector self._spikes_group = spikes_group self._sample_index_search = None @@ -792,73 +777,6 @@ def get_default_zarr_compressor(clevel: int = 5): return BloscCodec(cname="zstd", clevel=clevel, shuffle=BloscShuffle.bitshuffle) -def build_codec_pipeline(filters=None, compressors=None): - """ - Build zarr v3 codec kwargs from filters and compressors. - - Classifies codecs into the three slots accepted by ``zarr.Group.create_array()``: - 1. ``filters`` — ArrayArrayCodec (e.g. Delta) - 2. ``serializer`` — ArrayBytesCodec (e.g. WavPack, BytesCodec) - 3. ``compressors``— BytesBytesCodec (e.g. BloscCodec, ZstdCodec) - - This allows callers to pass an ArrayBytesCodec (e.g. WavPack) as a - compressor and have it placed in the correct serializer slot automatically. - - Parameters - ---------- - filters : ArrayArrayCodec or list of ArrayArrayCodec or None - Codec(s) applied before serialization. - compressors : codec or list of codecs or None - Can be a mix of ArrayBytesCodec (serializer) and BytesBytesCodec - (byte-level compressors). At most one ArrayBytesCodec is allowed. - - Returns - ------- - dict - Keyword arguments to unpack into ``zarr.Group.create_array()``. - Only keys with explicit values are included; omitted keys let zarr - use its defaults. - - Raises - ------ - ValueError - If filters contain non-ArrayArrayCodec instances, if more than one - ArrayBytesCodec is provided, or if an unrecognised codec type is passed. - """ - from zarr.abc.codec import ArrayArrayCodec, ArrayBytesCodec, BytesBytesCodec - - if filters is None: - filters = [] - if not isinstance(filters, (list, tuple)): - filters = [filters] - - if compressors is None: - compressors = [] - if not isinstance(compressors, (list, tuple)): - compressors = [compressors] - - for f in filters: - if not isinstance(f, ArrayArrayCodec): - raise ValueError(f"All filters must be ArrayArrayCodec instances, got {type(f)}") - - serializers = [c for c in compressors if isinstance(c, ArrayBytesCodec)] - byte_compressors = [c for c in compressors if isinstance(c, BytesBytesCodec)] - invalid = [c for c in compressors if not isinstance(c, (ArrayBytesCodec, BytesBytesCodec))] - - if invalid: - raise ValueError( - f"Compressors must be ArrayBytesCodec or BytesBytesCodec instances, got {[type(c) for c in invalid]}" - ) - if len(serializers) > 1: - raise ValueError("Only one ArrayBytesCodec (serializer) is allowed in the codec pipeline.") - - codec_kwargs = {} - codec_kwargs["filters"] = filters - codec_kwargs["serializer"] = serializers[0] if len(serializers) == 1 else "auto" - codec_kwargs["compressors"] = byte_compressors - return codec_kwargs - - def _has_string_fields(dtype: np.dtype) -> bool: """Return True if dtype is or contains fixed-length unicode (U) sub-fields.""" if dtype.names: @@ -908,9 +826,15 @@ def add_sorting_to_zarr_group( kwargs : dict Other arguments passed to the zarr writer: - * "compressors" : zarr compressor or None, default: None - * "target_chunk_size_bytes" : tuple or None, default: None - The target chunk size for the zarr datasets. + * "target_chunk_size_bytes" : int or None, default: 10 * 1024 * 1024 (10 MB) + The target chunk size for each spike field in the zarr arrays. + If explicitly set to None, no chunking is used. + * "compressors" : zarr compressors or str or None, default: "default" + The compressor used for each spike field in the zarr arrays. + If "default", the SpikeInterface default zarr compressor is used (Blosc-zstd, level 5, bitshuffle). + * "filters" : list of zarr filters or None, default: None + The list of zarr filters to apply to each spike field in the zarr arrays. + If None, a Delta filter is applied. * "shard_factor" : int or None, default: None If given, a shard will be shard_factor * chunk size. """ @@ -926,9 +850,9 @@ def add_sorting_to_zarr_group( zarr_group.attrs["num_segments"] = int(num_segments) zarr_group.create_array(name="unit_ids", data=sorting.unit_ids, compressors=None) - compressor = kwargs.get("compressors") or kwargs.get("compressor") - if compressor is None: - compressor = get_default_zarr_compressor() + compressors = kwargs.get("compressors", "default") + if isinstance(compressors, str) and compressors == "default": + compressors = get_default_zarr_compressor() # Save the full structured spike array. The "segment_index" field is additionally stored as # "segment_slices" (an attribute on the spikes array) which contains the start and end indices @@ -937,27 +861,7 @@ def add_sorting_to_zarr_group( chunks = None shards = None - target_chunk_size_bytes = kwargs.get("target_chunk_size_bytes") - - # We need https://github.com/zarr-developers/zarr-python/pull/3996 released before being able to - # access the structured array lazily. Until then, we always materialise it. - # For now, let's keep the old by field implementation - # if target_chunk_size_bytes is not None: - # spike_num_bytes = spikes.dtype.itemsize - # target_chunk_size = target_chunk_size_bytes // spike_num_bytes - # chunks = (target_chunk_size,) - # shard_factor = kwargs.get("shard_factor") - # if shard_factor is not None: - # shards = (target_chunk_size * shard_factor,) - # else: - # chunks = (len(spikes),) - # spikes = zarr_group.create_array( - # name="spikes", - # data=spikes, - # chunks=chunks, - # shards=shards, - # **codec_kwargs - # ) + target_chunk_size_bytes = kwargs.get("target_chunk_size_bytes", 10 * 1024 * 1024) # Save sub fields: in this case chunks and shards are set per field spikes_group = zarr_group.create_group(name="spikes") @@ -973,7 +877,10 @@ def add_sorting_to_zarr_group( shards = (target_chunk_size * shard_factor,) else: chunks = (len(spikes),) - codec_kwargs = build_codec_pipeline(filters=Delta(dtype=spikes[field].dtype.str), compressors=compressor) + filters = kwargs.get("filters", None) + if filters is None: + filters = Delta(dtype=spikes[field].dtype.str) + codec_kwargs = build_codec_pipeline(filters=filters, compressors=compressors) spikes_group.create_array(name=field, data=spikes[field], chunks=chunks, shards=shards, **codec_kwargs) segment_slices = [] @@ -985,7 +892,6 @@ def add_sorting_to_zarr_group( segment_slices.append([i0, i1]) segment_slices = np.array(segment_slices, dtype="int64") spikes_group.create_array(name="segment_slices", data=segment_slices, compressors=None) - # spikes.attrs["segment_slices"] = segment_slices.tolist() # first sample_index of every zarr chunk: lets a lazy reader search sample_index # one chunk at a time (see ZarrSampleIndexSearch) instead of materialising it @@ -993,7 +899,7 @@ def add_sorting_to_zarr_group( spikes_group.create_array( name="sample_index_chunk_firsts", data=np.asarray(spikes["sample_index"][::chunk_length], dtype="int64"), - compressor=None, + compressors=None, ) add_properties_and_annotations(zarr_group, sorting) From c0e71ad063aba4dcaf4742d99a8d8c1b08469429 Mon Sep 17 00:00:00 2001 From: Chris Halcrow <57948917+chrishalcrow@users.noreply.github.com> Date: Fri, 23 Oct 2026 16:21:36 +0100 Subject: [PATCH 58/59] Apply suggestion from @chrishalcrow --- pyproject.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 0fdf06d34f..05c1e5e851 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,8 +28,8 @@ dependencies = [ # for release we need pypi, so this needs to be commented "neo @ git+https://github.com/NeuralEnsemble/python-neo.git", # for testing we need probeinterface Zarr v3 branch, so this needs to be commented - #"probeinterface @ git+https://github.com/SpikeInterface/probeinterface.git", - "probeinterface @ git+https://github.com/alejoe91/probeinterface.git@zarrv3", + "probeinterface @ git+https://github.com/SpikeInterface/probeinterface.git", + # "neo>=0.14.5", # "probeinterface>=0.4.0", "packaging", From 75b9839723ed08e94414157e844de62cf03aec4a Mon Sep 17 00:00:00 2001 From: Chris Halcrow <57948917+chrishalcrow@users.noreply.github.com> Date: Fri, 23 Oct 2026 16:22:51 +0100 Subject: [PATCH 59/59] tidy pyproject.toml Co-authored-by: Chris Halcrow <57948917+chrishalcrow@users.noreply.github.com> --- pyproject.toml | 2 -- 1 file changed, 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 05c1e5e851..4265be1622 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,9 +27,7 @@ dependencies = [ "zarr>=3,<4", # for release we need pypi, so this needs to be commented "neo @ git+https://github.com/NeuralEnsemble/python-neo.git", - # for testing we need probeinterface Zarr v3 branch, so this needs to be commented "probeinterface @ git+https://github.com/SpikeInterface/probeinterface.git", - # "neo>=0.14.5", # "probeinterface>=0.4.0", "packaging",