Repository: SpikeInterface/spikeinterface
Pull request #143 (2021-05-12 09:38:48)

Implement rec.frame_slice()


Patch:
diff --git a/spikeinterface/core/__init__.py b/spikeinterface/core/__init__.py
index 2302c9a57..02c53163e 100644
--- a/spikeinterface/core/__init__.py
+++ b/spikeinterface/core/__init__.py
@@ -21,6 +21,7 @@ from .numpyextractors import NumpyRecording , NumpySorting
 # utility extractors (equivalent to OLD subrecording/subsorting)
 from .channelslicerecording import ChannelSliceRecording
 from .unitsselectionsorting import UnitsSelectionSorting
+from .frameslicerecording import FrameSliceRecording
 
 # utils to append and concatenate segment (equivalent to OLD MultiRecordingTimeExtractor)
 from .segmentutils import (
diff --git a/spikeinterface/core/baserecording.py b/spikeinterface/core/baserecording.py
index 3c5d1afcc..ae5cd830b 100644
--- a/spikeinterface/core/baserecording.py
+++ b/spikeinterface/core/baserecording.py
@@ -395,6 +395,11 @@ class BaseRecording(BaseExtractor):
         from spikeinterface import ChannelSliceRecording
         sub_recording = ChannelSliceRecording(self, channel_ids, renamed_channel_ids=renamed_channel_ids)
         return sub_recording
+    
+    def frame_slice(self, start_frame, end_frame):
+        from spikeinterface import FrameSliceRecording
+        sub_recording = FrameSliceRecording(self, start_frame=start_frame, end_frame=end_frame)
+        return sub_recording
 
     def split_by(self, property='group', outputs='list'):
         assert outputs in ('list', 'dict')
diff --git a/spikeinterface/core/frameslicerecording.py b/spikeinterface/core/frameslicerecording.py
new file mode 100644
index 000000000..815db7ca4
--- /dev/null
+++ b/spikeinterface/core/frameslicerecording.py
@@ -0,0 +1,60 @@
+import numpy as np
+
+from .baserecording import BaseRecording, BaseRecordingSegment
+
+class FrameSliceRecording(BaseRecording):
+    """
+    Class to get a lazy frame slice.
+    Work only with mono segment recording.
+    """
+
+    def __init__(self, parent_recording, start_frame=None, end_frame=None):
+        channel_ids = parent_recording.get_channel_ids()
+    
+        assert parent_recording.get_num_segments() == 1, 'FrameSliceRecording work only with one segment'
+        
+        parent_size = parent_recording.get_num_samples(0)
+        if start_frame is None:
+            start_frame = 0
+        else:
+            assert 0 <= start_frame < parent_size
+
+        if end_frame is None:
+            end_frame = parent_size
+        else:
+            assert 0 <= start_frame <= parent_size
+
+        BaseRecording.__init__(self,
+                               sampling_frequency=parent_recording.get_sampling_frequency(),
+                               channel_ids=channel_ids,
+                               dtype=parent_recording.get_dtype())
+
+        # link recording segment
+        parent_segment = parent_recording._recording_segments[0]
+        sub_segment = FrameSliceRecordingSegment(parent_segment, start_frame, end_frame)
+        self.add_recording_segment(sub_segment)
+
+        # update dump dict
+        self._kwargs = {'parent_recording': parent_recording.to_dict(), 'start_frame': int(start_frame), 'end_frame': int(end_frame)}
+
+
+class FrameSliceRecordingSegment(BaseRecordingSegment):
+    def __init__(self, parent_recording_segment, start_frame, end_frame):
+        BaseRecordingSegment.__init__(self)
+        self._parent_recording_segment = parent_recording_segment
+        self.start_frame = start_frame
+        self.end_frame = end_frame
+        
+
+    def get_num_samples(self):
+        return self.end_frame - self.start_frame
+
+    def get_traces(self, start_frame, end_frame, channel_indices):
+        if start_frame is None:
+            start_frame = 0
+        if end_frame is None:
+            end_frame = self.get_num_samples()
+        parent_start = self.start_frame + start_frame
+        parent_end = self.start_frame + end_frame
+        traces = self._parent_recording_segment.get_traces(parent_start, parent_end, channel_indices)
+        return traces
diff --git a/spikeinterface/sorters/tridesclous/tridesclous.py b/spikeinterface/sorters/tridesclous/tridesclous.py
index d91450488..2c1f35f35 100644
--- a/spikeinterface/sorters/tridesclous/tridesclous.py
+++ b/spikeinterface/sorters/tridesclous/tridesclous.py
@@ -102,7 +102,7 @@ class TridesclousSorter(BaseSorter):
             # save binary file (chunk by hcunk) into a new file
             num_chan = recording.get_num_channels()
             dtype = recording.get_dtype().str
-            files_path = [str(output_folder / f'raw_signals_{i}.raw') for i in num_seg]
+            files_path = [str(output_folder / f'raw_signals_{i}.raw') for i in range(num_seg)]
             BinaryRecordingExtractor.write_recording(recording, files_path=files_path,
                                                                     dtype=dtype, total_memory="500M", n_jobs=-1,
                                                                     verbose=False, progress_bar=verbose)
@@ -175,7 +175,7 @@ def make_nested_tdc_params(tdc_dataio, chan_grp, **new_params):
         params['preprocessor']['highpass_freq'] = new_params['freq_min']
 
     if 'freq_max' in new_params:
-        params['preprocessor']['highpass_freq'] = new_params['freq_max']
+        params['preprocessor']['lowpass_freq'] = new_params['freq_max']
 
     if 'common_ref_removal' in new_params:
         params['preprocessor']['common_ref_removal'] = new_params['common_ref_removal']


Test patch:
diff --git a/spikeinterface/core/tests/test_frameslicerecording.py b/spikeinterface/core/tests/test_frameslicerecording.py
new file mode 100644
index 000000000..f671786e5
--- /dev/null
+++ b/spikeinterface/core/tests/test_frameslicerecording.py
@@ -0,0 +1,38 @@
+import pytest
+import numpy as np
+
+from spikeinterface.core import FrameSliceRecording, NumpyRecording
+
+
+def test_FrameSliceRecording():
+    traces = np.zeros((1000, 5), dtype='float64')
+    traces[:] = np.arange(1000)[:, None]
+    sampling_frequency = 30000
+    rec = NumpyRecording([traces], sampling_frequency)
+    
+    
+    sub_rec = rec.frame_slice(None, None)
+    assert sub_rec.get_num_samples(0) == 1000
+    traces = sub_rec.get_traces()
+    assert np.array_equal(traces[:, 0], np.arange(0, 1000, dtype='float64'))
+
+    sub_rec = rec.frame_slice(None, 10)
+    assert sub_rec.get_num_samples(0) == 10
+    traces = sub_rec.get_traces()
+    assert np.array_equal(traces[:, 0], np.arange(0, 10, dtype='float64'))
+
+
+    sub_rec = rec.frame_slice(900, 1000)
+    assert sub_rec.get_num_samples(0) == 100
+    traces = sub_rec.get_traces()
+    assert np.array_equal(traces[:, 0], np.arange(900, 1000, dtype='float64'))
+
+    sub_rec = rec.frame_slice(10, 85)
+    assert sub_rec.get_num_samples(0) == 75
+    traces = sub_rec.get_traces()
+    assert np.array_equal(traces[:, 0], np.arange(10, 85, dtype='float64'))
+
+    
+
+if __name__ == '__main__':
+    test_FrameSliceRecording()
\ No newline at end of file