Repository: SpikeInterface/spikeinterface (python)

Issue:
Unit aggregation should preserve unit ids.
Like https://github.com/SpikeInterface/spikeinterface/issues/2789 but for UnitAggregation. Currently it does not:

```python
from spikeinterface.core import generate_sorting


sorting1 = generate_sorting(num_units=3)
sorting1 = sorting1.rename_units(new_unit_ids=["unit1", "unit2", "unit3"])

sorting2 = generate_sorting(num_units=3)
sorting2 = sorting2.rename_units(new_unit_ids=["unit4", "unit5", "unit6"])

from spikeinterface.core import aggregate_units

sorting = aggregate_units([sorting1, sorting2])
sorting.get_unit_ids()
array([0, 1, 2, 3, 4, 5])
```

Pull request:
Units aggregation preserve unit ids of aggregated sorters
This should close #3152

Patch:
diff --git a/src/spikeinterface/core/base.py b/src/spikeinterface/core/base.py
index 8b037ad10..f3e300fff 100644
--- a/src/spikeinterface/core/base.py
+++ b/src/spikeinterface/core/base.py
@@ -7,6 +7,7 @@ import warnings
 import weakref
 import json
 import pickle
+import os
 import random
 import string
 from packaging.version import parse
@@ -950,14 +951,13 @@ class BaseExtractor:
         folder.mkdir(parents=True, exist_ok=False)
 
         # dump provenance
+        provenance_file = folder / f"provenance.json"
         if self.check_serializability("json"):
-            provenance_file = folder / f"provenance.json"
-            self.dump(provenance_file)
-        elif self.check_serializability("pickle"):
-            provenance_file = folder / f"provenance.pkl"
             self.dump(provenance_file)
         else:
-            warnings.warn("The extractor is not serializable to file. The provenance will not be saved.")
+            provenance_file.write_text(
+                json.dumps({"warning": "the provenace is not json serializable!!!"}), encoding="utf8"
+            )
 
         self.save_metadata_to_folder(folder)
 
@@ -1024,6 +1024,7 @@ class BaseExtractor:
         cached: ZarrExtractor
             Saved copy of the extractor.
         """
+        import zarr
         from .zarrextractors import read_zarr
 
         save_kwargs.pop("format", None)
diff --git a/src/spikeinterface/core/generate.py b/src/spikeinterface/core/generate.py
index 6ce94114c..f5312f9c4 100644
--- a/src/spikeinterface/core/generate.py
+++ b/src/spikeinterface/core/generate.py
@@ -1803,8 +1803,6 @@ class InjectTemplatesRecording(BaseRecording):
             )
             self.add_recording_segment(recording_segment)
 
-        # to discuss: maybe we could set json serializability to False always
-        # because templates could be large!
         if not sorting.check_serializability("json"):
             self._serializability["json"] = False
         if parent_recording is not None:
diff --git a/src/spikeinterface/core/unitsaggregationsorting.py b/src/spikeinterface/core/unitsaggregationsorting.py
index ea019268f..9eb37e31e 100644
--- a/src/spikeinterface/core/unitsaggregationsorting.py
+++ b/src/spikeinterface/core/unitsaggregationsorting.py
@@ -34,7 +34,21 @@ class UnitsAggregationSorting(BaseSorting):
             )
             unit_ids = list(renamed_unit_ids)
         else:
-            unit_ids = list(np.arange(num_all_units))
+            unit_ids_dtypes = [sort.get_unit_ids().dtype for sort in sorting_list]
+            all_ids_are_same_type = np.unique(unit_ids_dtypes).size == 1
+            all_units_ids_are_unique = False
+            if all_ids_are_same_type:
+                combined_ids = np.concatenate([sort.get_unit_ids() for sort in sorting_list])
+                all_units_ids_are_unique = np.unique(combined_ids).size == num_all_units
+
+            if all_ids_are_same_type and all_units_ids_are_unique:
+                unit_ids = combined_ids
+            else:
+                default_unit_ids = [str(i) for i in range(num_all_units)]
+                if all_ids_are_same_type and np.issubdtype(unit_ids_dtypes[0], np.integer):
+                    unit_ids = np.arange(num_all_units, dtype=np.uint64)
+                else:
+                    unit_ids = default_unit_ids
 
         # unit map maps unit ids that are used to get spike trains
         u_id = 0
diff --git a/src/spikeinterface/generation/drift_tools.py b/src/spikeinterface/generation/drift_tools.py
index 70e13160f..cce2e08b5 100644
--- a/src/spikeinterface/generation/drift_tools.py
+++ b/src/spikeinterface/generation/drift_tools.py
@@ -458,9 +458,6 @@ class InjectDriftingTemplatesRecording(BaseRecording):
 
         self.set_probe(drifting_templates.probe, in_place=True)
 
-        # templates are too large, we don't serialize them to JSON
-        self._serializability["json"] = False
-
         self._kwargs = {
             "sorting": sorting,
             "drifting_templates": drifting_templates,


Test patch:
diff --git a/src/spikeinterface/core/tests/test_unitsaggregationsorting.py b/src/spikeinterface/core/tests/test_unitsaggregationsorting.py
index b6cb479c7..09b8affe5 100644
--- a/src/spikeinterface/core/tests/test_unitsaggregationsorting.py
+++ b/src/spikeinterface/core/tests/test_unitsaggregationsorting.py
@@ -5,6 +5,7 @@ from spikeinterface.core import aggregate_units
 
 from spikeinterface.core import NpzSortingExtractor
 from spikeinterface.core import create_sorting_npz
+from spikeinterface.core import generate_sorting
 
 
 def test_unitsaggregationsorting(create_cache_folder):
@@ -92,5 +93,42 @@ def test_unitsaggregationsorting(create_cache_folder):
     print(sorting_agg_prop.get_property("brain_area"))
 
 
+def test_unit_aggregation_preserve_ids():
+
+    sorting1 = generate_sorting(num_units=3)
+    sorting1 = sorting1.rename_units(new_unit_ids=["unit1", "unit2", "unit3"])
+
+    sorting2 = generate_sorting(num_units=3)
+    sorting2 = sorting2.rename_units(new_unit_ids=["unit4", "unit5", "unit6"])
+
+    aggregated_sorting = aggregate_units([sorting1, sorting2])
+    assert aggregated_sorting.get_num_units() == 6
+    assert list(aggregated_sorting.get_unit_ids()) == ["unit1", "unit2", "unit3", "unit4", "unit5", "unit6"]
+
+
+def test_unit_aggregation_does_not_preserve_ids_if_not_unique():
+    sorting1 = generate_sorting(num_units=3)
+    sorting1 = sorting1.rename_units(new_unit_ids=["unit1", "unit2", "unit3"])
+
+    sorting2 = generate_sorting(num_units=3)
+    sorting2 = sorting2.rename_units(new_unit_ids=["unit1", "unit2", "unit3"])
+
+    aggregated_sorting = aggregate_units([sorting1, sorting2])
+    assert aggregated_sorting.get_num_units() == 6
+    assert list(aggregated_sorting.get_unit_ids()) == ["0", "1", "2", "3", "4", "5"]
+
+
+def test_unit_aggregation_does_not_preserve_ids_not_the_same_type():
+    sorting1 = generate_sorting(num_units=3)
+    sorting1 = sorting1.rename_units(new_unit_ids=["unit1", "unit2", "unit3"])
+
+    sorting2 = generate_sorting(num_units=2)
+    sorting2 = sorting2.rename_units(new_unit_ids=[1, 2])
+
+    aggregated_sorting = aggregate_units([sorting1, sorting2])
+    assert aggregated_sorting.get_num_units() == 5
+    assert list(aggregated_sorting.get_unit_ids()) == ["0", "1", "2", "3", "4"]
+
+
 if __name__ == "__main__":
     test_unitsaggregationsorting()