Repository: SpikeInterface/spikeinterface (python)

Issue:
Make `TransformSorting` class
Functions that take an existing `BaseSorting` and modify it, e.g., adding synchrony or removing spikes will return a `TransformedSorting` object.

Internally, this new class will have useful references to:
- `added_spikes_mask`: indices or mask of added spikes in the new spike vector
- `removed_spikes_mask`: indices or mask of removed spikes in the original spike vector

Pull request:
Creation of a `TransformSorting` object to track modifications and bencharmk
Draft prototype for a TransformedSorting object. I need feedback, because not sure exactly of what we really want. This object should take any sorting, and dict of units to remove/add spikes (dict of spikes (np array with sepcified dtype obtained via to_spike_vector?)), and then be sorting itself, with some additional attribute: 
- the sorting it has been obtained from
- an array of the removed spikes from the original sorting (as a plain binary mask)
- an array of the added spikes (as a plain binary mask)
I have concerns about the input signature of such an objects. Should we provide spike_vectors? Dicts?


Closes #2028 

Patch:
diff --git a/src/spikeinterface/comparison/__init__.py b/src/spikeinterface/comparison/__init__.py
index bff85dde4..648ef4ed7 100644
--- a/src/spikeinterface/comparison/__init__.py
+++ b/src/spikeinterface/comparison/__init__.py
@@ -36,7 +36,7 @@ from .correlogram import CorrelogramGTComparison, CorrelogramGTStudy
 from .hybrid import (
     HybridSpikesRecording,
     HybridUnitsRecording,
-    generate_injected_sorting,
+    generate_sorting_to_inject,
     create_hybrid_units_recording,
     create_hybrid_spikes_recording,
 )
diff --git a/src/spikeinterface/comparison/hybrid.py b/src/spikeinterface/comparison/hybrid.py
index 8af64536f..54230a438 100644
--- a/src/spikeinterface/comparison/hybrid.py
+++ b/src/spikeinterface/comparison/hybrid.py
@@ -8,7 +8,12 @@ from spikeinterface.core import (
     NumpySorting,
 )
 from spikeinterface.core.core_tools import define_function_from_class
-from spikeinterface.core.generate import generate_sorting, InjectTemplatesRecording, _ensure_seed
+from spikeinterface.core.generate import (
+    generate_sorting,
+    InjectTemplatesRecording,
+    _ensure_seed,
+    generate_sorting_to_inject,
+)
 
 
 class HybridUnitsRecording(InjectTemplatesRecording):
@@ -174,7 +179,7 @@ class HybridSpikesRecording(InjectTemplatesRecording):
             num_samples = [
                 target_recording.get_num_frames(seg_index) for seg_index in range(target_recording.get_num_segments())
             ]
-            self.injected_sorting = generate_injected_sorting(
+            self.injected_sorting = generate_sorting_to_inject(
                 target_sorting, num_samples, max_injected_per_unit, injected_rate, refractory_period_ms
             )
         else:
@@ -201,44 +206,6 @@ class HybridSpikesRecording(InjectTemplatesRecording):
         )
 
 
-def generate_injected_sorting(
-    sorting: BaseSorting,
-    num_samples: List[int],
-    max_injected_per_unit: int = 1000,
-    injected_rate: float = 0.05,
-    refractory_period_ms: float = 1.5,
-) -> NumpySorting:
-    injected_spike_trains = [{} for seg_index in range(sorting.get_num_segments())]
-    t_r = int(round(refractory_period_ms * sorting.get_sampling_frequency() * 1e-3))
-
-    for segment_index in range(sorting.get_num_segments()):
-        for unit_id in sorting.unit_ids:
-            spike_train = sorting.get_unit_spike_train(unit_id, segment_index=segment_index)
-            n_injection = min(max_injected_per_unit, int(round(injected_rate * len(spike_train))))
-            # Inject more, then take out all that violate the refractory period.
-            n = int(n_injection + 10 * np.sqrt(n_injection))
-            injected_spike_train = np.sort(
-                np.random.uniform(low=0, high=num_samples[segment_index], size=n).astype(np.int64)
-            )
-
-            # Remove spikes that are in the refractory period.
-            violations = np.where(np.diff(injected_spike_train) < t_r)[0]
-            injected_spike_train = np.delete(injected_spike_train, violations)
-
-            # Remove spikes that violate the refractory period of the real spikes.
-            # TODO: Need a better & faster way than this.
-            min_diff = np.min(np.abs(injected_spike_train[:, None] - spike_train[None, :]), axis=1)
-            violations = min_diff < t_r
-            injected_spike_train = injected_spike_train[~violations]
-
-            if len(injected_spike_train) > n_injection:
-                injected_spike_train = np.sort(np.random.choice(injected_spike_train, n_injection, replace=False))
-
-            injected_spike_trains[segment_index][unit_id] = injected_spike_train
-
-    return NumpySorting.from_unit_dict(injected_spike_trains, sorting.get_sampling_frequency())
-
-
 create_hybrid_units_recording = define_function_from_class(
     source_class=HybridUnitsRecording, name="create_hybrid_units_recording"
 )
diff --git a/src/spikeinterface/core/generate.py b/src/spikeinterface/core/generate.py
index 3e56467ca..aadfe0ba2 100644
--- a/src/spikeinterface/core/generate.py
+++ b/src/spikeinterface/core/generate.py
@@ -4,8 +4,8 @@ import numpy as np
 from typing import Union, Optional, List, Literal
 import warnings
 
-
-from .numpyextractors import NumpyRecording, NumpySorting
+from .basesorting import SpikeVectorSortingSegment
+from .numpyextractors import NumpySorting
 from .basesorting import minimum_spike_dtype
 
 from probeinterface import Probe, generate_linear_probe, generate_multi_columns_probe
@@ -85,7 +85,7 @@ def generate_recording(
 def generate_sorting(
     num_units=5,
     sampling_frequency=30000.0,  # in Hz
-    durations=[10.325, 3.5],  #  in s for 2 segments
+    durations=[10.325, 3.5],  # in s for 2 segments
     firing_rates=3.0,
     empty_units=None,
     refractory_period_ms=4.0,  # in ms
@@ -195,8 +195,8 @@ def add_synchrony_to_sorting(sorting, sync_event_ratio=0.3, seed=None):
 
     Returns
     -------
-    sorting : NumpySorting
-        The sorting object
+    sorting : TransformSorting
+        The sorting object, keeping track of added spikes
 
     """
     rng = np.random.default_rng(seed)
@@ -220,16 +220,338 @@ def add_synchrony_to_sorting(sorting, sync_event_ratio=0.3, seed=None):
             continue
         new_unit_indices[i] = rng.choice(units_not_used)
         units_used_for_spike[sample_index] = np.append(units_used_for_spike[sample_index], new_unit_indices[i])
+
     spikes_duplicated["unit_index"] = new_unit_indices
-    spikes_all = np.concatenate((spikes, spikes_duplicated))
-    sort_idxs = np.lexsort([spikes_all["sample_index"], spikes_all["segment_index"]])
-    spikes_all = spikes_all[sort_idxs]
+    sort_idxs = np.lexsort([spikes_duplicated["sample_index"], spikes_duplicated["segment_index"]])
+    spikes_duplicated = spikes_duplicated[sort_idxs]
 
-    sorting = NumpySorting(spikes=spikes_all, sampling_frequency=sorting.sampling_frequency, unit_ids=unit_ids)
+    synchronous_spikes = NumpySorting(spikes_duplicated, sorting.get_sampling_frequency(), unit_ids)
+    sorting = TransformSorting.add_from_sorting(sorting, synchronous_spikes)
 
     return sorting
 
 
+def generate_sorting_to_inject(
+    sorting: BaseSorting,
+    num_samples: List[int],
+    max_injected_per_unit: int = 1000,
+    injected_rate: float = 0.05,
+    refractory_period_ms: float = 1.5,
+    seed=None,
+) -> NumpySorting:
+    """
+    Generates a sorting with spikes that are can be injected into the already existing sorting without violating
+    the refractory period.
+
+    Parameters
+    ----------
+    sorting : BaseSorting
+        The sorting object
+    num_samples: list of size num_segments.
+        The number of samples in all the segments of the sorting, to generate spike times
+        covering entire the entire duration of the segments
+    max_injected_per_unit: int, default 1000
+        The maximal number of spikes injected per units
+    injected_rate: float, default 0.05
+        The rate at which spikes are injected
+    refractory_period_ms: float, default 1.5
+        The refractory period that should not be violated while injecting new spikes
+    seed: int, default None
+        The random seed
+
+    Returns
+    -------
+    sorting : NumpySorting
+        The sorting object with the spikes to inject
+
+    """
+
+    injected_spike_trains = [{} for seg_index in range(sorting.get_num_segments())]
+    t_r = int(round(refractory_period_ms * sorting.get_sampling_frequency() * 1e-3))
+
+    rng = np.random.default_rng(seed=seed)
+
+    for segment_index in range(sorting.get_num_segments()):
+        for unit_id in sorting.unit_ids:
+            spike_train = sorting.get_unit_spike_train(unit_id, segment_index=segment_index)
+            n_injection = min(max_injected_per_unit, int(round(injected_rate * len(spike_train))))
+            # Inject more, then take out all that violate the refractory period.
+            n = int(n_injection + 10 * np.sqrt(n_injection))
+            injected_spike_train = np.sort(
+                np.random.uniform(low=0, high=num_samples[segment_index], size=n).astype(np.int64)
+            )
+
+            # Remove spikes that are in the refractory period.
+            violations = np.where(np.diff(injected_spike_train) < t_r)[0]
+            injected_spike_train = np.delete(injected_spike_train, violations)
+
+            # Remove spikes that violate the refractory period of the real spikes.
+            # TODO: Need a better & faster way than this.
+            min_diff = np.min(np.abs(injected_spike_train[:, None] - spike_train[None, :]), axis=1)
+            violations = min_diff < t_r
+            injected_spike_train = injected_spike_train[~violations]
+
+            if len(injected_spike_train) > n_injection:
+                injected_spike_train = np.sort(np.random.choice(injected_spike_train, n_injection, replace=False))
+
+            injected_spike_trains[segment_index][unit_id] = injected_spike_train
+
+    return NumpySorting.from_unit_dict(injected_spike_trains, sorting.get_sampling_frequency())
+
+
+class TransformSorting(BaseSorting):
+    """
+    Generates a sorting object keeping track of added spikes/units from an external spike_vector.
+    More precisely, the TransformSorting objects keeps two internal arrays added_spikes_from_existing_units and
+    added_spikes_from_new_units as boolean mask to track (in the representation as a spike vector) where
+    modifications have been made
+
+    Parameters
+    ----------
+    sorting : BaseSorting
+        The sorting object
+    added_spikes_existing_units : np.array (spike_vector)
+        The spikes that should be added to the sorting object, for existing units
+    added_spikes_new_units: np.array (spike_vector)
+        The spikes that should be added to the sorting object, for new units
+    new_units_ids: list
+        The unit_ids that should be added if spikes for new units are added
+    refractory_period_ms : float, default None
+        The refractory period violation to prevent duplicates and/or unphysiological addition
+        of spikes. Any spike times in added_spikes violating the refractory period will be
+        discarded
+
+    Returns
+    -------
+    sorting : TransformSorting
+        The sorting object with the added spikes and/or units
+    """
+
+    def __init__(
+        self,
+        sorting: BaseSorting,
+        added_spikes_existing_units=None,
+        added_spikes_new_units=None,
+        new_unit_ids: Optional[List[Union[str, int]]] = None,
+        refractory_period_ms: Optional[float] = None,
+    ):
+        sampling_frequency = sorting.get_sampling_frequency()
+        unit_ids = list(sorting.get_unit_ids())
+
+        if new_unit_ids is not None:
+            new_unit_ids = list(new_unit_ids)
+            assert ~np.any(
+                np.isin(new_unit_ids, sorting.unit_ids)
+            ), "some units ids are already present. Consider using added_spikes_existing_units"
+            if len(new_unit_ids) > 0:
+                assert type(unit_ids[0]) == type(new_unit_ids[0]), "unit_ids should have the same type"
+                unit_ids = unit_ids + list(new_unit_ids)
+
+        BaseSorting.__init__(self, sampling_frequency, unit_ids)
+
+        self.parent_unit_ids = sorting.unit_ids
+        self._cached_spike_vector = sorting.to_spike_vector().copy()
+        self.refractory_period_ms = refractory_period_ms
+
+        self.added_spikes_from_existing_mask = np.zeros(len(self._cached_spike_vector), dtype=bool)
+        self.added_spikes_from_new_mask = np.zeros(len(self._cached_spike_vector), dtype=bool)
+
+        if added_spikes_existing_units is not None and len(added_spikes_existing_units) > 0:
+            assert (
+                added_spikes_existing_units.dtype == minimum_spike_dtype
+            ), "added_spikes_existing_units should be a spike vector"
+            added_unit_indices = np.arange(len(self.parent_unit_ids))
+            self._cached_spike_vector = np.concatenate((self._cached_spike_vector, added_spikes_existing_units))
+            self.added_spikes_from_existing_mask = np.concatenate(
+                (self.added_spikes_from_existing_mask, np.ones(len(added_spikes_existing_units), dtype=bool))
+            )
+            self.added_spikes_from_new_mask = np.concatenate(
+                (self.added_spikes_from_new_mask, np.zeros(len(added_spikes_existing_units), dtype=bool))
+            )
+
+        if added_spikes_new_units is not None and len(added_spikes_new_units) > 0:
+            assert (
+                added_spikes_new_units.dtype == minimum_spike_dtype
+            ), "added_spikes_new_units should be a spike vector"
+            self._cached_spike_vector = np.concatenate((self._cached_spike_vector, added_spikes_new_units))
+            self.added_spikes_from_existing_mask = np.concatenate(
+                (self.added_spikes_from_existing_mask, np.zeros(len(added_spikes_new_units), dtype=bool))
+            )
+            self.added_spikes_from_new_mask = np.concatenate(
+                (self.added_spikes_from_new_mask, np.ones(len(added_spikes_new_units), dtype=bool))
+            )
+
+        sort_idxs = np.lexsort([self._cached_spike_vector["sample_index"], self._cached_spike_vector["segment_index"]])
+        self._cached_spike_vector = self._cached_spike_vector[sort_idxs]
+        self.added_spikes_from_existing_mask = self.added_spikes_from_existing_mask[sort_idxs]
+        self.added_spikes_from_new_mask = self.added_spikes_from_new_mask[sort_idxs]
+
+        # We need to add the sorting segments
+        for segment_index in range(sorting.get_num_segments()):
+            segment = SpikeVectorSortingSegment(self._cached_spike_vector, segment_index, unit_ids=self.unit_ids)
+            self.add_sorting_segment(segment)
+
+        if self.refractory_period_ms is not None:
+            self.clean_refractory_period()
+
+        self._kwargs = dict(
+            sorting=sorting,
+            added_spikes_existing_units=added_spikes_existing_units,
+            added_spikes_new_units=added_spikes_new_units,
+            new_unit_ids=new_unit_ids,
+            refractory_period_ms=refractory_period_ms,
+        )
+
+    @property
+    def added_spikes_mask(self):
+        return np.logical_or(self.added_spikes_from_existing_mask, self.added_spikes_from_new_mask)
+
+    def get_added_spikes_indices(self):
+        return np.nonzero(self.added_spikes_mask)[0]
+
+    def get_added_spikes_from_existing_indices(self):
+        return np.nonzero(self.added_spikes_from_existing_mask)[0]
+
+    def get_added_spikes_from_new_indices(self):
+        return np.nonzero(self.added_spikes_from_new_mask)[0]
+
+    def get_added_units_inds(self):
+        return self.unit_ids[len(self.parent_unit_ids) :]
+
+    @staticmethod
+    def add_from_sorting(sorting1: BaseSorting, sorting2: BaseSorting, refractory_period_ms=None) -> "TransformSorting":
+        """
+        Construct TransformSorting by adding one sorting to one other.
+
+        Parameters
+        ----------
+        sorting1: the first sorting
+        sorting2: the second sorting
+        refractory_period_ms : float, default None
+            The refractory period violation to prevent duplicates and/or unphysiological addition
+            of spikes. Any spike times in added_spikes violating the refractory period will be
+            discarded
+        """
+        assert (
+            sorting1.get_sampling_frequency() == sorting2.get_sampling_frequency()
+        ), "sampling_frequency should be the same"
+        assert type(sorting1.unit_ids[0]) == type(sorting2.unit_ids[0]), "unit_ids should have the same type"
+        # We detect the indices that are shared by the two sortings
+        mask1 = np.isin(sorting2.unit_ids, sorting1.unit_ids)
+        common_ids = sorting2.unit_ids[mask1]
+        exclusive_ids = sorting2.unit_ids[~mask1]
+
+        # We detect the indicies in the spike_vectors
+        idx1 = sorting1.ids_to_indices(common_ids)
+        idx2 = sorting2.ids_to_indices(common_ids)
+
+        spike_vector_2 = sorting2.to_spike_vector()
+        from_existing_units = np.isin(spike_vector_2["unit_index"], idx2)
+        common = spike_vector_2[from_existing_units].copy()
+
+        # If indices are not the same, we need to remap
+        if not np.all(idx1 == idx2):
+            old_indices = common["unit_index"].copy()
+            for i, j in zip(idx1, idx2):
+                mask = old_indices == j
+                common["unit_index"][mask] = i
+
+        idx1 = len(sorting1.unit_ids) + np.arange(len(exclusive_ids), dtype=int)
+        idx2 = sorting2.ids_to_indices(exclusive_ids)
+
+        not_common = spike_vector_2[~from_existing_units].copy()
+
+        # If indices are not the same, we need to remap
+        if not np.all(idx1 == idx2):
+            old_indices = not_common["unit_index"].copy()
+            for i, j in zip(idx1, idx2):
+                mask = old_indices == j
+                not_common["unit_index"][mask] = i
+
+        sorting = TransformSorting(
+            sorting1,
+            added_spikes_existing_units=common,
+            added_spikes_new_units=not_common,
+            new_unit_ids=exclusive_ids,
+            refractory_period_ms=refractory_period_ms,
+        )
+        return sorting
+
+    @staticmethod
+    def add_from_unit_dict(
+        sorting1: BaseSorting, units_dict_list: dict, refractory_period_ms=None
+    ) -> "TransformSorting":
+        """
+        Construct TransformSorting by adding one sorting with a
+        list of dict. The list length is the segment count.
+        Each dict have unit_ids as keys and spike times as values.
+
+        Parameters
+        ----------
+
+        sorting1: the first sorting
+        dict_list: list of dict
+        refractory_period_ms : float, default None
+            The refractory period violation to prevent duplicates and/or unphysiological addition
+            of spikes. Any spike times in added_spikes violating the refractory period will be
+            discarded
+        """
+        sorting2 = NumpySorting.from_unit_dict(units_dict_list, sorting1.get_sampling_frequency())
+        sorting = TransformSorting.add_from_sorting(sorting1, sorting2, refractory_period_ms)
+        return sorting
+
+    @staticmethod
+    def from_times_labels(
+        sorting1, times_list, labels_list, sampling_frequency, unit_ids=None, refractory_period_ms=None
+    ) -> "NumpySorting":
+        """
+        Construct TransformSorting from:
+          * an array of spike times (in frames)
+          * an array of spike labels and adds all the
+        In case of multisegment, it is a list of array.
+
+        Parameters
+        ----------
+        sorting1: the first sorting
+        times_list: list of array (or array)
+            An array of spike times (in frames)
+        labels_list: list of array (or array)
+            An array of spike labels corresponding to the given times
+        unit_ids: list or None, default: None
+            The explicit list of unit_ids that should be extracted from labels_list
+            If None, then it will be np.unique(labels_list)
+        refractory_period_ms : float, default None
+            The refractory period violation to prevent duplicates and/or unphysiological addition
+            of spikes. Any spike times in added_spikes violating the refractory period will be
+            discarded
+        """
+
+        sorting2 = NumpySorting.from_times_labels(times_list, labels_list, sampling_frequency, unit_ids)
+        sorting = TransformSorting.add_from_sorting(sorting1, sorting2, refractory_period_ms)
+        return sorting
+
+    def clean_refractory_period(self):
+        ## This function will remove the added spikes that will violate RPV, but does not affect the
+        ## spikes in the original sorting. So if some RPV violation are present in this sorting,
+        ## they will be left untouched
+        unit_indices = np.unique(self._cached_spike_vector["unit_index"])
+        rpv = int(self.get_sampling_frequency() * self.refractory_period_ms / 1000)
+        to_keep = ~self.added_spikes_from_existing_mask.copy()
+        for segment_index in range(self.get_num_segments()):
+            for unit_ind in unit_indices:
+                (indices,) = np.nonzero(
+                    (self._cached_spike_vector["unit_index"] == unit_ind)
+                    * (self._cached_spike_vector["segment_index"] == segment_index)
+                )
+                to_keep[indices[1:]] = np.logical_or(
+                    to_keep[indices[1:]], np.diff(self._cached_spike_vector[indices]["sample_index"]) > rpv
+                )
+
+        self._cached_spike_vector = self._cached_spike_vector[to_keep]
+        self.added_spikes_from_existing_mask = self.added_spikes_from_existing_mask[to_keep]
+        self.added_spikes_from_new_mask = self.added_spikes_from_new_mask[to_keep]
+
+
 def create_sorting_npz(num_seg, file_path):
     # create a NPZ sorting file
     d = {}
@@ -413,7 +735,7 @@ def inject_some_duplicate_units(sorting, num=4, max_shift=5, ratio=None, seed=No
 
     Parameters
     ----------
-    soring :
+    sorting :
         Original sorting
     num : int
         Number of injected units
@@ -444,6 +766,8 @@ def inject_some_duplicate_units(sorting, num=4, max_shift=5, ratio=None, seed=No
             unit_id: sorting.get_unit_spike_train(unit_id, segment_index=segment_index) for unit_id in sorting.unit_ids
         }
 
+        r = {}
+
         # inject some duplicate
         for i, unit_id in enumerate(other_ids):
             original_times = d[sorting.unit_ids[i]]
@@ -457,10 +781,11 @@ def inject_some_duplicate_units(sorting, num=4, max_shift=5, ratio=None, seed=No
             # clip inside 0 and last spike
             times = np.clip(times, 0, original_times[-1])
             times = np.sort(times)
-            d[unit_id] = times
-        spiketrains.append(d)
+            r[unit_id] = times
+        spiketrains.append(r)
 
-    sorting_with_dup = NumpySorting.from_unit_dict(spiketrains, sampling_frequency=sorting.get_sampling_frequency())
+    sorting_new_units = NumpySorting.from_unit_dict(spiketrains, sampling_frequency=sorting.get_sampling_frequency())
+    sorting_with_dup = TransformSorting.add_from_sorting(sorting, sorting_new_units)
 
     return sorting_with_dup
 


Test patch:
diff --git a/src/spikeinterface/comparison/tests/test_hybrid.py b/src/spikeinterface/comparison/tests/test_hybrid.py
index 7528848f3..144e7aacd 100644
--- a/src/spikeinterface/comparison/tests/test_hybrid.py
+++ b/src/spikeinterface/comparison/tests/test_hybrid.py
@@ -6,7 +6,6 @@ from spikeinterface.core.testing import check_recordings_equal
 from spikeinterface.comparison import (
     create_hybrid_units_recording,
     create_hybrid_spikes_recording,
-    generate_injected_sorting,
 )
 from spikeinterface.extractors import toy_example
 from spikeinterface.preprocessing import bandpass_filter
@@ -89,16 +88,8 @@ def test_hybrid_spikes_recording():
     check_recordings_equal(hybrid_spikes_recording, saved_2job, return_scaled=False)
 
 
-def test_generate_injected_sorting():
-    recording = load_extractor(cache_folder / "recording")
-    sorting = load_extractor(cache_folder / "sorting")
-    injected_sorting = generate_injected_sorting(
-        sorting, [recording.get_num_frames(seg_index) for seg_index in range(recording.get_num_segments())]
-    )
-
-
 if __name__ == "__main__":
     setup_module()
-    test_generate_injected_sorting()
+    test_generate_sorting_to_inject()
     test_hybrid_units_recording()
     test_hybrid_spikes_recording()
diff --git a/src/spikeinterface/core/tests/test_generate.py b/src/spikeinterface/core/tests/test_generate.py
index c1b7230b9..79d61a5ff 100644
--- a/src/spikeinterface/core/tests/test_generate.py
+++ b/src/spikeinterface/core/tests/test_generate.py
@@ -10,6 +10,7 @@ from spikeinterface.core.generate import (
     generate_recording,
     generate_sorting,
     NoiseGeneratorRecording,
+    TransformSorting,
     generate_recording_by_size,
     InjectTemplatesRecording,
     generate_single_fake_waveform,
@@ -17,8 +18,10 @@ from spikeinterface.core.generate import (
     generate_channel_locations,
     generate_unit_locations,
     generate_ground_truth_recording,
+    generate_sorting_to_inject,
 )
 
+from spikeinterface.core.numpyextractors import NumpySorting
 
 from spikeinterface.core.core_tools import convert_bytes_to_str
 
@@ -452,6 +455,60 @@ def test_inject_templates():
         check_recordings_equal(rec, saved_loaded, return_scaled=False)
 
 
+def test_transformsorting():
+    sorting_1 = generate_sorting(seed=0)
+    sorting_2 = generate_sorting(seed=1)
+    sorting_3 = generate_sorting(num_units=50, seed=1)
+
+    transformed_1 = TransformSorting(sorting_1, sorting_2.to_spike_vector())
+    n_spikes_1 = len(sorting_1.to_spike_vector())
+    n_spikes_2 = len(sorting_2.to_spike_vector())
+    n_spikes_added_1 = len(transformed_1.to_spike_vector())
+    assert n_spikes_added_1 == n_spikes_1 + n_spikes_2
+
+    transformed = TransformSorting.add_from_sorting(sorting_1, sorting_3)
+    assert len(transformed.unit_ids) == 50
+
+    sorting_1 = NumpySorting.from_unit_dict({46: np.array([0, 150], dtype=int)}, sampling_frequency=20000.0)
+    sorting_2 = NumpySorting.from_unit_dict(
+        {0: np.array([100, 2000], dtype=int), 3: np.array([200, 4000], dtype=int)}, sampling_frequency=20000.0
+    )
+    transformed = TransformSorting.add_from_sorting(sorting_1, sorting_2)
+    assert len(transformed.unit_ids) == 3
+    assert np.all(np.array([k for k in transformed.count_num_spikes_per_unit(outputs="array")]) == 2)
+
+    sorting_1 = NumpySorting.from_unit_dict({0: np.array([12], dtype=int)}, sampling_frequency=20000.0)
+    sorting_2 = NumpySorting.from_unit_dict(
+        {0: np.array([150], dtype=int), 3: np.array([12, 150], dtype=int)}, sampling_frequency=20000.0
+    )
+    transformed = TransformSorting.add_from_sorting(sorting_1, sorting_2)
+    assert len(transformed.unit_ids) == 2
+    target_array = np.array([2, 2])
+    source_array = np.array([k for k in transformed.count_num_spikes_per_unit(outputs="array")])
+    assert np.array_equal(source_array, target_array)
+
+    assert transformed.get_added_spikes_from_existing_indices().size == 1
+    assert transformed.get_added_spikes_from_new_indices().size == 2
+    assert transformed.get_added_units_inds() == [3]
+
+    transformed = TransformSorting.add_from_unit_dict(sorting_1, {46: np.array([12, 150], dtype=int)})
+
+    sorting_1 = generate_sorting(seed=0)
+    transformed = TransformSorting(sorting_1, sorting_1.to_spike_vector(), refractory_period_ms=0)
+    assert len(sorting_1.to_spike_vector()) == len(transformed.to_spike_vector())
+
+    transformed = TransformSorting(sorting_1, sorting_1.to_spike_vector(), refractory_period_ms=5)
+    assert 2 * len(sorting_1.to_spike_vector()) > len(transformed.to_spike_vector())
+
+    transformed_2 = TransformSorting(sorting_1, transformed.to_spike_vector())
+    assert len(transformed_2.to_spike_vector()) > len(transformed.to_spike_vector())
+
+    assert np.sum(transformed_2.get_added_spikes_indices()) >= np.sum(transformed_2.get_added_spikes_from_new_indices())
+    assert np.sum(transformed_2.get_added_spikes_indices()) >= np.sum(
+        transformed_2.get_added_spikes_from_existing_indices()
+    )
+
+
 def test_generate_ground_truth_recording():
     rec, sorting = generate_ground_truth_recording(upsample_factor=None)
     assert rec.templates.ndim == 3
@@ -460,6 +517,19 @@ def test_generate_ground_truth_recording():
     assert rec.templates.ndim == 4
 
 
+def test_generate_sorting_to_inject():
+    durations = [10.0, 20.0]
+    sorting = generate_sorting(num_units=10, durations=durations, sampling_frequency=30000, firing_rates=1.0)
+    injected_sorting = generate_sorting_to_inject(
+        sorting, [int(duration * sorting.sampling_frequency) for duration in durations]
+    )
+    num_spikes = sorting.count_num_spikes_per_unit()
+    num_injected_spikes = injected_sorting.count_num_spikes_per_unit()
+    # injected spikes should be less than original spikes
+    for unit_id in num_spikes:
+        assert num_injected_spikes[unit_id] <= num_spikes[unit_id]
+
+
 if __name__ == "__main__":
     strategy = "tile_pregenerated"
     # strategy = "on_the_fly"
@@ -471,7 +541,7 @@ if __name__ == "__main__":
     # test_noise_generator_consistency_after_dump(strategy, None)
     # test_generate_recording()
     # test_generate_single_fake_waveform()
-    test_generate_unit_locations()
+    test_transformsorting()
     # test_generate_templates()
     # test_inject_templates()
     # test_generate_ground_truth_recording()