Merge pull request #3694 from alejoe91/analyzer-set-unit-property

Add `SortingAnalyzer.set_sorting_property()/get_sorting_property()` functions

Author: alejoe91

src/spikeinterface/core/base.py

@@ -145,7 +145,7 @@ def ids_to_indices(
             non_existent_ids = [id for id in ids if id not in self._main_ids]
             if non_existent_ids:
                 error_msg = (
-                    f"IDs {non_existent_ids} are not channel ids of the extractor. \n"
+                    f"IDs {non_existent_ids} are not ids of the extractor. \n"
                     f"Available ids are {self._main_ids} with dtype {self._main_ids.dtype}"
                 )
                 raise ValueError(error_msg)
@@ -231,10 +231,10 @@ def set_property(
         ids : list/np.array, default: None
             List of subset of ids to set the values, default: None
             if None which is the default all the ids are set or changed
-        missing_value : object, default: None
+        missing_value : Any, default: None
             In case the property is set on a subset of values ("ids" not None),
-            it specifies the how the missing values should be filled.
-            The missing_value has to be specified for types int and unsigned int.
+            This argument specifies how to fill missing values.
+            The `missing_value` is required for types int and unsigned int.
         """
 
         if values is None:
@@ -293,6 +293,11 @@ def set_property(
                     ), f"Mismatch between existing property dtype {existing_property.kind} and provided values dtype {dtype_kind}."
 
                 indices = self.ids_to_indices(ids)
+                if dtype_kind == "U":
+                    # re-adjust the size of the property
+                    existing_unicode_size = max(len(s) for s in self._properties[key])
+                    new_unicode_size = max(max(len(s) for s in values), existing_unicode_size)
+                    self._properties[key] = self._properties[key].astype(f"<U{new_unicode_size}")
                 self._properties[key][indices] = values
             else:
                 indices = self.ids_to_indices(ids)

src/spikeinterface/core/sortinganalyzer.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import Literal, Optional
+from typing import Literal, Optional, Any
 
 from pathlib import Path
 from itertools import chain
@@ -223,13 +223,13 @@ class SortingAnalyzer:
 
     def __init__(
         self,
-        sorting=None,
-        recording=None,
-        rec_attributes=None,
-        format=None,
-        sparsity=None,
-        return_scaled=True,
-        backend_options=None,
+        sorting: BaseSorting,
+        recording: BaseRecording | None = None,
+        rec_attributes: dict | None = None,
+        format: str | None = None,
+        sparsity: ChannelSparsity | None = None,
+        return_scaled: bool = True,
+        backend_options: dict | None = None,
     ):
         # very fast init because checks are done in load and create
         self.sorting = sorting
@@ -239,6 +239,7 @@ def __init__(
         self.format = format
         self.sparsity = sparsity
         self.return_scaled = return_scaled
+        self.folder: str | Path | None = None
 
         # this is used to store temporary recording
         self._temporary_recording = None
@@ -742,6 +743,75 @@ def set_temporary_recording(self, recording: BaseRecording, check_dtype: bool =
             warnings.warn("SortingAnalyzer recording is already set. The current recording is temporarily replaced.")
         self._temporary_recording = recording
 
+    def set_sorting_property(
+        self,
+        key,
+        values: list | np.ndarray | tuple,
+        ids: list | np.ndarray | tuple | None = None,
+        missing_value: Any = None,
+        save: bool = True,
+    ) -> None:
+        """
+        Set property vector for unit ids.
+
+        If the SortingAnalyzer backend is in memory, the property will be only set in memory.
+        If the SortingAnalyzer backend is `binary_folder` or `zarr`, the property will also
+        be saved to to the backend.
+
+        Parameters
+        ----------
+        key : str
+            The property name
+        values : np.array
+            Array of values for the property
+        ids : list/np.array, default: None
+            List of subset of ids to set the values.
+            if None all the ids are set or changed
+        missing_value : Any, default: None
+            In case the property is set on a subset of values ("ids" not None),
+            This argument specifies how to fill missing values.
+            The `missing_value` is required for types int and unsigned int.
+        save : bool, default: True
+            If True, the property is saved to the backend if possible.
+        """
+        self.sorting.set_property(key, values, ids=ids, missing_value=missing_value)
+        if not self.is_read_only() and save:
+            if self.format == "binary_folder":
+                np.save(self.folder / "sorting" / "properties" / f"{key}.npy", self.sorting.get_property(key))
+            elif self.format == "zarr":
+                import zarr
+
+                zarr_root = self._get_zarr_root(mode="r+")
+                prop_values = self.sorting.get_property(key)
+                if prop_values.dtype.kind == "O":
+                    warnings.warn(f"Property {key} not saved because it is a python Object type")
+                else:
+                    if key in zarr_root["sorting"]["properties"]:
+                        zarr_root["sorting"]["properties"][key][:] = prop_values
+                    else:
+                        zarr_root["sorting"]["properties"].create_dataset(name=key, data=prop_values, compressor=None)
+                    # IMPORTANT: we need to re-consolidate the zarr store!
+                    zarr.consolidate_metadata(zarr_root.store)
+
+    def get_sorting_property(self, key: str, ids: Optional[Iterable] = None) -> np.ndarray:
+        """
+        Get property vector for unit ids.
+
+        Parameters
+        ----------
+        key : str
+            The property name
+        ids : list/np.array, default: None
+            List of subset of ids to get the values.
+            if None all the ids are returned
+
+        Returns
+        -------
+        values : np.array
+            Array of values for the property
+        """
+        return self.sorting.get_property(key, ids=ids)
+
     def _save_or_select_or_merge(
         self,
         format="binary_folder",
@@ -1177,6 +1247,9 @@ def get_sorting_provenance(self):
                 filename = self.folder / f"sorting_provenance.{type}"
                 sorting_provenance = None
                 if filename.exists():
+                    # try-except here is because it's not required to be able
+                    # to load the sorting provenance, as the user might have deleted
+                    # the original sorting folder
                     try:
                         sorting_provenance = load(filename, base_folder=self.folder)
                         break
@@ -1186,11 +1259,16 @@ def get_sorting_provenance(self):
 
         elif self.format == "zarr":
             zarr_root = self._get_zarr_root(mode="r")
+            sorting_provenance = None
             if "sorting_provenance" in zarr_root.keys():
-                sort_dict = zarr_root["sorting_provenance"][0]
-                sorting_provenance = load(sort_dict, base_folder=self.folder)
-            else:
-                sorting_provenance = None
+                # try-except here is because it's not required to be able
+                # to load the sorting provenance, as the user might have deleted
+                # the original sorting folder
+                try:
+                    sort_dict = zarr_root["sorting_provenance"][0]
+                    sorting_provenance = load(sort_dict, base_folder=self.folder)
+                except:
+                    pass
 
         return sorting_provenance
 

src/spikeinterface/core/tests/test_sortinganalyzer.py

@@ -66,6 +66,12 @@ def test_SortingAnalyzer_memory(tmp_path, dataset):
     )
     assert not sorting_analyzer.return_scaled
 
+    # test set_sorting_property
+    sorting_analyzer.set_sorting_property(key="quality", values=["good"] * len(sorting_analyzer.unit_ids))
+    sorting_analyzer.set_sorting_property(key="number", values=np.arange(len(sorting_analyzer.unit_ids)))
+    assert "quality" in sorting_analyzer.sorting.get_property_keys()
+    assert "number" in sorting_analyzer.sorting.get_property_keys()
+
 
 def test_SortingAnalyzer_binary_folder(tmp_path, dataset):
     recording, sorting = dataset
@@ -103,6 +109,15 @@ def test_SortingAnalyzer_binary_folder(tmp_path, dataset):
     assert not sorting_analyzer.return_scaled
     _check_sorting_analyzers(sorting_analyzer, sorting, cache_folder=tmp_path)
 
+    # test set_sorting_property
+    sorting_analyzer.set_sorting_property(key="quality", values=["good"] * len(sorting_analyzer.unit_ids))
+    sorting_analyzer.set_sorting_property(key="number", values=np.arange(len(sorting_analyzer.unit_ids)))
+    assert "quality" in sorting_analyzer.sorting.get_property_keys()
+    assert "number" in sorting_analyzer.sorting.get_property_keys()
+    sorting_analyzer_reloded = load_sorting_analyzer(folder, format="auto")
+    assert "quality" in sorting_analyzer_reloded.sorting.get_property_keys()
+    assert "number" in sorting_analyzer.sorting.get_property_keys()
+
 
 def test_SortingAnalyzer_zarr(tmp_path, dataset):
     recording, sorting = dataset
@@ -176,6 +191,15 @@ def test_SortingAnalyzer_zarr(tmp_path, dataset):
         == LZMA.codec_id
     )
 
+    # test set_sorting_property
+    sorting_analyzer.set_sorting_property(key="quality", values=["good"] * len(sorting_analyzer.unit_ids))
+    sorting_analyzer.set_sorting_property(key="number", values=np.arange(len(sorting_analyzer.unit_ids)))
+    assert "quality" in sorting_analyzer.sorting.get_property_keys()
+    assert "number" in sorting_analyzer.sorting.get_property_keys()
+    sorting_analyzer_reloded = load_sorting_analyzer(sorting_analyzer.folder, format="auto")
+    assert "quality" in sorting_analyzer_reloded.sorting.get_property_keys()
+    assert "number" in sorting_analyzer.sorting.get_property_keys()
+
 
 def test_load_without_runtime_info(tmp_path, dataset):
     import zarr
@@ -262,9 +286,6 @@ def _check_sorting_analyzers(sorting_analyzer, original_sorting, cache_folder):
     assert "sampling_frequency" in sorting_analyzer.rec_attributes
     assert "num_samples" in sorting_analyzer.rec_attributes
 
-    probe = sorting_analyzer.get_probe()
-    sparsity = sorting_analyzer.sparsity
-
     # compute
     sorting_analyzer.compute("dummy", param1=5.5)
     # equivalent
@@ -367,6 +388,9 @@ def _check_sorting_analyzers(sorting_analyzer, original_sorting, cache_folder):
         else:
             folder = None
         sorting_analyzer4 = sorting_analyzer.merge_units(merge_unit_groups=[[0, 1]], format=format, folder=folder)
+        assert 0 not in sorting_analyzer4.unit_ids
+        assert 1 not in sorting_analyzer4.unit_ids
+        assert len(sorting_analyzer4.unit_ids) == len(sorting_analyzer.unit_ids) - 1
 
         if format != "memory":
             if format == "zarr":
@@ -380,6 +404,10 @@ def _check_sorting_analyzers(sorting_analyzer, original_sorting, cache_folder):
         sorting_analyzer5 = sorting_analyzer.merge_units(
             merge_unit_groups=[[0, 1]], new_unit_ids=[50], format=format, folder=folder, merging_mode="hard"
         )
+        assert 0 not in sorting_analyzer5.unit_ids
+        assert 1 not in sorting_analyzer5.unit_ids
+        assert len(sorting_analyzer5.unit_ids) == len(sorting_analyzer.unit_ids) - 1
+        assert 50 in sorting_analyzer5.unit_ids
 
     # test compute with extension-specific params
     sorting_analyzer.compute(["dummy"], extension_params={"dummy": {"param1": 5.5}})