Implement feedback

Author: alejoe91

src/spikeinterface/curation/curation_format.py

@@ -128,7 +128,7 @@ def apply_curation_labels(
         sorting.set_property(key, all_values)
 
     for new_unit_id, merge in zip(new_unit_ids, curation_model.merges):
-        old_group_ids = merge.merge_unit_group
+        old_group_ids = merge.unit_ids
         for label_key, label_def in curation_model.label_definitions.items():
             if label_def.exclusive:
                 group_values = []
@@ -221,7 +221,7 @@ def apply_curation(
         if len(curation_model.merges) > 0:
             sorting, _, new_unit_ids = apply_merges_to_sorting(
                 sorting,
-                merge_unit_groups=[m.merge_unit_group for m in curation_model.merges],
+                merge_unit_groups=[m.unit_ids for m in curation_model.merges],
                 censor_ms=censor_ms,
                 return_extra=True,
                 new_id_strategy=new_id_strategy,
@@ -237,7 +237,7 @@ def apply_curation(
             analyzer = analyzer.remove_units(curation_model.removed)
         if len(curation_model.removed) > 0:
             analyzer, new_unit_ids = analyzer.merge_units(
-                merge_unit_groups=[m.merge_unit_group for m in curation_model.merges],
+                merge_unit_groups=[m.unit_ids for m in curation_model.merges],
                 censor_ms=censor_ms,
                 merging_mode=merging_mode,
                 sparsity_overlap=sparsity_overlap,

src/spikeinterface/curation/curation_model.py

@@ -16,24 +16,31 @@ class ManualLabel(BaseModel):
 
 
 class Merge(BaseModel):
-    merge_unit_group: List[Union[int, str]] = Field(..., description="List of groups of units to be merged")
-    merge_new_unit_id: Optional[Union[int, str]] = Field(default=None, description="New unit IDs for the merge group")
+    unit_ids: List[Union[int, str]] = Field(..., description="List of unit ids to be merged")
+    new_unit_id: Optional[Union[int, str]] = Field(default=None, description="New unit IDs for the merge group")
 
 
 class Split(BaseModel):
     unit_id: Union[int, str] = Field(..., description="ID of the unit")
-    split_mode: Literal["indices", "labels"] = Field(
+    mode: Literal["indices", "labels"] = Field(
         default="indices",
         description=(
             "Mode of the split. The split can be defined by indices or labels. "
             "If indices, the split is defined by the a list of lists of indices of spikes within spikes "
-            "belonging to the unit (`split_indices`). "
-            "If labels, the split is defined by a list of labels for each spike (`split_labels`). "
+            "belonging to the unit (`indices`). "
+            "If labels, the split is defined by a list of labels for each spike (`labels`). "
         ),
     )
-    split_indices: Optional[Union[List[List[int]]]] = Field(default=None, description="List of indices for the split")
-    split_labels: Optional[List[int]] = Field(default=None, description="List of labels for the split")
-    split_new_unit_ids: Optional[List[Union[int, str]]] = Field(
+    indices: Optional[Union[List[int], List[List[int]]]] = Field(
+        default=None,
+        description=(
+            "List of indices for the split. If a list of indices, the unit is splt in 2 (provided indices/others). "
+            "If a list of lists, the unit is split in multiple groups (one for each list of indices), plus an optional "
+            "extra if the spike train has more spikes than the sum of the indices in the lists."
+        ),
+    )
+    labels: Optional[List[int]] = Field(default=None, description="List of labels for the split")
+    new_unit_ids: Optional[List[Union[int, str]]] = Field(
         default=None, description="List of new unit IDs for each split"
     )
 
@@ -129,25 +136,36 @@ def check_merges(cls, values):
         # Validate merges
         for merge in merges:
             # Check unit ids exist
-            for unit_id in merge.merge_unit_group:
+            for unit_id in merge.unit_ids:
                 if unit_id not in unit_ids:
                     raise ValueError(f"Merge unit group unit_id {unit_id} is not in the unit list")
 
             # Check minimum group size
-            if len(merge.merge_unit_group) < 2:
+            if len(merge.unit_ids) < 2:
                 raise ValueError("Merge unit groups must have at least 2 elements")
 
             # Check new unit id not already used
-            if merge.merge_new_unit_id is not None:
-                if merge.merge_new_unit_id in unit_ids:
-                    raise ValueError(f"New unit ID {merge.merge_new_unit_id} is already in the unit list")
+            if merge.new_unit_id is not None:
+                if merge.new_unit_id in unit_ids:
+                    raise ValueError(f"New unit ID {merge.new_unit_id} is already in the unit list")
 
         values["merges"] = merges
         return values
 
     @classmethod
     def check_splits(cls, values):
-
+        """
+        Checks and validates the splits in the curation model.
+        If `splits` is a dictionary with unit_id as key and split data as values,
+        it converts it to a list of Split objects.
+        Each Split object is then validated:
+        - Checks if the unit_id exists in the unit_ids list.
+        - Validates the mode (indices or labels).
+        - If mode is indices, checks that indices are defined and not empty, and that there are no duplicate indices.
+        - If mode is labels, checks that labels are defined and not empty.
+        - Validates new unit IDs if provided, ensuring they are not already in the unit_ids list and match the
+          number of splits.
+        """
         unit_ids = list(values["unit_ids"])
         splits = values.get("splits")
         if splits is None:
@@ -162,12 +180,12 @@ def check_splits(cls, values):
                     split_list.append(
                         {
                             "unit_id": unit_id,
-                            "split_mode": "indices",
-                            "split_indices": [list(indices) for indices in split_data],
+                            "mode": "indices",
+                            "indices": [list(indices) for indices in split_data],
                         }
                     )
                 else:
-                    split_list.append({"unit_id": unit_id, "split_mode": "labels", "split_labels": list(split_data)})
+                    split_list.append({"unit_id": unit_id, "mode": "labels", "labels": list(split_data)})
             splits = split_list
 
         # Make a copy of the list
@@ -177,12 +195,12 @@ def check_splits(cls, values):
         for i, split in enumerate(splits):
             if isinstance(split, dict):
                 split = dict(split)
-                if "split_indices" in split:
-                    split["split_indices"] = [list(indices) for indices in split["split_indices"]]
-                if "split_labels" in split:
-                    split["split_labels"] = list(split["split_labels"])
-                if "split_new_unit_ids" in split:
-                    split["split_new_unit_ids"] = list(split["split_new_unit_ids"])
+                if "indices" in split:
+                    split["indices"] = [list(indices) for indices in split["indices"]]
+                if "labels" in split:
+                    split["labels"] = list(split["labels"])
+                if "new_unit_ids" in split:
+                    split["new_unit_ids"] = list(split["new_unit_ids"])
                 splits[i] = Split(**split)
 
         # Validate splits
@@ -192,36 +210,36 @@ def check_splits(cls, values):
                 raise ValueError(f"Split unit_id {split.unit_id} is not in the unit list")
 
             # Validate based on mode
-            if split.split_mode == "indices":
-                if split.split_indices is None:
-                    raise ValueError(f"Split unit {split.unit_id} has no split_indices defined")
-                if len(split.split_indices) < 1:
-                    raise ValueError(f"Split unit {split.unit_id} has empty split_indices")
+            if split.mode == "indices":
+                if split.indices is None:
+                    raise ValueError(f"Split unit {split.unit_id} has no indices defined")
+                if len(split.indices) < 1:
+                    raise ValueError(f"Split unit {split.unit_id} has empty indices")
                 # Check no duplicate indices
-                all_indices = list(chain.from_iterable(split.split_indices))
+                all_indices = list(chain.from_iterable(split.indices))
                 if len(all_indices) != len(set(all_indices)):
                     raise ValueError(f"Split unit {split.unit_id} has duplicate indices")
 
-            elif split.split_mode == "labels":
-                if split.split_labels is None:
-                    raise ValueError(f"Split unit {split.unit_id} has no split_labels defined")
-                if len(split.split_labels) == 0:
-                    raise ValueError(f"Split unit {split.unit_id} has empty split_labels")
+            elif split.mode == "labels":
+                if split.labels is None:
+                    raise ValueError(f"Split unit {split.unit_id} has no labels defined")
+                if len(split.labels) == 0:
+                    raise ValueError(f"Split unit {split.unit_id} has empty labels")
 
             # Validate new unit IDs
-            if split.split_new_unit_ids is not None:
-                if split.split_mode == "indices":
-                    if len(split.split_new_unit_ids) != len(split.split_indices):
+            if split.new_unit_ids is not None:
+                if split.mode == "indices":
+                    if len(split.new_unit_ids) != len(split.indices):
                         raise ValueError(
                             f"Number of new unit IDs does not match number of splits for unit {split.unit_id}"
                         )
-                elif split.split_mode == "labels":
-                    if len(split.split_new_unit_ids) != len(set(split.split_labels)):
+                elif split.mode == "labels":
+                    if len(split.new_unit_ids) != len(set(split.labels)):
                         raise ValueError(
                             f"Number of new unit IDs does not match number of unique labels for unit {split.unit_id}"
                         )
 
-                for new_id in split.split_new_unit_ids:
+                for new_id in split.new_unit_ids:
                     if new_id in unit_ids:
                         raise ValueError(f"New unit ID {new_id} is already in the unit list")
 
@@ -312,7 +330,7 @@ def validate_curation_dict(cls, values):
 
         labeled_unit_set = set([lbl.unit_id for lbl in values.manual_labels]) if values.manual_labels else set()
         merged_units_set = (
-            set(chain.from_iterable(merge.merge_unit_group for merge in values.merges)) if values.merges else set()
+            set(chain.from_iterable(merge.unit_ids for merge in values.merges)) if values.merges else set()
         )
         split_units_set = set(split.unit_id for split in values.splits) if values.splits else set()
         removed_set = set(values.removed) if values.removed else set()
@@ -329,7 +347,7 @@ def validate_curation_dict(cls, values):
             raise ValueError("Curation format: some removed units are not in the unit list")
 
         # Check for units being merged multiple times
-        all_merging_groups = [set(merge.merge_unit_group) for merge in values.merges] if values.merges else []
+        all_merging_groups = [set(merge.unit_ids) for merge in values.merges] if values.merges else []
         for gp_1, gp_2 in combinations(all_merging_groups, 2):
             if len(gp_1.intersection(gp_2)) != 0:
                 raise ValueError("Curation format: some units belong to multiple merge groups")

src/spikeinterface/curation/sortingview_curation.py

@@ -112,7 +112,7 @@ def apply_sortingview_curation(
         clean_merges = []
         for merge in curation_model.merges:
             clean_merge = []
-            for unit_id in merge.merge_unit_group:
+            for unit_id in merge.unit_ids:
                 if unit_id not in curation_model.removed:
                     clean_merge.append(unit_id)
             if len(clean_merge) > 1:

src/spikeinterface/curation/tests/test_curation_model.py

@@ -98,15 +98,15 @@ def test_merge_units():
 
     model = CurationModel(**valid_merge)
     assert len(model.merges) == 2
-    assert model.merges[0].merge_new_unit_id == 5
-    assert model.merges[1].merge_new_unit_id == 6
+    assert model.merges[0].new_unit_id == 5
+    assert model.merges[1].new_unit_id == 6
 
     # Test dictionary format
     valid_merge_dict = {"format_version": "2", "unit_ids": [1, 2, 3, 4], "merges": {5: [1, 2], 6: [3, 4]}}
 
     model = CurationModel(**valid_merge_dict)
     assert len(model.merges) == 2
-    merge_new_ids = {merge.merge_new_unit_id for merge in model.merges}
+    merge_new_ids = {merge.new_unit_id for merge in model.merges}
     assert merge_new_ids == {5, 6}
 
     # Test list format
@@ -158,8 +158,8 @@ def test_split_units():
 
     model = CurationModel(**valid_split_indices)
     assert len(model.splits) == 1
-    assert model.splits[0].split_mode == "indices"
-    assert len(model.splits[0].split_indices) == 2
+    assert model.splits[0].mode == "indices"
+    assert len(model.splits[0].indices) == 2
 
     # Test labels mode with list format
     valid_split_labels = {
@@ -172,8 +172,8 @@ def test_split_units():
 
     model = CurationModel(**valid_split_labels)
     assert len(model.splits) == 1
-    assert model.splits[0].split_mode == "labels"
-    assert len(set(model.splits[0].split_labels)) == 3
+    assert model.splits[0].mode == "labels"
+    assert len(set(model.splits[0].labels)) == 3
 
     # Test dictionary format with indices
     valid_split_dict = {
@@ -187,7 +187,7 @@ def test_split_units():
 
     model = CurationModel(**valid_split_dict)
     assert len(model.splits) == 2
-    assert all(split.split_mode == "indices" for split in model.splits)
+    assert all(split.mode == "indices" for split in model.splits)
 
     # Test invalid unit ID
     invalid_unit_id = {