Repository: SpikeInterface/spikeinterface
Pull request #2403 (2024-01-12 10:42:06)

Improve ZarrSortingExtractor
Add option for ZarrSortingExtractor to read in sub group.


Patch:
diff --git a/src/spikeinterface/comparison/basecomparison.py b/src/spikeinterface/comparison/basecomparison.py
index 4133167a0..7a231f3cb 100644
--- a/src/spikeinterface/comparison/basecomparison.py
+++ b/src/spikeinterface/comparison/basecomparison.py
@@ -91,7 +91,7 @@ class BaseMultiComparison(BaseComparison):
     ):
         # do pairwise matching
         if self._verbose:
-            print("Multicomparison step 1: pairwise comparison")
+            print("Multicomaprison step 1: pairwise comparison")
 
         self.comparisons = {}
         for i in range(len(self.object_list)):
@@ -133,7 +133,7 @@ class BaseMultiComparison(BaseComparison):
 
     def _clean_graph(self):
         if self._verbose:
-            print("Multicomparison step 3: clean graph")
+            print("Multicomaprison step 3: clean graph")
         clean_graph = self.graph.copy()
         import networkx as nx
 
diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py
index c9540ffe4..a6aa4a2d4 100644
--- a/src/spikeinterface/core/zarrextractors.py
+++ b/src/spikeinterface/core/zarrextractors.py
@@ -154,11 +154,12 @@ class ZarrSortingExtractor(BaseSorting):
 
     Parameters
     ----------
-    root_path: str or Path
+    folder_path: str or Path
         Path to the zarr root file
     storage_options: dict or None
         Storage options for zarr `store`. E.g., if "s3://" or "gcs://" they can provide authentication methods, etc.
-
+    zarr_group: str or None, default: None
+        Optional zarr group path to load the sorting from. This can be used when the sorting is not stored at the root, but in sub group.
     Returns
     -------
     sorting: ZarrSortingExtractor
@@ -171,12 +172,16 @@ class ZarrSortingExtractor(BaseSorting):
     installation_mesg = ""
     name = "zarr"
 
-    def __init__(self, folder_path: Path | str, storage_options: dict | None = None):
+    def __init__(self, folder_path: Path | str, storage_options: dict | None = None, zarr_group: str | None = None):
         assert self.installed, self.installation_mesg
 
         folder_path, folder_path_kwarg = resolve_zarr_path(folder_path)
 
-        self._root = zarr.open(str(folder_path), mode="r", storage_options=storage_options)
+        zarr_root = self._root = zarr.open(str(folder_path), mode="r", storage_options=storage_options)
+        if zarr_group is None:
+            self._root = zarr_root
+        else:
+            self._root = zarr_root[zarr_group]
 
         sampling_frequency = self._root.attrs.get("sampling_frequency", None)
         num_segments = self._root.attrs.get("num_segments", None)
@@ -215,7 +220,7 @@ class ZarrSortingExtractor(BaseSorting):
         if annotations is not None:
             self.annotate(**annotations)
 
-        self._kwargs = {"root_path": folder_path_kwarg, "storage_options": storage_options}
+        self._kwargs = {"folder_path": folder_path_kwarg, "storage_options": storage_options, "zarr_group": zarr_group}
 
     @staticmethod
     def write_sorting(sorting: BaseSorting, folder_path: str | Path, storage_options: dict | None = None, **kwargs):


Test patch:
diff --git a/src/spikeinterface/core/tests/test_zarrextractors.py b/src/spikeinterface/core/tests/test_zarrextractors.py
new file mode 100644
index 000000000..72247cb42
--- /dev/null
+++ b/src/spikeinterface/core/tests/test_zarrextractors.py
@@ -0,0 +1,46 @@
+import pytest
+from pathlib import Path
+
+import shutil
+
+import zarr
+
+from spikeinterface.core import (
+    ZarrRecordingExtractor,
+    ZarrSortingExtractor,
+    generate_sorting,
+    load_extractor,
+)
+from spikeinterface.core.zarrextractors import add_sorting_to_zarr_group
+
+if hasattr(pytest, "global_test_folder"):
+    cache_folder = pytest.global_test_folder / "core"
+else:
+    cache_folder = Path("cache_folder") / "core"
+
+
+def test_ZarrSortingExtractor():
+    np_sorting = generate_sorting()
+
+    # store in root standard normal way
+    folder = cache_folder / "zarr_sorting"
+    if folder.is_dir():
+        shutil.rmtree(folder)
+    ZarrSortingExtractor.write_sorting(np_sorting, folder)
+    sorting = ZarrSortingExtractor(folder)
+    sorting = load_extractor(sorting.to_dict())
+
+    # store the sorting in a sub group (for instance SortingResult)
+    folder = cache_folder / "zarr_sorting_sub_group"
+    if folder.is_dir():
+        shutil.rmtree(folder)
+    zarr_root = zarr.open(folder, mode="w")
+    zarr_sorting_group = zarr_root.create_group("sorting")
+    add_sorting_to_zarr_group(sorting, zarr_sorting_group)
+    sorting = ZarrSortingExtractor(folder, zarr_group="sorting")
+    # and reaload
+    sorting = load_extractor(sorting.to_dict())
+
+
+if __name__ == "__main__":
+    test_ZarrSortingExtractor()