Merge pull request #30 from int-brain-lab/olive-saturation

saturation application with apodization

Author: alejoe91

src/spikeinterface/preprocessing/detect_artifacts.py

@@ -717,6 +717,7 @@ def __init__(
         job_kwargs: dict | None = None,
         mode: Literal["zeros", "noise"] = "zeros",
         noise_levels_kwargs: dict | None = None,
+        apodization: int = 7,
         seed: int | None = None,
         artifact_periods=None,
     ) -> None:
@@ -729,7 +730,7 @@ def __init__(
                 recording_to_detect, method=method, method_kwargs=method_kwargs, job_kwargs=job_kwargs
             )
         super().__init__(
-            recording, periods=artifact_periods, mode=mode, noise_levels_kwargs=noise_levels_kwargs, seed=seed
+            recording, periods=artifact_periods, mode=mode, noise_levels_kwargs=noise_levels_kwargs, seed=seed, apodization=apodization
         )
 
         self._kwargs = dict(
@@ -742,6 +743,7 @@ def __init__(
             noise_levels_kwargs=noise_levels_kwargs,
             seed=seed,
             artifact_periods=artifact_periods,
+            apodization=apodization,
         )
 
 

src/spikeinterface/preprocessing/silence_periods.py

@@ -1,4 +1,5 @@
 import numpy as np
+import scipy.signal
 
 from spikeinterface.core.core_tools import define_function_handling_dict_from_class
 from .basepreprocessor import BasePreprocessor, BasePreprocessorSegment
@@ -48,14 +49,15 @@ def __init__(
         self,
         recording,
         periods=None,
-        # this is keep for backward compatibility
+        # this is kept for backward compatibility
         list_periods=None,
         mode="zeros",
         noise_levels=None,
+        apodization=7,
         seed=None,
         **noise_levels_kwargs,
     ):
-        available_modes = ("zeros", "noise")
+        available_modes = ("zeros", "noise", "apodization")
         num_seg = recording.get_num_segments()
 
         # handle backward compatibility with previous version
@@ -108,11 +110,11 @@ def __init__(
             i1 = seg_limits[seg_index + 1]
             periods_in_seg = periods[i0:i1]
             rec_segment = SilencedPeriodsRecordingSegment(
-                parent_segment, periods_in_seg, mode, noise_generator, seg_index
+                parent_segment, periods_in_seg, mode, noise_generator, seg_index, apodization=apodization,
             )
             self.add_recording_segment(rec_segment)
 
-        self._kwargs = dict(recording=recording, periods=periods, mode=mode, seed=seed, noise_levels=noise_levels)
+        self._kwargs = dict(recording=recording, periods=periods, mode=mode, seed=seed, noise_levels=noise_levels, apodization=apodization)
 
 
 def _all_period_list_to_periods_vec(list_periods, num_seg):
@@ -154,12 +156,13 @@ def _check_periods(periods, num_seg):
 
 
 class SilencedPeriodsRecordingSegment(BasePreprocessorSegment):
-    def __init__(self, parent_recording_segment, periods, mode, noise_generator, seg_index):
+    def __init__(self, parent_recording_segment, periods, mode, noise_generator, seg_index, apodization=7):
         BasePreprocessorSegment.__init__(self, parent_recording_segment)
         self.periods = periods
         self.mode = mode
         self.seg_index = seg_index
         self.noise_generator = noise_generator
+        self.apodization = apodization
 
     def get_traces(self, start_frame, end_frame, channel_indices):
         traces = self.parent_recording_segment.get_traces(start_frame, end_frame, channel_indices)
@@ -185,7 +188,13 @@ def get_traces(self, start_frame, end_frame, channel_indices):
                             :, channel_indices
                         ]
                         traces[onset:offset, :] = noise[onset:offset]
-
+                    elif self.mode == "apodization":
+                        # apply a cosine taper to the saturation to create a mute function
+                        mute = np.zeros(traces.shape[0], dtype=np.float32)
+                        mute[onset:offset] = 1
+                        win = scipy.signal.windows.cosine(self.apodization)
+                        mute = np.maximum(0, 1 - scipy.signal.convolve(mute, win, mode="same"))
+                        traces = (traces.astype(np.float32) * mute[:, np.newaxis]).astype(traces.dtype)
         return traces