Repository: SpikeInterface/spikeinterface
Pull request #598 (2022-04-14 10:08:50)

Updates to zarr, and save sorting to memory


Patch:
diff --git a/spikeinterface/core/base.py b/spikeinterface/core/base.py
index 690cf260e..3d9e662b1 100644
--- a/spikeinterface/core/base.py
+++ b/spikeinterface/core/base.py
@@ -554,7 +554,7 @@ class BaseExtractor:
         Parameters
         ----------
         kwargs: Keyword arguments for saving.
-            * format: "memory" or "binary" (for recording) / "memory" or "npz" for sorting.
+            * format: "memory", "zarr", or "binary" (for recording) / "memory" or "npz" for sorting.
                 In case format is not memory, the recording is saved to a folder
             * folder: if provided, the folder path where the object is saved
             * name: if provided and folder is not given, the name of the folder in the global temporary
@@ -728,6 +728,7 @@ class BaseExtractor:
         save_kwargs['zarr_root'] = zarr_root
         save_kwargs['zarr_path'] = zarr_path
         cached = self._save(folder=None, verbose=verbose, **save_kwargs)
+        cached_annotations = deepcopy(cached._annotations)
 
         # save properties
         prop_group = zarr_root.create_group('properties')
@@ -740,6 +741,8 @@ class BaseExtractor:
 
         # copy properties/
         self.copy_metadata(cached)
+        # append annotations on compression
+        cached._annotations.update(cached_annotations)
 
         return cached
 
diff --git a/spikeinterface/core/baserecording.py b/spikeinterface/core/baserecording.py
index ea30942fc..5aba1ec2f 100644
--- a/spikeinterface/core/baserecording.py
+++ b/spikeinterface/core/baserecording.py
@@ -211,6 +211,8 @@ class BaseRecording(BaseExtractor):
                                     channel_ids=self.channel_ids)
 
         elif format == 'zarr':
+            from .zarrrecordingextractor import get_default_zarr_compressor, ZarrRecordingExtractor
+            
             zarr_root = save_kwargs.get('zarr_root', None)
             zarr_path = save_kwargs.get('zarr_path', None)
 
@@ -224,6 +226,12 @@ class BaseRecording(BaseExtractor):
             if dtype is None:
                 dtype = self.get_dtype()
             compressor = save_kwargs.get('compressor', None)
+            
+            if compressor is None:
+                compressor = get_default_zarr_compressor()
+                print(f"Using default zarr compressor: {compressor}. To use a different compressor, use the "
+                      f"'compressor' argument")
+            
             filters = save_kwargs.get('filters', None)
 
             job_kwargs = {k: save_kwargs[k]
@@ -252,7 +260,6 @@ class BaseRecording(BaseExtractor):
                 zarr_root.create_dataset(name="t_starts", data=t_starts,
                                          compressor=None)
 
-            from .zarrrecordingextractor import ZarrRecordingExtractor
             cached = ZarrRecordingExtractor(zarr_path)
 
         elif format == 'nwb':
diff --git a/spikeinterface/core/basesorting.py b/spikeinterface/core/basesorting.py
index edb14bf38..388ac66bc 100644
--- a/spikeinterface/core/basesorting.py
+++ b/spikeinterface/core/basesorting.py
@@ -119,8 +119,9 @@ class BaseSorting(BaseExtractor):
                     "The registered recording will not be persistent on disk, but only available in memory")
                 cached.register_recording(self._recording)
         else:
-            raise ValueError(f'format {format} not supported')
-
+            from .numpyextractors import NumpySorting
+            cached = NumpySorting.from_extractor(self)
+            
         return cached
 
     def get_unit_property(self, unit_id, key):
diff --git a/spikeinterface/core/zarrrecordingextractor.py b/spikeinterface/core/zarrrecordingextractor.py
index 635972402..00933689c 100644
--- a/spikeinterface/core/zarrrecordingextractor.py
+++ b/spikeinterface/core/zarrrecordingextractor.py
@@ -69,6 +69,9 @@ class ZarrRecordingExtractor(BaseRecording):
         dtype = np.dtype(dtype)
         t_starts = self._root.get("t_starts", None)
 
+        total_nbytes = 0
+        total_nbytes_stored = 0
+        cr_by_segment = {}
         for segment_index in range(num_segments):
             trace_name = f"traces_seg{segment_index}"
             assert len(channel_ids) == self._root[trace_name].shape[1], \
@@ -89,6 +92,13 @@ class ZarrRecordingExtractor(BaseRecording):
                 time_kwargs["sampling_frequency"] = sampling_frequency
 
             rec_segment = ZarrRecordingSegment(self._root, trace_name, **time_kwargs)
+            
+            nbytes_segment = self._root[trace_name].nbytes
+            nbytes_stored_segment = self._root[trace_name].nbytes_stored
+            cr_by_segment[segment_index] = nbytes_segment / nbytes_stored_segment
+            
+            total_nbytes += nbytes_segment
+            total_nbytes_stored += nbytes_stored_segment
             self.add_recording_segment(rec_segment)
 
         # load probe
@@ -108,6 +118,9 @@ class ZarrRecordingExtractor(BaseRecording):
         annotations = self._root.attrs.get("annotations", None)
         if annotations is not None:
             self.annotate(**annotations)
+        # annotate compression ratios
+        cr = total_nbytes / total_nbytes_stored
+        self.annotate(compression_ratio=cr, compression_ratio_segments=cr_by_segment)
         
         self._kwargs = {'root_path': str(root_path.absolute())}
 


Test patch:
diff --git a/spikeinterface/core/tests/test_basesorting.py b/spikeinterface/core/tests/test_basesorting.py
index 3e4ceb59b..9fe35ab4a 100644
--- a/spikeinterface/core/tests/test_basesorting.py
+++ b/spikeinterface/core/tests/test_basesorting.py
@@ -59,6 +59,9 @@ def test_BaseSorting():
     sorting2 = BaseExtractor.load_from_folder(folder)
     # but also possible
     sorting3 = BaseExtractor.load(folder)
+    
+    # save to memory
+    sorting4 = sorting.save(format="memory")
 
     spikes = sorting.get_all_spike_trains()
     # print(spikes)