Merge pull request #1495 from alejoe91/sync-973-to-main

Sync #973 to main

Author: samuelgarcia

doc/api.rst

@@ -149,6 +149,7 @@ spikeinterface.preprocessing
     .. autofunction:: correct_lsb
     .. autofunction:: detect_bad_channels
     .. autofunction:: filter
+    .. autofunction:: gaussian_bandpass_filter
     .. autofunction:: highpass_filter
     .. autofunction:: highpass_spatial_filter
     .. autofunction:: interpolate_bad_channels

src/spikeinterface/preprocessing/filter_gaussian.py

@@ -0,0 +1,98 @@
+from typing import Iterable, Union
+import numpy as np
+from scipy.stats import norm
+from spikeinterface.core import BaseRecording, BaseRecordingSegment, get_chunk_with_margin
+from spikeinterface.core.core_tools import define_function_from_class
+from .basepreprocessor import BasePreprocessor, BasePreprocessorSegment
+
+
+class GaussianBandpassFilterRecording(BasePreprocessor):
+    """
+    Class for performing a bandpass gaussian filtering/smoothing on a recording.
+    This is done by a convolution with a Gaussian kernel, which acts as a lowpass-filter.
+    The highpass-filter can be computed by subtracting the result.
+
+    Here, the bandpass is computed in the Fourier domain to accelerate the computation.
+
+    Parameters
+    ----------
+    recording: BaseRecording
+        The recording extractor to be filtered.
+    freq_min: float
+        The lower frequency cutoff for the bandpass filter.
+    freq_max: float
+        The higher frequency cutoff for the bandpass filter.
+
+    Returns
+    -------
+    gaussian_bandpass_filtered_recording: GaussianBandpassFilterRecording
+        The filtered recording extractor object.
+    """
+    name = 'gaussian_bandpass_filter'
+
+    def __init__(self, recording: BaseRecording, freq_min: float = 300., freq_max: float = 5000.):
+        sf = recording.sampling_frequency
+        BasePreprocessor.__init__(self, recording)
+        self.annotate(is_filtered=True)
+
+        for parent_segment in recording._recording_segments:
+            self.add_recording_segment(GaussianFilterRecordingSegment(parent_segment, freq_min, freq_max))
+
+        self._kwargs = {'recording': recording.to_dict(), 'freq_min': freq_min, 'freq_max': freq_max}
+
+
+class GaussianFilterRecordingSegment(BasePreprocessorSegment):
+
+    def __init__(self, parent_recording_segment: BaseRecordingSegment, freq_min: float, freq_max: float):
+        BasePreprocessorSegment.__init__(self, parent_recording_segment)
+
+        self.freq_min = freq_min
+        self.freq_max = freq_max
+        self.cached_gaussian  = dict()
+
+        sf = parent_recording_segment.sampling_frequency
+        low_sigma  = sf / (2*np.pi * freq_min)
+        high_sigma = sf / (2*np.pi * freq_max)
+        self.margin = int(max(low_sigma, high_sigma) * 6. + 1)
+
+    def get_traces(self, start_frame: Union[int, None] = None, end_frame: Union[int, None] = None,
+                   channel_indices: Union[Iterable, None] = None):
+        traces, left_margin, right_margin = get_chunk_with_margin(self.parent_recording_segment, start_frame,
+                                                                  end_frame, channel_indices, self.margin)
+        dtype = traces.dtype
+        
+        traces_fft = np.fft.fft(traces, axis=0)
+        gauss_low  = self._create_gaussian(traces.shape[0], self.freq_min)
+        gauss_high = self._create_gaussian(traces.shape[0], self.freq_max)
+
+        filtered_fft = traces_fft * (gauss_high - gauss_low)[:, None]
+        filtered_traces = np.real(np.fft.ifft(filtered_fft, axis=0))
+
+        if right_margin > 0:
+            return filtered_traces[left_margin : -right_margin, :].astype(dtype)
+        else:
+            return filtered_traces[left_margin:, :].astype(dtype)
+
+    def _create_gaussian(self, N: int, cutoff_f: float):
+        if cutoff_f in self.cached_gaussian and N in self.cached_gaussian[cutoff_f]:
+            return self.cached_gaussian[cutoff_f][N]
+
+        sf = self.parent_recording_segment.sampling_frequency
+        faxis = np.fft.fftfreq(N, d=1/sf)
+
+        if cutoff_f > sf / 8:  # The Fourier transform of a Gaussian with a very low sigma isn't a Gaussian.
+            sigma = sf / (2*np.pi * cutoff_f)
+            limit = int(round(6*sigma)) + 1
+            xaxis = np.arange(-limit, limit+1) / sigma
+            gaussian = norm.pdf(xaxis) / sigma
+            gaussian = np.abs(np.fft.fft(gaussian, n=N))
+        else:
+            gaussian = norm.pdf(faxis / cutoff_f) * np.sqrt(2*np.pi)
+
+        if cutoff_f not in self.cached_gaussian:
+            self.cached_gaussian[cutoff_f] = dict()
+        self.cached_gaussian[cutoff_f][N] = gaussian
+
+        return gaussian
+
+gaussian_bandpass_filter = define_function_from_class(source_class=GaussianBandpassFilterRecording, name="gaussian_filter")

src/spikeinterface/preprocessing/preprocessinglist.py

@@ -5,6 +5,7 @@
                      NotchFilterRecording, notch_filter,
                      HighpassFilterRecording, highpass_filter,
                      )
+from .filter_gaussian import (GaussianBandpassFilterRecording, gaussian_bandpass_filter)
 from .normalize_scale import (
     NormalizeByQuantileRecording, normalize_by_quantile,
     ScaleRecording, scale,
@@ -34,6 +35,7 @@
     BandpassFilterRecording,
     HighpassFilterRecording,
     NotchFilterRecording,
+    GaussianBandpassFilterRecording,
 
     # gain offset stuff
     NormalizeByQuantileRecording,

src/spikeinterface/preprocessing/tests/test_filter_gaussian.py

@@ -0,0 +1,45 @@
+import numpy as np
+import pytest
+from pathlib import Path
+from spikeinterface.core import load_extractor, set_global_tmp_folder
+from spikeinterface.core.testing import check_recordings_equal
+from spikeinterface.core.generate import generate_recording
+from spikeinterface.preprocessing import gaussian_bandpass_filter
+
+
+if hasattr(pytest, "global_test_folder"):
+    cache_folder = pytest.global_test_folder / "preprocessing" / "gaussian_bandpass_filter"
+else:
+    cache_folder = Path("cache_folder") / "preprocessing" / "gaussian_bandpass_filter"
+
+set_global_tmp_folder(cache_folder)
+cache_folder.mkdir(parents=True, exist_ok=True)
+
+
+def test_filter_gaussian():
+    recording = generate_recording(num_channels=3)
+    recording.annotate(is_filtered=True)
+    recording = recording.save(folder=cache_folder / "recording")
+
+    rec_filtered = gaussian_bandpass_filter(recording)
+
+    assert rec_filtered.dtype == recording.dtype
+    assert rec_filtered.get_traces(segment_index=0, end_frame=100).dtype == rec_filtered.dtype
+    assert rec_filtered.get_traces(segment_index=0, end_frame=600).shape == (600, 3)
+    assert rec_filtered.get_traces(segment_index=0, start_frame=100, end_frame=600).shape == (500, 3)
+    assert rec_filtered.get_traces(segment_index=1, start_frame=rec_filtered.get_num_frames(1) - 200).shape == (200, 3)
+
+    # Check dumpability
+    saved_loaded = load_extractor(rec_filtered.to_dict())
+    check_recordings_equal(rec_filtered, saved_loaded, return_scaled=False)
+
+    saved_1job = rec_filtered.save(folder=cache_folder / "1job")
+    saved_2job = rec_filtered.save(folder=cache_folder / "2job", n_jobs=2, chunk_duration='1s')
+
+    for seg_idx in range(rec_filtered.get_num_segments()):
+        original_trace = rec_filtered.get_traces(seg_idx)
+        saved1_trace = saved_1job.get_traces(seg_idx)
+        saved2_trace = saved_2job.get_traces(seg_idx)
+
+        assert np.allclose(original_trace[60:-60], saved1_trace[60:-60], rtol=1e-3, atol=1e-3)
+        assert np.allclose(original_trace[60:-60], saved2_trace[60:-60], rtol=1e-3, atol=1e-3)