Add fir_filter step and FIR highpass and bandpass designs

doc/api.rst

@@ -189,6 +189,8 @@ spikeinterface.preprocessing
     .. autofunction:: save_motion_info
     .. autofunction:: decimate
     .. autofunction:: depth_order
+    .. autofunction:: design_fir_bandpass
+    .. autofunction:: design_fir_highpass
     .. autofunction:: detect_bad_channels
     .. autofunction:: detect_and_interpolate_bad_channels
     .. autofunction:: detect_and_remove_bad_channels
@@ -198,6 +200,7 @@ spikeinterface.preprocessing
     .. autofunction:: detect_and_remove_artifacts
     .. autofunction:: directional_derivative
     .. autofunction:: filter
+    .. autofunction:: fir_filter
     .. autofunction:: gaussian_filter
     .. autofunction:: highpass_filter
     .. autofunction:: highpass_spatial_filter

src/spikeinterface/core/testing.py

@@ -91,6 +91,35 @@ def check_recordings_equal(
         check_extractor_properties_equal(RX1, RX2)
 
 
+def check_chunk_invariance(
+    recording: BaseRecording, chunk_size: int, atol: float = 0.0, segment_index: int = 0
+) -> float:
+    """
+    Check that traces read in chunks of `chunk_size` and concatenated equal traces read in one call.
+
+    A preprocessor whose margin covers its full support should pass with `atol` at float rounding; one that
+    truncates an unbounded response, such as an IIR filter, should pass with `atol` at its stated seam error.
+
+    Returns
+    -------
+    max_error : float
+        The largest absolute difference between the chunked and the whole traces.
+    """
+    whole = recording.get_traces(segment_index=segment_index).astype(np.float64)
+    num_samples = whole.shape[0]
+    chunked = np.concatenate(
+        [
+            recording.get_traces(
+                segment_index=segment_index, start_frame=start, end_frame=min(start + chunk_size, num_samples)
+            )
+            for start in range(0, num_samples, chunk_size)
+        ]
+    )
+    max_error = float(np.abs(chunked - whole).max())
+    assert max_error <= atol, f"chunked traces differ from whole traces by {max_error}, above {atol}"
+    return max_error
+
+
 def check_sortings_equal(
     SX1: BaseSorting,
     SX2: BaseSorting,

src/spikeinterface/preprocessing/__init__.py

@@ -12,6 +12,7 @@
 from .preprocessing_tools import get_spatial_interpolation_kernel
 from .detect_bad_channels import detect_bad_channels
 from .correct_lsb import correct_lsb
+from .filter_design import design_fir_bandpass, design_fir_highpass
 
 from .pipeline import (
     apply_preprocessing_pipeline,

src/spikeinterface/preprocessing/filter_design.py

@@ -0,0 +1,116 @@
+"""Filter design, separate from filter application, so kernels can be inspected and composed before use."""
+
+from __future__ import annotations
+
+import numpy as np
+
+
+def _stopband_db(taps, fs, edges):
+    """Smallest attenuation, in dB, of `taps` over the frequency intervals `edges`, sampled 16 times per ripple."""
+    import scipy.signal
+
+    worst = 0.0
+    for low, high in edges:
+        freqs = np.linspace(low, high, max(64, int(16 * (high - low) / fs * taps.size)))
+        _, response = scipy.signal.freqz(taps, worN=freqs, fs=fs)
+        worst = max(worst, np.abs(response).max())
+    return -20 * np.log10(max(worst, 1e-300))
+
+
+def design_fir_bandpass(freq_min, freq_max, fs, transition_low_hz=175.0, transition_high_hz=2000.0, stopband_db=45.0):
+    """
+    Design a linear-phase bandpass FIR kernel by the Kaiser window method.
+
+    The response is -6 dB at `freq_min` and `freq_max`, and at least `stopband_db` down below
+    ``freq_min - transition_low_hz / 2`` and above ``freq_max + transition_high_hz / 2``. Passband ripple is at
+    most about ``2 * 10 ** (-stopband_db / 20)`` in amplitude. The kernel convolves a highpass and a lowpass,
+    each windowed to its own transition: one window would give both edges the narrower width, and a sharp
+    upper edge rings for longer.
+
+    Parameters
+    ----------
+    freq_min, freq_max : float
+        Lower and upper -6 dB cutoffs, in Hz.
+    fs : float
+        Sampling frequency, in Hz.
+    transition_low_hz, transition_high_hz : float, default: 175.0, 2000.0
+        Width of the transition band centred on each cutoff, in Hz.
+    stopband_db : float, default: 45.0
+        Least attenuation in both stopbands, in dB.
+
+    Returns
+    -------
+    taps : np.ndarray
+        Symmetric, odd-length float64 kernel, for ``fir_filter``.
+    """
+    import scipy.signal
+
+    stop_low = freq_min - transition_low_hz / 2
+    stop_high = freq_max + transition_high_hz / 2
+    if transition_low_hz <= 0 or transition_high_hz <= 0 or stopband_db <= 0:
+        raise ValueError("transition_low_hz, transition_high_hz and stopband_db must be positive")
+    if not 0 < stop_low < freq_min < freq_max < stop_high < fs / 2:
+        raise ValueError(
+            f"The transition bands, [{stop_low}, {freq_min + transition_low_hz / 2}] and "
+            f"[{freq_max - transition_high_hz / 2}, {stop_high}] Hz, must lie in (0, {fs / 2}) Hz in that order"
+        )
+
+    def edge(cutoff, width, pass_zero):
+        num_taps, beta = scipy.signal.kaiserord(stopband_db, width / (fs / 2))
+        return num_taps | 1, beta, cutoff, pass_zero
+
+    edges = [edge(freq_min, transition_low_hz, False), edge(freq_max, transition_high_hz, True)]
+    stopbands = [(0.0, stop_low), (stop_high, fs / 2)]
+    while True:
+        highpass, lowpass = (
+            scipy.signal.firwin(n, cutoff, window=("kaiser", beta), pass_zero=pass_zero, fs=fs)
+            for n, beta, cutoff, pass_zero in edges
+        )
+        taps = np.convolve(highpass, lowpass)
+        # The two kernels' ripples can add, so the window's nominal attenuation is checked and the failing edge
+        # lengthened.
+        misses = [_stopband_db(taps, fs, [band]) < stopband_db for band in stopbands]
+        if not any(misses):
+            return taps
+        edges = [(n + 2 * miss, beta, cutoff, pz) for (n, beta, cutoff, pz), miss in zip(edges, misses)]
+
+
+def design_fir_highpass(freq_min, fs, transition_low_hz=175.0, stopband_db=45.0):
+    """
+    Design a linear-phase highpass FIR kernel by the Kaiser window method.
+
+    The response is -6 dB at `freq_min` and at least `stopband_db` down below ``freq_min - transition_low_hz / 2``,
+    with passband ripple of at most about ``10 ** (-stopband_db / 20)`` in amplitude.
+
+    Parameters
+    ----------
+    freq_min : float
+        The -6 dB cutoff, in Hz.
+    fs : float
+        Sampling frequency, in Hz.
+    transition_low_hz : float, default: 175.0
+        Width of the transition band centred on the cutoff, in Hz.
+    stopband_db : float, default: 45.0
+        Least attenuation in the stopband, in dB.
+
+    Returns
+    -------
+    taps : np.ndarray
+        Symmetric, odd-length float64 kernel, for ``fir_filter``.
+    """
+    import scipy.signal
+
+    stop = freq_min - transition_low_hz / 2
+    if transition_low_hz <= 0 or stopband_db <= 0:
+        raise ValueError("transition_low_hz and stopband_db must be positive")
+    if not 0 < stop < freq_min + transition_low_hz / 2 < fs / 2:
+        raise ValueError(
+            f"The transition band, [{stop}, {freq_min + transition_low_hz / 2}] Hz, must lie in (0, {fs / 2}) Hz"
+        )
+    num_taps, beta = scipy.signal.kaiserord(stopband_db, transition_low_hz / (fs / 2))
+    num_taps |= 1  # a highpass needs a type I kernel
+    while True:
+        taps = scipy.signal.firwin(num_taps, freq_min, window=("kaiser", beta), pass_zero=False, fs=fs)
+        if _stopband_db(taps, fs, [(0.0, stop)]) >= stopband_db:
+            return taps
+        num_taps += 2

src/spikeinterface/preprocessing/fir_filter.py

@@ -0,0 +1,218 @@
+from __future__ import annotations
+
+import math
+import threading
+
+import numpy as np
+
+from spikeinterface.core.core_tools import define_function_handling_dict_from_class
+
+from .basepreprocessor import BasePreprocessor, BasePreprocessorSegment
+from .filter import fix_dtype
+
+
+def _detect_l2_bytes():
+    """Per-core L2 size from Linux sysfs, else 1 MB, which fits in any current core's L2."""
+    for index in range(5):
+        cache = f"/sys/devices/system/cpu/cpu0/cache/index{index}"
+        try:
+            with open(f"{cache}/level") as f:
+                level = int(f.read())
+            with open(f"{cache}/type") as f:
+                cache_type = f.read().strip()
+            if level != 2 or cache_type not in ("Unified", "Data"):
+                continue
+            with open(f"{cache}/size") as f:
+                size = f.read().strip()
+        except (OSError, ValueError):
+            continue
+        scale = {"K": 1024, "M": 1024**2}.get(size[-1:], 1)
+        return int(size.rstrip("KM")) * scale
+    return 1_000_000
+
+
+_L2_BYTES = _detect_l2_bytes()
+
+
+def _overlap_save_fft_length(signal_length, num_taps):
+    """FFT length minimising FFT work per output sample, as ``scipy.signal.oaconvolve`` chooses it.
+
+    Unlike ``oaconvolve``, the length is rounded up to one that is fast for a real FFT: lengths with a factor of 7
+    or 11 are fast for a complex FFT but slower for ``rfft``.
+    """
+    import scipy.fft
+    from scipy.special import lambertw
+
+    whole = int(scipy.fft.next_fast_len(signal_length + num_taps - 1, real=True))
+    if num_taps == 1 or num_taps >= signal_length / 2:
+        return whole
+    overlap = num_taps - 1
+    optimal = -overlap * lambertw(-1 / (2 * math.e * overlap), k=-1).real
+    return min(whole, int(scipy.fft.next_fast_len(math.ceil(optimal), real=True)))
+
+
+def _channels_per_block(num_channels, fft_length, bytes_per_sample, pool_size=1, headroom=0.6):
+    """
+    Channels per overlap-save block: as many as keep the block's buffers in L2, but no more than gives each of
+    `pool_size` workers a block.
+    """
+    fit_l2 = max(1, int(headroom * _L2_BYTES / (bytes_per_sample * fft_length)))
+    return min(fit_l2, max(1, math.ceil(num_channels / pool_size)))
+
+
+class FIRFilterRecording(BasePreprocessor):
+    """
+    Applies a FIR filter to each channel, by overlap-save convolution.
+
+    Output sample ``n`` is ``sum_j taps[j] * x[n + center - j]``, so with the default `center` a symmetric
+    kernel is zero-phase. Every output sample depends only on the ``len(taps) - 1`` input samples around it,
+    which the filter reads from the parent, so the output does not depend on how it is chunked.
+
+    Parameters
+    ----------
+    recording : Recording
+        The recording to filter.
+    taps : array-like
+        Kernel of shape ``(num_taps,)``, applied to every channel, or ``(num_kernels, num_taps)``, with each
+        channel's row given by `channel_groups`. Designed with, for example, ``design_fir_bandpass`` or
+        ``scipy.signal.firwin``.
+    channel_groups : array-like | None, default: None
+        For a 2D `taps`, the row of `taps` for each channel. None requires one row per channel.
+    center : int | None, default: None
+        Index of the tap aligned with the output sample. None means ``(num_taps - 1) // 2``, zero-phase for a
+        symmetric kernel; 0 makes the filter causal.
+    dtype : dtype | None, default: None
+        The dtype of the output. None uses the parent's. Computation is in float32, and integer output is
+        rounded.
+    add_reflect_padding : bool, default: False
+        If True, the input is reflected beyond the recording's edges, otherwise it is zero there.
+
+    Returns
+    -------
+    filtered_recording : FIRFilterRecording
+        The filtered recording.
+    """
+
+    def __init__(self, recording, taps, channel_groups=None, center=None, dtype=None, add_reflect_padding=False):
+        taps = np.asarray(taps, dtype=np.float64)
+        if taps.ndim not in (1, 2) or taps.shape[-1] == 0:
+            raise ValueError(f"taps must have shape (num_taps,) or (num_kernels, num_taps), got {taps.shape}")
+        kernels = np.atleast_2d(taps)
+        num_kernels, num_taps = kernels.shape
+        num_channels = recording.get_num_channels()
+        if channel_groups is None:
+            if num_kernels not in (1, num_channels):
+                raise ValueError(f"{num_kernels} kernels for {num_channels} channels need channel_groups")
+            groups = np.zeros(num_channels, dtype=np.intp) if num_kernels == 1 else np.arange(num_channels)
+        else:
+            groups = np.asarray(channel_groups, dtype=np.intp)
+            if groups.shape != (num_channels,) or groups.min() < 0 or groups.max() >= num_kernels:
+                raise ValueError(f"channel_groups must give a row of taps, 0 to {num_kernels - 1}, for each channel")
+        center = (num_taps - 1) // 2 if center is None else int(center)
+        if not 0 <= center < num_taps:
+            raise ValueError(f"center must be in [0, {num_taps}), got {center}")
+        dtype = fix_dtype(recording, dtype)
+
+        BasePreprocessor.__init__(self, recording, dtype=dtype)
+        self.annotate(is_filtered=True)
+        if "offset_to_uV" in self.get_property_keys():
+            # A constant offset passes through scaled by each kernel's DC gain.
+            self.set_channel_offsets(self.get_channel_offsets() * kernels.sum(axis=1)[groups])
+
+        for parent_segment in recording.segments:
+            self.add_recording_segment(
+                FIRFilterRecordingSegment(parent_segment, kernels, groups, center, dtype, add_reflect_padding)
+            )
+
+        self._kwargs = dict(
+            recording=recording,
+            taps=taps.tolist(),
+            channel_groups=None if channel_groups is None else groups.tolist(),
+            center=center,
+            dtype=dtype.str,
+            add_reflect_padding=add_reflect_padding,
+        )
+
+
+class FIRFilterRecordingSegment(BasePreprocessorSegment):
+    def __init__(self, parent_recording_segment, kernels, channel_groups, center, dtype, add_reflect_padding):
+        BasePreprocessorSegment.__init__(self, parent_recording_segment)
+        self.kernels = kernels
+        self.num_taps = kernels.shape[1]
+        self.channel_groups = channel_groups
+        self.center = center
+        self.dtype = dtype
+        self.add_reflect_padding = add_reflect_padding
+        # Kernel spectra depend on the FFT length, which depends on the chunk length.
+        self._kernel_spectra = {}
+        self._kernel_spectra_lock = threading.Lock()
+
+    def _get_kernel_spectra(self, fft_length):
+        import scipy.fft
+
+        with self._kernel_spectra_lock:
+            spectra = self._kernel_spectra.get(fft_length)
+            if spectra is None:
+                spectra = scipy.fft.rfft(self.kernels.T, n=fft_length, axis=0).astype(np.complex64)
+                self._kernel_spectra[fft_length] = spectra
+        return spectra
+
+    def _get_input(self, start_frame, end_frame, channel_indices):
+        """Parent traces for every input sample the kernel reaches, padded beyond the recording's edges."""
+        left = self.num_taps - 1 - self.center
+        right = self.center
+        fetch_start = max(0, start_frame - left)
+        fetch_end = min(self.parent_recording_segment.get_num_samples(), end_frame + right)
+        traces = self.parent_recording_segment.get_traces(fetch_start, fetch_end, channel_indices)
+        pad = ((fetch_start - (start_frame - left), end_frame + right - fetch_end), (0, 0))
+        if pad[0] != (0, 0):
+            traces = np.pad(traces, pad, mode="reflect" if self.add_reflect_padding else "constant")
+        return traces
+
+    def get_traces(self, start_frame, end_frame, channel_indices):
+        import scipy.fft
+
+        if start_frame is None:
+            start_frame = 0
+        if end_frame is None:
+            end_frame = self.get_num_samples()
+        if channel_indices is None:
+            channel_indices = slice(None)
+        traces = self._get_input(start_frame, end_frame, channel_indices)
+        groups = self.channel_groups[channel_indices]
+        num_samples, num_channels = end_frame - start_frame, traces.shape[1]
+        num_taps = self.num_taps
+
+        fft_length = _overlap_save_fft_length(num_samples + num_taps - 1, num_taps)
+        step = fft_length - num_taps + 1
+        spectra = self._get_kernel_spectra(fft_length)
+        one_kernel = spectra.shape[1] == 1
+        # The float32 block and its complex64 spectrum take 8 bytes per sample, and per-channel kernel spectra 4 more.
+        slab = _channels_per_block(num_channels, fft_length, bytes_per_sample=8 if one_kernel else 12)
+
+        # Each block is rounded and cast as it is written, so no full float32 copy of the chunk is held.
+        out = np.empty((num_samples, num_channels), dtype=self.dtype)
+        round_to_integer = np.issubdtype(self.dtype, np.integer)
+        buffer = np.empty((fft_length, slab), dtype=np.float32)
+        for c0 in range(0, num_channels, slab):
+            c1 = min(num_channels, c0 + slab)
+            tile = buffer[:, : c1 - c0]
+            tile_spectra = spectra if one_kernel else spectra[:, groups[c0:c1]]
+            for t0 in range(0, num_samples, step):
+                num_out = min(step, num_samples - t0)
+                num_in = num_out + num_taps - 1
+                tile[:num_in] = traces[t0 : t0 + num_in, c0:c1]
+                tile[num_in:] = 0
+                spectrum = scipy.fft.rfft(tile, axis=0)
+                spectrum *= tile_spectra
+                # The first num_taps - 1 samples of each circular convolution wrap around and are discarded.
+                result = scipy.fft.irfft(spectrum, n=fft_length, axis=0, overwrite_x=True)
+                valid = result[num_taps - 1 : num_taps - 1 + num_out]
+                if round_to_integer:
+                    np.rint(valid, out=valid)
+                out[t0 : t0 + num_out, c0:c1] = valid
+
+        return out
+
+
+fir_filter = define_function_handling_dict_from_class(source_class=FIRFilterRecording, name="fir_filter")