Repository: SpikeInterface/spikeinterface
Pull request #464 (2022-01-17 09:21:34)

Add recording link to sorting
Added `register_recording` to link a recording to a sorting in memory. This enables to propagate the `get_times` mechanism to the sorting object.

* The `BaseSorter` automatically registers the recording to the output sorting object
* recording object is propagated to `UnitsSelectionSorting`, but it cannot be propagated to `UnitsAggregationSorting` and `FrameSliceSorting` (as they would require more complicated times management)
* Added a `return_times` option to return spike trains as times in second (either using the registered recording or the sampling frequency)

- [x] add tests
- [x] clear examples and documentation


Patch:
diff --git a/doc/module_core.rst b/doc/module_core.rst
index 995533073..42f12eb95 100644
--- a/doc/module_core.rst
+++ b/doc/module_core.rst
@@ -25,6 +25,7 @@ It represents an extracellular recording and has the following features:
 * store object annotations
 * enable grouping, splitting, and slicing
 * handle segment operations (e.g. concatenation)
+* handle time information
 
 
 Sorting
@@ -38,6 +39,7 @@ It represents a spike sorted output and has the following features:
 * store channel properties
 * store object annotations
 * enable selection of sub-units
+* handle time information
 
 
 Event
diff --git a/examples/modules/core/plot_6_handle_times.py b/examples/modules/core/plot_6_handle_times.py
new file mode 100644
index 000000000..87ce34611
--- /dev/null
+++ b/examples/modules/core/plot_6_handle_times.py
@@ -0,0 +1,52 @@
+"""
+Handle time information
+=======================
+
+By default, :code:`spikeinterface` assumes that a recording is uniformly sampled and it starts at 0 seconds.
+However, in some cases there could be a different start time or even some missing frames in the recording.
+
+This notebook shows how to handle time informaiton in :code:`spikeinterface` recording and sorting objects.
+"""
+from spikeinterface.extractors import toy_example
+
+##############################################################################
+# First let's generate toy example with a single segment:
+
+rec, sort = toy_example(num_segments=1)
+
+
+##############################################################################
+# Generally, the time information would be automaticall loaded when reading a 
+# recording. 
+# However, sometimes we might need to add a time vector externally.
+# For example, now let's create a time vector by getting the default times and 
+# adding 5 s:
+
+default_times = rec.get_times()
+print(default_times[:10])
+new_times = default_times + 5
+
+##############################################################################
+# We can now set the new time vector with the :code:`set_times()` function.
+# Additionally, we can register to recording object to the sorting one so that 
+# time information can be accessed by the sorting object as well (note that this 
+# link is lost in case the sorting object is saved to disk!):
+
+rec.set_times(new_times)
+sort.register_recording(rec)
+
+# print new times
+print(rec.get_times()[:10])
+
+# print spike times (internally uses registered recording times)
+spike_times0 = sort.get_unit_spike_train(sort.unit_ids[0], return_times=True)
+print(spike_times0[:10])
+
+
+##############################################################################
+# While here we have shown how to set times only for a mono-segment recording,
+# times can also be handled in multi-segment recordings (using the 
+# :code:`segment_index` argument when calling :code:`set_times()`).
+#
+# Finally, you you run spike sorting through :code:`spikeinterface`, the recording
+# is automatically registered to the output sorting object!
diff --git a/spikeinterface/core/basesorting.py b/spikeinterface/core/basesorting.py
index 12ba0bf25..eb061a808 100644
--- a/spikeinterface/core/basesorting.py
+++ b/spikeinterface/core/basesorting.py
@@ -1,6 +1,6 @@
 from typing import List, Union
-
 import numpy as np
+import warnings
 
 from .base import BaseExtractor, BaseSegment
 
@@ -15,6 +15,8 @@ class BaseSorting(BaseExtractor):
         BaseExtractor.__init__(self, unit_ids)
         self._sampling_frequency = sampling_frequency
         self._sorting_segments: List[BaseSortingSegment] = []
+        # this weak link is to handle times from a recording object
+        self._recording = None
 
     def __repr__(self):
         clsname = self.__class__.__name__
@@ -51,10 +53,54 @@ class BaseSorting(BaseExtractor):
                              segment_index: Union[int, None] = None,
                              start_frame: Union[int, None] = None,
                              end_frame: Union[int, None] = None,
+                             return_times: bool = False
                              ):
         segment_index = self._check_segment_index(segment_index)
-        S = self._sorting_segments[segment_index]
-        return S.get_unit_spike_train(unit_id=unit_id, start_frame=start_frame, end_frame=end_frame).astype("int64")
+        segment = self._sorting_segments[segment_index]
+        spike_train = segment.get_unit_spike_train(
+            unit_id=unit_id, start_frame=start_frame, end_frame=end_frame).astype("int64")
+        if return_times:
+            if self.has_recording():
+                times = self.get_times(segment_index=segment_index)
+                return times[spike_train]
+            else:
+                return spike_train / self.get_sampling_frequency()
+        else:
+            return spike_train
+
+    def register_recording(self, recording):
+        assert np.isclose(self.get_sampling_frequency(),
+                          recording.get_sampling_frequency(),
+                          atol=0.1), "The recording has a different sampling frequency than the sorting!"
+        self._recording = recording
+
+    def has_recording(self):
+        return self._recording is not None
+
+    def has_time_vector(self, segment_index=None):
+        """
+        Check if the segment of the registered recording has a time vector.
+        """
+        segment_index = self._check_segment_index(segment_index)
+        if self.has_recording():
+            return self._recording.has_time_vector(segment_index=segment_index)
+        else:
+            return False
+
+    def get_times(self, segment_index=None):
+        """
+        Get time vector for a registered recording segment.
+
+        If a recording is registered:
+            * if the segment has a time_vector, then it is returned
+            * if not, a time_vector is constructed on the fly with sampling frequency
+        If there is no registered recording it returns None
+        """
+        segment_index = self._check_segment_index(segment_index)
+        if self.has_recording():
+            return self._recording.get_times(segment_index=segment_index)
+        else:
+            return None
 
     def _save(self, format='npz', **save_kwargs):
         """
@@ -68,6 +114,9 @@ class BaseSorting(BaseExtractor):
             save_path = folder / 'sorting_cached.npz'
             NpzSortingExtractor.write_sorting(self, save_path)
             cached = NpzSortingExtractor(save_path)
+            if self.has_recording():
+                warnings.warn("The registered recording will not be persistent on disk, but only available in memory")
+                cached.register_recording(self._recording)
         else:
             raise ValueError(f'format {format} not supported')
 
diff --git a/spikeinterface/core/frameslicesorting.py b/spikeinterface/core/frameslicesorting.py
index f0dbb6ea2..8f0be4320 100644
--- a/spikeinterface/core/frameslicesorting.py
+++ b/spikeinterface/core/frameslicesorting.py
@@ -1,4 +1,5 @@
 import numpy as np
+import warnings
 
 from .basesorting import BaseSorting, BaseSortingSegment
 
@@ -44,6 +45,10 @@ class FrameSliceSorting(BaseSorting):
         # copy properties and annotations
         parent_sorting.copy_metadata(self)
 
+        if parent_sorting.has_recording():
+            self.register_recording(parent_sorting._recording.frame_slice(start_frame=start_frame,
+                                                                          end_frame=end_frame))
+
         # update dump dict
         self._kwargs = {'parent_sorting': parent_sorting.to_dict(), 'start_frame': int(start_frame),
                         'end_frame': int(end_frame)}
diff --git a/spikeinterface/core/unitsaggregationsorting.py b/spikeinterface/core/unitsaggregationsorting.py
index e0c2e7cc5..551c019df 100644
--- a/spikeinterface/core/unitsaggregationsorting.py
+++ b/spikeinterface/core/unitsaggregationsorting.py
@@ -1,5 +1,5 @@
 from typing import List, Union
-
+import warnings
 import numpy as np
 
 from .basesorting import BaseSorting, BaseSortingSegment
@@ -67,6 +67,10 @@ class UnitsAggregationSorting(BaseSorting):
             sub_segment = UnitsAggregationSortingSegment(unit_map, parent_segments)
             self.add_sorting_segment(sub_segment)
 
+        if np.any([sort.has_recording() for sort in sorting_list]):
+            warnings.warn(
+                "Cannot propagate registered recording to UnitsAggregationSorting")
+
         self._sortings = sorting_list
         self._kwargs = {'sorting_list': [sort.to_dict() for sort in sorting_list],
                         'renamed_unit_ids': renamed_unit_ids}
diff --git a/spikeinterface/core/unitsselectionsorting.py b/spikeinterface/core/unitsselectionsorting.py
index 3139ad00c..cdf6c78b8 100644
--- a/spikeinterface/core/unitsselectionsorting.py
+++ b/spikeinterface/core/unitsselectionsorting.py
@@ -40,6 +40,9 @@ class UnitsSelectionSorting(BaseSorting):
 
         parent_sorting.copy_metadata(self, only_main=False, ids=self._unit_ids)
 
+        if parent_sorting.has_recording():
+            self.register_recording(parent_sorting._recording)
+
         self._kwargs = dict(parent_sorting=parent_sorting.to_dict(), unit_ids=unit_ids,
                             renamed_unit_ids=renamed_unit_ids)
 
diff --git a/spikeinterface/sorters/basesorter.py b/spikeinterface/sorters/basesorter.py
index ce63d7e63..c8a74a248 100644
--- a/spikeinterface/sorters/basesorter.py
+++ b/spikeinterface/sorters/basesorter.py
@@ -253,7 +253,9 @@ class BaseSorter:
             raise SpikeSortingError(
                 "Spike sorting failed. You can inspect the runtime trace in spikeinterface_log.json")
 
+        recording = load_extractor(output_folder / 'spikeinterface_recording.json')
         sorting = cls._get_result_from_folder(output_folder)
+        sorting.register_recording(recording)
         return sorting
 
     #############################################


Test patch:
diff --git a/spikeinterface/core/tests/test_time_handling.py b/spikeinterface/core/tests/test_time_handling.py
new file mode 100644
index 000000000..b300ef40f
--- /dev/null
+++ b/spikeinterface/core/tests/test_time_handling.py
@@ -0,0 +1,93 @@
+"""
+test for BaseSorting are done with NpzSortingExtractor.
+but check only for BaseRecording general methods.
+"""
+import shutil
+from pathlib import Path
+import pytest
+import numpy as np
+
+from spikeinterface.core import NpzSortingExtractor, load_extractor
+from spikeinterface.core.base import BaseExtractor
+
+from spikeinterface.core.testing_tools import generate_recording, generate_sorting
+
+
+def _clean_all():
+    cache_folder = './my_cache_folder'
+    if Path(cache_folder).exists():
+        shutil.rmtree(cache_folder)
+
+
+def setup_module():
+    _clean_all()
+
+
+def teardown_module():
+    _clean_all()
+
+
+def test_time_handling():
+    cache_folder = Path('./my_cache_folder')
+    durations = [[10], [10, 5]]
+
+    # test multi-segment
+    for i, dur in enumerate(durations):
+        rec = generate_recording(num_channels=4, durations=dur)
+        sort = generate_sorting(num_units=10, durations=dur)
+
+        for segment_index in range(rec.get_num_segments()):
+            original_times = rec.get_times(segment_index=segment_index)
+            new_times = original_times + 5
+            rec.set_times(new_times, segment_index=segment_index)
+
+        sort.register_recording(rec)
+        assert sort.has_recording()
+
+        rec_cache = rec.save(folder=cache_folder / f"rec{i}")
+
+        for segment_index in range(sort.get_num_segments()):
+            assert rec.has_time_vector(segment_index=segment_index)
+            assert sort.has_time_vector(segment_index=segment_index)
+
+            # times are correctly saved by the recording
+            assert np.allclose(rec.get_times(segment_index=segment_index), 
+                               rec_cache.get_times(segment_index=segment_index))
+
+            # spike times are correctly adjusted
+            for u in sort.get_unit_ids():
+                spike_times = sort.get_unit_spike_train(u, segment_index=segment_index,
+                                                        return_times=True)
+                rec_times = rec.get_times(segment_index=segment_index)
+                assert np.all(spike_times >= rec_times[0])
+                assert np.all(spike_times < rec_times[-1])
+
+
+def test_frame_slicing():
+    duration = [10]
+
+    rec = generate_recording(num_channels=4, durations=duration)
+    sort = generate_sorting(num_units=10, durations=duration)
+
+    original_times = rec.get_times()
+    new_times = original_times + 5
+    rec.set_times(new_times)
+
+    sort.register_recording(rec)
+
+    start_frame = 3 * rec.get_sampling_frequency()
+    end_frame = 7 * rec.get_sampling_frequency()
+
+    rec_slice = rec.frame_slice(start_frame=start_frame, end_frame=end_frame)
+    sort_slice = sort.frame_slice(start_frame=start_frame, end_frame=end_frame)
+
+    for u in sort_slice.get_unit_ids():
+        spike_times = sort_slice.get_unit_spike_train(u, return_times=True)
+        rec_times = rec_slice.get_times()
+        assert np.all(spike_times >= rec_times[0])
+        assert np.all(spike_times < rec_times[-1])
+
+
+if __name__ == '__main__':
+    _clean_all()
+    test_frame_slicing()