Repository: SpikeInterface/spikeinterface
Pull request #3817 (2025-03-28 22:21:57)

Add splitting functionality to curation and `SortingAnalyzer`
Depends on #3760 

Implements splitting functionality in:
- `sorting_tools`: [`apply_splits_to_sorting`](https://github.com/alejoe91/spikeinterface/blob/splitting-curation/src/spikeinterface/core/sorting_tools.py#L452)
- `SortingAnalyzer`: [`.split_units()`](https://github.com/alejoe91/spikeinterface/blob/splitting-curation/src/spikeinterface/core/sortinganalyzer.py#L1299)
- `AnalyzerExtension`: [`_split_extension_data`](https://github.com/alejoe91/spikeinterface/blob/splitting-curation/src/spikeinterface/core/sortinganalyzer.py#L2182) (implemented function for all extensions and [extended tests](https://github.com/alejoe91/spikeinterface/blob/splitting-curation/src/spikeinterface/postprocessing/tests/test_multi_extensions.py#L224))


Patch:
diff --git a/src/spikeinterface/core/__init__.py b/src/spikeinterface/core/__init__.py
index fb2e173b3..fdcfc73c2 100644
--- a/src/spikeinterface/core/__init__.py
+++ b/src/spikeinterface/core/__init__.py
@@ -109,7 +109,12 @@ from .recording_tools import (
     get_chunk_with_margin,
     order_channels_by_depth,
 )
-from .sorting_tools import spike_vector_to_spike_trains, random_spikes_selection, apply_merges_to_sorting
+from .sorting_tools import (
+    spike_vector_to_spike_trains,
+    random_spikes_selection,
+    apply_merges_to_sorting,
+    apply_splits_to_sorting,
+)
 
 from .waveform_tools import extract_waveforms_to_buffers, estimate_templates, estimate_templates_with_accumulator
 from .snippets_tools import snippets_from_sorting
diff --git a/src/spikeinterface/core/analyzer_extension_core.py b/src/spikeinterface/core/analyzer_extension_core.py
index e237342f5..4a045c255 100644
--- a/src/spikeinterface/core/analyzer_extension_core.py
+++ b/src/spikeinterface/core/analyzer_extension_core.py
@@ -9,6 +9,7 @@ It also implements:
   * ComputeNoiseLevels which is very convenient to have
 """
 
+import warnings
 import numpy as np
 
 from .sortinganalyzer import AnalyzerExtension, register_result_extension
@@ -92,6 +93,11 @@ class ComputeRandomSpikes(AnalyzerExtension):
             new_data["random_spikes_indices"] = np.flatnonzero(selected_mask[keep_mask])
         return new_data
 
+    def _split_extension_data(self, split_units, new_unit_ids, new_sorting_analyzer, verbose=False, **job_kwargs):
+        new_data = dict()
+        new_data["random_spikes_indices"] = self.data["random_spikes_indices"].copy()
+        return new_data
+
     def _get_data(self):
         return self.data["random_spikes_indices"]
 
@@ -243,8 +249,6 @@ class ComputeWaveforms(AnalyzerExtension):
     def _merge_extension_data(
         self, merge_unit_groups, new_unit_ids, new_sorting_analyzer, keep_mask=None, verbose=False, **job_kwargs
     ):
-        new_data = dict()
-
         waveforms = self.data["waveforms"]
         some_spikes = self.sorting_analyzer.get_extension("random_spikes").get_random_spikes()
         if keep_mask is not None:
@@ -275,6 +279,11 @@ class ComputeWaveforms(AnalyzerExtension):
 
         return dict(waveforms=waveforms)
 
+    def _split_extension_data(self, split_units, new_unit_ids, new_sorting_analyzer, verbose=False, **job_kwargs):
+        # splitting only affects random spikes, not waveforms
+        new_data = dict(waveforms=self.data["waveforms"].copy())
+        return new_data
+
     def get_waveforms_one_unit(self, unit_id, force_dense: bool = False):
         """
         Returns the waveforms of a unit id.
@@ -554,6 +563,49 @@ class ComputeTemplates(AnalyzerExtension):
 
         return new_data
 
+    def _split_extension_data(self, split_units, new_unit_ids, new_sorting_analyzer, verbose=False, **job_kwargs):
+        if not new_sorting_analyzer.has_extension("waveforms"):
+            warnings.warn(
+                "Splitting templates without the 'waveforms' extension will simply copy the template of the unit that "
+                "was split to the new split units. This is not recommended and may lead to incorrect results. It is "
+                "recommended to compute the 'waveforms' extension before splitting, or to use 'hard' splitting mode.",
+            )
+        new_data = dict()
+        for operator, arr in self.data.items():
+            # we first copy the unsplit units
+            new_array = np.zeros((len(new_sorting_analyzer.unit_ids), arr.shape[1], arr.shape[2]), dtype=arr.dtype)
+            new_analyzer_unit_ids = list(new_sorting_analyzer.unit_ids)
+            unsplit_unit_ids = [unit_id for unit_id in self.sorting_analyzer.unit_ids if unit_id not in split_units]
+            new_indices = np.array([new_analyzer_unit_ids.index(unit_id) for unit_id in unsplit_unit_ids])
+            old_indices = self.sorting_analyzer.sorting.ids_to_indices(unsplit_unit_ids)
+            new_array[new_indices, ...] = arr[old_indices, ...]
+
+            for split_unit_id, new_splits in zip(split_units, new_unit_ids):
+                if new_sorting_analyzer.has_extension("waveforms"):
+                    for new_unit_id in new_splits:
+                        split_unit_index = new_sorting_analyzer.sorting.id_to_index(new_unit_id)
+                        wfs = new_sorting_analyzer.get_extension("waveforms").get_waveforms_one_unit(
+                            new_unit_id, force_dense=True
+                        )
+
+                        if operator == "average":
+                            arr = np.average(wfs, axis=0)
+                        elif operator == "std":
+                            arr = np.std(wfs, axis=0)
+                        elif operator == "median":
+                            arr = np.median(wfs, axis=0)
+                        elif "percentile" in operator:
+                            _, percentile = operator.splot("_")
+                            arr = np.percentile(wfs, float(percentile), axis=0)
+                        new_array[split_unit_index, ...] = arr
+                else:
+                    split_unit_index = self.sorting_analyzer.sorting.id_to_index(split_unit_id)
+                    old_template = arr[split_unit_index, ...]
+                    new_indices = new_sorting_analyzer.sorting.ids_to_indices(new_splits)
+                    new_array[new_indices, ...] = np.tile(old_template, (len(new_splits), 1, 1))
+            new_data[operator] = new_array
+        return new_data
+
     def _get_data(self, operator="average", percentile=None, outputs="numpy"):
         if operator != "percentile":
             key = operator
@@ -727,6 +779,10 @@ class ComputeNoiseLevels(AnalyzerExtension):
         # this does not depend on units
         return self.data.copy()
 
+    def _split_extension_data(self, split_units, new_unit_ids, new_sorting_analyzer, verbose=False, **job_kwargs):
+        # this does not depend on units
+        return self.data.copy()
+
     def _run(self, verbose=False, **job_kwargs):
         self.data["noise_levels"] = get_noise_levels(
             self.sorting_analyzer.recording,
diff --git a/src/spikeinterface/core/sorting_tools.py b/src/spikeinterface/core/sorting_tools.py
index d5a899b9e..c74c76d01 100644
--- a/src/spikeinterface/core/sorting_tools.py
+++ b/src/spikeinterface/core/sorting_tools.py
@@ -1,11 +1,13 @@
 from __future__ import annotations
 
+import warnings
 import importlib.util
 
 import numpy as np
 
-from .basesorting import BaseSorting
-from .numpyextractors import NumpySorting
+from spikeinterface.core.base import BaseExtractor
+from spikeinterface.core.basesorting import BaseSorting
+from spikeinterface.core.numpyextractors import NumpySorting
 
 numba_spec = importlib.util.find_spec("numba")
 if numba_spec is not None:
@@ -226,6 +228,7 @@ def random_spikes_selection(
     return random_spikes_indices
 
 
+### MERGING ZONE ###
 def apply_merges_to_sorting(
     sorting: BaseSorting,
     merge_unit_groups: list[list[int | str]] | list[tuple[int | str]],
@@ -319,12 +322,69 @@ def apply_merges_to_sorting(
                 keep_mask[group_indices[inds + 1]] = False
 
     spikes = spikes[keep_mask]
-    sorting = NumpySorting(spikes, sorting.sampling_frequency, all_unit_ids)
+    merge_sorting = NumpySorting(spikes, sorting.sampling_frequency, all_unit_ids)
+    set_properties_after_merging(merge_sorting, sorting, merge_unit_groups, new_unit_ids=new_unit_ids)
 
     if return_extra:
-        return sorting, keep_mask, new_unit_ids
+        return merge_sorting, keep_mask, new_unit_ids
     else:
-        return sorting
+        return merge_sorting
+
+
+def set_properties_after_merging(
+    sorting_post_merge: BaseSorting,
+    sorting_pre_merge: BaseSorting,
+    merge_unit_groups: list[list[int | str]],
+    new_unit_ids: list[int | str],
+):
+    """
+    Add properties to the merge sorting object after merging units.
+    The properties of the merged units are propagated only if they are the same
+    for all units in the merge group.
+
+    Parameters
+    ----------
+    sorting_post_merge : BaseSorting
+        The Sorting object after merging units.
+    sorting_pre_merge : BaseSorting
+        The Sorting object before merging units.
+    merge_unit_groups : list
+        The groups of unit ids that were merged.
+    new_unit_ids : list
+        A list of new unit_ids for each merge.
+    """
+    prop_keys = sorting_pre_merge.get_property_keys()
+    pre_unit_ids = sorting_pre_merge.unit_ids
+    post_unit_ids = sorting_post_merge.unit_ids
+
+    kept_unit_ids = post_unit_ids[np.isin(post_unit_ids, pre_unit_ids)]
+    keep_pre_inds = sorting_pre_merge.ids_to_indices(kept_unit_ids)
+    keep_post_inds = sorting_post_merge.ids_to_indices(kept_unit_ids)
+
+    for key in prop_keys:
+        parent_values = sorting_pre_merge.get_property(key)
+
+        # propagate keep values
+        shape = (len(sorting_post_merge.unit_ids),) + parent_values.shape[1:]
+        new_values = np.empty(shape=shape, dtype=parent_values.dtype)
+        new_values[keep_post_inds] = parent_values[keep_pre_inds]
+        for new_id, merge_group in zip(new_unit_ids, merge_unit_groups):
+            merged_indices = sorting_pre_merge.ids_to_indices(merge_group)
+            merge_values = parent_values[merged_indices]
+            same_property_values = np.all([np.array_equal(m, merge_values[0]) for m in merge_values[1:]])
+            new_index = sorting_post_merge.id_to_index(new_id)
+            if same_property_values:
+                # and new values only if they are all similar
+                new_values[new_index] = merge_values[0]
+            else:
+                default_missing_values = BaseExtractor.default_missing_property_values
+                new_values[new_index] = default_missing_values[parent_values.dtype.kind]
+        sorting_post_merge.set_property(key, new_values)
+
+    # set is_merged property
+    is_merged = np.ones(len(sorting_post_merge.unit_ids), dtype=bool)
+    is_merged[keep_post_inds] = False
+    sorting_post_merge.set_property("is_merged", is_merged)
 
 
 def _get_ids_after_merging(old_unit_ids, merge_unit_groups, new_unit_ids):
@@ -375,7 +435,7 @@ def _get_ids_after_merging(old_unit_ids, merge_unit_groups, new_unit_ids):
 
 def generate_unit_ids_for_merge_group(old_unit_ids, merge_unit_groups, new_unit_ids=None, new_id_strategy="append"):
     """
-    Function to generate new units ids during a merging procedure. If new_units_ids
+    Function to generate new units ids during a merging procedure. If `new_units_ids`
     are provided, it will return these unit ids, checking that they have the the same
     length as `merge_unit_groups`.
 
@@ -440,3 +500,301 @@ def generate_unit_ids_for_merge_group(old_unit_ids, merge_unit_groups, new_unit_
             raise ValueError("wrong new_id_strategy")
 
     return new_unit_ids
+
+
+### SPLITTING ZONE ###
+def apply_splits_to_sorting(
+    sorting: BaseSorting,
+    unit_splits: dict[int | str, list[list[int | str]]],
+    new_unit_ids: list[list[int | str]] | None = None,
+    return_extra: bool = False,
+    new_id_strategy: str = "append",
+):
+    """
+    Apply a the splits to a sorting object.
+
+    This function is not lazy and creates a new NumpySorting with a compact spike_vector as fast as possible.
+    The `unit_splits` should be a dict with the unit ids as keys and a list of lists of spike indices as values.
+    For each split, the list of spike indices should contain the indices of the spikes to be assigned to each split and
+    it should be complete (i.e. the sum of the lengths of the sublists must equal the number of spikes in the unit).
+    If `new_unit_ids` is not None, it will use these new unit ids for the split units.
+    If `new_unit_ids` is None, it will generate new unit ids according to `new_id_strategy`.
+
+    Parameters
+    ----------
+    sorting : BaseSorting
+        The Sorting object to apply splits.
+    unit_splits : dict
+        A dictionary with the split unit id as key and a list of lists of spike indices for each split.
+        The split indices for each unit MUST be a list of lists, where each sublist (at least two) contains the
+        indices of the spikes to be assigned to the each split. The sum of the lengths of the sublists must equal
+        the number of spikes in the unit.
+    new_unit_ids : list | None, default: None
+        List of new unit_ids for each split. If given, it needs to have the same length as `unit_splits`.
+        and each element must have the same length as the corresponding list of split indices.
+        If None, new ids will be generated.
+    return_extra : bool, default: False
+        If True, also return the new_unit_ids.
+    new_id_strategy : "append" | "split", default: "append"
+        The strategy that should be used, if `new_unit_ids` is None, to create new unit_ids.
+
+            * "append" : new_units_ids will be added at the end of max(sorging.unit_ids)
+            * "split" : new_unit_ids will be the created as {split_unit_id]-{split_number}
+                        (e.g. when splitting unit "13" in 2: "13-0" / "13-1").
+                        Only works if unit_ids are str otherwise switch to "append"
+
+    Returns
+    -------
+    sorting : NumpySorting
+        The newly create sorting with the split units.
+    """
+    check_unit_splits_consistency(unit_splits, sorting)
+    spikes = sorting.to_spike_vector().copy()
+
+    # here we assume that unit_splits split_indices are already full.
+    # this is true when running via apply_curation
+
+    new_unit_ids = generate_unit_ids_for_split(
+        sorting.unit_ids, unit_splits, new_unit_ids=new_unit_ids, new_id_strategy=new_id_strategy
+    )
+    all_unit_ids = _get_ids_after_splitting(sorting.unit_ids, unit_splits, new_unit_ids)
+    all_unit_ids = list(all_unit_ids)
+
+    num_seg = sorting.get_num_segments()
+    seg_lims = np.searchsorted(spikes["segment_index"], np.arange(0, num_seg + 2))
+    segment_slices = [(seg_lims[i], seg_lims[i + 1]) for i in range(num_seg)]
+
+    # using this function vaoid to use the mask approach and simplify a lot the algo
+    spike_vector_list = [spikes[s0:s1] for s0, s1 in segment_slices]
+    spike_indices = spike_vector_to_indices(spike_vector_list, sorting.unit_ids, absolute_index=True)
+
+    for unit_id in sorting.unit_ids:
+        if unit_id in unit_splits:
+            split_indices = unit_splits[unit_id]
+            new_split_ids = new_unit_ids[list(unit_splits.keys()).index(unit_id)]
+
+            for split, new_unit_id in zip(split_indices, new_split_ids):
+                new_unit_index = all_unit_ids.index(new_unit_id)
+                # split_indices are a concatenation across segments with absolute indices
+                # so we need to concatenate the spike indices across segments
+                spike_indices_unit = np.concatenate(
+                    [spike_indices[segment_index][unit_id] for segment_index in range(num_seg)]
+                )
+                spikes["unit_index"][spike_indices_unit[split]] = new_unit_index
+        else:
+            new_unit_index = all_unit_ids.index(unit_id)
+            for segment_index in range(num_seg):
+                spike_inds = spike_indices[segment_index][unit_id]
+                spikes["unit_index"][spike_inds] = new_unit_index
+    split_sorting = NumpySorting(spikes, sorting.sampling_frequency, all_unit_ids)
+    set_properties_after_splits(
+        split_sorting,
+        sorting,
+        list(unit_splits.keys()),
+        new_unit_ids=new_unit_ids,
+    )
+
+    if return_extra:
+        return split_sorting, new_unit_ids
+    else:
+        return split_sorting
+
+
+def set_properties_after_splits(
+    sorting_post_split: BaseSorting,
+    sorting_pre_split: BaseSorting,
+    split_unit_ids: list[int | str],
+    new_unit_ids: list[list[int | str]],
+):
+    """
+    Add properties to the split sorting object after splitting units.
+    The properties of the split units are propagated to the new split units.
+
+    Parameters
+    ----------
+    sorting_post_split : BaseSorting
+        The Sorting object after splitting units.
+    sorting_pre_split : BaseSorting
+        The Sorting object before splitting units.
+    split_unit_ids : list
+        The unit ids that were split.
+    new_unit_ids : list
+        A list of new unit_ids for each split.
+    """
+    prop_keys = sorting_pre_split.get_property_keys()
+    pre_unit_ids = sorting_pre_split.unit_ids
+    post_unit_ids = sorting_post_split.unit_ids
+
+    kept_unit_ids = post_unit_ids[np.isin(post_unit_ids, pre_unit_ids)]
+    keep_pre_inds = sorting_pre_split.ids_to_indices(kept_unit_ids)
+    keep_post_inds = sorting_post_split.ids_to_indices(kept_unit_ids)
+
+    for key in prop_keys:
+        parent_values = sorting_pre_split.get_property(key)
+
+        # propagate keep values
+        shape = (len(sorting_post_split.unit_ids),) + parent_values.shape[1:]
+        new_values = np.empty(shape=shape, dtype=parent_values.dtype)
+        new_values[keep_post_inds] = parent_values[keep_pre_inds]
+        for split_unit, new_split_ids in zip(split_unit_ids, new_unit_ids):
+            split_index = sorting_pre_split.id_to_index(split_unit)
+            split_value = parent_values[split_index]
+            # propagate the split value to all new unit ids
+            new_unit_indices = sorting_post_split.ids_to_indices(new_split_ids)
+            new_values[new_unit_indices] = split_value
+        sorting_post_split.set_property(key, new_values)
+
+    # set is_merged property
+    is_split = np.ones(len(sorting_post_split.unit_ids), dtype=bool)
+    is_split[keep_post_inds] = False
+    sorting_post_split.set_property("is_split", is_split)
+
+
+def generate_unit_ids_for_split(old_unit_ids, unit_splits, new_unit_ids=None, new_id_strategy="append"):
+    """
+    Function to generate new units ids during a splitting procedure. If `new_units_ids`
+    are provided, it will return these unit ids, checking that they are consistent with
+    `unit_splits`.
+
+    Parameters
+    ----------
+    old_unit_ids : np.array
+        The old unit_ids.
+    unit_splits : dict
+
+    new_unit_ids : list | None, default: None
+        Optional new unit_ids for split units. If given, it needs to have the same length as `merge_unit_groups`.
+        If None, new ids will be generated.
+    new_id_strategy : "append" | "split", default: "append"
+        The strategy that should be used, if `new_unit_ids` is None, to create new unit_ids.
+
+            * "append" : new_units_ids will be added at the end of max(sorging.unit_ids)
+            * "split" : new_unit_ids will be the created as {split_unit_id]-{split_number}
+                        (e.g. when splitting unit "13" in 2: "13-0" / "13-1").
+                        Only works if unit_ids are str otherwise switch to "append"
+
+    Returns
+    -------
+    new_unit_ids : list of lists
+        The new units_ids associated with the merges.
+    """
+    assert new_id_strategy in ["append", "split"], "new_id_strategy should be 'append' or 'split'"
+    old_unit_ids = np.asarray(old_unit_ids)
+
+    if new_unit_ids is not None:
+        for split_unit, new_split_ids in zip(unit_splits.values(), new_unit_ids):
+            # then only doing a consistency check
+            assert len(split_unit) == len(new_split_ids), "new_unit_ids should have the same len as unit_splits.values"
+            # new_unit_ids can also be part of old_unit_ids only inside the same group:
+            assert all(
+                new_split_id not in old_unit_ids for new_split_id in new_split_ids
+            ), "new_unit_ids already exists but outside the split groups"
+    else:
+        dtype = old_unit_ids.dtype
+        if np.issubdtype(dtype, np.integer) and new_id_strategy == "split":
+            warnings.warn("new_id_strategy 'split' is not compatible with integer unit_ids. Switching to 'append'.")
+            new_id_strategy = "append"
+
+        new_unit_ids = []
+        current_unit_ids = old_unit_ids.copy()
+        for unit_to_split, split_indices in unit_splits.items():
+            num_splits = len(split_indices)
+            # select new_unit_ids greater that the max id, event greater than the numerical str ids
+            if new_id_strategy == "append":
+                if np.issubdtype(dtype, np.character):
+                    # dtype str
+                    if all(p.isdigit() for p in current_unit_ids):
+                        # All str are digit : we can generate a max
+                        m = max(int(p) for p in current_unit_ids) + 1
+                        new_units_for_split = [str(m + i) for i in range(num_splits)]
+                    else:
+                        # we cannot automatically find new names
+                        new_units_for_split = [f"{unit_to_split}-split{i}" for i in range(num_splits)]
+                else:
+                    # dtype int
+                    new_units_for_split = list(max(current_unit_ids) + 1 + np.arange(num_splits, dtype=dtype))
+                # we append the new split unit ids to continue to increment the max id
+                current_unit_ids = np.concatenate([current_unit_ids, new_units_for_split])
+            elif new_id_strategy == "split":
+                # we made sure that dtype is not integer
+                new_units_for_split = [f"{unit_to_split}-{i}" for i in np.arange(len(split_indices))]
+            new_unit_ids.append(new_units_for_split)
+
+    return new_unit_ids
+
+
+def check_unit_splits_consistency(unit_splits, sorting):
+    """
+    Function to check the consistency of unit_splits indices with the sorting object.
+    It checks that the split indices for each unit are a list of lists, where each sublist (at least two)
+    contains the indices of the spikes to be assigned to each split. The sum of the lengths
+    of the sublists must equal the number of spikes in the unit.
+
+    Parameters
+    ----------
+    unit_splits : dict
+        A dictionary with the split unit id as key and a list of numpy arrays or lists of spike indices for each split.
+    sorting : BaseSorting
+        The sorting object containing spike information.
+
+    Raises
+    ------
+    ValueError
+        If the unit_splits are not in the expected format or if the total number of spikes in the splits does not match
+        the number of spikes in the unit.
+    """
+    num_spikes = sorting.count_num_spikes_per_unit()
+    for unit_id, split_indices in unit_splits.items():
+        if not isinstance(split_indices, (list, np.ndarray)):
+            raise ValueError(f"unit_splits[{unit_id}] should be a list or numpy array, got {type(split_indices)}")
+        if not all(isinstance(indices, (list, np.ndarray)) for indices in split_indices):
+            raise ValueError(f"unit_splits[{unit_id}] should be a list of lists or numpy arrays")
+        if len(split_indices) < 2:
+            raise ValueError(f"unit_splits[{unit_id}] should have at least two splits")
+        total_spikes_in_split = sum(len(indices) for indices in split_indices)
+        if total_spikes_in_split != num_spikes[unit_id]:
+            raise ValueError(
+                f"Total spikes in unit {unit_id} split ({total_spikes_in_split}) does not match the number of spikes in the unit ({num_spikes[unit_id]})"
+            )
+
+
+def _get_ids_after_splitting(old_unit_ids, split_units, new_unit_ids):
+    """
+    Function to get the list of unique unit_ids after some splits, with given new_units_ids would
+    be provided.
+
+    Every new unit_id will be added at the end if not already present.
+
+    Parameters
+    ----------
+    old_unit_ids : np.array
+        The old unit_ids.
+    split_units : dict
+        A dict of split units. Each element needs to have at least two elements (two units to split).
+    new_unit_ids : list | None
+        A new unit_ids for split units. If given, it needs to have the same length as `split_units` values.
+
+    Returns
+    -------
+
+    all_unit_ids :  The unit ids in the split sorting
+        The units_ids that will be present after splits
+
+    """
+    old_unit_ids = np.asarray(old_unit_ids)
+    dtype = old_unit_ids.dtype
+    if dtype.kind == "U":
+        # the new dtype can be longer
+        dtype = "U"
+
+    assert len(new_unit_ids) == len(split_units), "new_unit_ids should have the same len as merge_unit_groups"
+    for new_unit_in_split, unit_to_split in zip(new_unit_ids, split_units.keys()):
+        assert len(new_unit_in_split) == len(
+            split_units[unit_to_split]
+        ), "new_unit_ids should have the same len as split_units values"
+
+    all_unit_ids = list(old_unit_ids.copy())
+    for split_unit, split_new_units in zip(split_units, new_unit_ids):
+        all_unit_ids.remove(split_unit)
+        all_unit_ids.extend(split_new_units)
+    return np.array(all_unit_ids, dtype=dtype)
diff --git a/src/spikeinterface/core/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py
index 5339f8b7b..f8d11dd15 100644
--- a/src/spikeinterface/core/sortinganalyzer.py
+++ b/src/spikeinterface/core/sortinganalyzer.py
@@ -31,7 +31,13 @@ from .core_tools import (
     is_path_remote,
     clean_zarr_folder_name,
 )
-from .sorting_tools import generate_unit_ids_for_merge_group, _get_ids_after_merging
+from .sorting_tools import (
+    generate_unit_ids_for_merge_group,
+    generate_unit_ids_for_split,
+    check_unit_splits_consistency,
+    _get_ids_after_merging,
+    _get_ids_after_splitting,
+)
 from .job_tools import split_job_kwargs
 from .numpyextractors import NumpySorting
 from .sparsity import ChannelSparsity, estimate_sparsity
@@ -888,7 +894,7 @@ class SortingAnalyzer:
         else:
             return mergeable
 
-    def _save_or_select_or_merge(
+    def _save_or_select_or_merge_or_split(
         self,
         format="binary_folder",
         folder=None,
@@ -897,9 +903,12 @@ class SortingAnalyzer:
         censor_ms=None,
         merging_mode="soft",
         sparsity_overlap=0.75,
-        verbose=False,
-        new_unit_ids=None,
+        merge_new_unit_ids=None,
+        split_units=None,
+        splitting_mode="soft",
+        split_new_unit_ids=None,
         backend_options=None,
+        verbose=False,
         **job_kwargs,
     ) -> "SortingAnalyzer":
         """
@@ -921,12 +930,19 @@ class SortingAnalyzer:
             When merging units, any spikes violating this refractory period will be discarded.
         merging_mode : "soft" | "hard", default: "soft"
             How merges are performed. In the "soft" mode, merges will be approximated, with no smart merging
-            of the extension data.
+            of the extension data. In the "hard" mode, the extensions for merged units will be recomputed.
         sparsity_overlap : float, default 0.75
             The percentage of overlap that units should share in order to accept merges. If this criteria is not
             achieved, soft merging will not be performed.
-        new_unit_ids : list or None, default: None
+        merge_new_unit_ids : list or None, default: None
             The new unit ids for merged units. Required if `merge_unit_groups` is not None.
+        split_units : dict or None, default: None
+            A dictionary with the keys being the unit ids to split and the values being the split indices.
+        splitting_mode : "soft" | "hard", default: "soft"
+            How splits are performed. In the "soft" mode, splits will be approximated, with no smart splitting.
+            If `splitting_mode` is "hard", the extensons for split units willbe recomputed.
+        split_new_unit_ids : list or None, default: None
+            The new unit ids for split units. Required if `split_units` is not None.
         verbose : bool, default: False
             If True, output is verbose.
         backend_options : dict | None, default: None
@@ -949,38 +965,64 @@ class SortingAnalyzer:
         else:
             recording = None
 
-        if self.sparsity is not None and unit_ids is None and merge_unit_groups is None:
-            sparsity = self.sparsity
-        elif self.sparsity is not None and unit_ids is not None and merge_unit_groups is None:
-            sparsity_mask = self.sparsity.mask[np.isin(self.unit_ids, unit_ids), :]
-            sparsity = ChannelSparsity(sparsity_mask, unit_ids, self.channel_ids)
-        elif self.sparsity is not None and merge_unit_groups is not None:
-            all_unit_ids = unit_ids
-            sparsity_mask = np.zeros((len(all_unit_ids), self.sparsity.mask.shape[1]), dtype=bool)
-            mergeable, masks = self.are_units_mergeable(
-                merge_unit_groups,
-                sparsity_overlap=sparsity_overlap,
-                merging_mode=merging_mode,
-                return_masks=True,
-            )
+        has_removed = unit_ids is not None
+        has_merges = merge_unit_groups is not None
+        has_splits = split_units is not None
+        assert not has_merges if has_splits else True, "Cannot merge and split at the same time"
+
+        if self.sparsity is not None:
+            if not has_removed and not has_merges and not has_splits:
+                # no changes in units
+                sparsity = self.sparsity
+            elif has_removed and not has_merges and not has_splits:
+                # remove units
+                sparsity_mask = self.sparsity.mask[np.isin(self.unit_ids, unit_ids), :]
+                sparsity = ChannelSparsity(sparsity_mask, unit_ids, self.channel_ids)
+            elif has_merges:
+                # merge units
+                all_unit_ids = unit_ids
+                sparsity_mask = np.zeros((len(all_unit_ids), self.sparsity.mask.shape[1]), dtype=bool)
+                mergeable, masks = self.are_units_mergeable(
+                    merge_unit_groups,
+                    sparsity_overlap=sparsity_overlap,
+                    merging_mode=merging_mode,
+                    return_masks=True,
+                )
 
-            for unit_index, unit_id in enumerate(all_unit_ids):
-                if unit_id in new_unit_ids:
-                    merge_unit_group = tuple(merge_unit_groups[new_unit_ids.index(unit_id)])
-                    if not mergeable[merge_unit_group]:
-                        # unit groups can be not mergeable only in "soft" mode
-                        # see are_units_mergeable() function
-                        raise Exception(
-                            f"The sparsity of {merge_unit_group} do not overlap enough for a soft merge using "
-                            f"a sparsity threshold of {sparsity_overlap}. You can either lower the threshold "
-                            "or use a hard merge."
-                        )
-                    sparsity_mask[unit_index] = masks[merge_unit_group]
-                else:
-                    # This means that the unit is already in the previous sorting
-                    index = self.sorting.id_to_index(unit_id)
-                    sparsity_mask[unit_index] = self.sparsity.mask[index]
-            sparsity = ChannelSparsity(sparsity_mask, list(all_unit_ids), self.channel_ids)
+                for unit_index, unit_id in enumerate(all_unit_ids):
+                    if unit_id in merge_new_unit_ids:
+                        merge_unit_group = tuple(merge_unit_groups[merge_new_unit_ids.index(unit_id)])
+                        if not mergeable[merge_unit_group]:
+                            raise Exception(
+                                f"The sparsity of {merge_unit_group} do not overlap enough for a soft merge using "
+                                f"a sparsity threshold of {sparsity_overlap}. You can either lower the threshold or use "
+                                "a hard merge."
+                            )
+                        else:
+                            sparsity_mask[unit_index] = masks[merge_unit_group]
+                    else:
+                        # This means that the unit is already in the previous sorting
+                        index = self.sorting.id_to_index(unit_id)
+                        sparsity_mask[unit_index] = self.sparsity.mask[index]
+                sparsity = ChannelSparsity(sparsity_mask, list(all_unit_ids), self.channel_ids)
+            elif has_splits:
+                # split units
+                all_unit_ids = unit_ids
+                original_unit_ids = self.unit_ids
+                sparsity_mask = np.zeros((len(all_unit_ids), self.sparsity.mask.shape[1]), dtype=bool)
+                for unit_index, unit_id in enumerate(all_unit_ids):
+                    if unit_id not in original_unit_ids:
+                        # then it is a new unit
+                        # we assign the original sparsity
+                        for split_unit, new_unit_ids in zip(split_units, split_new_unit_ids):
+                            if unit_id in new_unit_ids:
+                                original_unit_index = self.sorting.id_to_index(split_unit)
+                                sparsity_mask[unit_index] = self.sparsity.mask[original_unit_index]
+                                break
+                    else:
+                        original_unit_index = self.sorting.id_to_index(unit_id)
+                        sparsity_mask[unit_index] = self.sparsity.mask[original_unit_index]
+                sparsity = ChannelSparsity(sparsity_mask, list(all_unit_ids), self.channel_ids)
         else:
             sparsity = None
 
@@ -990,25 +1032,33 @@ class SortingAnalyzer:
             # if the original sorting object is not available anymore (kilosort folder deleted, ....), take the copy
             sorting_provenance = self.sorting
 
-        if merge_unit_groups is None:
+        if merge_unit_groups is None and split_units is None:
             # when only some unit_ids then the sorting must be sliced
             # TODO check that unit_ids are in same order otherwise many extension do handle it properly!!!!
             sorting_provenance = sorting_provenance.select_units(unit_ids)
-        else:
+        elif merge_unit_groups is not None:
+            assert split_units is None, "split_units must be None when merge_unit_groups is None"
             from spikeinterface.core.sorting_tools import apply_merges_to_sorting
 
             sorting_provenance, keep_mask, _ = apply_merges_to_sorting(
                 sorting=sorting_provenance,
                 merge_unit_groups=merge_unit_groups,
-                new_unit_ids=new_unit_ids,
+                new_unit_ids=merge_new_unit_ids,
                 censor_ms=censor_ms,
                 return_extra=True,
             )
             if censor_ms is None:
                 # in this case having keep_mask None is faster instead of having a vector of ones
                 keep_mask = None
-            # TODO: sam/pierre would create a curation field / curation.json with the applied merges.
-            # What do you think?
+        elif split_units is not None:
+            assert merge_unit_groups is None, "merge_unit_groups must be None when split_units is not None"
+            from spikeinterface.core.sorting_tools import apply_splits_to_sorting
+
+            sorting_provenance = apply_splits_to_sorting(
+                sorting=sorting_provenance,
+                unit_splits=split_units,
+                new_unit_ids=split_new_unit_ids,
+            )
 
         backend_options = {} if backend_options is None else backend_options
 
@@ -1057,26 +1107,34 @@ class SortingAnalyzer:
         recompute_dict = {}
 
         for extension_name, extension in sorted_extensions.items():
-            if merge_unit_groups is None:
+            if merge_unit_groups is None and split_units is None:
                 # copy full or select
                 new_sorting_analyzer.extensions[extension_name] = extension.copy(
                     new_sorting_analyzer, unit_ids=unit_ids
                 )
-            else:
+            elif merge_unit_groups is not None:
                 # merge
                 if merging_mode == "soft":
                     new_sorting_analyzer.extensions[extension_name] = extension.merge(
                         new_sorting_analyzer,
                         merge_unit_groups=merge_unit_groups,
-                        new_unit_ids=new_unit_ids,
+                        new_unit_ids=merge_new_unit_ids,
                         keep_mask=keep_mask,
                         verbose=verbose,
                         **job_kwargs,
                     )
                 elif merging_mode == "hard":
                     recompute_dict[extension_name] = extension.params
+            else:
+                # split
+                if splitting_mode == "soft":
+                    new_sorting_analyzer.extensions[extension_name] = extension.split(
+                        new_sorting_analyzer, split_units=split_units, new_unit_ids=split_new_unit_ids, verbose=verbose
+                    )
+                elif splitting_mode == "hard":
+                    recompute_dict[extension_name] = extension.params
 
-        if merge_unit_groups is not None and merging_mode == "hard" and len(recompute_dict) > 0:
+        if len(recompute_dict) > 0:
             new_sorting_analyzer.compute_several_extensions(recompute_dict, save=True, verbose=verbose, **job_kwargs)
 
         return new_sorting_analyzer
@@ -1104,7 +1162,7 @@ class SortingAnalyzer:
         """
         if format == "zarr":
             folder = clean_zarr_folder_name(folder)
-        return self._save_or_select_or_merge(format=format, folder=folder, backend_options=backend_options)
+        return self._save_or_select_or_merge_or_split(format=format, folder=folder, backend_options=backend_options)
 
     def select_units(self, unit_ids, format="memory", folder=None) -> "SortingAnalyzer":
         """
@@ -1131,7 +1189,7 @@ class SortingAnalyzer:
         # TODO check that unit_ids are in same order otherwise many extension do handle it properly!!!!
         if format == "zarr":
             folder = clean_zarr_folder_name(folder)
-        return self._save_or_select_or_merge(format=format, folder=folder, unit_ids=unit_ids)
+        return self._save_or_select_or_merge_or_split(format=format, folder=folder, unit_ids=unit_ids)
 
     def remove_units(self, remove_unit_ids, format="memory", folder=None) -> "SortingAnalyzer":
         """
@@ -1159,7 +1217,7 @@ class SortingAnalyzer:
         unit_ids = self.unit_ids[~np.isin(self.unit_ids, remove_unit_ids)]
         if format == "zarr":
             folder = clean_zarr_folder_name(folder)
-        return self._save_or_select_or_merge(format=format, folder=folder, unit_ids=unit_ids)
+        return self._save_or_select_or_merge_or_split(format=format, folder=folder, unit_ids=unit_ids)
 
     def merge_units(
         self,
@@ -1224,28 +1282,18 @@ class SortingAnalyzer:
         assert merging_mode in ["soft", "hard"], "Merging mode should be either soft or hard"
 
         if len(merge_unit_groups) == 0:
-            # TODO I think we should raise an error or at least make a copy and not return itself
-            if return_new_unit_ids:
-                return self, []
-            else:
-                return self
+            raise ValueError("Merging requires at least one group of units to merge")
 
         for units in merge_unit_groups:
-            # TODO more checks like one units is only in one group
             if len(units) < 2:
                 raise ValueError("Merging requires at least two units to merge")
 
-        # TODO : no this function did not exists before
-        if not isinstance(merge_unit_groups[0], (list, tuple)):
-            # keep backward compatibility : the previous behavior was only one merge
-            merge_unit_groups = [merge_unit_groups]
-
         new_unit_ids = generate_unit_ids_for_merge_group(
             self.unit_ids, merge_unit_groups, new_unit_ids, new_id_strategy
         )
         all_unit_ids = _get_ids_after_merging(self.unit_ids, merge_unit_groups, new_unit_ids=new_unit_ids)
 
-        new_analyzer = self._save_or_select_or_merge(
+        new_analyzer = self._save_or_select_or_merge_or_split(
             format=format,
             folder=folder,
             merge_unit_groups=merge_unit_groups,
@@ -1254,7 +1302,79 @@ class SortingAnalyzer:
             merging_mode=merging_mode,
             sparsity_overlap=sparsity_overlap,
             verbose=verbose,
-            new_unit_ids=new_unit_ids,
+            merge_new_unit_ids=new_unit_ids,
+            **job_kwargs,
+        )
+        if return_new_unit_ids:
+            return new_analyzer, new_unit_ids
+        else:
+            return new_analyzer
+
+    def split_units(
+        self,
+        split_units: dict[list[str | int], list[int] | list[list[int]]],
+        new_unit_ids: list[list[int | str]] | None = None,
+        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 | tuple[SortingAnalyzer, list[int | str]]":
+        """
+        This method is equivalent to `save_as()` but with a list of splits that have to be achieved.
+        Split units by creating a new SortingAnalyzer object with the appropriate splits
+
+        Extensions are also updated to display the split `unit_ids`.
+
+        Parameters
+        ----------
+        split_units : dict
+            A dictionary with the keys being the unit ids to split and the values being the split indices.
+            The split indices for each unit MUST be a list of lists, where each sublist (at least two) contains the
+            indices of the spikes to be assigned to the each split. The sum of the lengths of the sublists must equal
+            the number of spikes in the unit.
+        new_unit_ids : None | list, default: None
+            A new unit_ids for split units. If given, it needs to have the same length as `merge_unit_groups`. If None,
+            merged units will have the first unit_id of every lists of merges
+        new_id_strategy : "append" | "split", default: "append"
+            The strategy that should be used, if `new_unit_ids` is None, to create new unit_ids.
+
+                * "append" : new_units_ids will be added at the end of max(sorting.unit_ids)
+                * "split" : new_unit_ids will be the original unit_id to split with -{subsplit}
+        return_new_unit_ids : bool, default False
+            Alse return new_unit_ids which are the ids of the new units.
+        folder : Path | None, default: None
+            The new folder where the analyzer with merged units is copied for `format` "binary_folder" or "zarr"
+        format : "memory" | "binary_folder" | "zarr", default: "memory"
+            The format of SortingAnalyzer
+        verbose : bool, default: False
+            Whether to display calculations (such as sparsity estimation)
+
+        Returns
+        -------
+        analyzer :  SortingAnalyzer
+            The newly create `SortingAnalyzer` with the selected units
+        """
+
+        if format == "zarr":
+            folder = clean_zarr_folder_name(folder)
+
+        if len(split_units) == 0:
+            raise ValueError("Splitting requires at least one unit to split")
+
+        check_unit_splits_consistency(split_units, self.sorting)
+
+        new_unit_ids = generate_unit_ids_for_split(self.unit_ids, split_units, new_unit_ids, new_id_strategy)
+        all_unit_ids = _get_ids_after_splitting(self.unit_ids, split_units, new_unit_ids=new_unit_ids)
+
+        new_analyzer = self._save_or_select_or_merge_or_split(
+            format=format,
+            folder=folder,
+            split_units=split_units,
+            unit_ids=all_unit_ids,
+            verbose=verbose,
+            split_new_unit_ids=new_unit_ids,
             **job_kwargs,
         )
         if return_new_unit_ids:
@@ -1266,7 +1386,7 @@ class SortingAnalyzer:
         """
         Create a a copy of SortingAnalyzer with format "memory".
         """
-        return self._save_or_select_or_merge(format="memory", folder=None)
+        return self._save_or_select_or_merge_or_split(format="memory", folder=None)
 
     def is_read_only(self) -> bool:
         if self.format == "memory":
@@ -2071,6 +2191,10 @@ class AnalyzerExtension:
         # must be implemented in subclass
         raise NotImplementedError
 
+    def _split_extension_data(self, split_units, new_unit_ids, new_sorting_analyzer, verbose=False, **job_kwargs):
+        # must be implemented in subclass
+        raise NotImplementedError
+
     def _get_pipeline_nodes(self):
         # must be implemented in subclass only if use_nodepipeline=True
         raise NotImplementedError
@@ -2307,6 +2431,23 @@ class AnalyzerExtension:
         new_extension.save()
         return new_extension
 
+    def split(
+        self,
+        new_sorting_analyzer,
+        split_units,
+        new_unit_ids,
+        verbose=False,
+        **job_kwargs,
+    ):
+        new_extension = self.__class__(new_sorting_analyzer)
+        new_extension.params = self.params.copy()
+        new_extension.data = self._split_extension_data(
+            split_units, new_unit_ids, new_sorting_analyzer, verbose=verbose, **job_kwargs
+        )
+        new_extension.run_info = copy(self.run_info)
+        new_extension.save()
+        return new_extension
+
     def run(self, save=True, **kwargs):
         if save and not self.sorting_analyzer.is_read_only():
             # NB: this call to _save_params() also resets the folder or zarr group
diff --git a/src/spikeinterface/curation/curation_format.py b/src/spikeinterface/curation/curation_format.py
index bcd6bd9a4..baea4399f 100644
--- a/src/spikeinterface/curation/curation_format.py
+++ b/src/spikeinterface/curation/curation_format.py
@@ -1,8 +1,9 @@
 from __future__ import annotations
 
 import numpy as np
+from itertools import chain
 
-from spikeinterface.core import BaseSorting, SortingAnalyzer, apply_merges_to_sorting
+from spikeinterface.core import BaseSorting, SortingAnalyzer, apply_merges_to_sorting, apply_splits_to_sorting
 from spikeinterface.curation.curation_model import CurationModel
 
 
@@ -100,34 +101,45 @@ def curation_label_to_dataframe(curation_dict_or_model: dict | CurationModel):
 
 
 def apply_curation_labels(
-    sorting: BaseSorting, new_unit_ids: list[int, str], curation_dict_or_model: dict | CurationModel
+    sorting_or_analyzer: BaseSorting | SortingAnalyzer, curation_dict_or_model: dict | CurationModel
 ):
     """
-    Apply manual labels after merges.
+    Apply manual labels after merges/splits.
 
     Rules:
-      * label for non merge is applied first
+      * label for non merged units is applied first
       * for merged group, when exclusive=True, if all have the same label then this label is applied
       * for merged group, when exclusive=False, if one unit has the label then the new one have also it
+      * for split units, the original label is applied to all split units
     """
     if isinstance(curation_dict_or_model, dict):
         curation_model = CurationModel(**curation_dict_or_model)
     else:
         curation_model = curation_dict_or_model
 
+    if isinstance(sorting_or_analyzer, BaseSorting):
+        sorting = sorting_or_analyzer
+    else:
+        sorting = sorting_or_analyzer.sorting
+
     # Please note that manual_labels is done on the unit_ids before the merge!!!
     manual_labels = curation_label_to_vectors(curation_model)
 
-    # apply on non merged
+    # apply on non merged / split
+    merge_new_unit_ids = [m.new_unit_id for m in curation_model.merges]
+    split_new_unit_ids = [m.new_unit_ids for m in curation_model.splits]
+    split_new_unit_ids = list(chain(*split_new_unit_ids))
+
+    merged_split_units = merge_new_unit_ids + split_new_unit_ids
     for key, values in manual_labels.items():
         all_values = np.zeros(sorting.unit_ids.size, dtype=values.dtype)
         for unit_ind, unit_id in enumerate(sorting.unit_ids):
-            if unit_id not in new_unit_ids:
+            if unit_id not in merged_split_units:
                 ind = list(curation_model.unit_ids).index(unit_id)
                 all_values[unit_ind] = values[ind]
         sorting.set_property(key, all_values)
 
-    for new_unit_id, merge in zip(new_unit_ids, curation_model.merges):
+    for new_unit_id, merge in zip(merge_new_unit_ids, curation_model.merges):
         old_group_ids = merge.unit_ids
         for label_key, label_def in curation_model.label_definitions.items():
             if label_def.exclusive:
@@ -141,7 +153,6 @@ def apply_curation_labels(
                     # all group has the same label or empty
                     sorting.set_property(key, values=group_values[:1], ids=[new_unit_id])
             else:
-
                 for key in label_def.label_options:
                     group_values = []
                     for unit_id in old_group_ids:
@@ -151,6 +162,23 @@ def apply_curation_labels(
                     new_value = np.any(group_values)
                     sorting.set_property(key, values=[new_value], ids=[new_unit_id])
 
+    # splits
+    for split in curation_model.splits:
+        # propagate property of splut unit to new units
+        old_unit = split.unit_id
+        new_unit_ids = split.new_unit_ids
+        for label_key, label_def in curation_model.label_definitions.items():
+            if label_def.exclusive:
+                ind = list(curation_model.unit_ids).index(old_unit)
+                value = manual_labels[label_key][ind]
+                if value != "":
+                    sorting.set_property(label_key, values=[value] * len(new_unit_ids), ids=new_unit_ids)
+            else:
+                for key in label_def.label_options:
+                    ind = list(curation_model.unit_ids).index(old_unit)
+                    value = manual_labels[key][ind]
+                    sorting.set_property(key, values=[value] * len(new_unit_ids), ids=new_unit_ids)
+
 
 def apply_curation(
     sorting_or_analyzer: BaseSorting | SortingAnalyzer,
@@ -168,7 +196,8 @@ def apply_curation(
     Steps are done in this order:
       1. Apply removal using curation_dict["removed"]
       2. Apply merges using curation_dict["merges"]
-      3. Set labels using curation_dict["manual_labels"]
+      3. Apply splits using curation_dict["splits"]
+      4. Set labels using curation_dict["manual_labels"]
 
     A new Sorting or SortingAnalyzer (in memory) is returned.
     The user (an adult) has the responsability to save it somewhere (or not).
@@ -202,11 +231,15 @@ def apply_curation(
 
     Returns
     -------
-    sorting_or_analyzer : Sorting | SortingAnalyzer
+    curated_sorting_or_analyzer : Sorting | SortingAnalyzer
         The curated object.
-
-
     """
+    assert isinstance(
+        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)}"
     if isinstance(curation_dict_or_model, dict):
         curation_model = CurationModel(**curation_dict_or_model)
     else:
@@ -215,29 +248,29 @@ def apply_curation(
     if not np.array_equal(np.asarray(curation_model.unit_ids), sorting_or_analyzer.unit_ids):
         raise ValueError("unit_ids from the curation_dict do not match the one from Sorting or SortingAnalyzer")
 
-    if isinstance(sorting_or_analyzer, BaseSorting):
-        sorting = sorting_or_analyzer
-        sorting = sorting.remove_units(curation_model.removed)
-        if len(curation_model.merges) > 0:
-            sorting, _, new_unit_ids = apply_merges_to_sorting(
-                sorting,
-                merge_unit_groups=[m.unit_ids for m in curation_model.merges],
+    # 1. Remove units
+    if len(curation_model.removed) > 0:
+        curated_sorting_or_analyzer = sorting_or_analyzer.remove_units(curation_model.removed)
+    else:
+        curated_sorting_or_analyzer = sorting_or_analyzer
+
+    # 2. Merge units
+    if len(curation_model.merges) > 0:
+        merge_unit_groups = [m.unit_ids for m in curation_model.merges]
+        merge_new_unit_ids = [m.new_unit_id for m in curation_model.merges if m.new_unit_id is not None]
+        if len(merge_new_unit_ids) == 0:
+            merge_new_unit_ids = None
+        if isinstance(sorting_or_analyzer, BaseSorting):
+            curated_sorting_or_analyzer, _, new_unit_ids = apply_merges_to_sorting(
+                curated_sorting_or_analyzer,
+                merge_unit_groups=merge_unit_groups,
                 censor_ms=censor_ms,
-                return_extra=True,
                 new_id_strategy=new_id_strategy,
+                return_extra=True,
             )
         else:
-            new_unit_ids = []
-        apply_curation_labels(sorting, new_unit_ids, curation_model)
-        return sorting
-
-    elif isinstance(sorting_or_analyzer, SortingAnalyzer):
-        analyzer = sorting_or_analyzer
-        if len(curation_model.removed) > 0:
-            analyzer = analyzer.remove_units(curation_model.removed)
-        if len(curation_model.removed) > 0:
-            analyzer, new_unit_ids = analyzer.merge_units(
-                merge_unit_groups=[m.unit_ids for m in curation_model.merges],
+            curated_sorting_or_analyzer, new_unit_ids = curated_sorting_or_analyzer.merge_units(
+                merge_unit_groups=merge_unit_groups,
                 censor_ms=censor_ms,
                 merging_mode=merging_mode,
                 sparsity_overlap=sparsity_overlap,
@@ -247,11 +280,43 @@ def apply_curation(
                 verbose=verbose,
                 **job_kwargs,
             )
+        for i, merge_unit_id in enumerate(new_unit_ids):
+            curation_model.merges[i].new_unit_id = merge_unit_id
+
+    # 3. Split units
+    if len(curation_model.splits) > 0:
+        split_units = {}
+        for split in curation_model.splits:
+            sorting = (
+                curated_sorting_or_analyzer
+                if isinstance(sorting_or_analyzer, BaseSorting)
+                else sorting_or_analyzer.sorting
+            )
+            split_units[split.unit_id] = split.get_full_spike_indices(sorting)
+        split_new_unit_ids = [s.new_unit_ids for s in curation_model.splits if s.new_unit_ids is not None]
+        if len(split_new_unit_ids) == 0:
+            split_new_unit_ids = None
+        if isinstance(sorting_or_analyzer, BaseSorting):
+            curated_sorting_or_analyzer, new_unit_ids = apply_splits_to_sorting(
+                curated_sorting_or_analyzer,
+                split_units,
+                new_unit_ids=split_new_unit_ids,
+                new_id_strategy=new_id_strategy,
+                return_extra=True,
+            )
         else:
-            new_unit_ids = []
-        apply_curation_labels(analyzer.sorting, new_unit_ids, curation_model)
-        return analyzer
-    else:
-        raise TypeError(
-            f"`sorting_or_analyzer` must be a Sorting or a SortingAnalyzer, not an object of type {type(sorting_or_analyzer)}"
-        )
+            curated_sorting_or_analyzer, new_unit_ids = curated_sorting_or_analyzer.split_units(
+                split_units,
+                new_id_strategy=new_id_strategy,
+                return_new_unit_ids=True,
+                new_unit_ids=split_new_unit_ids,
+                format="memory",
+                verbose=verbose,
+            )
+        for i, split_unit_ids in enumerate(new_unit_ids):
+            curation_model.splits[i].new_unit_ids = split_unit_ids
+
+    # 4. Apply labels
+    apply_curation_labels(curated_sorting_or_analyzer, curation_model)
+
+    return curated_sorting_or_analyzer
diff --git a/src/spikeinterface/curation/curation_model.py b/src/spikeinterface/curation/curation_model.py
index cc5bc8e4c..5eb025ace 100644
--- a/src/spikeinterface/curation/curation_model.py
+++ b/src/spikeinterface/curation/curation_model.py
@@ -3,6 +3,8 @@ from typing import List, Dict, Union, Optional, Literal, Tuple
 from itertools import chain, combinations
 import numpy as np
 
+from spikeinterface import BaseSorting
+
 
 class LabelDefinition(BaseModel):
     name: str = Field(description="Name of the label")
@@ -31,12 +33,11 @@ class Split(BaseModel):
             "If labels, the split is defined by a list of labels for each spike (`labels`). "
         ),
     )
-    indices: Optional[Union[List[int], List[List[int]]]] = Field(
+    indices: Optional[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."
+            "List of indices for the split. 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")
@@ -44,6 +45,35 @@ class Split(BaseModel):
         default=None, description="List of new unit IDs for each split"
     )
 
+    def get_full_spike_indices(self, sorting: BaseSorting):
+        """
+        Get the full indices of the spikes in the split for different split modes.
+        """
+        num_spikes = sorting.count_num_spikes_per_unit()[self.unit_id]
+        if self.mode == "indices":
+            # check the sum of split_indices is equal to num_spikes
+            num_spikes_in_split = sum(len(indices) for indices in self.indices)
+            if num_spikes_in_split != num_spikes:
+                # add remaining spike indices
+                full_spike_indices = list(self.indices)
+                existing_indices = np.concatenate(self.indices)
+                remaining_indices = np.setdiff1d(np.arange(num_spikes), existing_indices)
+                full_spike_indices.append(remaining_indices)
+            else:
+                full_spike_indices = self.indices
+        elif self.mode == "labels":
+            assert len(self.labels) == num_spikes, (
+                f"In 'labels' mode, the number of.labels ({len(self.labels)}) "
+                f"must match the number of spikes in the unit ({num_spikes})"
+            )
+            # convert to spike indices
+            full_spike_indices = []
+            for label in np.unique(self.labels):
+                label_indices = np.where(self.labels == label)[0]
+                full_spike_indices.append(label_indices)
+
+        return full_spike_indices
+
 
 class CurationModel(BaseModel):
     supported_versions: Tuple[Literal["1"], Literal["2"]] = Field(
@@ -74,7 +104,6 @@ class CurationModel(BaseModel):
 
     @classmethod
     def check_manual_labels(cls, values):
-
         unit_ids = list(values["unit_ids"])
         manual_labels = values.get("manual_labels")
         if manual_labels is None:
@@ -106,7 +135,6 @@ class CurationModel(BaseModel):
 
     @classmethod
     def check_merges(cls, values):
-
         unit_ids = list(values["unit_ids"])
         merges = values.get("merges")
         if merges is None:
@@ -229,7 +257,10 @@ class CurationModel(BaseModel):
             # Validate new unit IDs
             if split.new_unit_ids is not None:
                 if split.mode == "indices":
-                    if len(split.new_unit_ids) != len(split.indices):
+                    if (
+                        len(split.new_unit_ids) != len(split.indices)
+                        and len(split.new_unit_ids) != len(split.indices) + 1
+                    ):
                         raise ValueError(
                             f"Number of new unit IDs does not match number of splits for unit {split.unit_id}"
                         )
diff --git a/src/spikeinterface/postprocessing/amplitude_scalings.py b/src/spikeinterface/postprocessing/amplitude_scalings.py
index 278151a93..ff926a998 100644
--- a/src/spikeinterface/postprocessing/amplitude_scalings.py
+++ b/src/spikeinterface/postprocessing/amplitude_scalings.py
@@ -128,6 +128,9 @@ class ComputeAmplitudeScalings(AnalyzerExtension):
 
         return new_data
 
+    def _split_extension_data(self, split_units, new_unit_ids, new_sorting_analyzer, verbose=False, **job_kwargs):
+        return self.data.copy()
+
     def _get_pipeline_nodes(self):
 
         recording = self.sorting_analyzer.recording
diff --git a/src/spikeinterface/postprocessing/correlograms.py b/src/spikeinterface/postprocessing/correlograms.py
index dd05bc1bd..fee5cb4c6 100644
--- a/src/spikeinterface/postprocessing/correlograms.py
+++ b/src/spikeinterface/postprocessing/correlograms.py
@@ -166,9 +166,6 @@ class ComputeCorrelograms(AnalyzerExtension):
                 if unit_involved_in_merge is False:
                     old_to_new_unit_index_map[old_unit_index] = new_sorting_analyzer.sorting.id_to_index(old_unit)
 
-            need_to_append = False
-            delete_from = 1
-
             correlograms, new_bins = deepcopy(self.get_data())
 
             for new_unit_id, merge_unit_group in zip(new_unit_ids, merge_unit_groups):
@@ -199,6 +196,12 @@ class ComputeCorrelograms(AnalyzerExtension):
             new_data = dict(ccgs=new_correlograms, bins=new_bins)
         return new_data
 
+    def _split_extension_data(self, split_units, new_unit_ids, new_sorting_analyzer, verbose=False, **job_kwargs):
+        # TODO: for now we just copy
+        new_ccgs, new_bins = _compute_correlograms_on_sorting(new_sorting_analyzer.sorting, **self.params)
+        new_data = dict(ccgs=new_ccgs, bins=new_bins)
+        return new_data
+
     def _run(self, verbose=False):
         ccgs, bins = _compute_correlograms_on_sorting(self.sorting_analyzer.sorting, **self.params)
         self.data["ccgs"] = ccgs
diff --git a/src/spikeinterface/postprocessing/isi.py b/src/spikeinterface/postprocessing/isi.py
index 6089d8eba..529dbf718 100644
--- a/src/spikeinterface/postprocessing/isi.py
+++ b/src/spikeinterface/postprocessing/isi.py
@@ -3,6 +3,7 @@ from __future__ import annotations
 import importlib.util
 
 import numpy as np
+from itertools import chain
 
 from spikeinterface.core.sortinganalyzer import register_result_extension, AnalyzerExtension
 
@@ -81,6 +82,29 @@ class ComputeISIHistograms(AnalyzerExtension):
         new_extension_data = dict(isi_histograms=new_isi_hists, bins=new_bins)
         return new_extension_data
 
+    def _split_extension_data(self, split_units, new_unit_ids, new_sorting_analyzer, verbose=False, **job_kwargs):
+        new_bins = self.data["bins"]
+        arr = self.data["isi_histograms"]
+        num_dims = arr.shape[1]
+        all_new_units = new_sorting_analyzer.unit_ids
+        new_isi_hists = np.zeros((len(all_new_units), num_dims), dtype=arr.dtype)
+
+        # compute all new isi at once
+        new_unit_ids_f = list(chain(*new_unit_ids))
+        new_sorting = new_sorting_analyzer.sorting.select_units(new_unit_ids_f)
+        only_new_hist, _ = _compute_isi_histograms(new_sorting, **self.params)
+
+        for unit_ind, unit_id in enumerate(all_new_units):
+            if unit_id not in new_unit_ids_f:
+                keep_unit_index = self.sorting_analyzer.sorting.id_to_index(unit_id)
+                new_isi_hists[unit_ind, :] = arr[keep_unit_index, :]
+            else:
+                new_unit_index = new_sorting.id_to_index(unit_id)
+                new_isi_hists[unit_ind, :] = only_new_hist[new_unit_index, :]
+
+        new_extension_data = dict(isi_histograms=new_isi_hists, bins=new_bins)
+        return new_extension_data
+
     def _run(self, verbose=False):
         isi_histograms, bins = _compute_isi_histograms(self.sorting_analyzer.sorting, **self.params)
         self.data["isi_histograms"] = isi_histograms
diff --git a/src/spikeinterface/postprocessing/principal_component.py b/src/spikeinterface/postprocessing/principal_component.py
index 8eb343c20..f5c1a7484 100644
--- a/src/spikeinterface/postprocessing/principal_component.py
+++ b/src/spikeinterface/postprocessing/principal_component.py
@@ -150,6 +150,10 @@ class ComputePrincipalComponents(AnalyzerExtension):
                 new_data[k] = v
         return new_data
 
+    def _split_extension_data(self, split_units, new_unit_ids, new_sorting_analyzer, verbose=False, **job_kwargs):
+        # splitting only changes random spikes assignments
+        return self.data.copy()
+
     def get_pca_model(self):
         """
         Returns the scikit-learn PCA model objects.
diff --git a/src/spikeinterface/postprocessing/spike_amplitudes.py b/src/spikeinterface/postprocessing/spike_amplitudes.py
index 61c09977b..8e95a28dd 100644
--- a/src/spikeinterface/postprocessing/spike_amplitudes.py
+++ b/src/spikeinterface/postprocessing/spike_amplitudes.py
@@ -72,6 +72,10 @@ class ComputeSpikeAmplitudes(AnalyzerExtension):
 
         return new_data
 
+    def _split_extension_data(self, split_units, new_unit_ids, new_sorting_analyzer, verbose=False, **job_kwargs):
+        # splitting only changes random spikes assignments
+        return self.data.copy()
+
     def _get_pipeline_nodes(self):
 
         recording = self.sorting_analyzer.recording
diff --git a/src/spikeinterface/postprocessing/spike_locations.py b/src/spikeinterface/postprocessing/spike_locations.py
index 214aabcf4..700abe232 100644
--- a/src/spikeinterface/postprocessing/spike_locations.py
+++ b/src/spikeinterface/postprocessing/spike_locations.py
@@ -103,6 +103,10 @@ class ComputeSpikeLocations(AnalyzerExtension):
         ### in a merged could be different. Should be discussed
         return dict(spike_locations=new_spike_locations)
 
+    def _split_extension_data(self, split_units, new_unit_ids, new_sorting_analyzer, verbose=False, **job_kwargs):
+        # splitting only changes random spikes assignments
+        return self.data.copy()
+
     def _get_pipeline_nodes(self):
         from spikeinterface.sortingcomponents.peak_localization import get_localization_pipeline_nodes
 
diff --git a/src/spikeinterface/postprocessing/template_metrics.py b/src/spikeinterface/postprocessing/template_metrics.py
index 66e315d12..7de9db964 100644
--- a/src/spikeinterface/postprocessing/template_metrics.py
+++ b/src/spikeinterface/postprocessing/template_metrics.py
@@ -8,6 +8,7 @@ from __future__ import annotations
 
 import numpy as np
 import warnings
+from itertools import chain
 from copy import deepcopy
 
 from spikeinterface.core.sortinganalyzer import register_result_extension, AnalyzerExtension
@@ -193,6 +194,26 @@ class ComputeTemplateMetrics(AnalyzerExtension):
         new_data = dict(metrics=metrics)
         return new_data
 
+    def _split_extension_data(self, split_units, new_unit_ids, new_sorting_analyzer, verbose=False, **job_kwargs):
+        import pandas as pd
+
+        metric_names = self.params["metric_names"]
+        old_metrics = self.data["metrics"]
+
+        all_unit_ids = new_sorting_analyzer.unit_ids
+        new_unit_ids_f = list(chain(*new_unit_ids))
+        not_new_ids = all_unit_ids[~np.isin(all_unit_ids, new_unit_ids_f)]
+
+        metrics = pd.DataFrame(index=all_unit_ids, columns=old_metrics.columns)
+
+        metrics.loc[not_new_ids, :] = old_metrics.loc[not_new_ids, :]
+        metrics.loc[new_unit_ids_f, :] = self._compute_metrics(
+            new_sorting_analyzer, new_unit_ids_f, verbose, metric_names, **job_kwargs
+        )
+
+        new_data = dict(metrics=metrics)
+        return new_data
+
     def _compute_metrics(self, sorting_analyzer, unit_ids=None, verbose=False, metric_names=None, **job_kwargs):
         """
         Compute template metrics.
diff --git a/src/spikeinterface/postprocessing/template_similarity.py b/src/spikeinterface/postprocessing/template_similarity.py
index 4f032dc78..2771d382b 100644
--- a/src/spikeinterface/postprocessing/template_similarity.py
+++ b/src/spikeinterface/postprocessing/template_similarity.py
@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import numpy as np
 import warnings
+from itertools import chain
 import importlib.util
 
 from spikeinterface.core.sortinganalyzer import register_result_extension, AnalyzerExtension
@@ -129,6 +130,58 @@ class ComputeTemplateSimilarity(AnalyzerExtension):
 
         return dict(similarity=similarity)
 
+    def _split_extension_data(self, split_units, new_unit_ids, new_sorting_analyzer, verbose=False, **job_kwargs):
+        num_shifts = int(self.params["max_lag_ms"] * self.sorting_analyzer.sampling_frequency / 1000)
+        all_templates_array = get_dense_templates_array(
+            new_sorting_analyzer, return_scaled=self.sorting_analyzer.return_scaled
+        )
+
+        new_unit_ids_f = list(chain(*new_unit_ids))
+        keep = np.isin(new_sorting_analyzer.unit_ids, new_unit_ids_f)
+        new_templates_array = all_templates_array[keep, :, :]
+        if new_sorting_analyzer.sparsity is None:
+            new_sparsity = None
+        else:
+            new_sparsity = ChannelSparsity(
+                new_sorting_analyzer.sparsity.mask[keep, :], new_unit_ids_f, new_sorting_analyzer.channel_ids
+            )
+
+        new_similarity = compute_similarity_with_templates_array(
+            new_templates_array,
+            all_templates_array,
+            method=self.params["method"],
+            num_shifts=num_shifts,
+            support=self.params["support"],
+            sparsity=new_sparsity,
+            other_sparsity=new_sorting_analyzer.sparsity,
+        )
+
+        old_similarity = self.data["similarity"]
+
+        all_new_unit_ids = new_sorting_analyzer.unit_ids
+        n = all_new_unit_ids.size
+        similarity = np.zeros((n, n), dtype=old_similarity.dtype)
+
+        # copy old similarity
+        for unit_ind1, unit_id1 in enumerate(all_new_unit_ids):
+            if unit_id1 not in new_unit_ids_f:
+                old_ind1 = self.sorting_analyzer.sorting.id_to_index(unit_id1)
+                for unit_ind2, unit_id2 in enumerate(all_new_unit_ids):
+                    if unit_id2 not in new_unit_ids_f:
+                        old_ind2 = self.sorting_analyzer.sorting.id_to_index(unit_id2)
+                        s = self.data["similarity"][old_ind1, old_ind2]
+                        similarity[unit_ind1, unit_ind2] = s
+                        similarity[unit_ind2, unit_ind1] = s
+
+        # insert new similarity both way
+        for unit_ind, unit_id in enumerate(all_new_unit_ids):
+            if unit_id in new_unit_ids_f:
+                new_index = list(new_unit_ids_f).index(unit_id)
+                similarity[unit_ind, :] = new_similarity[new_index, :]
+                similarity[:, unit_ind] = new_similarity[new_index, :]
+
+        return dict(similarity=similarity)
+
     def _run(self, verbose=False):
         num_shifts = int(self.params["max_lag_ms"] * self.sorting_analyzer.sampling_frequency / 1000)
         templates_array = get_dense_templates_array(
diff --git a/src/spikeinterface/postprocessing/unit_locations.py b/src/spikeinterface/postprocessing/unit_locations.py
index 53bebd52c..ea297f7b6 100644
--- a/src/spikeinterface/postprocessing/unit_locations.py
+++ b/src/spikeinterface/postprocessing/unit_locations.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import numpy as np
+from itertools import chain
 
 from spikeinterface.core.sortinganalyzer import register_result_extension, AnalyzerExtension
 from .localization_tools import _unit_location_methods
@@ -87,6 +88,30 @@ class ComputeUnitLocations(AnalyzerExtension):
 
         return dict(unit_locations=unit_location)
 
+    def _split_extension_data(self, split_units, new_unit_ids, new_sorting_analyzer, verbose=False, **job_kwargs):
+        old_unit_locations = self.data["unit_locations"]
+        num_dims = old_unit_locations.shape[1]
+
+        method = self.params.get("method")
+        method_kwargs = self.params.copy()
+        method_kwargs.pop("method")
+        func = _unit_location_methods[method]
+        new_unit_ids_f = list(chain(*new_unit_ids))
+        new_unit_locations = func(new_sorting_analyzer, unit_ids=new_unit_ids_f, **method_kwargs)
+        assert new_unit_locations.shape[0] == len(new_unit_ids_f)
+
+        all_new_unit_ids = new_sorting_analyzer.unit_ids
+        unit_location = np.zeros((len(all_new_unit_ids), num_dims), dtype=old_unit_locations.dtype)
+        for unit_index, unit_id in enumerate(all_new_unit_ids):
+            if unit_id not in new_unit_ids_f:
+                old_index = self.sorting_analyzer.sorting.id_to_index(unit_id)
+                unit_location[unit_index] = old_unit_locations[old_index]
+            else:
+                new_index = list(new_unit_ids_f).index(unit_id)
+                unit_location[unit_index] = new_unit_locations[new_index]
+
+        return dict(unit_locations=unit_location)
+
     def _run(self, verbose=False):
         method = self.params.get("method")
         method_kwargs = self.params.copy()
diff --git a/src/spikeinterface/qualitymetrics/quality_metric_calculator.py b/src/spikeinterface/qualitymetrics/quality_metric_calculator.py
index 134849e70..055fefc78 100644
--- a/src/spikeinterface/qualitymetrics/quality_metric_calculator.py
+++ b/src/spikeinterface/qualitymetrics/quality_metric_calculator.py
@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import warnings
+from itertools import chain
 from copy import deepcopy
 
 import numpy as np
@@ -158,6 +159,33 @@ class ComputeQualityMetrics(AnalyzerExtension):
         new_data = dict(metrics=metrics)
         return new_data
 
+    def _split_extension_data(self, split_units, new_unit_ids, new_sorting_analyzer, verbose=False, **job_kwargs):
+        import pandas as pd
+
+        metric_names = self.params["metric_names"]
+        old_metrics = self.data["metrics"]
+
+        all_unit_ids = new_sorting_analyzer.unit_ids
+        new_unit_ids_f = list(chain(*new_unit_ids))
+        not_new_ids = all_unit_ids[~np.isin(all_unit_ids, new_unit_ids_f)]
+
+        # this creates a new metrics dictionary, but the dtype for everything will be
+        # object. So we will need to fix this later after computing metrics
+        metrics = pd.DataFrame(index=all_unit_ids, columns=old_metrics.columns)
+        metrics.loc[not_new_ids, :] = old_metrics.loc[not_new_ids, :]
+        metrics.loc[new_unit_ids_f, :] = self._compute_metrics(
+            new_sorting_analyzer, new_unit_ids_f, verbose, metric_names, **job_kwargs
+        )
+
+        # we need to fix the dtypes after we compute everything because we have nans
+        # we can iterate through the columns and convert them back to the dtype
+        # of the original quality dataframe.
+        for column in old_metrics.columns:
+            metrics[column] = metrics[column].astype(old_metrics[column].dtype)
+
+        new_data = dict(metrics=metrics)
+        return new_data
+
     def _compute_metrics(self, sorting_analyzer, unit_ids=None, verbose=False, metric_names=None, **job_kwargs):
         """
         Compute quality metrics.


Test patch:
diff --git a/src/spikeinterface/core/tests/test_sortinganalyzer.py b/src/spikeinterface/core/tests/test_sortinganalyzer.py
index 7074c054b..ce248f00f 100644
--- a/src/spikeinterface/core/tests/test_sortinganalyzer.py
+++ b/src/spikeinterface/core/tests/test_sortinganalyzer.py
@@ -378,6 +378,7 @@ def _check_sorting_analyzers(sorting_analyzer, original_sorting, cache_folder):
         # unit 0, 2, ... should be removed
         assert np.all(~np.isin(data["result_two"], [0, 2]))
 
+        # test merges
         if format != "memory":
             if format == "zarr":
                 folder = cache_folder / f"test_SortingAnalyzer_merge_soft_with_{format}.zarr"
@@ -387,10 +388,14 @@ def _check_sorting_analyzers(sorting_analyzer, original_sorting, cache_folder):
                 shutil.rmtree(folder)
         else:
             folder = None
-        sorting_analyzer4 = sorting_analyzer.merge_units(merge_unit_groups=[[0, 1]], format=format, folder=folder)
+        sorting_analyzer4, new_unit_ids = sorting_analyzer.merge_units(
+            merge_unit_groups=[[0, 1]], format=format, folder=folder, return_new_unit_ids=True
+        )
         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
+        is_merged_values = sorting_analyzer4.sorting.get_property("is_merged")
+        assert is_merged_values[sorting_analyzer4.sorting.ids_to_indices(new_unit_ids)][0]
 
         if format != "memory":
             if format == "zarr":
@@ -401,13 +406,50 @@ def _check_sorting_analyzers(sorting_analyzer, original_sorting, cache_folder):
                 shutil.rmtree(folder)
         else:
             folder = None
-        sorting_analyzer5 = sorting_analyzer.merge_units(
-            merge_unit_groups=[[0, 1]], new_unit_ids=[50], format=format, folder=folder, merging_mode="hard"
+        sorting_analyzer5, new_unit_ids = sorting_analyzer.merge_units(
+            merge_unit_groups=[[0, 1]],
+            new_unit_ids=[50],
+            format=format,
+            folder=folder,
+            merging_mode="hard",
+            return_new_unit_ids=True,
         )
         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
+        is_merged_values = sorting_analyzer5.sorting.get_property("is_merged")
+        assert is_merged_values[sorting_analyzer5.sorting.id_to_index(50)]
+
+        # test splitting
+        if format != "memory":
+            if format == "zarr":
+                folder = cache_folder / f"test_SortingAnalyzer_split_with_{format}.zarr"
+            else:
+                folder = cache_folder / f"test_SortingAnalyzer_split_with_{format}"
+            if folder.exists():
+                shutil.rmtree(folder)
+        else:
+            folder = None
+        split_units = {}
+        num_spikes = sorting_analyzer.sorting.count_num_spikes_per_unit()
+        units_to_split = sorting_analyzer.unit_ids[:2]
+        for unit in units_to_split:
+            for unit in units_to_split:
+                split_units[unit] = [
+                    np.arange(num_spikes[unit] // 2),
+                    np.arange(num_spikes[unit] // 2, num_spikes[unit]),
+                ]
+
+        sorting_analyzer6, split_new_unit_ids = sorting_analyzer.split_units(
+            split_units=split_units, format=format, folder=folder, return_new_unit_ids=True
+        )
+        for unit_to_split in units_to_split:
+            assert unit_to_split not in sorting_analyzer6.unit_ids
+        assert len(sorting_analyzer6.unit_ids) == len(sorting_analyzer.unit_ids) + 2
+        is_split_values = sorting_analyzer6.sorting.get_property("is_split")
+        for new_unit_ids in split_new_unit_ids:
+            assert all(is_split_values[sorting_analyzer6.sorting.ids_to_indices(new_unit_ids)])
 
     # test compute with extension-specific params
     sorting_analyzer.compute(["dummy"], extension_params={"dummy": {"param1": 5.5}})
@@ -491,6 +533,14 @@ class DummyAnalyzerExtension(AnalyzerExtension):
 
         return new_data
 
+    def _split_extension_data(self, split_units, new_unit_ids, new_sorting_analyzer, verbose=False, **job_kwargs):
+        new_data = dict()
+        new_data["result_one"] = self.data["result_one"]
+        spikes = new_sorting_analyzer.sorting.to_spike_vector()
+        new_data["result_two"] = spikes["unit_index"].copy()
+        new_data["result_three"] = np.zeros((len(new_sorting_analyzer.unit_ids), 2))
+        return new_data
+
     def _get_data(self):
         return self.data["result_one"]
 
diff --git a/src/spikeinterface/curation/tests/test_curation_format.py b/src/spikeinterface/curation/tests/test_curation_format.py
index 587673b91..9707b0f08 100644
--- a/src/spikeinterface/curation/tests/test_curation_format.py
+++ b/src/spikeinterface/curation/tests/test_curation_format.py
@@ -127,10 +127,25 @@ curation_ids_str = {
 # Test dictionary format for merges with string IDs
 curation_ids_str_dict = {**curation_ids_str, "merges": {"u50": ["u3", "u6"], "u51": ["u10", "u14", "u20"]}}
 
+# This is a failure example with duplicated merge
+duplicate_merge = curation_ids_int.copy()
+duplicate_merge["merge_unit_groups"] = [[3, 6, 10], [10, 14, 20]]
+
 # Test with splits
 curation_with_splits = {
-    **curation_ids_int,
-    "splits": [{"unit_id": 2, "mode": "indices", "indices": [[0, 1, 2], [3, 4, 5]], "new_unit_ids": [50, 51]}],
+    "format_version": "2",
+    "unit_ids": [1, 2, 3, 6, 10, 14, 20, 31, 42],
+    "label_definitions": {
+        "quality": {"label_options": ["good", "noise", "MUA", "artifact"], "exclusive": True},
+        "putative_type": {
+            "label_options": ["excitatory", "inhibitory", "pyramidal", "mitral"],
+            "exclusive": False,
+        },
+    },
+    "manual_labels": [
+        {"unit_id": 2, "quality": ["good"], "putative_type": ["excitatory", "pyramidal"]},
+    ],
+    "splits": [{"unit_id": 2, "mode": "indices", "indices": [[0, 1, 2], [3, 4, 5]]}],
 }
 
 # Test dictionary format for splits
@@ -243,15 +258,144 @@ def test_apply_curation():
     assert sorting_curated.get_property("quality", ids=[2])[0] == "noise"
     assert sorting_curated.get_property("excitatory", ids=[2])[0]
 
-    # Test with splits
-    sorting_curated = apply_curation(sorting, curation_with_splits)
-    assert sorting_curated.get_property("quality", ids=[1])[0] == "good"
-
     # Test analyzer
     analyzer_curated = apply_curation(analyzer, curation_ids_int)
     assert "quality" in analyzer_curated.sorting.get_property_keys()
 
 
+def test_apply_curation_with_split():
+    recording, sorting = generate_ground_truth_recording(durations=[10.0], num_units=9, seed=2205)
+    sorting = sorting.rename_units(np.array([1, 2, 3, 6, 10, 14, 20, 31, 42]))
+    analyzer = create_sorting_analyzer(sorting, recording, sparse=False)
+
+    sorting_curated = apply_curation(sorting, curation_with_splits)
+    # the split indices are not complete, so an extra unit is added
+    assert len(sorting_curated.unit_ids) == len(sorting.unit_ids) + 2
+
+    assert 2 not in sorting_curated.unit_ids
+    split_unit_ids = [43, 44, 45]
+    for unit_id in split_unit_ids:
+        assert unit_id in sorting_curated.unit_ids
+        assert sorting_curated.get_property("quality", ids=[unit_id])[0] == "good"
+        assert sorting_curated.get_property("excitatory", ids=[unit_id])[0]
+        assert sorting_curated.get_property("pyramidal", ids=[unit_id])[0]
+
+    analyzer_curated = apply_curation(analyzer, curation_with_splits)
+    assert len(analyzer_curated.sorting.unit_ids) == len(analyzer.sorting.unit_ids) + 2
+
+    assert 2 not in analyzer_curated.unit_ids
+    for unit_id in split_unit_ids:
+        assert unit_id in analyzer_curated.unit_ids
+        assert analyzer_curated.sorting.get_property("quality", ids=[unit_id])[0] == "good"
+        assert analyzer_curated.sorting.get_property("excitatory", ids=[unit_id])[0]
+        assert analyzer_curated.sorting.get_property("pyramidal", ids=[unit_id])[0]
+
+
+def test_apply_curation_with_split_multi_segment():
+    recording, sorting = generate_ground_truth_recording(durations=[10.0, 10.0], num_units=9, seed=2205)
+    sorting = sorting.rename_units(np.array([1, 2, 3, 6, 10, 14, 20, 31, 42]))
+    analyzer = create_sorting_analyzer(sorting, recording, sparse=False)
+    num_segments = sorting.get_num_segments()
+
+    curation_with_splits_multi_segment = curation_with_splits.copy()
+
+    # we make a split so that each subsplit will have all spikes from different segments
+    split_unit_id = curation_with_splits_multi_segment["splits"][0]["unit_id"]
+    sv = sorting.to_spike_vector()
+    unit_index = sorting.id_to_index(split_unit_id)
+    spikes_from_split_unit = sv[sv["unit_index"] == unit_index]
+
+    split_indices = []
+    cum_spikes = 0
+    for segment_index in range(num_segments):
+        spikes_in_segment = spikes_from_split_unit[spikes_from_split_unit["segment_index"] == segment_index]
+        split_indices.append(np.arange(0, len(spikes_in_segment)) + cum_spikes)
+        cum_spikes += len(spikes_in_segment)
+
+    curation_with_splits_multi_segment["splits"][0]["indices"] = split_indices
+
+    sorting_curated = apply_curation(sorting, curation_with_splits_multi_segment)
+
+    assert len(sorting_curated.unit_ids) == len(sorting.unit_ids) + 1
+    assert 2 not in sorting_curated.unit_ids
+    assert 43 in sorting_curated.unit_ids
+    assert 44 in sorting_curated.unit_ids
+
+    # check that spike trains are correctly split across segments
+    for seg_index in range(num_segments):
+        st_43 = sorting_curated.get_unit_spike_train(43, segment_index=seg_index)
+        st_44 = sorting_curated.get_unit_spike_train(44, segment_index=seg_index)
+        if seg_index == 0:
+            assert len(st_43) > 0
+            assert len(st_44) == 0
+        else:
+            assert len(st_43) == 0
+            assert len(st_44) > 0
+
+
+def test_apply_curation_splits_with_mask():
+    recording, sorting = generate_ground_truth_recording(durations=[10.0], num_units=9, seed=2205)
+    sorting = sorting.rename_units(np.array([1, 2, 3, 6, 10, 14, 20, 31, 42]))
+    analyzer = create_sorting_analyzer(sorting, recording, sparse=False)
+
+    # Get number of spikes for unit 2
+    num_spikes = sorting.count_num_spikes_per_unit()[2]
+
+    # Create split labels that assign spikes to 3 different clusters
+    split_labels = np.zeros(num_spikes, dtype=int)
+    split_labels[: num_spikes // 3] = 0  # First third to cluster 0
+    split_labels[num_spikes // 3 : 2 * num_spikes // 3] = 1  # Second third to cluster 1
+    split_labels[2 * num_spikes // 3 :] = 2  # Last third to cluster 2
+
+    curation_with_mask_split = {
+        "format_version": "2",
+        "unit_ids": [1, 2, 3, 6, 10, 14, 20, 31, 42],
+        "label_definitions": {
+            "quality": {"label_options": ["good", "noise", "MUA", "artifact"], "exclusive": True},
+            "putative_type": {
+                "label_options": ["excitatory", "inhibitory", "pyramidal", "mitral"],
+                "exclusive": False,
+            },
+        },
+        "manual_labels": [
+            {"unit_id": 2, "quality": ["good"], "putative_type": ["excitatory", "pyramidal"]},
+        ],
+        "splits": [
+            {
+                "unit_id": 2,
+                "mode": "labels",
+                "labels": split_labels.tolist(),
+                "new_unit_ids": [43, 44, 45],
+            }
+        ],
+    }
+
+    sorting_curated = apply_curation(sorting, curation_with_mask_split)
+
+    # Check results
+    assert len(sorting_curated.unit_ids) == len(sorting.unit_ids) + 2  # Original units - 1 (split) + 3 (new)
+    assert 2 not in sorting_curated.unit_ids  # Original unit should be removed
+
+    # Check new split units
+    split_unit_ids = [43, 44, 45]
+    for unit_id in split_unit_ids:
+        assert unit_id in sorting_curated.unit_ids
+        # Check properties are propagated
+        assert sorting_curated.get_property("quality", ids=[unit_id])[0] == "good"
+        assert sorting_curated.get_property("excitatory", ids=[unit_id])[0]
+        assert sorting_curated.get_property("pyramidal", ids=[unit_id])[0]
+
+    # Check analyzer
+    analyzer_curated = apply_curation(analyzer, curation_with_mask_split)
+    assert len(analyzer_curated.sorting.unit_ids) == len(analyzer.sorting.unit_ids) + 2
+
+    # Verify split sizes
+    spike_counts = analyzer_curated.sorting.count_num_spikes_per_unit()
+    assert spike_counts[43] == num_spikes // 3  # First third
+    assert spike_counts[44] == num_spikes // 3  # Second third
+    assert spike_counts[45] == num_spikes - 2 * (num_spikes // 3)  # Remainder
+
+
 if __name__ == "__main__":
     test_curation_format_validation()
     test_to_from_json()
@@ -259,3 +403,5 @@ if __name__ == "__main__":
     test_curation_label_to_vectors()
     test_curation_label_to_dataframe()
     test_apply_curation()
+    test_apply_curation_with_split_multi_segment()
+    test_apply_curation_splits_with_mask()
diff --git a/src/spikeinterface/postprocessing/tests/test_multi_extensions.py b/src/spikeinterface/postprocessing/tests/test_multi_extensions.py
index be0070d94..87ff3cdeb 100644
--- a/src/spikeinterface/postprocessing/tests/test_multi_extensions.py
+++ b/src/spikeinterface/postprocessing/tests/test_multi_extensions.py
@@ -11,8 +11,59 @@ from spikeinterface import (
 )
 from spikeinterface.core.generate import inject_some_split_units
 
+# even if this is in postprocessing, we make an extension for quality metrics
+extension_dict = {
+    "noise_levels": dict(),
+    "random_spikes": dict(),
+    "waveforms": dict(),
+    "templates": dict(),
+    "principal_components": dict(),
+    "spike_amplitudes": dict(),
+    "template_similarity": dict(),
+    "correlograms": dict(),
+    "isi_histograms": dict(),
+    "amplitude_scalings": dict(handle_collisions=False),  # otherwise hard mode could fail due to dropped spikes
+    "spike_locations": dict(method="center_of_mass"),  # trick to avoid UserWarning
+    "unit_locations": dict(),
+    "template_metrics": dict(),
+    "quality_metrics": dict(metric_names=["firing_rate", "isi_violation", "snr"]),
+}
+extension_data_type = {
+    "noise_levels": None,
+    "templates": "unit",
+    "isi_histograms": "unit",
+    "unit_locations": "unit",
+    "spike_amplitudes": "spike",
+    "amplitude_scalings": "spike",
+    "spike_locations": "spike",
+    "quality_metrics": "pandas",
+    "template_metrics": "pandas",
+    "correlograms": "matrix",
+    "template_similarity": "matrix",
+    "principal_components": "random",
+    "waveforms": "random",
+    "random_spikes": "random_spikes",
+}
+data_with_miltiple_returns = ["isi_histograms", "correlograms"]
+# due to incremental PCA, hard computation could result in different results for PCA
+# the model is differents always
+random_computation = ["principal_components"]
+# for some extensions (templates, amplitude_scalings), since the templates slightly change for merges/splits
+# we allow a relative tolerance
+# (amplitud_scalings are the moste sensitive!)
+extensions_with_rel_tolerance_merge = {
+    "amplitude_scalings": 1e-1,
+    "templates": 1e-3,
+    "template_similarity": 1e-3,
+    "unit_locations": 1e-3,
+    "template_metrics": 1e-3,
+    "quality_metrics": 1e-3,
+}
+extensions_with_rel_tolerance_splits = {"amplitude_scalings": 1e-1}
 
-def get_dataset():
+
+def get_dataset_to_merge():
+    # generate a dataset with some split units to minimize merge errors
     recording, sorting = generate_ground_truth_recording(
         durations=[30.0],
         sampling_frequency=16000.0,
@@ -20,6 +71,7 @@ def get_dataset():
         num_units=10,
         generate_sorting_kwargs=dict(firing_rates=10.0, refractory_period_ms=4.0),
         noise_kwargs=dict(noise_levels=5.0, strategy="tile_pregenerated"),
+        generate_unit_locations_kwargs=dict(margin_um=10.0, minimum_z=2.0, maximum_z=15.0, minimum_distance=20),
         seed=2205,
     )
 
@@ -36,70 +88,65 @@ def get_dataset():
     sort_by_amp = np.argsort(list(get_template_extremum_amplitude(analyzer_raw).values()))[::-1]
     split_ids = sorting.unit_ids[sort_by_amp][:3]
 
-    sorting_with_splits, other_ids = inject_some_split_units(
+    sorting_with_splits, split_unit_ids = inject_some_split_units(
         sorting, num_split=3, split_ids=split_ids, output_ids=True, seed=0
     )
-    return recording, sorting_with_splits, other_ids
+    return recording, sorting_with_splits, split_unit_ids
+
+
+def get_dataset_to_split():
+    # generate a dataset and return large unit to split to minimize split errors
+    recording, sorting = generate_ground_truth_recording(
+        durations=[30.0],
+        sampling_frequency=16000.0,
+        num_channels=10,
+        num_units=10,
+        generate_sorting_kwargs=dict(firing_rates=10.0, refractory_period_ms=4.0),
+        noise_kwargs=dict(noise_levels=5.0, strategy="tile_pregenerated"),
+        seed=2205,
+    )
+
+    channel_ids_as_integers = [id for id in range(recording.get_num_channels())]
+    unit_ids_as_integers = [id for id in range(sorting.get_num_units())]
+    recording = recording.rename_channels(new_channel_ids=channel_ids_as_integers)
+    sorting = sorting.rename_units(new_unit_ids=unit_ids_as_integers)
+
+    # since templates are going to be averaged and this might be a problem for amplitude scaling
+    # we select the 3 units with the largest templates to split
+    analyzer_raw = create_sorting_analyzer(sorting, recording, format="memory", sparse=False)
+    analyzer_raw.compute(["random_spikes", "templates"])
+    # select 3 largest templates to split
+    sort_by_amp = np.argsort(list(get_template_extremum_amplitude(analyzer_raw).values()))[::-1]
+    large_units = sorting.unit_ids[sort_by_amp][:2]
+
+    return recording, sorting, large_units
 
 
 @pytest.fixture(scope="module")
-def dataset():
-    return get_dataset()
+def dataset_to_merge():
+    return get_dataset_to_merge()
+
+
+@pytest.fixture(scope="module")
+def dataset_to_split():
+    return get_dataset_to_split()
 
 
 @pytest.mark.parametrize("sparse", [False, True])
-def test_SortingAnalyzer_merge_all_extensions(dataset, sparse):
+def test_SortingAnalyzer_merge_all_extensions(dataset_to_merge, sparse):
     set_global_job_kwargs(n_jobs=1)
 
-    recording, sorting, other_ids = dataset
+    recording, sorting, other_ids = dataset_to_merge
 
     sorting_analyzer = create_sorting_analyzer(sorting, recording, format="memory", sparse=sparse)
+    extension_dict_merge = extension_dict.copy()
 
     # we apply the merges according to the artificial splits
     merges = [list(v) for v in other_ids.values()]
     split_unit_ids = np.ravel(merges)
     unmerged_unit_ids = sorting_analyzer.unit_ids[~np.isin(sorting_analyzer.unit_ids, split_unit_ids)]
 
-    # even if this is in postprocessing, we make an extension for quality metrics
-    extension_dict = {
-        "noise_levels": dict(),
-        "random_spikes": dict(),
-        "waveforms": dict(),
-        "templates": dict(),
-        "principal_components": dict(),
-        "spike_amplitudes": dict(),
-        "template_similarity": dict(),
-        "correlograms": dict(),
-        "isi_histograms": dict(),
-        "amplitude_scalings": dict(handle_collisions=False),  # otherwise hard mode could fail due to dropped spikes
-        "spike_locations": dict(method="center_of_mass"),  # trick to avoid UserWarning
-        "unit_locations": dict(),
-        "template_metrics": dict(),
-        "quality_metrics": dict(metric_names=["firing_rate", "isi_violation", "snr"]),
-    }
-    extension_data_type = {
-        "noise_levels": None,
-        "templates": "unit",
-        "isi_histograms": "unit",
-        "unit_locations": "unit",
-        "spike_amplitudes": "spike",
-        "amplitude_scalings": "spike",
-        "spike_locations": "spike",
-        "quality_metrics": "pandas",
-        "template_metrics": "pandas",
-        "correlograms": "matrix",
-        "template_similarity": "matrix",
-        "principal_components": "random",
-        "waveforms": "random",
-        "random_spikes": "random_spikes",
-    }
-    data_with_miltiple_returns = ["isi_histograms", "correlograms"]
-
-    # due to incremental PCA, hard computation could result in different results for PCA
-    # the model is differents always
-    random_computation = ["principal_components"]
-
-    sorting_analyzer.compute(extension_dict, n_jobs=1)
+    sorting_analyzer.compute(extension_dict_merge, n_jobs=1)
 
     # TODO: still some UserWarnings for n_jobs, where from?
     t0 = time.perf_counter()
@@ -155,14 +202,95 @@ def test_SortingAnalyzer_merge_all_extensions(dataset, sparse):
         )
 
         if ext not in random_computation:
+            if ext in extensions_with_rel_tolerance_merge:
+                rtol = extensions_with_rel_tolerance_merge[ext]
+            else:
+                rtol = 0
             if extension_data_type[ext] == "pandas":
                 data_hard_merged = data_hard_merged.dropna().to_numpy().astype("float")
                 data_soft_merged = data_soft_merged.dropna().to_numpy().astype("float")
             if data_hard_merged.dtype.fields is None:
-                assert np.allclose(data_hard_merged, data_soft_merged, rtol=0.1)
+                if not np.allclose(data_hard_merged, data_soft_merged, rtol=rtol):
+                    max_error = np.max(np.abs(data_hard_merged - data_soft_merged))
+                    raise Exception(f"Failed for {ext} - max error {max_error}")
             else:
                 for f in data_hard_merged.dtype.fields:
-                    assert np.allclose(data_hard_merged[f], data_soft_merged[f], rtol=0.1)
+                    if not np.allclose(data_hard_merged[f], data_soft_merged[f], rtol=rtol):
+                        max_error = np.max(np.abs(data_hard_merged[f] - data_soft_merged[f]))
+                        raise Exception(f"Failed for {ext} - field {f} - max error {max_error}")
+
+
+@pytest.mark.parametrize("sparse", [False, True])
+def test_SortingAnalyzer_split_all_extensions(dataset_to_split, sparse):
+    set_global_job_kwargs(n_jobs=1)
+
+    recording, sorting, units_to_split = dataset_to_split
+
+    sorting_analyzer = create_sorting_analyzer(sorting, recording, format="memory", sparse=sparse)
+    extension_dict_split = extension_dict.copy()
+    sorting_analyzer.compute(extension_dict, n_jobs=1)
+
+    # we randomly apply splits (at half of spiketrain)
+    num_spikes = sorting.count_num_spikes_per_unit()
+
+    unsplit_unit_ids = sorting_analyzer.unit_ids[~np.isin(sorting_analyzer.unit_ids, units_to_split)]
+    splits = {}
+    for unit in units_to_split:
+        splits[unit] = [np.arange(num_spikes[unit] // 2), np.arange(num_spikes[unit] // 2, num_spikes[unit])]
+
+    analyzer_split, split_unit_ids = sorting_analyzer.split_units(split_units=splits, return_new_unit_ids=True)
+    split_unit_ids = list(np.concatenate(split_unit_ids))
+
+    # also do a full recopute
+    analyzer_hard = create_sorting_analyzer(analyzer_split.sorting, recording, format="memory", sparse=sparse)
+    # we propagate random spikes to avoid random spikes to be recomputed
+    extension_dict_ = extension_dict_split.copy()
+    extension_dict_.pop("random_spikes")
+    analyzer_hard.extensions["random_spikes"] = analyzer_split.extensions["random_spikes"]
+    analyzer_hard.compute(extension_dict_, n_jobs=1)
+
+    for ext in extension_dict:
+        # 1. check that data are exactly the same for unchanged units between original/split
+        data_original = sorting_analyzer.get_extension(ext).get_data()
+        data_split = analyzer_split.get_extension(ext).get_data()
+        data_recompute = analyzer_hard.get_extension(ext).get_data()
+        if ext in data_with_miltiple_returns:
+            data_original = data_original[0]
+            data_split = data_split[0]
+            data_recompute = data_recompute[0]
+        data_original_unsplit = get_extension_data_for_units(
+            sorting_analyzer, data_original, unsplit_unit_ids, extension_data_type[ext]
+        )
+        data_split_unsplit = get_extension_data_for_units(
+            analyzer_split, data_split, unsplit_unit_ids, extension_data_type[ext]
+        )
+
+        np.testing.assert_array_equal(data_original_unsplit, data_split_unsplit)
+
+        # 2. check that split data are the same for extension split and recompute
+        data_split_soft = get_extension_data_for_units(
+            analyzer_split, data_split, split_unit_ids, extension_data_type[ext]
+        )
+        data_split_hard = get_extension_data_for_units(
+            analyzer_hard, data_recompute, split_unit_ids, extension_data_type[ext]
+        )
+        if ext not in random_computation:
+            if ext in extensions_with_rel_tolerance_splits:
+                rtol = extensions_with_rel_tolerance_splits[ext]
+            else:
+                rtol = 0
+            if extension_data_type[ext] == "pandas":
+                data_split_soft = data_split_soft.dropna().to_numpy().astype("float")
+                data_split_hard = data_split_hard.dropna().to_numpy().astype("float")
+            if data_split_hard.dtype.fields is None:
+                if not np.allclose(data_split_hard, data_split_soft, rtol=rtol):
+                    max_error = np.max(np.abs(data_split_hard - data_split_soft))
+                    raise Exception(f"Failed for {ext} - max error {max_error}")
+            else:
+                for f in data_split_hard.dtype.fields:
+                    if not np.allclose(data_split_hard[f], data_split_soft[f], rtol=rtol):
+                        max_error = np.max(np.abs(data_split_hard[f] - data_split_soft[f]))
+                        raise Exception(f"Failed for {ext} - field {f} - max error {max_error}")
 
 
 def get_extension_data_for_units(sorting_analyzer, data, unit_ids, ext_data_type):
@@ -191,5 +319,5 @@ def get_extension_data_for_units(sorting_analyzer, data, unit_ids, ext_data_type
 
 
 if __name__ == "__main__":
-    dataset = get_dataset()
+    dataset = get_dataset_to_merge()
     test_SortingAnalyzer_merge_all_extensions(dataset, False)