Repository: SpikeInterface/spikeinterface
Pull request #3403 (2024-09-12 16:50:45)

Add `get_channel_locations` to the base recording api
This will make the method appear on the BaseRecording API on the docs.

Plus, I added a docstring to the method. 

No changes on behavior with this diff.


Patch:
diff --git a/src/spikeinterface/core/baserecording.py b/src/spikeinterface/core/baserecording.py
index 082afd880..03001ae47 100644
--- a/src/spikeinterface/core/baserecording.py
+++ b/src/spikeinterface/core/baserecording.py
@@ -746,6 +746,30 @@ class BaseRecording(BaseRecordingSnippets):
 
         return SelectSegmentRecording(self, segment_indices=segment_indices)
 
+    def get_channel_locations(
+        self,
+        channel_ids: list | np.ndarray | tuple | None = None,
+        axes: "xy" | "yz" | "xz" | "xyz" = "xy",
+    ) -> np.ndarray:
+        """
+        Get the physical locations of specified channels.
+
+        Parameters
+        ----------
+        channel_ids : array-like, optional
+            The IDs of the channels for which to retrieve locations. If None, retrieves locations
+            for all available channels. Default is None.
+        axes : "xy" | "yz" | "xz" | "xyz", default: "xy"
+            The spatial axes to return, specified as a string (e.g., "xy", "xyz"). Default is "xy".
+
+        Returns
+        -------
+        np.ndarray
+            A 2D or 3D array of shape (n_channels, n_dimensions) containing the locations of the channels.
+            The number of dimensions depends on the `axes` argument (e.g., 2 for "xy", 3 for "xyz").
+        """
+        return super().get_channel_locations(channel_ids=channel_ids, axes=axes)
+
     def is_binary_compatible(self) -> bool:
         """
         Checks if the recording is "binary" compatible.
diff --git a/src/spikeinterface/core/baserecordingsnippets.py b/src/spikeinterface/core/baserecordingsnippets.py
index 428472bf9..3953c1f05 100644
--- a/src/spikeinterface/core/baserecordingsnippets.py
+++ b/src/spikeinterface/core/baserecordingsnippets.py
@@ -344,7 +344,7 @@ class BaseRecordingSnippets(BaseExtractor):
             raise ValueError("set_channel_locations(..) destroys the probe description, prefer _set_probes(..)")
         self.set_property("location", locations, ids=channel_ids)
 
-    def get_channel_locations(self, channel_ids=None, axes: str = "xy"):
+    def get_channel_locations(self, channel_ids=None, axes: str = "xy") -> np.ndarray:
         if channel_ids is None:
             channel_ids = self.get_channel_ids()
         channel_indices = self.ids_to_indices(channel_ids)
diff --git a/src/spikeinterface/core/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py
index 424fab7c5..49a31738e 100644
--- a/src/spikeinterface/core/sortinganalyzer.py
+++ b/src/spikeinterface/core/sortinganalyzer.py
@@ -77,9 +77,6 @@ def create_sorting_analyzer(
     return_scaled : bool, default: True
         All extensions that play with traces will use this global return_scaled : "waveforms", "noise_levels", "templates".
         This prevent return_scaled being differents from different extensions and having wrong snr for instance.
-    overwrite: bool, default: False
-        If True, overwrite the folder if it already exists.
-
 
     Returns
     -------
@@ -489,11 +486,7 @@ class SortingAnalyzer:
 
         if is_path_remote(str(self.folder)):
             mode = "r"
-        # we open_consolidated only if we are in read mode
-        if mode in ("r+", "a"):
-            zarr_root = zarr.open(str(self.folder), mode=mode, storage_options=self.storage_options)
-        else:
-            zarr_root = zarr.open_consolidated(self.folder, mode=mode, storage_options=self.storage_options)
+        zarr_root = zarr.open(self.folder, mode=mode, storage_options=self.storage_options)
         return zarr_root
 
     @classmethod
@@ -571,13 +564,11 @@ class SortingAnalyzer:
 
         recording_info = zarr_root.create_group("extensions")
 
-        zarr.consolidate_metadata(zarr_root.store)
-
     @classmethod
     def load_from_zarr(cls, folder, recording=None, storage_options=None):
         import zarr
 
-        zarr_root = zarr.open_consolidated(str(folder), mode="r", storage_options=storage_options)
+        zarr_root = zarr.open(str(folder), mode="r", storage_options=storage_options)
 
         # load internal sorting in memory
         sorting = NumpySorting.from_sorting(
@@ -623,7 +614,6 @@ class SortingAnalyzer:
             format="zarr",
             sparsity=sparsity,
             return_scaled=return_scaled,
-            storage_options=storage_options,
         )
 
         return sorting_analyzer
@@ -1472,7 +1462,7 @@ extension_params={"waveforms":{"ms_before":1.5, "ms_after": "2.5"}}\
         if self.format != "memory" and self.has_extension(extension_name):
             # need a reload to reset the folder
             ext = self.load_extension(extension_name)
-            ext.delete()
+            ext.reset()
 
         # remove from dict
         self.extensions.pop(extension_name, None)
@@ -2014,17 +2004,19 @@ class AnalyzerExtension:
             # NB: this call to _save_params() also resets the folder or zarr group
             self._save_params()
             self._save_importing_provenance()
+            self._save_run_info()
 
         t_start = perf_counter()
         self._run(**kwargs)
         t_end = perf_counter()
         self.run_info["runtime_s"] = t_end - t_start
-        self.run_info["run_completed"] = True
 
         if save and not self.sorting_analyzer.is_read_only():
-            self._save_run_info()
             self._save_data(**kwargs)
 
+        self.run_info["run_completed"] = True
+        self._save_run_info()
+
     def save(self, **kwargs):
         self._save_params()
         self._save_importing_provenance()
@@ -2070,7 +2062,7 @@ class AnalyzerExtension:
                     except:
                         raise Exception(f"Could not save {ext_data_name} as extension data")
         elif self.format == "zarr":
-            import zarr
+
             import numcodecs
 
             extension_group = self._get_zarr_extension_group(mode="r+")
@@ -2104,8 +2096,6 @@ class AnalyzerExtension:
                     except:
                         raise Exception(f"Could not save {ext_data_name} as extension data")
                     extension_group[ext_data_name].attrs["object"] = True
-            # we need to re-consolidate
-            zarr.consolidate_metadata(self.sorting_analyzer._get_zarr_root().store)
 
     def _reset_extension_folder(self):
         """
@@ -2120,35 +2110,8 @@ class AnalyzerExtension:
         elif self.format == "zarr":
             import zarr
 
-            zarr_root = self.sorting_analyzer._get_zarr_root(mode="r+")
-            _ = zarr_root["extensions"].create_group(self.extension_name, overwrite=True)
-            zarr.consolidate_metadata(zarr_root.store)
-
-    def _delete_extension_folder(self):
-        """
-        Delete the extension in a folder (binary or zarr).
-        """
-        if self.format == "binary_folder":
-            extension_folder = self._get_binary_extension_folder()
-            if extension_folder.is_dir():
-                shutil.rmtree(extension_folder)
-
-        elif self.format == "zarr":
-            import zarr
-
-            zarr_root = self.sorting_analyzer._get_zarr_root(mode="r+")
-            if self.extension_name in zarr_root["extensions"]:
-                del zarr_root["extensions"][self.extension_name]
-                zarr.consolidate_metadata(zarr_root.store)
-
-    def delete(self):
-        """
-        Delete the extension from the folder or zarr and from the dict.
-        """
-        self._delete_extension_folder()
-        self.params = None
-        self.run_info = self._default_run_info_dict()
-        self.data = dict()
+            zarr_root = zarr.open(self.folder, mode="r+")
+            extension_group = zarr_root["extensions"].create_group(self.extension_name, overwrite=True)
 
     def reset(self):
         """
@@ -2165,7 +2128,7 @@ class AnalyzerExtension:
         Set parameters for the extension and
         make it persistent in json.
         """
-        # this ensure data is also deleted and corresponds to params
+        # this ensure data is also deleted and corresponf to params
         # this also ensure the group is created
         self._reset_extension_folder()
 
@@ -2178,6 +2141,7 @@ class AnalyzerExtension:
         if save:
             self._save_params()
             self._save_importing_provenance()
+            self._save_run_info()
 
     def _save_params(self):
         params_to_save = self.params.copy()


Test patch:
diff --git a/src/spikeinterface/core/tests/test_analyzer_extension_core.py b/src/spikeinterface/core/tests/test_analyzer_extension_core.py
index 626899ab6..b4d96a339 100644
--- a/src/spikeinterface/core/tests/test_analyzer_extension_core.py
+++ b/src/spikeinterface/core/tests/test_analyzer_extension_core.py
@@ -79,20 +79,15 @@ def _check_result_extension(sorting_analyzer, extension_name, cache_folder):
 )
 def test_ComputeRandomSpikes(format, sparse, create_cache_folder):
     cache_folder = create_cache_folder
-    print("Creating analyzer")
     sorting_analyzer = get_sorting_analyzer(cache_folder, format=format, sparse=sparse)
 
-    print("Computing random spikes")
     ext = sorting_analyzer.compute("random_spikes", max_spikes_per_unit=10, seed=2205)
     indices = ext.data["random_spikes_indices"]
     assert indices.size == 10 * sorting_analyzer.sorting.unit_ids.size
 
-    print("Checking results")
     _check_result_extension(sorting_analyzer, "random_spikes", cache_folder)
-    print("Delering extension")
     sorting_analyzer.delete_extension("random_spikes")
 
-    print("Re-computing random spikes")
     ext = sorting_analyzer.compute("random_spikes", method="all")
     indices = ext.data["random_spikes_indices"]
     assert indices.size == len(sorting_analyzer.sorting.to_spike_vector())
diff --git a/src/spikeinterface/core/tests/test_sortinganalyzer.py b/src/spikeinterface/core/tests/test_sortinganalyzer.py
index 77b8f2c5b..3f45487f4 100644
--- a/src/spikeinterface/core/tests/test_sortinganalyzer.py
+++ b/src/spikeinterface/core/tests/test_sortinganalyzer.py
@@ -126,8 +126,6 @@ def test_SortingAnalyzer_zarr(tmp_path, dataset):
 
 
 def test_load_without_runtime_info(tmp_path, dataset):
-    import zarr
-
     recording, sorting = dataset
 
     folder = tmp_path / "test_SortingAnalyzer_run_info"
@@ -155,7 +153,6 @@ def test_load_without_runtime_info(tmp_path, dataset):
     root = sorting_analyzer._get_zarr_root(mode="r+")
     for ext in extensions:
         del root["extensions"][ext].attrs["run_info"]
-        zarr.consolidate_metadata(root.store)
     # should raise a warning for missing run_info
     with pytest.warns(UserWarning):
         sorting_analyzer = load_sorting_analyzer(folder, format="auto")