Sequential curation (#4298)

Co-authored-by: chrishalcrow <chrishalcrow@gmail.com>
Co-authored-by: Chris Halcrow <57948917+chrishalcrow@users.noreply.github.com>

Author: alejoe91

src/spikeinterface/core/analyzer_extension_core.py

@@ -9,6 +9,7 @@
   * ComputeNoiseLevels which is very convenient to have
 """
 
+from copy import copy
 import warnings
 import numpy as np
 from collections import namedtuple
@@ -387,6 +388,13 @@ def _handle_backward_compatibility_on_load(self):
             # compatibility february 2024 > july 2024
             self.params["ms_after"] = self.params["nafter"] * 1000.0 / self.sorting_analyzer.sampling_frequency
 
+        old_keys = copy(list(self.data.keys()))
+        for operator in old_keys:
+            if "pencentile" in operator:
+                fixed_operator = operator.replace("pencentile", "percentile")
+                self.data[fixed_operator] = self.data[operator]
+                del self.data[operator]
+
     def _set_params(self, ms_before: float = 1.0, ms_after: float = 2.0, operators=None):
         operators = operators or ["average", "std"]
         assert isinstance(operators, list)
@@ -485,7 +493,7 @@ def _compute_and_append_from_waveforms(self, operators):
             elif isinstance(operator, (list, tuple)):
                 operator, percentile = operator
                 assert operator == "percentile"
-                key = f"pencentile_{percentile}"
+                key = f"percentile_{percentile}"
             else:
                 raise ValueError(f"ComputeTemplates: wrong operator {operator}")
             self.data[key] = np.zeros((unit_ids.size, num_samples, channel_ids.size))
@@ -516,7 +524,7 @@ def _compute_and_append_from_waveforms(self, operators):
                 elif isinstance(operator, (list, tuple)):
                     operator, percentile = operator
                     arr = np.percentile(wfs, percentile, axis=0)
-                    key = f"pencentile_{percentile}"
+                    key = f"percentile_{percentile}"
 
                 if self.sparsity is None:
                     self.data[key][unit_index, :, :] = arr
@@ -606,7 +614,7 @@ def _split_extension_data(self, split_units, new_unit_ids, new_sorting_analyzer,
                         elif operator == "median":
                             arr = np.median(wfs, axis=0)
                         elif "percentile" in operator:
-                            _, percentile = operator.splot("_")
+                            _, percentile = operator.split("_")
                             arr = np.percentile(wfs, float(percentile), axis=0)
                         new_array[split_unit_index, ...] = arr
                 else:
@@ -676,7 +684,7 @@ def get_templates(self, unit_ids=None, operator="average", percentile=None, save
             key = operator
         else:
             assert percentile is not None, "You must provide percentile=... if `operator='percentile'`"
-            key = f"pencentile_{percentile}"
+            key = f"percentile_{percentile}"
 
         if key in self.data:
             templates_array = self.data[key]

src/spikeinterface/core/sortinganalyzer.py

@@ -2264,6 +2264,7 @@ class AnalyzerExtension:
       * _run()
       * _select_extension_data()
       * _merge_extension_data()
+      * _split_extension_data()
       * _get_data()
 
     The subclass must also set an `extension_name` class attribute which is not None by default.

src/spikeinterface/core/tests/test_analyzer_extension_core.py

@@ -160,7 +160,7 @@ def test_ComputeTemplates(format, sparse, create_cache_folder):
 
     # they all should be in data
     data = sorting_analyzer.get_extension("templates").data
-    for k in ["average", "std", "median", "pencentile_5.0", "pencentile_95.0"]:
+    for k in ["average", "std", "median", "percentile_5.0", "percentile_95.0"]:
         assert k in data.keys()
         assert data[k].shape[0] == sorting_analyzer.unit_ids.size
         assert data[k].shape[2] == sorting_analyzer.channel_ids.size

src/spikeinterface/core/waveforms_extractor_backwards_compatibility.py

@@ -324,7 +324,7 @@ def get_all_templates(
         ext = self.sorting_analyzer.get_extension("templates")
 
         if mode == "percentile":
-            key = f"pencentile_{percentile}"
+            key = f"percentile_{percentile}"
         else:
             key = mode
 

src/spikeinterface/curation/curation_format.py

@@ -6,7 +6,7 @@
 from itertools import chain
 
 from spikeinterface.core import BaseSorting, SortingAnalyzer, apply_merges_to_sorting, apply_splits_to_sorting
-from spikeinterface.curation.curation_model import CurationModel
+from spikeinterface.curation.curation_model import CurationModel, SequentialCuration
 
 
 def validate_curation_dict(curation_dict: dict):
@@ -138,7 +138,7 @@ def apply_curation_labels(
 
 def apply_curation(
     sorting_or_analyzer: BaseSorting | SortingAnalyzer,
-    curation_dict_or_model: dict | CurationModel,
+    curation_dict_or_model: dict | list | CurationModel | SequentialCuration,
     censor_ms: float | None = None,
     new_id_strategy: str = "append",
     merging_mode: str = "soft",
@@ -164,7 +164,7 @@ def apply_curation(
     ----------
     sorting_or_analyzer : Sorting | SortingAnalyzer
         The Sorting or SortingAnalyzer object to apply merges.
-    curation_dict : dict or CurationModel
+    curation_dict : dict | CurationModel | SequentialCuration
         The curation dict or model.
     censor_ms : float | None, default: None
         When applying the merges, any consecutive spikes within the `censor_ms` are removed. This can be thought of
@@ -199,14 +199,32 @@ def apply_curation(
         sorting_or_analyzer, (BaseSorting, SortingAnalyzer)
     ), f"`sorting_or_analyzer` must be a Sorting or a SortingAnalyzer, not an object of type {type(sorting_or_analyzer)}"
     assert isinstance(
-        curation_dict_or_model, (dict, CurationModel)
-    ), f"`curation_dict_or_model` must be a dict or a CurationModel, not an object of type {type(curation_dict_or_model)}"
+        curation_dict_or_model, (dict, list, CurationModel, SequentialCuration)
+    ), f"`curation_dict_or_model` must be a dict, CurationModel or a SequentialCuration not an object of type {type(curation_dict_or_model)}"
     if isinstance(curation_dict_or_model, dict):
         curation_model = CurationModel(**curation_dict_or_model)
+    elif isinstance(curation_dict_or_model, list):
+        curation_model = SequentialCuration(curation_steps=curation_dict_or_model)
     else:
         curation_model = curation_dict_or_model.model_copy(deep=True)
 
-    if not np.array_equal(np.asarray(curation_model.unit_ids), sorting_or_analyzer.unit_ids):
+    if isinstance(curation_model, SequentialCuration):
+        for c, single_curation_model in enumerate(curation_model.curation_steps):
+            if verbose:
+                print(f"Applying curation step: {c + 1} / {len(curation_model.curation_steps)}")
+            sorting_or_analyzer = apply_curation(
+                sorting_or_analyzer,
+                single_curation_model,
+                censor_ms=censor_ms,
+                merging_mode=merging_mode,
+                sparsity_overlap=sparsity_overlap,
+                raise_error_if_overlap_fails=raise_error_if_overlap_fails,
+                verbose=verbose,
+                job_kwargs=job_kwargs,
+            )
+        return sorting_or_analyzer
+
+    if not set(curation_model.unit_ids) == set(sorting_or_analyzer.unit_ids):
         raise ValueError("unit_ids from the curation_dict do not match the one from Sorting or SortingAnalyzer")
 
     # 1. Apply labels
@@ -228,13 +246,15 @@ def apply_curation(
             curated_sorting_or_analyzer, _, _ = apply_merges_to_sorting(
                 curated_sorting_or_analyzer,
                 merge_unit_groups=merge_unit_groups,
+                new_unit_ids=merge_new_unit_ids,
                 censor_ms=censor_ms,
                 new_id_strategy=new_id_strategy,
                 return_extra=True,
             )
         else:
             curated_sorting_or_analyzer, _ = curated_sorting_or_analyzer.merge_units(
                 merge_unit_groups=merge_unit_groups,
+                new_unit_ids=merge_new_unit_ids,
                 censor_ms=censor_ms,
                 merging_mode=merging_mode,
                 sparsity_overlap=sparsity_overlap,

src/spikeinterface/curation/curation_model.py

@@ -4,6 +4,7 @@
 import numpy as np
 
 from spikeinterface import BaseSorting
+from spikeinterface.core.sorting_tools import _get_ids_after_merging, _get_ids_after_splitting
 
 
 class LabelDefinition(BaseModel):
@@ -190,7 +191,7 @@ def check_merges(cls, values):
 
             # Check new unit id not already used
             if merge.new_unit_id is not None:
-                if merge.new_unit_id in unit_ids:
+                if merge.new_unit_id in unit_ids and merge.new_unit_id not in merge.unit_ids:
                     raise ValueError(f"New unit ID {merge.new_unit_id} is already in the unit list")
 
         values["merges"] = merges
@@ -366,6 +367,52 @@ def convert_old_format(cls, values):
                 values["removed"] = list(removed_units)
         return values
 
+    def get_final_ids_from_new_unit_ids(self) -> list:
+        """
+        Returns the final unit ids of the `curation_model`, when new unit ids are
+        given for each curation choice. Raises an error if new unit ids are missing
+        for any curation choice.
+
+        Returns
+        -------
+        final_ids : list
+            The ids of the sorting/analyzer after curation takes place
+        """
+        final_ids = list(self.unit_ids)
+        # 1. Remove units
+        for unit_id in self.removed:
+            if unit_id not in final_ids:
+                raise ValueError(f"Removed unit_id {unit_id} is not in the unit list")
+            final_ids.remove(unit_id)
+
+        # 2. Merge units
+        merge_unit_groups = []
+        new_merge_unit_ids = []
+        for merge in self.merges:
+            if merge.new_unit_id is None:
+                raise ValueError(
+                    f"The `new_unit_id` for the merge of units {merge.unit_ids} is `None`. This must be given."
+                )
+            merge_unit_groups.append(merge.unit_ids)
+            new_merge_unit_ids.append(merge.new_unit_id)
+        final_ids = _get_ids_after_merging(
+            final_ids, merge_unit_groups=merge_unit_groups, new_unit_ids=new_merge_unit_ids
+        )
+
+        # 3. Split units
+        split_units = {}
+        split_new_unit_ids = []
+        for split in self.splits:
+            if split.new_unit_ids is None:
+                raise ValueError(
+                    f"The `new_unit_ids` for the split of unit {split.unit_id} is `None`. These must be given."
+                )
+            # we only need the correct key and elements for the split, to mimic the output of the split function
+            split_units[split.unit_id] = [[]] * len(split.new_unit_ids)
+            split_new_unit_ids.append(split.new_unit_ids)
+        final_ids = _get_ids_after_splitting(final_ids, split_units=split_units, new_unit_ids=split_new_unit_ids)
+        return list(final_ids)
+
     @model_validator(mode="before")
     def validate_fields(cls, values):
         values = dict(values)
@@ -430,3 +477,43 @@ def validate_curation_dict(self):
                         )
 
         return self
+
+
+class SequentialCuration(BaseModel):
+    """
+    A Pydantic model which defines a sequence of curation steps. If using sequential curations,
+    we demand that each individual curation (except the final one) has manually defined new unit ids,
+    and that these match the unit ids of the following curation.
+    """
+
+    curation_steps: List[CurationModel] = Field(description="List of curation steps applied sequentially")
+
+    @model_validator(mode="after")
+    def validate_sequential_curation(self):
+
+        for curation in self.curation_steps[:-1]:
+            for merge in curation.merges:
+                if merge.new_unit_id is None:
+                    raise ValueError(
+                        "In a sequential curation, all curation decisions must have explicit `new_unit_id`s defined."
+                    )
+            for split in curation.splits:
+                if split.new_unit_ids is None:
+                    raiseValueError(
+                        "In a sequential curation, all curation decisions must have explicit `new_unit_id`s defined."
+                    )
+
+        for curation_index in range(len(self.curation_steps))[:-1]:
+
+            curation_1 = self.curation_steps[curation_index]
+            curation_2 = self.curation_steps[curation_index + 1]
+
+            previous_model_final_ids = curation_1.get_final_ids_from_new_unit_ids()
+            next_model_initial_ids = curation_2.unit_ids
+
+            if not (set(previous_model_final_ids) == set(next_model_initial_ids)):
+                raise ValueError(
+                    f"The initial unit_ids of curation {curation_index+1} do not match the final unit_ids of curation {curation_index}."
+                )
+
+        return self

src/spikeinterface/curation/tests/test_curation_format.py

@@ -171,6 +171,27 @@
 # This is a failure because unit 99 is not in the initial list
 unknown_removed_unit = {**curation_ids_int, "removed": [31, 42, 99]}
 
+# Sequential curation test data
+sequential_curation = [
+    {
+        "format_version": "2",
+        "unit_ids": [1, 2, 3, 4, 5],
+        "merges": [{"unit_ids": [3, 4], "new_unit_id": 34}],
+    },
+    {
+        "format_version": "2",
+        "unit_ids": [1, 2, 34, 5],
+        "splits": [{"unit_id": 34, "mode": "indices", "indices": [[0, 1, 2, 3]], "new_unit_ids": [340, 341]}],
+    },
+    {
+        "format_version": "2",
+        "unit_ids": [1, 2, 340, 341, 5],
+        "removed": [2, 5],
+        "merges": [{"unit_ids": [1, 340], "new_unit_id": 100}],
+        "splits": [{"unit_id": 341, "mode": "indices", "indices": [[0, 1, 2]], "new_unit_ids": [3410, 3411]}],
+    },
+]
+
 
 def test_curation_format_validation():
     # Test basic formats
@@ -412,6 +433,26 @@ def test_apply_curation_splits_with_mask():
     assert spike_counts[45] == num_spikes - 2 * (num_spikes // 3)  # Remainder
 
 
+def test_apply_sequential_curation():
+    recording, sorting = generate_ground_truth_recording(durations=[10.0], num_units=5, seed=2205)
+    sorting = sorting.rename_units([1, 2, 3, 4, 5])
+    analyzer = create_sorting_analyzer(sorting, recording, sparse=False)
+
+    # sequential curation steps:
+    # 1. merge 3 and 4 -> 34
+    # 2. split 34 -> 340, 341
+    # 3. remove 2, 5; merge 1 and 340 -> 100; split 341 -> 3410, 3411
+    analyzer_curated = apply_curation(analyzer, sequential_curation, verbose=True)
+    # initial -1(merge) +1(split) -2(remove) -1(merge) +1(split)
+    num_final_units = analyzer.get_num_units() - 1 + 1 - 2 - 1 + 1
+    assert analyzer_curated.get_num_units() == num_final_units
+
+    # check final unit ids
+    final_unit_ids = analyzer_curated.sorting.unit_ids
+    expected_final_unit_ids = [100, 3410, 3411]
+    assert set(final_unit_ids) == set(expected_final_unit_ids)
+
+
 if __name__ == "__main__":
     test_curation_format_validation()
     test_to_from_json()

src/spikeinterface/curation/tests/test_curation_model.py

@@ -3,7 +3,7 @@
 from pydantic import ValidationError
 import numpy as np
 
-from spikeinterface.curation.curation_model import CurationModel, LabelDefinition
+from spikeinterface.curation.curation_model import CurationModel, SequentialCuration, LabelDefinition
 
 
 # Test data for format version
@@ -282,3 +282,33 @@ def test_complete_model():
     assert len(model.merges) == 1
     assert len(model.splits) == 1
     assert len(model.removed) == 1
+
+
+def test_sequential_curation():
+    sequential_curation_steps_valid = [
+        {"format_version": "2", "unit_ids": [1, 2, 3, 4], "merges": [{"unit_ids": [1, 2], "new_unit_id": 22}]},
+        {
+            "format_version": "2",
+            "unit_ids": [3, 4, 22],
+            "splits": [
+                {"unit_id": 22, "mode": "indices", "indices": [[0, 1, 2], [3, 4, 5]], "new_unit_ids": [222, 223]}
+            ],
+        },
+        {"format_version": "2", "unit_ids": [3, 4, 222, 223], "removed": [223]},
+    ]
+
+    # this is valid
+    SequentialCuration(curation_steps=sequential_curation_steps_valid)
+
+    sequential_curation_steps_no_ids = sequential_curation_steps_valid.copy()
+    # remove new_unit_id in merge step
+    sequential_curation_steps_no_ids[0]["merges"][0]["new_unit_id"] = None
+
+    with pytest.raises(ValidationError):
+        SequentialCuration(curation_steps=sequential_curation_steps_no_ids)
+
+    sequential_curation_steps_invalid = sequential_curation_steps_valid.copy()
+    # invalid unit_ids in last step
+    sequential_curation_steps_invalid[2]["unit_ids"] = [3, 4, 222, 224]  # 224 should be 223
+    with pytest.raises(ValidationError):
+        SequentialCuration(curation_steps=sequential_curation_steps_invalid)

src/spikeinterface/sortingcomponents/clustering/itersplit_tools.py

@@ -74,9 +74,9 @@ def split_clusters(
     peak_labels = peak_labels.copy()
     split_count = np.zeros(peak_labels.size, dtype=int)
     recursion_level = 1
-    Executor = get_poolexecutor(n_jobs)
+    executor = get_poolexecutor(n_jobs)
 
-    with Executor(
+    with executor(
         max_workers=n_jobs,
         initializer=split_worker_init,
         mp_context=get_context(method=mp_context),
@@ -166,7 +166,6 @@ def split_worker_init(
     _ctx = {}
 
     _ctx["recording"] = recording
-    features_dict_or_folder
     _ctx["original_labels"] = original_labels
     _ctx["method"] = method
     _ctx["method_kwargs"] = method_kwargs