Merge pull request #2559 from zm711/zarr-json-serializability

Add `check_json` to `zarr` for `SortingAnalyzer` for sorting_provenance file writing

src/spikeinterface/core/__init__.py

@@ -146,6 +146,8 @@
 # SortingAnalyzer and AnalyzerExtension
 from .sortinganalyzer import SortingAnalyzer, AnalyzerExtension, create_sorting_analyzer, load_sorting_analyzer
 from .analyzer_extension_core import (
+    SelectRandomSpikes,
+    compute_select_random_spikes,
     ComputeWaveforms,
     compute_waveforms,
     ComputeTemplates,

src/spikeinterface/core/analyzer_extension_core.py

@@ -1,12 +1,13 @@
 """
 Implement AnalyzerExtension that are essential and imported in core
+  * SelectRandomSpikes
   * ComputeWaveforms
   * ComputeTemplates
 Theses two classes replace the WaveformExtractor
 
 It also implement:
   * ComputeFastTemplates which is equivalent but without extacting waveforms.
-  * ComputeNoiseLevels which is very convinient to have
+  * ComputeNoiseLevels which is very convenient to have
 """
 
 import numpy as np
@@ -105,6 +106,7 @@ def get_selected_indices_in_spike_train(self, unit_id, segment_index):
         return selected_spikes_in_spike_train
 
 
+compute_select_random_spikes = SelectRandomSpikes.function_factory()
 register_result_extension(SelectRandomSpikes)
 
 

src/spikeinterface/core/sortinganalyzer.py

@@ -431,7 +431,7 @@ def create_zarr(cls, folder, sorting, recording, sparsity, rec_attributes):
         # sorting provenance
         sort_dict = sorting.to_dict(relative_to=folder, recursive=True)
         if sorting.check_serializability("json"):
-            zarr_sort = np.array([sort_dict], dtype=object)
+            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)
@@ -507,7 +507,7 @@ def load_from_zarr(cls, folder, recording=None):
         # sparsity
         if "sparsity_mask" in zarr_root.attrs:
             # sparsity = zarr_root.attrs["sparsity"]
-            sparsity = ChannelSparsity(zarr_root["sparsity_mask"], self.unit_ids, rec_attributes["channel_ids"])
+            sparsity = ChannelSparsity(zarr_root["sparsity_mask"], cls.unit_ids, rec_attributes["channel_ids"])
         else:
             sparsity = None