Repository: SpikeInterface/spikeinterface
Pull request #1419 (2023-03-15 18:24:46)

Implement sorting concatenation
Hi,

Sorry the tests are not very well written, but it works !
The `ignore_times` kwarg is effectively forced to be True  in all cases in ConcatenateRecordings (I think?), but if behavior changes for the recording concatenation, `ignore_times=False` should work out of the box for the sorting concatenation.
Let me know if you have comments!
Best


Patch:
diff --git a/doc/modules/core.rst b/doc/modules/core.rst
index 5f3937f0d..3ac642ff1 100644
--- a/doc/modules/core.rst
+++ b/doc/modules/core.rst
@@ -639,9 +639,12 @@ We have several functions to manipulate segments of SpikeInterface objects. All
 
 
 
-The same functions, except for the :code:`concatenate_*` one, are also available for 
-:py:class:`~spikeinterface.core.BaseSorting` objects (:py:func:`~spikeinterface.core.append_sortings`, 
-:py:func:`~spikeinterface.core.split_sorting`, :py:func:`~spikeinterface.core.select_segment_sorting`).
+The same functions are also available for 
+:py:class:`~spikeinterface.core.BaseSorting` objects
+(:py:func:`~spikeinterface.core.append_sortings`,
+:py:func:`~spikeinterface.core.concatenate_sortings`, 
+:py:func:`~spikeinterface.core.split_sorting`,
+:py:func:`~spikeinterface.core.select_segment_sorting`).
 
 
 **Note** :py:func:`~spikeinterface.core.append_recordings` and:py:func:`~spikeinterface.core.concatenate_recordings`
diff --git a/spikeinterface/core/__init__.py b/spikeinterface/core/__init__.py
index ccb232cd0..c5cbfdedd 100644
--- a/spikeinterface/core/__init__.py
+++ b/spikeinterface/core/__init__.py
@@ -42,6 +42,8 @@ from .segmentutils import (
     SelectSegmentRecording,
     append_sortings,
     AppendSegmentSorting,
+    concatenate_sortings,
+    ConcatenateSegmentSorting,
     split_sorting,
     SplitSegmentSorting,
     select_segment_sorting,
diff --git a/spikeinterface/core/basesorting.py b/spikeinterface/core/basesorting.py
index 5381feb28..8232c3159 100644
--- a/spikeinterface/core/basesorting.py
+++ b/spikeinterface/core/basesorting.py
@@ -1,8 +1,10 @@
-from typing import List, Union, Optional
-import numpy as np
 import warnings
+from typing import List, Optional, Union
+
+import numpy as np
 
 from .base import BaseExtractor, BaseSegment
+from .waveform_tools import has_exceeding_spikes
 
 
 class BaseSorting(BaseExtractor):
@@ -53,6 +55,51 @@ class BaseSorting(BaseExtractor):
     def get_num_segments(self):
         return len(self._sorting_segments)
 
+    def get_num_samples(self, segment_index=None):
+        """Returns the number of samples of the associated recording for a segment.
+
+        Parameters
+        ----------
+        segment_index : int, optional
+            The segment index to retrieve the number of samples for. 
+            For multi-segment objects, it is required, by default None
+
+        Returns
+        -------
+        int
+            The number of samples
+        """
+        assert self.has_recording(), (
+            "This methods requires an associated recording. Call self.register_recording() first."
+        )
+        return self._recording.get_num_samples(segment_index=segment_index)
+
+    def get_total_samples(self):
+        """Returns the total number of samples of the associated recording.
+
+        Returns
+        -------
+        int
+            The total number of samples
+        """
+        s = 0
+        for segment_index in range(self.get_num_segments()):
+            s += self.get_num_samples(segment_index)
+        return s
+
+    def get_total_duration(self):
+        """Returns the total duration in s of the associated recording.
+
+        Returns
+        -------
+        float
+            The duration in seconds
+        """
+        assert self.has_recording(), (
+            "This methods requires an associated recording. Call self.register_recording() first."
+        )
+        return self._recording.get_total_duration()
+
     def get_unit_spike_train(
         self,
         unit_id,
@@ -76,10 +123,31 @@ class BaseSorting(BaseExtractor):
         else:
             return spike_frames
 
-    def register_recording(self, recording):
+    def register_recording(self, recording, check_spike_frames=True):
+        """Register a recording to the sorting.
+
+        Parameters
+        ----------
+        recording : BaseRecording
+            Recording with the same number of segments as current sorting.
+            Assigned to self._recording.
+        check_spike_frames : bool, optional
+            If True, assert for each segment that all spikes are within the recording's range.
+            By default True.
+        """
         assert np.isclose(self.get_sampling_frequency(),
                           recording.get_sampling_frequency(),
                           atol=0.1), "The recording has a different sampling frequency than the sorting!"
+        assert self.get_num_segments() == recording.get_num_segments(), (
+            "The recording has a different number of segments than the sorting!"
+        )
+        if check_spike_frames:
+            if has_exceeding_spikes(recording, self):
+                warnings.warn(
+                    "Some spikes are exceeding the recording's duration! "
+                    "Removing these excess spikes with `spikeinterface.curation.remove_excess_spikes()` "
+                    "Might be necessary for further postprocessing."
+                )
         self._recording = recording
 
     @property
diff --git a/spikeinterface/core/segmentutils.py b/spikeinterface/core/segmentutils.py
index 022b881f7..eb982d151 100644
--- a/spikeinterface/core/segmentutils.py
+++ b/spikeinterface/core/segmentutils.py
@@ -314,6 +314,184 @@ class ProxyAppendSortingSegment(BaseSortingSegment):
 append_sortings = define_function_from_class(source_class=AppendSegmentSorting, name='append_sortings')
 
 
+class ConcatenateSegmentSorting(BaseSorting):
+    """
+    Return a sorting that "concatenates" all segments from all sorting
+    into one sorting with a single segment. This operation is lazy.
+
+    For instance, given one recording with 2 segments and one recording with
+    3 segments, this class will give one recording with one large segment
+    made by concatenating the 5 segments. The returned spike times (originating
+    from each segment) are returned as relative to the start of the concatenated segment.
+
+    Time information is lost upon concatenation. By default `ignore_times` is True.
+    If it is False, you get an error unless:
+
+      * all segments DO NOT have times, AND
+      * all segment have t_start=None
+
+    Parameters
+    ----------
+    sorting_list : list of BaseSorting
+        A list of sortings. If `total_samples_list` is not provided, all
+        sortings should have an assigned recording.  Otherwise, all sortings
+        should be monosegments.
+    total_samples_list : list[int] or None
+        If the sortings have no assigned recording, the total number of samples
+        of each of the concatenated (monosegment) sortings is pulled from this
+        list.
+    ignore_times : bool
+        If True (default), time information (t_start, time_vector) is ignored
+        when concatenating the sortings' assigned recordings.
+    sampling_frequency_max_diff : float
+        Maximum allowed difference of sampling frequencies across sortings (default 0)
+    """
+
+    def __init__(self, sorting_list, total_samples_list=None, ignore_times=True, sampling_frequency_max_diff=0):
+
+        # Check that all sortings have a recording or that sortings' num_samples are provided
+        all_has_recording = all([sorting.has_recording() for sorting in sorting_list])
+        if not all_has_recording:
+            assert total_samples_list is not None, (
+                "Some concatenated sortings don't have a registered recording. "
+                "Call sorting.register_recording() or set `total_samples_list` kwarg."
+            )
+            assert len(total_samples_list) == len(sorting_list), (
+                "`total_samples_list` should have the same number of elements as `sorting_list`"
+            )
+            assert all([s.get_num_segments() == 1 for s in sorting_list]), (
+                "All sortings are expected to be monosegment."
+            )
+            assert ignore_times, (
+                "Concatenating sortings without registered recordings: "
+                "Use ignore_times=True to ignore time information."
+            )
+        else:
+            assert total_samples_list is None, "Sortings have registered recordings: Use `total_samples_list=None`"
+
+        # Pull metadata from AppendSorting object
+        one_sorting = append_sortings(sorting_list, sampling_frequency_max_diff=sampling_frequency_max_diff)
+        BaseSorting.__init__(self, one_sorting.sampling_frequency, one_sorting.unit_ids)
+        one_sorting.copy_metadata(self)
+
+        # Check and pull n_samples from each segment
+        parent_segments = []
+        parent_num_samples = []
+        for sorting_i, sorting in enumerate(sorting_list):
+            for segment_i, parent_segment in enumerate(sorting._sorting_segments):
+                # Check t_start is not assigned
+                segment_t_start = parent_segment._t_start
+                if not ignore_times:
+                    assert segment_t_start is None, ("ConcatenateSegmentSorting does not handle Sorting.t_start. "
+                                                     "Set time information only in the sortings' assigned recordings, "
+                                                     "or use ignore_times=True to ignore time information.")
+                # Pull num samples for each segment
+                if sorting.has_recording():
+                    segment_num_samples = sorting.get_num_samples(segment_index=segment_i)
+                else:
+                    segment_num_samples = total_samples_list[sorting_i]
+                # Check consistency between num samples and spike frames
+                for unit_id in sorting.unit_ids:
+                    unit_segment_spikes = parent_segment.get_unit_spike_train(
+                        unit_id=unit_id, start_frame=None, end_frame=None,
+                    )
+                    if any([spike_frame >= segment_num_samples for spike_frame in unit_segment_spikes]):
+                        raise ValueError(
+                            "Sortings' spike frames exceed the provided number of samples for some segment. "
+                            "If the sortings have registered recordings, you can remove these excess "
+                            "spikes with `spikeinterface.curation.remove_excess_spikes(sorting, sorting._recording)`. "
+                            "Otherwise, the `total_num_samples` argument may contain invalid values."
+                        )
+                parent_segments.append(parent_segment)
+                parent_num_samples.append(segment_num_samples)
+        self.parent_num_samples = parent_num_samples
+
+        # Add a single Concatenated segment
+        sorting_seg = ProxyConcatenateSortingSegment(parent_segments, parent_num_samples,
+                                                     one_sorting.get_sampling_frequency())
+        self.add_sorting_segment(sorting_seg)
+
+        # Assign concatenated recording if possible
+        if all_has_recording:
+            self.register_recording(
+                concatenate_recordings([s._recording for s in sorting_list], ignore_times=ignore_times)
+            )
+
+        self._kwargs = {'sorting_list': [sorting.to_dict() for sorting in sorting_list],
+                        'ignore_times': ignore_times}
+
+    def get_num_samples(self, segment_index=None):
+        """Overrides the BaseSorting method, which requires a recording."""
+        segment_index = self._check_segment_index(segment_index)
+        n_samples = self._sorting_segments[segment_index].get_num_samples()
+        if self.has_recording():  # Sanity check
+            assert n_samples == self._recording.get_num_samples(segment_index)
+        return n_samples
+
+
+class ProxyConcatenateSortingSegment(BaseSortingSegment):
+    def __init__(self, parent_segments, parent_num_samples, sampling_frequency):
+        BaseSortingSegment.__init__(self)
+        self.parent_segments = parent_segments
+        self.parent_num_samples = parent_num_samples
+        self.cumsum_length = np.cumsum([0] + self.parent_num_samples)
+        self.total_num_samples = np.sum(self.parent_num_samples)
+
+    def get_num_samples(self):
+        return self.total_num_samples
+
+    def get_unit_spike_train(
+        self,
+        unit_id,
+        start_frame,
+        end_frame,
+    ):
+        if start_frame is None:
+            start_frame = 0
+        if end_frame is None:
+            end_frame = self.get_num_samples()
+
+        i0 = np.searchsorted(self.cumsum_length, start_frame, side='right') - 1
+        i1 = np.searchsorted(self.cumsum_length, end_frame, side='right') - 1
+
+        # several case:
+        #  * come from one segment (i0 == i1)
+        #  * come from several segment (i0 < i1)
+        if i0 == i1:
+            #  one segment
+            sorting_seg = self.parent_segments[i0]
+            seg_start = self.cumsum_length[i0]
+            spike_frames = sorting_seg.get_unit_spike_train(unit_id, start_frame - seg_start, end_frame - seg_start) + seg_start
+        else:
+            #  several segments
+            all_spike_frames = []
+            for i in range(i0, i1 + 1):
+                if i == len(self.parent_segments):
+                    # limit case
+                    continue
+
+                sorting_seg = self.parent_segments[i]
+                seg_start = self.cumsum_length[i]
+                if i == i0:
+                    # first
+                    spike_frames_chunk = sorting_seg.get_unit_spike_train(unit_id, start_frame - seg_start, None) + seg_start
+                    all_spike_frames.append(spike_frames_chunk)
+                elif i == i1:
+                    # last
+                    if (end_frame - seg_start) > 0:
+                        spike_frames_chunk = sorting_seg.get_unit_spike_train(unit_id, None, end_frame - seg_start) + seg_start
+                        all_spike_frames.append(spike_frames_chunk)
+                else:
+                    # in between
+                    spike_frames_chunk = sorting_seg.get_unit_spike_train(unit_id, None, None) + seg_start
+                    all_spike_frames.append(spike_frames_chunk)
+            spike_frames = np.concatenate(all_spike_frames, axis=0)
+
+        return spike_frames
+
+concatenate_sortings = define_function_from_class(source_class=ConcatenateSegmentSorting, name='concatenate_sortings')
+
+
 class SplitSegmentSorting(BaseSorting):
     """Splits a sorting with a single segment to multiple segments
     based on the given list of recordings (must be in order)
@@ -400,4 +578,4 @@ class SelectSegmentSorting(BaseSorting):
 
 
 select_segment_sorting = define_function_from_class(source_class=SelectSegmentSorting,
-                                                    name='select_segment_sorting')
\ No newline at end of file
+                                                    name='select_segment_sorting')


Test patch:
diff --git a/spikeinterface/core/testing.py b/spikeinterface/core/testing.py
index 3cddc5e45..bcf600a1c 100644
--- a/spikeinterface/core/testing.py
+++ b/spikeinterface/core/testing.py
@@ -1,6 +1,13 @@
 import numpy as np
 from numpy.testing import assert_array_equal
 
+
+def check_sorted_arrays_equal(a1, a2):
+    a1 = np.sort(np.array(a1))
+    a2 = np.sort(np.array(a2))
+    assert_array_equal(a1, a2)
+
+
 def check_recordings_equal(RX1, RX2, return_scaled=True, force_dtype=None):
     assert RX1.get_num_segments() == RX2.get_num_segments()
 
diff --git a/spikeinterface/core/tests/test_basesorting.py b/spikeinterface/core/tests/test_basesorting.py
index 9fc5b468d..5f1d8dd71 100644
--- a/spikeinterface/core/tests/test_basesorting.py
+++ b/spikeinterface/core/tests/test_basesorting.py
@@ -4,13 +4,16 @@ but check only for BaseRecording general methods.
 """
 import shutil
 from pathlib import Path
-import pytest
+
 import numpy as np
+import pytest
+from numpy.testing import assert_raises
 
-from spikeinterface.core import NumpySorting, NpzSortingExtractor, load_extractor
+from spikeinterface.core import (NpzSortingExtractor, NumpyRecording,
+                                 NumpySorting, create_sorting_npz,
+                                 generate_sorting, load_extractor)
 from spikeinterface.core.base import BaseExtractor
-
-from spikeinterface.core import create_sorting_npz, generate_sorting
+from spikeinterface.core.testing import check_sorted_arrays_equal
 
 if hasattr(pytest, "global_test_folder"):
     cache_folder = pytest.global_test_folder / "core"
@@ -89,6 +92,44 @@ def test_BaseSorting():
         assert unit not in empty_units
 
 
+def test_npy_sorting():
+    sfreq = 10
+    spike_times_0 = {
+        "0": np.array([0, 1, 9]), # Max sample idx is 9 for a rec of length 10
+        "1": np.array([2, 5]),
+    }
+    spike_times_1 = {
+        "0": np.array([0, 1]),
+        "1": np.array([], dtype="int64"),
+    }
+    sorting = NumpySorting.from_dict(
+        [spike_times_0, spike_times_1], sfreq,
+    )
+
+    assert sorting.get_num_segments() == 2
+    assert set(sorting.get_unit_ids()) == set(["0", "1"])
+    check_sorted_arrays_equal(sorting.get_unit_spike_train(segment_index=0, unit_id="1"), [2, 5])
+
+    # Check registering a recording
+    seg_nframes = [10, 5]
+    rec = NumpyRecording([np.zeros((nframes, 10)) for nframes in seg_nframes], sampling_frequency=sfreq)
+    sorting.register_recording(rec)
+    assert sorting.get_num_samples(segment_index=0) == 10
+    assert sorting.get_num_samples(segment_index=1) == 5
+    assert sorting.get_total_samples() == 15
+
+    # Registering too short a recording raises a warning
+    seg_nframes = [9, 5]
+    rec = NumpyRecording([np.zeros((nframes, 10)) for nframes in seg_nframes], sampling_frequency=sfreq)
+    # assert_raises(Exception, sorting.register_recording, rec)
+    with pytest.warns():
+        sorting.register_recording(rec)
+
+    # Registering a rec with too many segments
+    seg_nframes = [9, 5, 10]
+    rec = NumpyRecording([np.zeros((nframes, 10)) for nframes in seg_nframes], sampling_frequency=sfreq)
+    assert_raises(Exception, sorting.register_recording, rec)
+
 def test_empty_sorting():
     sorting = NumpySorting.from_dict({}, 30000)
 
@@ -105,4 +146,5 @@ def test_empty_sorting():
 
 if __name__ == '__main__':
     test_BaseSorting()
+    test_npy_sorting()
     test_empty_sorting()
diff --git a/spikeinterface/core/tests/test_segmentutils.py b/spikeinterface/core/tests/test_segmentutils.py
index a600516df..4373b5dc6 100644
--- a/spikeinterface/core/tests/test_segmentutils.py
+++ b/spikeinterface/core/tests/test_segmentutils.py
@@ -1,21 +1,26 @@
-import pytest
-import numpy as np
+import warnings
 
-from spikeinterface.core import (
-    append_recordings, AppendSegmentRecording,
-    concatenate_recordings, ConcatenateSegmentRecording,
-    append_sortings, AppendSegmentSorting,
-    split_recording, select_segment_recording, split_sorting)
+import numpy as np
+import pytest
+from numpy.testing import assert_raises
 
-from spikeinterface.core import NumpyRecording, NumpySorting
+from spikeinterface.core import (AppendSegmentRecording, AppendSegmentSorting,
+                                 ConcatenateSegmentRecording,
+                                 ConcatenateSegmentSorting, NumpyRecording,
+                                 NumpySorting, append_recordings,
+                                 append_sortings, concatenate_recordings,
+                                 concatenate_sortings,
+                                 select_segment_recording, split_recording,
+                                 split_sorting)
+from spikeinterface.core.testing import check_sorted_arrays_equal
 
 
 def test_append_concatenate_recordings():
-    traces = np.zeros((1000, 5), dtype='float64')
-    traces[:] = np.arange(1000)[:, None]
+    base_traces = np.zeros((1000, 5), dtype='float64')
+    base_traces[:] = np.arange(1000)[:, None]
     sampling_frequency = 30000
-    rec0 = NumpyRecording([traces] * 3, sampling_frequency)
-    rec1 = NumpyRecording([traces] * 2, sampling_frequency)
+    rec0 = NumpyRecording([base_traces] * 3, sampling_frequency)
+    rec1 = NumpyRecording([base_traces] * 2, sampling_frequency)
 
     # append
     rec = append_recordings([rec0, rec1])
@@ -34,25 +39,32 @@ def test_append_concatenate_recordings():
 
     # case one segment
     traces = rec.get_traces(start_frame=0, end_frame=15)
-    assert np.array_equal(traces[:, 0], np.arange(0, 15))
+    assert np.array_equal(traces, base_traces[:15])
+    traces = rec.get_traces(start_frame=0, end_frame=1000)
+    assert np.array_equal(traces, base_traces)
     traces = rec.get_traces(start_frame=500, end_frame=750)
-    assert np.array_equal(traces[:, 0], np.arange(500, 750))
+    assert np.array_equal(traces, base_traces[500:750])
 
     # case on limit
     traces = rec.get_traces(start_frame=1000, end_frame=2000)
     assert traces.shape == (1000, 5)
-    assert np.array_equal(traces[:, 0], np.arange(0, 1000, dtype='float64'))
+    assert np.array_equal(traces, base_traces)
 
     # case total
     traces = rec.get_traces(start_frame=None, end_frame=None)
     assert traces.shape == (5000, 5)
-    assert traces[1000, 0] == 0
+    assert np.array_equal(traces[:1000], base_traces)
+    assert np.array_equal(traces[1000:2000], base_traces)
+    assert np.array_equal(traces[4000:5000], base_traces)
 
     # several segments
     traces = rec.get_traces(start_frame=50, end_frame=4500)
     assert traces.shape == (4450, 5)
-    assert traces[0, 0] == 50
-    assert traces[-1, 0] == 499
+    assert np.array_equal(traces[0:10], base_traces[50:60])
+    assert np.array_equal(traces[-10:], base_traces[490:500])
+    # across segments
+    assert np.array_equal(traces[949,:], base_traces[-1,:])
+    assert np.array_equal(traces[950,:], base_traces[0,:])
 
 
 def test_split_recordings():
@@ -78,23 +90,206 @@ def test_split_recordings():
     assert np.allclose(rec3.get_traces(), rec0.get_traces(segment_index=2))
 
 
-def test_append_sortings():
+def test_append_concatenate_sortings():
     sampling_frequency = 30000.
-    times = np.arange(0, 1000, 10)
-    labels = np.zeros(times.size, dtype='int64')
-    labels[0::3] = 0
-    labels[1::3] = 1
-    labels[2::3] = 2
+    nsamp0 = 1000
+    nsamp1 = 1500
+    times0 = np.arange(0, nsamp0)
+    times1 = np.arange(0, nsamp1-1) + 1
+    labels0 = np.zeros(times0.size, dtype='int64')
+    labels0[0::3] = 0
+    labels0[1::3] = 1
+    labels0[2::3] = 2
+    labels1 = np.zeros(times1.size, dtype='int64')
+    labels1[0::3] = 0
+    labels1[1::3] = 1
+    labels1[2::3] = 2
+
+    # Multisegment sortings
     sorting0 = NumpySorting.from_times_labels(
-        [times] * 3, [labels] * 3, sampling_frequency)
+        [times0] * 3, [labels0] * 3, sampling_frequency)
     sorting1 = NumpySorting.from_times_labels(
-        [times] * 2, [labels] * 2, sampling_frequency)
+        [times1] * 2, [labels1] * 2, sampling_frequency)
+    sorting_list = [sorting0, sorting1]
+    # Associated multisegment recordings
+    traces0 = np.zeros((nsamp0, 5), dtype='float64')
+    traces1 = np.zeros((nsamp1, 5), dtype='float64')
+    rec0 = NumpyRecording([traces0] * 3, sampling_frequency)
+    rec1 = NumpyRecording([traces1] * 2, sampling_frequency)
+    # Monosegemnt sortings
+    sorting0_mono = NumpySorting.from_times_labels(
+        [times0], [labels0], sampling_frequency)
+    sorting1_mono = NumpySorting.from_times_labels(
+        [times1], [labels1], sampling_frequency)
+    sorting_list_mono = [sorting0_mono, sorting1_mono]
+    # Associated multisegment recordings
+    traces0 = np.zeros((nsamp0, 5), dtype='float64')
+    traces1 = np.zeros((nsamp1, 5), dtype='float64')
+    rec0_mono = NumpyRecording([traces0], sampling_frequency)
+    rec1_mono = NumpyRecording([traces1], sampling_frequency)
+    # Recording too short
+    traces0_short = np.zeros((nsamp0-1, 5), dtype='float64')
+    rec0_mono_short = NumpyRecording([traces0_short], sampling_frequency)
 
-    sorting = append_sortings([sorting0, sorting1])
+    # Append
+    # Append multisegment
+    sorting = append_sortings(sorting_list)
     # print(sorting)
     assert sorting.get_num_segments() == 5
+    assert np.array_equal(
+        sorting.get_unit_spike_train(unit_id=0, segment_index=0),
+        times0[0::3],
+    )
+    assert np.array_equal(
+        sorting.get_unit_spike_train(unit_id=0, segment_index=3),
+        times1[0::3],
+    )
+    # Append monosegment
+    sorting_mono = append_sortings(sorting_list_mono)
+    assert sorting_mono.get_num_segments() == 2
+    assert np.array_equal(
+        sorting_mono.get_unit_spike_train(unit_id=0, segment_index=0),
+        times0[0::3],
+    )
+    assert np.array_equal(
+        sorting_mono.get_unit_spike_train(unit_id=0, segment_index=1),
+        times1[0::3],
+    )
+
+    # Concat
+    # Fails without registered recording
+    assert_raises(Exception, concatenate_sortings, sorting_list)
+    assert_raises(Exception, concatenate_sortings, sorting_list_mono)
+    # Fails with total_samples_list for multisegment
+    assert_raises(Exception, concatenate_sortings, sorting_list, total_samples_list="dummy")
+    # Succeeds without registered recording for mono
+    sorting_mono_norec = concatenate_sortings(sorting_list_mono, total_samples_list=[nsamp0, nsamp1])
+    assert sorting_mono_norec.get_num_segments() == 1
+    # Fails when excess spikes with total_samples_list
+    assert_raises(Exception, concatenate_sortings, sorting_list, total_samples_list=[nsamp0 - 1, nsamp1])
+    # Fails when excess spikes with registered recording
+    with warnings.catch_warnings():
+        warnings.filterwarnings('ignore')
+        sorting0_mono.register_recording(rec0_mono_short)
+    assert_raises(Exception, concatenate_sortings, [sorting0_mono])
+    
+    # Fails 
+
+    # With registered recording
+    sorting0.register_recording(rec0)
+    sorting1.register_recording(rec1)
+    sorting0_mono.register_recording(rec0_mono)
+    sorting1_mono.register_recording(rec1_mono)
+    sorting = concatenate_sortings(sorting_list)
+    sorting_mono = concatenate_sortings(sorting_list_mono)
+    # Check Segments
+    assert sorting.get_num_segments() == 1
+    assert sorting_mono.get_num_segments() == 1
+    # Check nsamps
+    assert sorting.get_total_samples() == 3 * nsamp0 + 2 * nsamp1
+    assert sorting_mono.get_total_samples() == nsamp0 + nsamp1
+    assert sorting.get_num_samples() == 3 * nsamp0 + 2 * nsamp1
+    assert sorting_mono.get_num_samples() == nsamp0 + nsamp1
+
+    # Check spike trains
+    unit_0_train = np.concatenate(
+        (
+            times0[::3],
+            nsamp0 + times0[::3],
+            2 * nsamp0 + times0[::3],
+            3 * nsamp0 + times1[::3],
+            (3 * nsamp0 + nsamp1) + times1[::3],
+        ),
+        axis=0,
+    )
+
+    # Within one segment
+    assert np.array_equal(
+        sorting.get_unit_spike_train(unit_id=0, segment_index=None, start_frame=0, end_frame=1000),
+        [t for t in unit_0_train if 0<=t<1000]
+    ) # Full first
+    assert np.array_equal(
+        sorting.get_unit_spike_train(unit_id=0, segment_index=None, start_frame=0, end_frame=300),
+        [t for t in unit_0_train if 0<=t<300]
+    ) # Part first
+    assert np.array_equal(
+        sorting.get_unit_spike_train(unit_id=0, segment_index=None, start_frame=500, end_frame=750),
+        [t for t in unit_0_train if 500<=t<750]
+    ) # Part first
+    assert np.array_equal(
+        sorting.get_unit_spike_train(unit_id=0, segment_index=None, start_frame=1000, end_frame=2000),
+        [t for t in unit_0_train if 1000<=t<2000]
+    ) # Full 2nd segment
+    assert np.array_equal(
+        sorting.get_unit_spike_train(unit_id=0, segment_index=None, start_frame=4500, end_frame=6000),
+        [t for t in unit_0_train if 4500<=t<6000]
+    ) # Full last segment
+    assert np.array_equal(
+        sorting.get_unit_spike_train(unit_id=0, segment_index=None, start_frame=4550, end_frame=5000),
+        [t for t in unit_0_train if 4550<=t<5000]
+    ) # Part of last segmetn
+
+    # Total
+    assert np.array_equal(sorting.get_unit_spike_train(unit_id=0), unit_0_train)
+
+    # Across segments
+    assert np.array_equal(
+        sorting.get_unit_spike_train(unit_id=0, segment_index=None, start_frame=950, end_frame=1050),
+        [t for t in unit_0_train if 950<=t<1050]
+    )
+
+    # Check same result with and without rec for mono segment
+    assert np.array_equal(
+        sorting_mono.get_unit_spike_train(unit_id=0),
+        sorting_mono_norec.get_unit_spike_train(unit_id=0),
+    )
+
+    # Slicing and concatenating back
+    # Full sorting
+    sorting_beginning = sorting.frame_slice(None, 1000)
+    sorting_end = sorting.frame_slice(1000, None)
+    sorting_concat = concatenate_sortings([sorting_beginning, sorting_end])
+    assert np.array_equal(
+        sorting.get_unit_spike_train(unit_id=0),
+        sorting_concat.get_unit_spike_train(unit_id=0),
+    )
+    # Slice in the middle
+    sorting_beginning = sorting.frame_slice(50, 1000)
+    sorting_end = sorting.frame_slice(1000, 2000)
+    sorting_concat = concatenate_sortings([sorting_beginning, sorting_end])
+    assert np.array_equal(
+        sorting.get_unit_spike_train(unit_id=0, start_frame=50, end_frame=2000) - 50,
+        sorting_concat.get_unit_spike_train(unit_id=0),
+    )
+
+    # Time conversion
+    # Basic (no time vector)
+    assert np.array_equal(
+        sorting.get_unit_spike_train(unit_id=0, return_times=True),
+        unit_0_train / sampling_frequency 
+    )
+
+    # Time conversion with non-trivial t_start/time_vector
+    sorting0 = NumpySorting.from_times_labels( [[0, 1]], [[0, 0]], sampling_frequency)
+    sorting1 = NumpySorting.from_times_labels( [[0, 1]], [[0, 0]], sampling_frequency)
+    traces = np.zeros((2, 5), dtype='float64')
+    rec0 = NumpyRecording([traces], sampling_frequency)
+    rec1 = NumpyRecording([traces], sampling_frequency)
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        rec0.set_times(np.array([1000, 1010]))
+        rec1.set_times(np.array([2000, 2010]))
+    sorting0.register_recording(rec0)
+    sorting1.register_recording(rec1)
+    # Actually, non-trivial time is not handled by concatenate_recordings
+    assert_raises(Exception, concatenate_sortings, [sorting0, sorting1], ignore_times=False)
+    # sorting = concatenate_sortings([sorting0, sorting1], ignore_times=False)
+    # assert np.array_equal(
+    #     sorting.get_unit_spike_train(unit_id=0, return_times=True),
+    #     [1000, 1010, 2000, 2010],
+    # )
 
 
 if __name__ == '__main__':
     test_append_concatenate_recordings()
-    test_append_sortings()
+    test_append_concatenate_sortings()