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/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
 

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

@@ -21,21 +21,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,14 +53,15 @@ 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
@@ -108,11 +114,23 @@ 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)
+        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,12 +172,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_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)
@@ -185,7 +204,15 @@ def get_traces(self, start_frame, end_frame, channel_indices):
                             :, channel_indices
                         ]
                         traces[onset:offset, :] = noise[onset:offset]
-
+                    elif self.mode == "apodization":
+                        import scipy.signal
+
+                        # 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_samples)
+                        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
 
 

src/spikeinterface/preprocessing/tests/test_detect_artifacts.py

@@ -5,6 +5,7 @@
     detect_artifact_periods,
     detect_saturation_periods,
     detect_artifact_periods_by_envelope,
+    detect_and_remove_artifacts,
 )
 
 
@@ -238,6 +239,119 @@ def test_detect_saturation_periods(debug_plots):
     assert np.array_equal(periods, periods_entry_with_annotation)
 
 
+def test_detect_saturation_signed():
+    import scipy.signal
+
+    num_chans = 32
+    sampling_frequency = 30000
+    chunk_size = 30000
+    job_kwargs = {"chunk_size": chunk_size}
+
+    sat_value = 1200
+    noise_level = 10
+    rng = np.random.default_rng(0)
+    data = noise_level * rng.uniform(low=-0.5, high=0.5, size=(90000, num_chans)) * 10
+
+    sos = scipy.signal.butter(N=3, Wn=8000 / (sampling_frequency / 2), btype="low", output="sos")
+    data = scipy.signal.sosfiltfilt(sos, data, axis=0)
+
+    # Inject positive saturation in first third, negative in second third
+    pos_start, pos_stop = 15000, 15500
+    neg_start, neg_stop = 45000, 45500
+    data[pos_start:pos_stop, :] = sat_value
+    data[neg_start:neg_stop, :] = -sat_value
+
+    gain = 2.34
+    offset = 0
+    data_int16 = np.clip(np.rint((data - offset) / gain), -32768, 32767).astype(np.int16)
+
+    recording = NumpyRecording(data_int16, sampling_frequency)
+    recording.set_channel_gains(gain)
+    recording.set_channel_offsets([offset] * num_chans)
+
+    periods = detect_saturation_periods(
+        recording,
+        saturation_threshold_uV=sat_value * 0.98,
+        signed=True,
+        job_kwargs=job_kwargs,
+    )
+
+    # Output dtype must include the "sign" field
+    assert "sign" in periods.dtype.names
+
+    pos_periods = periods[periods["sign"] == "positive"]
+    neg_periods = periods[periods["sign"] == "negative"]
+    assert len(pos_periods) > 0, "No positive saturation periods detected"
+    assert len(neg_periods) > 0, "No negative saturation periods detected"
+
+    # Positive period should be near the injected positive saturation
+    tolerance = 1
+    assert np.any(np.abs(pos_periods["start_sample_index"] - pos_start) <= tolerance)
+    assert np.any(np.abs(pos_periods["end_sample_index"] - pos_stop) <= tolerance)
+
+    # Negative period should be near the injected negative saturation
+    assert np.any(np.abs(neg_periods["start_sample_index"] - neg_start) <= tolerance)
+    assert np.any(np.abs(neg_periods["end_sample_index"] - neg_stop) <= tolerance)
+
+    # Positive periods must not contain any sample indices from the negative injection
+    for p in pos_periods:
+        assert not (p["start_sample_index"] < neg_stop and p["end_sample_index"] > neg_start)
+
+    # Negative periods must not contain any sample indices from the positive injection
+    for p in neg_periods:
+        assert not (p["start_sample_index"] < pos_stop and p["end_sample_index"] > pos_start)
+
+
+def test_detect_and_remove_artifacts():
+    import scipy.signal
+
+    num_chans = 32
+    sampling_frequency = 30000
+    chunk_size = 30000
+    job_kwargs = {"chunk_size": chunk_size}
+
+    sat_value = 1200
+    noise_level = 10
+    rng = np.random.default_rng(0)
+    data = noise_level * rng.uniform(low=-0.5, high=0.5, size=(90000, num_chans)) * 10
+
+    sos = scipy.signal.butter(N=3, Wn=8000 / (sampling_frequency / 2), btype="low", output="sos")
+    data = scipy.signal.sosfiltfilt(sos, data, axis=0)
+
+    sat_start, sat_stop = 15000, 15500
+    data[sat_start:sat_stop, :] = sat_value
+
+    gain = 2.34
+    offset = 0
+    data_int16 = np.clip(np.rint((data - offset) / gain), -32768, 32767).astype(np.int16)
+
+    recording = NumpyRecording(data_int16, sampling_frequency)
+    recording.set_channel_gains(gain)
+    recording.set_channel_offsets([offset] * num_chans)
+
+    # Basic usage: detect and zero out saturation in one step
+    cleaned = detect_and_remove_artifacts(
+        recording,
+        method="saturation",
+        method_kwargs=dict(saturation_threshold_uV=sat_value * 0.98),
+        job_kwargs=job_kwargs,
+    )
+    traces = cleaned.get_traces(segment_index=0)
+    assert traces[sat_start + 100, 0] == 0, "Saturated samples should be zeroed"
+    assert traces[0, 0] != 0, "Non-saturated samples should not be zeroed"
+
+    # recording_to_detect: detect on raw recording, silence a separate (processed) recording
+    # We use the same recording here just to exercise the code path
+    cleaned_with_detect = detect_and_remove_artifacts(
+        recording,
+        recording_to_detect=recording,
+        method="saturation",
+        method_kwargs=dict(saturation_threshold_uV=sat_value * 0.98),
+        job_kwargs=job_kwargs,
+    )
+    assert np.array_equal(cleaned.get_traces(), cleaned_with_detect.get_traces())
+
+
 if __name__ == "__main__":
     # test_detect_artifact_by_envelope(True)
     test_detect_saturation_periods(False)

src/spikeinterface/preprocessing/tests/test_pipeline.py

@@ -1,3 +1,5 @@
+import pytest
+
 from spikeinterface.generation import generate_recording, generate_ground_truth_recording
 from spikeinterface.preprocessing import (
     apply_preprocessing_pipeline,
@@ -212,6 +214,51 @@ def test_loading_from_analyzer(create_cache_folder):
     check_recordings_equal(pp_recording, pp_recording_from_zarr)
 
 
+def test_pipeline_recording_arg_substitution():
+    """
+    Tests that if a preprocessing step in the pipeline has an argument that is a string of the form "pipeline[preprocessor_name]",
+    then this string is replaced by the recording output by the preprocessor with name "preprocessor_name". This allows users to
+    use outputs of previous preprocessors as arguments for later preprocessors in the same pipeline.
+    """
+    from spikeinterface.preprocessing.filter import BandpassFilterRecording
+    from spikeinterface.preprocessing.common_reference import CommonReferenceRecording
+    from spikeinterface.preprocessing.detect_artifacts import DetectAndRemoveArtifactsRecording
+
+    rec = generate_recording(durations=[1])
+
+    # "recording" argument is protected, as it is the default argument for the recording to preprocess
+    pipeline_dict_wrong = {
+        "common_reference": {},
+        "bandpass_filter": {"recording": "pipeline[raw]"},
+    }
+    with pytest.raises(ValueError):
+        pp_rec_from_pipeline = apply_preprocessing_pipeline(rec, pipeline_dict_wrong)
+
+    # The argument using the pipeline substitution must be a string with "recording" as substring
+    pipeline_dict_wrong2 = {
+        "common_reference": {},
+        "bandpass_filter": {"freq_min": "pipeline[raw]"},
+    }
+    with pytest.raises(ValueError):
+        pp_rec_from_pipeline = apply_preprocessing_pipeline(rec, pipeline_dict_wrong2)
+
+    # Correct usage: the "recording_to_detect" argument for the "detect_and_remove_artifacts" step is set to be the
+    # output of the "bandpass_filter" step, which is correctly substituted when applying the pipeline.
+    # The "recording" argument for the "detect_and_remove_artifacts" step should be set to the output of the
+    # "common_reference" step, as this is the last preprocessor in the pipeline before it.
+    pipeline_dict_correct = {
+        "bandpass_filter": {},
+        "common_reference": {},
+        "detect_and_remove_artifacts": {"recording_to_detect": "pipeline[bandpass_filter]"},
+    }
+    pp_rec_from_pipeline = apply_preprocessing_pipeline(rec, pipeline_dict_correct)
+    # Check that the recording argument for detect step is common ref,
+    # and that the recording_to_detect argument for detect_and_remove_artifacts is also the output of bandpass_filter
+    assert isinstance(pp_rec_from_pipeline._kwargs["recording_to_detect"], BandpassFilterRecording)
+    assert isinstance(pp_rec_from_pipeline._kwargs["recording"], CommonReferenceRecording)
+    assert isinstance(pp_rec_from_pipeline, DetectAndRemoveArtifactsRecording)
+
+
 if __name__ == "__main__":
     import tempfile
     from pathlib import Path