Repository: SpikeInterface/spikeinterface

Issue:
Zarr `load_sorting_analyzer` trying to use non-existent sorting folder
Made a sorting analyzer; and moved it. It now tries to find the sorting output folder, but can't and errors. So

``` python
import spikeinterface.full as si
sa = si.load_sorting_analyzer("/Users/chris/zarr_sorting_analyzer.zarr", load_extensions=False)
sa.select_units(unit_ids=[3,4])
```

gives the error:

```
---------------------------------------------------------------------------
FileNotFoundError                         Traceback (most recent call last)
Cell In[1], line 3
      1 import spikeinterface.full as si
      2 sa = si.load_sorting_analyzer("/Users/chris/ks_3.zarr", load_extensions=False)
----> 3 sa.select_units(unit_ids=[3,4])

File ~/Work/Edinburgh/Spike/si_dev/spikeinterface_dev/spikeinterface/src/spikeinterface/core/sortinganalyzer.py:1057, in SortingAnalyzer.select_units(self, unit_ids, format, folder)
   1055 if format == "zarr":
   1056     folder = clean_zarr_folder_name(folder)
-> 1057 return self._save_or_select_or_merge(format=format, folder=folder, unit_ids=unit_ids)

File ~/Work/Edinburgh/Spike/si_dev/spikeinterface_dev/spikeinterface/src/spikeinterface/core/sortinganalyzer.py:911, in SortingAnalyzer._save_or_select_or_merge(self, format, folder, unit_ids, merge_unit_groups, censor_ms, merging_mode, sparsity_overlap, verbose, new_unit_ids, backend_options, **job_kwargs)
    908     sparsity = None
    910 # Note that the sorting is a copy we need to go back to the orginal sorting (if available)
--> 911 sorting_provenance = self.get_sorting_provenance()
    912 if sorting_provenance is None:
    913     # if the original sorting object is not available anymore (kilosort folder deleted, ....), take the copy
    914     sorting_provenance = self.sorting

File ~/Work/Edinburgh/Spike/si_dev/spikeinterface_dev/spikeinterface/src/spikeinterface/core/sortinganalyzer.py:1261, in SortingAnalyzer.get_sorting_provenance(self)
   1259 if "sorting_provenance" in zarr_root.keys():
   1260     sort_dict = zarr_root["sorting_provenance"][0]
-> 1261     sorting_provenance = load(sort_dict, base_folder=self.folder)
   1262 else:
   1263     sorting_provenance = None

File ~/Work/Edinburgh/Spike/si_dev/spikeinterface_dev/spikeinterface/src/spikeinterface/core/loading.py:73, in load(file_or_folder_or_dict, **kwargs)
     71     if object_type is None:
     72         raise ValueError(_error_msg.format(file_path="provided dictionary"))
---> 73     return _load_object_from_dict(d, object_type, base_folder=base_folder)
     75 is_local = not is_path_remote(file_or_folder_or_dict)
     76 if is_local:

File ~/Work/Edinburgh/Spike/si_dev/spikeinterface_dev/spikeinterface/src/spikeinterface/core/loading.py:175, in _load_object_from_dict(d, object_type, base_folder)
    173 def _load_object_from_dict(d, object_type, base_folder=None):
    174     if object_type in ("Recording", "Sorting", "Recording|Sorting"):
--> 175         return BaseExtractor.from_dict(d, base_folder=base_folder)
    177     elif object_type == "Templates":
    178         from spikeinterface.core import Templates

File ~/Work/Edinburgh/Spike/si_dev/spikeinterface_dev/spikeinterface/src/spikeinterface/core/base.py:534, in BaseExtractor.from_dict(dictionary, base_folder)
    532     assert base_folder is not None, "When  relative_paths=True, need to provide base_folder"
    533     dictionary = make_paths_absolute(dictionary, base_folder)
--> 534 extractor = _load_extractor_from_dict(dictionary)
    535 folder_metadata = dictionary.get("folder_metadata", None)
    536 if folder_metadata is not None:

File ~/Work/Edinburgh/Spike/si_dev/spikeinterface_dev/spikeinterface/src/spikeinterface/core/base.py:1075, in _load_extractor_from_dict(dic)
   1073     new_kwargs[name] = {k: transform_dict_to_extractor(v) for k, v in value.items()}
   1074 elif isinstance(value, list):
-> 1075     new_kwargs[name] = [transform_dict_to_extractor(e) for e in value]
   1076 else:
   1077     new_kwargs[name] = value

File ~/Work/Edinburgh/Spike/si_dev/spikeinterface_dev/spikeinterface/src/spikeinterface/core/base.py:1068, in _load_extractor_from_dict.<locals>.<lambda>(x)
   1066 # Create new kwargs to avoid modifying the original dict["kwargs"]
   1067 new_kwargs = dict()
-> 1068 transform_dict_to_extractor = lambda x: _load_extractor_from_dict(x) if is_dict_extractor(x) else x
   1069 for name, value in dic["kwargs"].items():
   1070     if is_dict_extractor(value):

File ~/Work/Edinburgh/Spike/si_dev/spikeinterface_dev/spikeinterface/src/spikeinterface/core/base.py:1090, in _load_extractor_from_dict(dic)
   1084     warnings.warn(
   1085         f"Versions are not the same. This might lead to compatibility errors. "
   1086         f"Using {class_name.split('.')[0]}=={dic['version']} is recommended"
   1087     )
   1089 # Initialize the extractor
-> 1090 extractor = extractor_class(**new_kwargs)
   1092 extractor._annotations.update(dic["annotations"])
   1093 for k, v in dic["properties"].items():

File ~/Work/Edinburgh/Spike/si_dev/spikeinterface_dev/spikeinterface/src/spikeinterface/extractors/phykilosortextractors.py:265, in KiloSortSortingExtractor.__init__(self, folder_path, keep_good_only, remove_empty_units)
    264 def __init__(self, folder_path: Path | str, keep_good_only: bool = False, remove_empty_units: bool = True):
--> 265     BasePhyKilosortSortingExtractor.__init__(
    266         self,
    267         folder_path,
    268         exclude_cluster_groups=None,
    269         keep_good_only=keep_good_only,
    270         remove_empty_units=remove_empty_units,
    271     )
    273     self._kwargs = {"folder_path": str(Path(folder_path).absolute()), "keep_good_only": keep_good_only}

File ~/Work/Edinburgh/Spike/si_dev/spikeinterface_dev/spikeinterface/src/spikeinterface/extractors/phykilosortextractors.py:47, in BasePhyKilosortSortingExtractor.__init__(self, folder_path, exclude_cluster_groups, keep_good_only, remove_empty_units, load_all_cluster_properties)
     44     raise ImportError(self.installation_mesg)
     46 phy_folder = Path(folder_path)
---> 47 spike_times = np.load(phy_folder / "spike_times.npy").astype(int)
     49 if (phy_folder / "spike_clusters.npy").is_file():
     50     spike_clusters = np.load(phy_folder / "spike_clusters.npy")

File ~/.my_venvs/intro_to_si/lib/python3.12/site-packages/numpy/lib/npyio.py:427, in load(file, mmap_mode, allow_pickle, fix_imports, encoding, max_header_size)
    425     own_fid = False
    426 else:
--> 427     fid = stack.enter_context(open(os_fspath(file), "rb"))
    428     own_fid = True
    430 # Code to distinguish from NumPy binary files and pickles.

FileNotFoundError: [Errno 2] No such file or directory: '../kilosort4_sort_3/0/sorter_output/spike_times.npy'
```

Replacing https://github.com/SpikeInterface/spikeinterface/blob/dd0c6d33f98f071362fa06abb631a7629b97de58/src/spikeinterface/core/sortinganalyzer.py#L1191
with a try/except works, but maybe there's a clever-er thing to do @alejoe91 ?? I'll do a PR if my simple fix is ok!

Hints from issue comments:
@chrishalcrow: Also, does this mean I'm the only one using zarr'd sorting analyzers??? They're really really small, everyone!@alejoe91: I use them all the time! Probably a problem of the select units!

The problem seems to be that the `select_units`..I'll take a look :)@alejoe91: yeah, we should have try-except as in the binary folder :) this is totally acceptable since the analyzer makes is own copy of the sorting, so a user could easily delete one. If you're ok, I'll just update #3694 to fix it


Solution patch:
diff --git a/src/spikeinterface/core/base.py b/src/spikeinterface/core/base.py
index 2dc7e0e9b..39eacdbef 100644
--- a/src/spikeinterface/core/base.py
+++ b/src/spikeinterface/core/base.py
@@ -145,7 +145,7 @@ class BaseExtractor:
             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 @@ class BaseExtractor:
         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 @@ class BaseExtractor:
                     ), 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)
diff --git a/src/spikeinterface/core/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py
index c284b9d6a..3b41625db 100644
--- a/src/spikeinterface/core/sortinganalyzer.py
+++ b/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 @@ class SortingAnalyzer:
         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 @@ class SortingAnalyzer:
             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 @@ class SortingAnalyzer:
                 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 @@ class SortingAnalyzer:
 
         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
 


Test patch:
diff --git a/src/spikeinterface/core/tests/test_sortinganalyzer.py b/src/spikeinterface/core/tests/test_sortinganalyzer.py
index 15f089f78..a76540f41 100644
--- a/src/spikeinterface/core/tests/test_sortinganalyzer.py
+++ b/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}})