Add pydantic curation model, improve merging rules, and add splitting model

pyproject.toml

@@ -28,6 +28,7 @@ dependencies = [
     "neo>=0.14.0",
     "probeinterface>=0.2.23",
     "packaging",
+    "pydantic",
 ]
 
 [build-system]

src/spikeinterface/core/sorting_tools.py

@@ -232,8 +232,13 @@ def random_spikes_selection(
 
 
 def apply_merges_to_sorting(
-    sorting, merge_unit_groups, new_unit_ids=None, censor_ms=None, return_extra=False, new_id_strategy="append"
-):
+    sorting: BaseSorting,
+    merge_unit_groups: list[list[int | str]] | list[tuple[int | str]],
+    new_unit_ids: list[int | str] | None = None,
+    censor_ms: float | None = None,
+    return_extra: bool = False,
+    new_id_strategy: str = "append",
+) -> NumpySorting | tuple[NumpySorting, np.ndarray, list[int | str]]:
     """
     Apply a resolved representation of the merges to a sorting object.
 
@@ -245,9 +250,9 @@ def apply_merges_to_sorting(
 
     Parameters
     ----------
-    sorting : Sorting
+    sorting : BaseSorting
         The Sorting object to apply merges.
-    merge_unit_groups : list/tuple of lists/tuples
+    merge_unit_groups : list of lists/tuples
         A list of lists for every merge group. Each element needs to have at least two elements (two units to merge),
         but it can also have more (merge multiple units at once).
     new_unit_ids : list | None, default: None

src/spikeinterface/core/sortinganalyzer.py

@@ -1140,18 +1140,18 @@ def remove_units(self, remove_unit_ids, format="memory", folder=None) -> "Sortin
 
     def merge_units(
         self,
-        merge_unit_groups,
-        new_unit_ids=None,
-        censor_ms=None,
-        merging_mode="soft",
-        sparsity_overlap=0.75,
-        new_id_strategy="append",
-        return_new_unit_ids=False,
-        format="memory",
-        folder=None,
-        verbose=False,
+        merge_unit_groups: list[list[str | int]] | list[tuple[str | int]],
+        new_unit_ids: list[int | str] | None = None,
+        censor_ms: float | None = None,
+        merging_mode: str = "soft",
+        sparsity_overlap: float = 0.75,
+        new_id_strategy: str = "append",
+        return_new_unit_ids: bool = False,
+        format: str = "memory",
+        folder: Path | str | None = None,
+        verbose: bool = False,
         **job_kwargs,
-    ) -> "SortingAnalyzer":
+    ) -> "SortingAnalyzer | tuple[SortingAnalyzer, list[int | str]]":
         """
         This method is equivalent to `save_as()` but with a list of merges that have to be achieved.
         Merges units by creating a new SortingAnalyzer object with the appropriate merges