Extend artifact detection: DetectAndRemoveArtifacts preprocessor and signed saturation

doc/modules/preprocessing.rst

@@ -149,6 +149,37 @@ can also be obtained from the pipeline object directly:
     dict_used_to_make_pipeline = preprocessing_pipeline.preprocessor_dict
 
 
+Some preprocessing steps, such as :code:`detect_and_remove_artifacts`, allow you to specify an input recording
+and optionally another recording to perform some computation (e.g., detect artifacts on the output of a previous
+preprocessor, but remove them on the the parent preprocessor). In this case, the string "pipeline[preprocessor_name]"
+can be used in the dictionary to specify that the recording argument for this step should be the output of a previous
+preprocessor in the same pipeline. For example, if we want to use the output of the "bandpass_filter" step as the
+recording to detect artifacts, we can specify it as follows:
+
+.. code-block:: python
+
+    preprocessing_dict = {
+        'bandpass_filter': {'freq_min': 250},
+        'common_reference': {'operator': 'median', 'reference': 'global'},
+        'detect_and_remove_artifacts': {'recording_to_detect': 'pipeline[bandpass_filter]'},
+    }
+
+This will detect artifacts on the output of the "bandpass_filter" step, but the artifacts will be removed on the output
+of the "common_reference" step (since the parent recording for "detect_and_remove_artifacts" is by default the output of
+the previous step in the pipeline, which is "common_reference" in this case).
+To specify the "raw" recording, i.e., the input to the pipeline, we can use "pipeline[raw]".
+For example, if we want to detect artifacts on the raw recording, we can specify it as follows:
+
+
+.. code-block:: python
+
+    preprocessing_dict = {
+        'bandpass_filter': {'freq_min': 250},
+        'common_reference': {'operator': 'median', 'reference': 'global'},
+        'detect_and_remove_artifacts': {'recording_to_detect': 'pipeline[raw]'},
+    }
+
+
 Impact on recording dtype
 -------------------------
 

src/spikeinterface/preprocessing/detect_bad_channels.py

@@ -308,6 +308,9 @@ def detect_bad_channels(
         if channel_filters is None:
             channel_filters = allowed_filters
 
+        if isinstance(channel_filters, list):
+            channel_filters = set(channel_filters)
+
         if not isinstance(channel_filters, set):
             raise ValueError(f"channel_filters must be None or a set of the following values : {allowed_filters} ")
 

src/spikeinterface/preprocessing/pipeline.py

@@ -99,11 +99,29 @@ def _apply(self, recording, apply_precomputed_kwargs=False):
             Preprocessed recording
 
         """
-
-        for preprocessor_name, kwargs in self.preprocessor_dict.items():
-
+        instantiated_recordings = {"raw": recording}
+        for preprocessor_name, kwargs_ in self.preprocessor_dict.items():
+            kwargs = kwargs_.copy()
             dont_apply_kwargs = ["recording", "parent_recording"]
 
+            for k, v in kwargs.items():
+                if isinstance(v, str) and "pipeline[" in v:
+                    if "recording" not in k:
+                        raise ValueError(
+                            f"Cannot substitute recording for argument '{k}' of preprocessor '{preprocessor_name}' "
+                            f"because this argument is not meant to be a recording object."
+                        )
+                    if k in dont_apply_kwargs:
+                        raise ValueError(
+                            f"Cannot substitute recording for argument '{k}' of preprocessor '{preprocessor_name}' "
+                            f"because this argument is reserved for the recording to be preprocessed."
+                        )
+                    rec_name = v.split("pipeline[")[-1].split("]")[0]
+                    substituted_recording = instantiated_recordings.get(rec_name)
+                    if substituted_recording is None:
+                        raise ValueError(f"Cannot find recording '{rec_name}' from previous steps in the pipeline.")
+                    kwargs[k] = substituted_recording
+
             if not apply_precomputed_kwargs:
                 preprocessor_class = pp_names_to_classes[preprocessor_name]
                 precomputable_kwarg_names = preprocessor_class._precomputable_kwarg_names
@@ -112,6 +130,7 @@ def _apply(self, recording, apply_precomputed_kwargs=False):
             non_rec_kwargs = {key: value for key, value in kwargs.items() if key not in dont_apply_kwargs}
             pp_output = pp_names_to_functions[preprocessor_name](recording, **non_rec_kwargs)
             recording = pp_output
+            instantiated_recordings[preprocessor_name] = recording
 
         return recording
 
@@ -305,6 +324,12 @@ def _load_pp_from_dict(prov_dict, kwargs_dict):
     for name, value in prov_dict["kwargs"].items():
         if is_dict_extractor(value):
             this_level_kwargs[name] = _load_pp_from_dict(value, kwargs_dict)
+        elif isinstance(value, BaseRecording):
+            extractor_as_dict = value.to_dict()
+            if name in ["recording", "parent_recording"]:
+                this_level_kwargs[name] = _load_pp_from_dict(extractor_as_dict, kwargs_dict)
+            else:  # this branch takes care of other arguments being a recording, e.g., `recording_to_detect`
+                this_level_kwargs[name] = value
         elif isinstance(value, dict):
             this_level_kwargs[name] = {k: prov_dict_to_kwargs_dict(v) for k, v in value.items()}
         elif isinstance(value, list):

src/spikeinterface/preprocessing/preprocessing_classes.py

@@ -48,6 +48,7 @@
 from .depth_order import DepthOrderRecording, depth_order
 from .astype import AstypeRecording, astype
 from .unsigned_to_signed import UnsignedToSignedRecording, unsigned_to_signed
+from .detect_artifacts import DetectAndRemoveArtifactsRecording, detect_and_remove_artifacts
 
 # from .silence_artifacts import SilencedArtifactsRecording, silence_artifacts
 
@@ -72,6 +73,8 @@
     # bad channel detection/interpolation
     DetectAndRemoveBadChannelsRecording: detect_and_remove_bad_channels,
     DetectAndInterpolateBadChannelsRecording: detect_and_interpolate_bad_channels,
+    # artifact/saturation handling
+    DetectAndRemoveArtifactsRecording: detect_and_remove_artifacts,
     # misc
     RectifyRecording: rectify,
     ClipRecording: clip,

src/spikeinterface/preprocessing/silence_periods.py

@@ -1,12 +1,13 @@
 import numpy as np
 
-from spikeinterface.core.core_tools import define_function_handling_dict_from_class
-from .basepreprocessor import BasePreprocessor, BasePreprocessorSegment
 
-from spikeinterface.core import get_noise_levels
+from spikeinterface.core.base import base_period_dtype
+from spikeinterface.core.core_tools import define_function_handling_dict_from_class
+from spikeinterface.core.recording_tools import get_noise_levels, get_chunk_with_margin
 from spikeinterface.core.generate import NoiseGeneratorRecording
 from spikeinterface.core.job_tools import split_job_kwargs
-from spikeinterface.core.base import base_period_dtype
+
+from .basepreprocessor import BasePreprocessor, BasePreprocessorSegment
 
 
 class SilencedPeriodsRecording(BasePreprocessor):
@@ -21,21 +22,26 @@ class SilencedPeriodsRecording(BasePreprocessor):
     ----------
     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
-    noise_levels : array
-        Noise levels if already computed
-    seed : int | None, default: None
-        Random seed for `get_noise_levels` and `NoiseGeneratorRecording`.
-        If none, `get_noise_levels` uses `seed=0` and `NoiseGeneratorRecording` generates a random seed using `numpy.random.default_rng`.
-    mode : "zeros" | "noise, default: "zeros"
+    periods : np.array
+        A numpy array with dtype `base_period_dtype` and fields
+        "segment_index", "start_sample_index", "end_sample_index".
+        Each row corresponds to a period to silence.
+    mode : "zeros" | "noise" | "apodization", default: "zeros"
         Determines what periods are replaced by. Can be one of the following:
 
         - "zeros": 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
+        - "apodization": The periods zeroed, but are apodized with a cosine taper (using `apodization_samples`)
+    apodization_samples : int, default: 7
+        The factor used for the cosine taper when mode is "apodization". Higher values create a wider taper.
+    noise_levels : array
+        Noise levels if already computed
+    seed : int | None, default: None
+        Random seed for `get_noise_levels` and `NoiseGeneratorRecording`.
+        If none, `get_noise_levels` uses `seed=0` and `NoiseGeneratorRecording` generates a random seed using `numpy.random.default_rng`.
     **noise_levels_kwargs : Keyword arguments for `spikeinterface.core.get_noise_levels()` function
 
     Returns
@@ -48,22 +54,29 @@ def __init__(
         self,
         recording,
         periods=None,
-        # this is keep for backward compatibility
+        # this is kept for backward compatibility
         list_periods=None,
         mode="zeros",
+        apodization_samples=7,
         noise_levels=None,
         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
         if list_periods is not None:
-            assert periods is None
+            assert periods is None, (
+                "You cannot specify both list_periods and periods. "
+                f"Please specify only periods, which should be a np.array with dtype {base_period_dtype}"
+            )
             periods = _all_period_list_to_periods_vec(list_periods, num_seg)
         else:
-            assert list_periods is None
+            assert list_periods is None, (
+                "list_periods is deprecated. Please specify periods, which should be a np.array with "
+                f"dtype {base_period_dtype}"
+            )
             if not isinstance(periods, np.ndarray):
                 raise ValueError(f"periods must be a np.array with dtype {base_period_dtype}")
 
@@ -108,11 +121,26 @@ 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_samples=apodization_samples,
             )
             self.add_recording_segment(rec_segment)
 
-        self._kwargs = dict(recording=recording, periods=periods, mode=mode, seed=seed, noise_levels=noise_levels)
+        # the base_period_dtype is a structured dtype, which is not json serializable
+        self._serializability["json"] = False
+
+        self._kwargs = dict(
+            recording=recording,
+            periods=periods,
+            mode=mode,
+            seed=seed,
+            noise_levels=noise_levels,
+            apodization_samples=apodization_samples,
+        )
 
 
 def _all_period_list_to_periods_vec(list_periods, num_seg):
@@ -154,18 +182,28 @@ 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_samples=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_samples = apodization_samples
 
     def get_traces(self, start_frame, end_frame, channel_indices):
-        traces = self.parent_recording_segment.get_traces(start_frame, end_frame, channel_indices)
+        if self.mode in ("zeros", "noise"):
+            margin = 0
+        elif self.mode == "apodization":
+            margin = self.apodization_samples
+        else:
+            raise ValueError(f"Unknown method {self.mode}")
+
+        traces, left_margin, right_margin = get_chunk_with_margin(
+            self.parent_recording_segment, start_frame, end_frame, channel_indices, margin=margin
+        )
 
         if self.periods.size > 0:
-            new_interval = np.array([start_frame, end_frame])
+            new_interval = np.array([start_frame - margin, end_frame + margin])
 
             lower_index = np.searchsorted(self.periods["end_sample_index"], new_interval[0])
             upper_index = np.searchsorted(self.periods["start_sample_index"], new_interval[1])
@@ -174,9 +212,14 @@ def get_traces(self, start_frame, end_frame, channel_indices):
                 traces = traces.copy()
 
                 periods_in_interval = self.periods[lower_index:upper_index]
+
+                # For apodization, we pre-allocate the mute function and cosine window
+                if self.mode == "apodization":
+                    mute_mask = np.zeros(traces.shape[0], dtype=np.float32)
+
                 for period in periods_in_interval:
-                    onset = max(0, period["start_sample_index"] - start_frame)
-                    offset = min(period["end_sample_index"] - start_frame, end_frame)
+                    onset = max(0, period["start_sample_index"] - start_frame - margin)
+                    offset = min(period["end_sample_index"] - start_frame + margin, end_frame + margin)
 
                     if self.mode == "zeros":
                         traces[onset:offset, :] = 0
@@ -185,8 +228,20 @@ def get_traces(self, start_frame, end_frame, channel_indices):
                             :, channel_indices
                         ]
                         traces[onset:offset, :] = noise[onset:offset]
-
-        return traces
+                    elif self.mode == "apodization":
+                        # apply a cosine taper to the saturation to create a mute function
+                        mute_mask[onset:offset] = 1
+
+                # For apodization, we apply the mute function including all periods to the whole trace,
+                # so that the edges of the silenced periods are smoothly tapered
+                if self.mode == "apodization":
+                    import scipy.signal
+
+                    win = scipy.signal.windows.cosine(self.apodization_samples)
+                    mute = np.maximum(0, 1 - scipy.signal.convolve(mute_mask, win, mode="same"))
+                    traces = (traces.astype(np.float32) * mute[:, np.newaxis]).astype(traces.dtype)
+        # discard margin
+        return traces[left_margin : traces.shape[0] - right_margin, :]
 
 
 # function for API