Repository: SpikeInterface/spikeinterface
Pull request #1706 (2023-06-07 11:08:50)

Add preprocessing for padding traces
Adding a preprocessing step that pads traces to the left and to the right.


Patch:
diff --git a/src/spikeinterface/preprocessing/zero_channel_pad.py b/src/spikeinterface/preprocessing/zero_channel_pad.py
index 4488ce0f0..124b2b080 100644
--- a/src/spikeinterface/preprocessing/zero_channel_pad.py
+++ b/src/spikeinterface/preprocessing/zero_channel_pad.py
@@ -1,10 +1,137 @@
 from typing import Union
+
 import numpy as np
+
 from spikeinterface.core import BaseRecording, BaseRecordingSegment
 from .basepreprocessor import BasePreprocessor, BasePreprocessorSegment
 from spikeinterface.core.core_tools import define_function_from_class
 
 
+class TracePaddedRecording(BasePreprocessor):
+    """
+    A Pre-processor class to lazily pad recordings.
+
+    The class retrieves traces from the parent recording segment and pads them with zeros at both ends.
+
+    Parameters
+    ----------
+    parent_recording_segment : BaseRecording
+        The parent recording segment from which the traces are to be retrieved.
+    padding_start : int
+        The amount of padding to add to the left of the traces. Default is 0. It has to be non-negative
+    padding_end : int
+        The amount of padding to add to the right of the traces. Default is 0. It has to be non-negative
+    fill_value: float
+        The value to pad with. Default is 0.
+    """
+
+    def __init__(
+        self, parent_recording: BaseRecording, padding_start: int = 0, padding_end: int = 0, fill_value: float = 0.0
+    ):
+        assert padding_end >= 0 and padding_start >= 0, "Paddings must be >= 0"
+        super().__init__(recording=parent_recording)
+
+        self.padding_start = padding_start
+        self.padding_end = padding_end
+        self.fill_value = fill_value
+        for segment in parent_recording._recording_segments:
+            recording_segment = TracePaddedRecordingSegment(
+                segment,
+                parent_recording.get_num_channels(),
+                self.dtype,
+                self.padding_start,
+                self.padding_end,
+                self.fill_value,
+            )
+            self.add_recording_segment(recording_segment)
+
+        self._kwargs = dict(
+            parent_recording=parent_recording,
+            padding_start=padding_start,
+            padding_end=padding_end,
+            fill_value=fill_value,
+        )
+
+
+class TracePaddedRecordingSegment(BasePreprocessorSegment):
+    def __init__(
+        self,
+        parent_recording_segment: BaseRecordingSegment,
+        num_channels,
+        dtype,
+        paddign_left,
+        padding_end,
+        fill_value,
+    ):
+        self.padding_start = paddign_left
+        self.padding_end = padding_end
+        self.fill_value = fill_value
+        self.num_channels = num_channels
+        self.num_samples_in_original_segment = parent_recording_segment.get_num_samples()
+        self.dtype = dtype
+
+        super().__init__(parent_recording_segment=parent_recording_segment)
+
+    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()
+
+        # This contains the padded elements by default and we add the original traces if necessary
+        trace_size = end_frame - start_frame
+        if isinstance(channel_indices, (np.ndarray, list)):
+            num_channels = len(channel_indices)
+        elif channel_indices == slice(None):
+            num_channels = self.num_channels
+        else:
+            raise ValueError(f"Unsupported channel_indices type: {type(channel_indices)} raise an issue on github ")
+
+        # This avoids an extra memory allocation if we are within the confines of the old traces
+        if start_frame > self.padding_start and end_frame < self.num_samples_in_original_segment + self.padding_start:
+            return self.get_original_traces_shifted(start_frame, end_frame, channel_indices)
+
+        # Else, we start with the full padded traces and allocate the original traces in the middle
+        output_traces = np.full(shape=(trace_size, num_channels), fill_value=self.fill_value, dtype=self.dtype)
+
+        # After the padding, the original traces are placed in the middle until the end of the original traces
+        if end_frame >= self.padding_start:
+            original_traces = self.get_original_traces_shifted(
+                start_frame=start_frame,
+                end_frame=end_frame,
+                channel_indices=channel_indices,
+            )
+
+            padded_frames_at_the_start = max(self.padding_start - start_frame, 0)
+            end_of_original_traces = padded_frames_at_the_start + original_traces.shape[0]
+            output_traces[padded_frames_at_the_start:end_of_original_traces, :] = original_traces
+
+        return output_traces
+
+    def get_original_traces_shifted(self, start_frame, end_frame, channel_indices):
+        """
+        Retrieves the corresponding original (unpadded) traces from the parent recording segment.
+
+        This method calculates the start and end frames of the original traces, taking into account
+        the start padding. It then fetches the original traces using these frames from the parent
+        recording segment.
+
+        """
+        original_start_frame = max(start_frame - self.padding_start, 0)
+        original_end_frame = min(end_frame - self.padding_start, self.num_samples_in_original_segment)
+        original_traces = self.parent_recording_segment.get_traces(
+            start_frame=original_start_frame,
+            end_frame=original_end_frame,
+            channel_indices=channel_indices,
+        )
+
+        return original_traces
+
+    def get_num_samples(self):
+        "Overide the parent method to return the padded number of samples"
+        return self.num_samples_in_original_segment + self.padding_start + self.padding_end
+
+
 class ZeroChannelPaddedRecording(BaseRecording):
     name = "zero_channel_pad"
     installed = True
@@ -85,3 +212,4 @@ class ZeroChannelPaddedRecordingSegment(BasePreprocessorSegment):
 
 # function for API
 zero_channel_pad = define_function_from_class(source_class=ZeroChannelPaddedRecording, name="zero_channel_pad")
+pad_traces = define_function_from_class(source_class=TracePaddedRecording, name="pad_traces")


Test patch:
diff --git a/src/spikeinterface/preprocessing/tests/test_zero_padding.py b/src/spikeinterface/preprocessing/tests/test_zero_padding.py
index 3c90bcabb..75d64b008 100644
--- a/src/spikeinterface/preprocessing/tests/test_zero_padding.py
+++ b/src/spikeinterface/preprocessing/tests/test_zero_padding.py
@@ -4,8 +4,10 @@ import numpy as np
 
 from spikeinterface import set_global_tmp_folder
 from spikeinterface.core import generate_recording
+from spikeinterface.core.numpyextractors import NumpyRecording
 
 from spikeinterface.preprocessing import zero_channel_pad
+from spikeinterface.preprocessing.zero_channel_pad import TracePaddedRecording
 
 if hasattr(pytest, "global_test_folder"):
     cache_folder = pytest.global_test_folder / "preprocessing"
@@ -15,7 +17,7 @@ else:
 set_global_tmp_folder(cache_folder)
 
 
-def test_zero_padding():
+def test_zero_padding_channel():
     num_original_channels = 4
     num_padded_channels = num_original_channels + 8
     rec = generate_recording(num_channels=num_original_channels, durations=[10])
@@ -34,5 +36,227 @@ def test_zero_padding():
     assert np.allclose(tr[:, :num_original_channels], rec.get_traces())
 
 
+@pytest.fixture
+def recording():
+    num_channels = 4
+    num_samples = 10
+    rng = np.random.default_rng(seed=0)
+    traces = rng.random(size=(num_samples, num_channels))
+    traces_list = [traces]
+    recording = NumpyRecording(traces_list=traces_list, sampling_frequency=30_000)
+    return recording
+
+
+@pytest.mark.parametrize("padding_start, padding_end", [(5, 5), (0, 5), (5, 0), (0, 0)])
+def test_trace_padded_recording_full_trace(recording, padding_start, padding_end):
+    num_channels = recording.get_num_channels()
+    num_samples = recording.get_num_samples()
+
+    padded_recording = TracePaddedRecording(
+        parent_recording=recording,
+        padding_start=padding_start,
+        padding_end=padding_end,
+    )
+    padded_traces = padded_recording.get_traces()
+
+    # Until padding_start the traces should be filled with zeros
+    padded_traces_start = padded_traces[:padding_start, :]
+    expected_padding = np.zeros((padding_start, num_channels))
+    assert np.allclose(padded_traces_start, expected_padding)
+
+    # Then from padding_start to the number of samples plus padding_start it should be the original traces
+    original_traces = recording.get_traces()
+    last_frame_of_original_trace = num_samples + padding_start
+    original_traces_from_padding = padded_traces[padding_start:last_frame_of_original_trace, :]
+    assert np.allclose(original_traces_from_padding, original_traces)
+
+    # After the original trace is over it should have zeros until the end of the padding right
+    padded_traces_end = padded_traces[last_frame_of_original_trace:, :]
+    expected_padding = np.zeros((padding_end, num_channels))
+    assert np.allclose(padded_traces_end, expected_padding)
+
+
+@pytest.mark.parametrize("padding_start, padding_end", [(5, 5), (0, 5), (5, 0), (0, 0)])
+def test_trace_padded_recording_full_trace_with_channel_indices(recording, padding_start, padding_end):
+    num_samples = recording.get_num_samples()
+
+    padded_recording = TracePaddedRecording(
+        parent_recording=recording,
+        padding_start=padding_start,
+        padding_end=padding_end,
+    )
+
+    channel_indices = [2, 0]
+    num_channels_sub_straces = len(channel_indices)
+    padded_traces = padded_recording.get_traces(channel_ids=channel_indices)
+
+    # Until padding_start the traces should be filled with zeros
+    assert np.allclose(padded_traces[:padding_start, :], np.zeros((padding_start, num_channels_sub_straces)))
+    last_frame_of_original_trace = num_samples + padding_start
+
+    # Then from padding_start to the number of samples plus padding_start it should be the original traces
+    original_traces = recording.get_traces(channel_ids=channel_indices)
+    original_traces_from_padding = padded_traces[padding_start:last_frame_of_original_trace, :]
+    assert np.allclose(original_traces_from_padding, original_traces)
+
+    # After the original trace is over it should have zeros until the end of the padding right
+    padded_traces_end = padded_traces[last_frame_of_original_trace:, :]
+    expected_padding = np.zeros((padding_end, num_channels_sub_straces))
+    assert np.allclose(padded_traces_end, expected_padding)
+
+
+@pytest.mark.parametrize("padding_start, padding_end", [(5, 5), (0, 5), (5, 0), (0, 0)])
+def test_trace_padded_recording_retrieve_original_trace(recording, padding_start, padding_end):
+    num_samples = recording.get_num_samples()
+
+    padded_recording = TracePaddedRecording(
+        parent_recording=recording,
+        padding_start=padding_start,
+        padding_end=padding_end,
+    )
+
+    # These are the limits of the original trace
+    start_frame = padding_start
+    end_frame = num_samples + padding_start
+    padded_traces = padded_recording.get_traces(start_frame=start_frame, end_frame=end_frame)
+
+    original_traces = recording.get_traces()
+    assert np.allclose(padded_traces, original_traces)
+
+
+@pytest.mark.parametrize("padding_start, padding_end", [(5, 5), (0, 5), (5, 0), (0, 0)])
+def test_trace_padded_recording_retrieve_partial_original_trace(recording, padding_start, padding_end):
+    num_samples = recording.get_num_samples()
+
+    padded_recording = TracePaddedRecording(
+        parent_recording=recording,
+        padding_start=padding_start,
+        padding_end=padding_end,
+    )
+
+    # We extract a trace that is smaller than the original trace
+    start_frame_original_traces = 2
+    end_frame_original_traces = num_samples - 3
+
+    # Adjust the start and end frame to take into account the padding in the padded recoderer
+    start_frame = padding_start + start_frame_original_traces
+    end_frame = padding_start + end_frame_original_traces
+    padded_traces = padded_recording.get_traces(start_frame=start_frame, end_frame=end_frame)
+
+    # Test for a match
+    original_traces = recording.get_traces(start_frame=start_frame_original_traces, end_frame=end_frame_original_traces)
+    assert np.allclose(padded_traces, original_traces)
+
+
+@pytest.mark.parametrize("padding_start, padding_end", [(5, 5), (5, 0)])
+def test_trace_padded_recording_retrieve_start_padding_and_partial_original_trace(
+    recording, padding_start, padding_end
+):
+    num_samples = recording.get_num_samples()
+    num_channels = recording.get_num_channels()
+
+    padded_recording = TracePaddedRecording(
+        parent_recording=recording,
+        padding_start=padding_start,
+        padding_end=padding_end,
+    )
+
+    # We extract a traces with padding at the beginning and smaller than original trace at the end
+    start_padding_to_retrieve = 2
+    end_frame_original_traces = num_samples - 3
+
+    # Adjust the start and end frame to take into account the padding in the padded recoderer
+    start_frame = padding_start - start_padding_to_retrieve
+    end_frame = padding_start + end_frame_original_traces
+    padded_traces = padded_recording.get_traces(start_frame=start_frame, end_frame=end_frame)
+
+    # Check the the beginning of the trace is is padded with zeros
+    start_padding = padded_traces[:start_padding_to_retrieve, :]
+    expected_padding = np.zeros((start_padding_to_retrieve, num_channels))
+    assert np.allclose(start_padding, expected_padding)
+
+    # Check the partial retrieval of the original series
+    partial_original_traces = recording.get_traces(start_frame=0, end_frame=end_frame_original_traces)
+    partial_orignal_traces_from_padded_traces = padded_traces[start_padding_to_retrieve:, :]
+
+    assert np.allclose(partial_original_traces, partial_orignal_traces_from_padded_traces)
+
+
+@pytest.mark.parametrize("padding_start, padding_end", [(5, 5), (0, 5)])
+def test_trace_padded_recording_retrieve_end_padding_and_partial_original_trace(recording, padding_start, padding_end):
+    num_samples = recording.get_num_samples()
+    num_channels = recording.get_num_channels()
+
+    padded_recording = TracePaddedRecording(
+        parent_recording=recording,
+        padding_start=padding_start,
+        padding_end=padding_end,
+    )
+
+    # We extract traces that are smaller than original at the start and then padded at the end
+    start_frame_original_traces = 2
+    end_padding_to_retrieve = 3
+    end_frame_original_traces = num_samples
+
+    # Adjust the start and end frame to take into account the padding in the padded recoderer
+    start_frame = padding_start + start_frame_original_traces
+    end_frame = padding_start + end_frame_original_traces + end_padding_to_retrieve
+    padded_traces = padded_recording.get_traces(start_frame=start_frame, end_frame=end_frame)
+
+    # Check the partial retrieval of the original series at the start
+    partial_original_traces = recording.get_traces(
+        start_frame=start_frame_original_traces, end_frame=end_frame_original_traces
+    )
+    partial_orignal_traces_from_padded_traces = padded_traces[:-end_padding_to_retrieve, :]
+    assert np.allclose(partial_orignal_traces_from_padded_traces, partial_original_traces)
+
+    # Check that the end of the padded traces is padded with zeros
+    last_padding = padded_traces[end_frame_original_traces:, :]
+    expected_padding = np.zeros((end_padding_to_retrieve, num_channels))
+    assert np.allclose(last_padding, expected_padding)
+
+
+@pytest.mark.parametrize("padding_start, padding_end", [(5, 5), (0, 5), (5, 0)])
+def test_trace_padded_recording_retrieve_traces_with_partial_padding(recording, padding_start, padding_end):
+    num_samples = recording.get_num_samples()
+    num_channels = recording.get_num_channels()
+
+    padded_recording = TracePaddedRecording(
+        parent_recording=recording,
+        padding_start=padding_start,
+        padding_end=padding_end,
+    )
+
+    # Define how many frames are padded at the start and end (take into account they could be 0)
+    number_of_padded_frames_at_start = 2 if padding_start > 0 else 0
+    number_of_paded_frames_at_end = 3 if padding_end > 0 else 0
+
+    # Start before the padding starts and ensure the end includes padding by going
+    # further than the number of samples of the original recording
+    start_frame = padding_start - number_of_padded_frames_at_start
+    end_frame = padding_start + num_samples + number_of_paded_frames_at_end
+
+    # Retrieve the traces which should have partial padding at the start and end
+    padded_traces = padded_recording.get_traces(start_frame=start_frame, end_frame=end_frame)
+
+    # Find that there as many frames padded a the start as expected
+    padded_traces_start = padded_traces[:number_of_padded_frames_at_start, :]
+    expected_zeros = np.zeros((number_of_padded_frames_at_start, num_channels))
+    assert np.allclose(padded_traces_start, expected_zeros)
+
+    # Then from padding_start to the number of samples plus padding_start this should be the original traces
+    frame_where_original_trace_ends = number_of_padded_frames_at_start + num_samples
+    original_traces_from_padded_traces = padded_traces[
+        number_of_padded_frames_at_start:frame_where_original_trace_ends, :
+    ]
+    original_traces = recording.get_traces()
+    assert np.allclose(original_traces_from_padded_traces, original_traces)
+
+    # Find that there as many frames padded at the end as expected
+    padded_traces_end = padded_traces[frame_where_original_trace_ends:, :]
+    expected_zeros = np.zeros((number_of_paded_frames_at_end, num_channels))
+    assert np.allclose(padded_traces_end, expected_zeros)
+
+
 if __name__ == "__main__":
-    test_zero_padding()
+    test_zero_padding_channel()