From 6f47bf04dc6baea8bb71378d3ce5c258be5ec81d Mon Sep 17 00:00:00 2001 From: Heberto Mayorquin Date: Fri, 16 Oct 2026 19:53:56 -0300 Subject: [PATCH 1/3] Make motion-corrected and silenced recordings JSON-serializable --- src/spikeinterface/core/motion.py | 7 ++++--- src/spikeinterface/preprocessing/silence_periods.py | 6 +++--- .../motion/motion_interpolation.py | 13 ++++++++----- 3 files changed, 15 insertions(+), 11 deletions(-) diff --git a/src/spikeinterface/core/motion.py b/src/spikeinterface/core/motion.py index 6a99ca29c8..50622b2e7d 100644 --- a/src/spikeinterface/core/motion.py +++ b/src/spikeinterface/core/motion.py @@ -166,10 +166,11 @@ def to_dict(self): @staticmethod def from_dict(d): + # arrays are lists when the dict comes from json return Motion( - d["displacement"], - d["temporal_bins_s"], - d["spatial_bins_um"], + [np.asarray(displacement) for displacement in d["displacement"]], + [np.asarray(temporal_bins_s) for temporal_bins_s in d["temporal_bins_s"]], + np.asarray(d["spatial_bins_um"]), direction=d["direction"], interpolation_method=d["interpolation_method"], ) diff --git a/src/spikeinterface/preprocessing/silence_periods.py b/src/spikeinterface/preprocessing/silence_periods.py index 430348bbb3..912b606174 100644 --- a/src/spikeinterface/preprocessing/silence_periods.py +++ b/src/spikeinterface/preprocessing/silence_periods.py @@ -79,6 +79,9 @@ def __init__( "list_periods is deprecated. Please specify periods, which should be a np.array with " f"dtype {base_period_dtype}" ) + if isinstance(periods, list): + # periods reloaded from json are a list of [segment_index, start_sample_index, end_sample_index] + periods = np.array([tuple(period) for period in periods], dtype=base_period_dtype) if not isinstance(periods, np.ndarray): raise ValueError(f"periods must be a np.array with dtype {base_period_dtype}") @@ -132,9 +135,6 @@ def __init__( ) self.add_recording_segment(rec_segment) - # the base_period_dtype is a structured dtype, which is not json serializable - self._serializability["json"] = False - self._kwargs = dict( recording=recording, periods=periods, diff --git a/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py b/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py index a144940853..0eb63b8b6b 100644 --- a/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py +++ b/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py @@ -1,6 +1,6 @@ import numpy as np from spikeinterface.core.core_tools import define_function_handling_dict_from_class -from spikeinterface.core.motion import ensure_time_bins +from spikeinterface.core.motion import Motion, ensure_time_bins from spikeinterface.preprocessing import get_spatial_interpolation_kernel from spikeinterface.preprocessing.basepreprocessor import BasePreprocessor, BasePreprocessorSegment from spikeinterface.preprocessing.filter import fix_dtype @@ -353,6 +353,10 @@ def __init__( ): # assert recording.get_num_segments() == 1, "correct_motion() is only available for single-segment recordings" + if isinstance(motion, dict): + # a Motion reloaded from json is a dict + motion = Motion.from_dict(motion) + channel_locations = recording.get_channel_locations() assert channel_locations.ndim >= motion.dim, ( f"'direction' {motion.direction} not available. " @@ -419,6 +423,9 @@ def __init__( interpolation_time_bin_centers_s = motion.temporal_bins_s interpolation_time_bin_edges_s = motion.temporal_bin_edges_s else: + if isinstance(interpolation_time_bin_centers_s, list): + # time bins reloaded from json are lists + interpolation_time_bin_centers_s = [np.asarray(bins) for bins in interpolation_time_bin_centers_s] interpolation_time_bin_centers_s, interpolation_time_bin_edges_s = ensure_time_bins( interpolation_time_bin_centers_s, interpolation_time_bin_edges_s ) @@ -453,10 +460,6 @@ def __init__( ) self.add_recording_segment(rec_segment) - # this object is currently not JSON-serializable because the Motion obejct cannot be reloaded properly - # see issue #3313 - self._serializability["json"] = False - self._kwargs = dict( recording=recording, motion=motion, From e20102420aa180e1718e412544877f929eb08372 Mon Sep 17 00:00:00 2001 From: Heberto Mayorquin Date: Mon, 19 Oct 2026 12:34:39 -0600 Subject: [PATCH 2/3] Store JSON data in zarr SortingAnalyzers as byte arrays instead of object arrays --- src/spikeinterface/core/sortinganalyzer.py | 31 +++++----- src/spikeinterface/core/zarr_tools.py | 61 ++++++++++++++++++++ src/spikeinterface/preprocessing/pipeline.py | 3 +- 3 files changed, 79 insertions(+), 16 deletions(-) create mode 100644 src/spikeinterface/core/zarr_tools.py diff --git a/src/spikeinterface/core/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py index b702237445..4b4377b159 100644 --- a/src/spikeinterface/core/sortinganalyzer.py +++ b/src/spikeinterface/core/sortinganalyzer.py @@ -41,6 +41,7 @@ from .sparsity import ChannelSparsity, estimate_sparsity from .sortingfolder import NumpyFolderSorting from .zarrextractors import get_default_zarr_compressor, ZarrSortingExtractor, super_zarr_open, _write_object_array +from .zarr_tools import write_json_array, read_encoded_array from .node_pipeline import run_node_pipeline # Typing hints @@ -1076,17 +1077,14 @@ def create_zarr( if recording is not None: rec_dict = recording.to_dict(relative_to=relative_to, recursive=True) if recording.check_serializability("json"): - _write_object_array(zarr_root, "recording", check_json(rec_dict), codec="json") - elif recording.check_serializability("pickle"): - try: - _write_object_array(zarr_root, "recording", rec_dict, codec="pickle") - except: - warnings.warn( - "Failed to serialize recording with Pickle Codec! " - "The recording link will be lost for future load" - ) + compressor = saving_options.get("compressor", get_default_zarr_compressor()) + write_json_array(zarr_root, "recording", check_json(rec_dict), compressor=compressor) else: - warnings.warn("The Recording is not serializable! The recording link will be lost for future load") + warnings.warn( + "The recording is not JSON-serializable, so it is not saved in the zarr SortingAnalyzer, which " + "will load without a recording. Use `set_temporary_recording()` after loading to attach it " + "again, or save the SortingAnalyzer in 'binary_folder' format, which keeps it." + ) else: assert rec_attributes is not None, "recording or rec_attributes must be provided" warnings.warn("Recording not provided, instantiating SortingAnalyzer in recordingless mode.") @@ -1187,7 +1185,7 @@ def load_from_zarr( if recording is None: rec_field = zarr_root.get("recording") if rec_field is not None: - rec_dict = rec_field[0] + rec_dict = read_encoded_array(rec_field) try: recording = load(rec_dict, base_folder=folder) except: @@ -3371,7 +3369,10 @@ def load_data(self, lazy=False): 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: + if "encoding" in ext_data_.attrs: + ext_data = read_encoded_array(ext_data_) + elif "dict" in ext_data_.attrs: + # legacy json object array (spikeinterface<=0.105) ext_data = ext_data_[0] elif "dataframe" in ext_data_.attrs: import pandas as pd @@ -3552,9 +3553,9 @@ def _save_data(self): if ext_data_name in extension_group: del extension_group[ext_data_name] if isinstance(ext_data, (dict, list)): - ext_data_ = check_json(ext_data) - _write_object_array(extension_group, ext_data_name, ext_data_, codec="json") - extension_group[ext_data_name].attrs["dict"] = True + write_json_array( + extension_group, ext_data_name, check_json(ext_data), compressor=saving_options["compressor"] + ) elif isinstance(ext_data, np.ndarray): # only save the array if the dataset does not already exist, since it is created directly # by the run_node_pipeline() function in the case of nodepipeline extensions diff --git a/src/spikeinterface/core/zarr_tools.py b/src/spikeinterface/core/zarr_tools.py new file mode 100644 index 0000000000..023de78dbe --- /dev/null +++ b/src/spikeinterface/core/zarr_tools.py @@ -0,0 +1,61 @@ +import json + +import numpy as np + + +def write_json_array(group, name: str, data, compressor=None): + """ + Save json-serializable data as the utf-8 bytes of its json text in a uint8 array. + + This replaces the length-1 object arrays with a json codec, which zarr-python 3 cannot write. + The data is stored as array data rather than as an attribute, so it is only read when needed: + attributes are read every time the group is opened, which is slow for large data such as a + recording dict. The array is marked with an "encoding" attribute, so that other encodings can be + added later. See `read_encoded_array` for the reverse operation. + + Parameters + ---------- + group : zarr.Group + The zarr group in which the array is created. + name : str + The name of the array. + data : dict | list + The json-serializable data to save. + compressor : numcodecs.Codec | None, default: None + The compressor of the array. + + Returns + ------- + array : zarr.Array + 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.attrs["encoding"] = "json" + return array + + +def read_encoded_array(array): + """ + Read data saved by `write_json_array`, or a legacy length-1 object array (json or pickle codec) + saved by spikeinterface<=0.105. + + Parameters + ---------- + array : zarr.Array + The array to read. + + Returns + ------- + data : Any + The saved data. + """ + encoding = array.attrs.get("encoding", None) + if encoding is None: + # legacy length-1 object array + return array[0] + if encoding == "json": + return json.loads(np.asarray(array[:]).tobytes().decode("utf-8")) + raise ValueError( + f"Unknown encoding {encoding!r} for {array.name}. It was probably saved by a newer version of spikeinterface." + ) diff --git a/src/spikeinterface/preprocessing/pipeline.py b/src/spikeinterface/preprocessing/pipeline.py index 0fe320681c..b75cae49f8 100644 --- a/src/spikeinterface/preprocessing/pipeline.py +++ b/src/spikeinterface/preprocessing/pipeline.py @@ -4,6 +4,7 @@ from spikeinterface.core import BaseRecording from spikeinterface.core.core_tools import is_dict_extractor, is_path_remote from spikeinterface.core.zarrextractors import super_zarr_open +from spikeinterface.core.zarr_tools import read_encoded_array from spikeinterface.preprocessing.preprocessing_classes import preprocessor_dict, _all_preprocesser_dict pp_names_to_functions = {preprocessor.__name__: preprocessor for preprocessor in preprocessor_dict.values()} @@ -265,7 +266,7 @@ def get_preprocessing_list_from_analyzer(analyzer_folder, format="auto", backend rec_field = zarr_root.get("recording") if rec_field is not None: - recording_dict = rec_field[0] + recording_dict = read_encoded_array(rec_field) else: recording_dict = {} From 15049610fd2ce5e8e32841900a5eabc84ee015b3 Mon Sep 17 00:00:00 2001 From: Heberto Mayorquin Date: Mon, 19 Oct 2026 14:58:53 -0600 Subject: [PATCH 3/3] Store pickled extension data in zarr SortingAnalyzers as byte arrays instead of object arrays --- src/spikeinterface/core/sortinganalyzer.py | 16 +++++--- src/spikeinterface/core/zarr_tools.py | 38 +++++++++++++++++- src/spikeinterface/core/zarrextractors.py | 45 ---------------------- 3 files changed, 46 insertions(+), 53 deletions(-) diff --git a/src/spikeinterface/core/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py index 4b4377b159..bd03f0b55e 100644 --- a/src/spikeinterface/core/sortinganalyzer.py +++ b/src/spikeinterface/core/sortinganalyzer.py @@ -40,8 +40,8 @@ from .numpyextractors import NumpySorting from .sparsity import ChannelSparsity, estimate_sparsity from .sortingfolder import NumpyFolderSorting -from .zarrextractors import get_default_zarr_compressor, ZarrSortingExtractor, super_zarr_open, _write_object_array -from .zarr_tools import write_json_array, read_encoded_array +from .zarrextractors import get_default_zarr_compressor, ZarrSortingExtractor, super_zarr_open +from .zarr_tools import write_json_array, write_pickle_array, read_encoded_array from .node_pipeline import run_node_pipeline # Typing hints @@ -3370,10 +3370,12 @@ def load_data(self, lazy=False): for ext_data_name in extension_group.keys(): ext_data_ = extension_group[ext_data_name] if "encoding" in ext_data_.attrs: + if lazy and ext_data_.attrs["encoding"] == "pickle": + continue ext_data = read_encoded_array(ext_data_) elif "dict" in ext_data_.attrs: # legacy json object array (spikeinterface<=0.105) - ext_data = ext_data_[0] + ext_data = read_encoded_array(ext_data_) elif "dataframe" in ext_data_.attrs: import pandas as pd @@ -3384,9 +3386,10 @@ def load_data(self, lazy=False): ext_data.loc[:, col] = ext_data_[col][:] ext_data = ext_data.convert_dtypes() elif "object" in ext_data_.attrs: + # legacy pickle object array (spikeinterface<=0.105) if lazy: continue - ext_data = ext_data_[0] + ext_data = read_encoded_array(ext_data_) else: ext_data = ext_data_ if lazy else np.array(ext_data_[:]) self.set_data(ext_data_name, ext_data) @@ -3577,10 +3580,11 @@ def _save_data(self): else: # any object try: - _write_object_array(extension_group, ext_data_name, ext_data, codec="pickle") + write_pickle_array( + extension_group, ext_data_name, ext_data, compressor=saving_options["compressor"] + ) except: raise Exception(f"Could not save {ext_data_name} as extension data") - extension_group[ext_data_name].attrs["object"] = True def _reset_extension_folder(self): """ diff --git a/src/spikeinterface/core/zarr_tools.py b/src/spikeinterface/core/zarr_tools.py index 023de78dbe..22931e9723 100644 --- a/src/spikeinterface/core/zarr_tools.py +++ b/src/spikeinterface/core/zarr_tools.py @@ -1,4 +1,5 @@ import json +import pickle import numpy as np @@ -35,10 +36,40 @@ def write_json_array(group, name: str, data, compressor=None): return array +def write_pickle_array(group, name: str, data, compressor=None): + """ + Save a python object as its pickle bytes in a uint8 array. + + This replaces the length-1 object arrays with a pickle codec, which zarr-python 3 cannot write. + It is used for extension data that is not json-serializable, such as the PCA models of the + "principal_components" extension. See `read_encoded_array` for the reverse operation. + + Parameters + ---------- + group : zarr.Group + The zarr group in which the array is created. + name : str + The name of the array. + data : Any + The picklable object to save. + compressor : numcodecs.Codec | None, default: None + The compressor of the array. + + Returns + ------- + array : zarr.Array + The created array. + """ + payload = np.frombuffer(pickle.dumps(data), dtype="uint8") + array = group.create_dataset(name=name, data=payload, compressor=compressor) + array.attrs["encoding"] = "pickle" + return array + + def read_encoded_array(array): """ - Read data saved by `write_json_array`, or a legacy length-1 object array (json or pickle codec) - saved by spikeinterface<=0.105. + Read data saved by `write_json_array` or `write_pickle_array`, or a legacy length-1 object array + (json or pickle codec) saved by spikeinterface<=0.105. Parameters ---------- @@ -52,10 +83,13 @@ def read_encoded_array(array): """ encoding = array.attrs.get("encoding", None) if encoding is None: + # With love for Alessio # legacy length-1 object array return array[0] if encoding == "json": return json.loads(np.asarray(array[:]).tobytes().decode("utf-8")) + if encoding == "pickle": + return pickle.loads(np.asarray(array[:]).tobytes()) raise ValueError( f"Unknown encoding {encoding!r} for {array.name}. It was probably saved by a newer version of spikeinterface." ) diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py index d2a604e38b..4f478e425c 100644 --- a/src/spikeinterface/core/zarrextractors.py +++ b/src/spikeinterface/core/zarrextractors.py @@ -744,51 +744,6 @@ def create_zarr_path_for_write(folder_path: str | Path, overwrite: bool = False) return folder_path -def _write_object_array( - group, - name: str, - data, - codec: str = "json", - overwrite: bool = True, -): - """ - Write a length-1 object-dtype array holding a Python dict/list/object. - - Centralizes the v2/v3 codec-placement difference for object blobs: under zarr-v2 - the object codec goes in ``object_codec=``; under zarr-v3 it goes in ``filters=`` - (wrapped via ``numcodecs.zarr3.*``). The helper picks the right path automatically. - - Parameters - ---------- - group : zarr.Group - The zarr group to write into. - name : str - Name of the array inside ``group``. - data : Any - The Python object to store. Wrapped into ``np.array([data], dtype=object)``. - codec : {"json", "pickle"}, default: "json" - Which object codec to use. - overwrite : bool, default: True - Whether to overwrite an existing array with the same name. - """ - import numcodecs - - if codec == "json": - codec_instance = numcodecs.JSON() - elif codec == "pickle": - codec_instance = numcodecs.Pickle() - else: - raise ValueError(f"codec must be 'json' or 'pickle', got {codec!r}") - - arr = np.array([data], dtype=object) - return group.create_dataset( - name=name, - data=arr, - object_codec=codec_instance, - overwrite=overwrite, - ) - - def get_default_zarr_compressor(clevel: int = 5): """ Return default Zarr compressor object for good preformance in int16