Merge pull request #1310 from yger/silence_interval

Silence interval

Author: alejoe91

spikeinterface/preprocessing/preprocessinglist.py

@@ -17,6 +17,7 @@
     ClipRecording, clip)
 from .common_reference import CommonReferenceRecording, common_reference
 from .remove_artifacts import RemoveArtifactsRecording, remove_artifacts
+from .silence_periods import SilencedPeriodsRecording, silence_periods
 from .phase_shift import PhaseShiftRecording, phase_shift
 from .zero_channel_pad import ZeroChannelPaddedRecording, zero_channel_pad
 from .deepinterpolation import DeepInterpolatedRecording, deepinterpolate
@@ -51,6 +52,7 @@
     RectifyRecording,
     ClipRecording,
     BlankSaturationRecording,
+    SilencedPeriodsRecording,
     RemoveArtifactsRecording,
     ZeroChannelPaddedRecording,
     DeepInterpolatedRecording,

spikeinterface/preprocessing/silence_periods.py

@@ -0,0 +1,127 @@
+import numpy as np
+import scipy.interpolate
+
+from spikeinterface.core.core_tools import define_function_from_class
+from .basepreprocessor import BasePreprocessor, BasePreprocessorSegment
+
+from ..core import get_random_data_chunks, get_noise_levels
+
+class SilencedPeriodsRecording(BasePreprocessor):
+    """
+    Silence user-defined periods from recording extractor traces. By default, 
+    periods are zeroed-out (mode = 'zeros'). You can also fill the periods with noise.
+    Note that both methods assume that traces that are centered around zero. 
+    If this is not the case, make sure you apply a filter or center function prior to 
+    silencing periods.
+
+    Parameters
+    ----------
+    recording: RecordingExtractor
+        The recording extractor to silance periods
+    list_periods: list of lists/arrays
+        One list per segment of tuples (start_frame, end_frame) to silence
+
+    mode: str
+        Determines what periods are replaced by. Can be one of the following:
+            
+        - 'zeros' (default): Artifacts are replaced by zeros.
+
+        - 'noise': The periods are filled with a gaussion noise that has the
+                   same variance that the one in the recordings, on a per channel
+                   basis
+    **random_chunk_kwargs: Keyword arguments for `spikeinterface.core.get_random_data_chunk()` function
+
+    Returns
+    -------
+    silence_recording: SilencedPeriodsRecording
+        The recording extractor after silencing some periods    
+    """
+    name = 'silence_periods'
+
+    def __init__(self, recording, list_periods, mode='zeros', 
+                 **random_chunk_kwargs):
+
+        available_modes = ('zeros', 'noise')
+        num_seg = recording.get_num_segments()
+
+
+        if num_seg == 1:
+            if isinstance(list_periods, (list, np.ndarray)) and not np.isscalar(list_periods[0]):
+                # when unique segment accept list instead of of list of list/arrays
+                list_periods = [list_periods]
+
+        # some checks
+        assert mode in available_modes, f"mode {mode} is not an available mode: {available_modes}"
+
+        assert isinstance(list_periods, list), "'list_periods' must be a list (one per segment)"
+        assert len(list_periods) == num_seg, "'list_periods' must have the same length as the number of segments"
+        assert all(isinstance(list_periods[i], (list, np.ndarray)) for i in range(num_seg)), \
+            "Each element of 'list_periods' must be array-like"
+
+        for periods in list_periods:
+            if len(periods) > 0:
+                assert np.all(np.diff(np.array(periods), axis=1) > 0), "t_stops should be larger than t_starts"
+                assert np.all(periods[i][1] < periods[i + 1][0] for i in np.arange(len(periods) - 1)), \
+                    "Intervals should not overlap"
+
+        if mode in ['noise']:
+            noise_levels = get_noise_levels(recording, return_scaled=False, concatenated=True, **random_chunk_kwargs)
+        else:
+            noise_levels = None
+
+        BasePreprocessor.__init__(self, recording)
+        for seg_index, parent_segment in enumerate(recording._recording_segments):
+            periods = list_periods[seg_index]
+            periods = np.asarray(periods, dtype='int64')
+            periods = np.sort(periods, axis=0)
+            rec_segment = SilencedPeriodsRecordingSegment(parent_segment, periods, mode, noise_levels)
+            self.add_recording_segment(rec_segment)
+        
+        self._kwargs = dict(recording=recording.to_dict(), list_periods=list_periods,
+                            mode=mode, noise_levels=noise_levels)
+
+
+class SilencedPeriodsRecordingSegment(BasePreprocessorSegment):
+
+    def __init__(self, parent_recording_segment, periods, mode, noise_levels):
+        BasePreprocessorSegment.__init__(self, parent_recording_segment)
+        self.periods = periods
+        self.mode = mode
+        self.noise_levels = noise_levels
+
+    def get_traces(self, start_frame, end_frame, channel_indices):
+
+        traces = self.parent_recording_segment.get_traces(start_frame, end_frame, channel_indices)
+        traces = traces.copy()
+        num_channels = traces.shape[1]
+
+        if start_frame is None:
+            start_frame = 0
+        if end_frame is None:
+            end_frame = self.get_num_samples()
+
+        if len(self.periods) > 0:
+            new_interval = np.array([start_frame, end_frame])
+            lower_index = np.searchsorted(self.periods[:, 1], new_interval[0])
+            upper_index = np.searchsorted(self.periods[:, 0], new_interval[1])
+
+            if upper_index > lower_index:
+
+                periods_in_interval = self.periods[lower_index:upper_index]
+
+                for period in periods_in_interval:
+
+                    onset = max(0, period[0] - start_frame)
+                    offset = min(period[1] - start_frame, end_frame)
+
+                    if self.mode == 'zeros':
+                        traces[onset:offset, :] = 0
+                    elif self.mode == 'noise':
+                        traces[onset:offset, :] = self.noise_levels[channel_indices] * \
+                            np.random.randn(offset - onset, num_channels)
+
+        return traces
+
+
+# function for API
+silence_periods = define_function_from_class(source_class=SilencedPeriodsRecording, name="silence_periods")

spikeinterface/preprocessing/tests/test_silence.py

@@ -0,0 +1,53 @@
+import pytest
+from pathlib import Path
+import shutil
+
+from spikeinterface import set_global_tmp_folder
+from spikeinterface.core import generate_recording
+
+from spikeinterface.preprocessing import silence_periods
+
+
+from spikeinterface.core import get_noise_levels
+
+import numpy as np
+
+
+if hasattr(pytest, "global_test_folder"):
+    cache_folder = pytest.global_test_folder / "preprocessing"
+else:
+    cache_folder = Path("cache_folder") / "preprocessing"
+
+set_global_tmp_folder(cache_folder)
+
+def test_silence():
+    rec = generate_recording()
+
+    rec0 = silence_periods(rec, list_periods=[[[0, 1000], [5000, 6000]], []], mode="zeros")
+    rec0.save(verbose=False)
+    traces_in0 = rec0.get_traces(segment_index=0, start_frame=0, end_frame=1000)
+    traces_in1 = rec0.get_traces(segment_index=0, start_frame=5000, end_frame=6000)
+    traces_out0 = rec0.get_traces(segment_index=0, start_frame=2000, end_frame=3000)
+    assert np.all(traces_in0 == 0)
+    assert np.all(traces_in1 == 0)
+    assert not np.all(traces_out0 == 0)
+
+    rec1 = silence_periods(rec, list_periods=[[[0, 1000], [5000, 6000]], []], mode="noise")
+    rec1 = rec1.save(folder=cache_folder / "rec_w_noise", verbose=False)
+    noise_levels = get_noise_levels(rec, return_scaled=False)
+    traces_in0 = rec1.get_traces(segment_index=0, start_frame=0, end_frame=1000)
+    traces_in1 = rec1.get_traces(segment_index=0, start_frame=5000, end_frame=6000)
+    assert np.abs((np.std(traces_in0, axis=0) - noise_levels) < 0.1).sum()
+    assert np.abs((np.std(traces_in1, axis=0) - noise_levels)).sum() < 0.1
+
+    traces_mix = rec0.get_traces(segment_index=0, start_frame=900, end_frame=5100)
+    traces_original = rec.get_traces(segment_index=0, start_frame=900, end_frame=5100)
+    assert np.all(traces_original[100:-100] == traces_mix[100:-100])
+    assert np.all(traces_mix[:100] == 0)
+    assert np.all(traces_mix[-100:] == 0)
+    assert not np.all(traces_mix[:200] == 0)
+    assert not np.all(traces_mix[:-200] == 0)
+
+
+if __name__ == '__main__':
+    test_silence()