[EHN] Filterbank as module

[bruAristimunha]

In this PR, we've implemented a FilterBank inside a Pytorch Module, we have a very strong dependency on MNE and we haven't implemented a way to apply the IIR filter.

Basically we use mne to create the filter and use fftconvolve to apply the filter. The implementation of fftconvolve comes from pytorch audio.

What I need help with here, as not an expert in filtering:

• Check that I haven't made any fundamental mistakes, such as additional dimensions or related details.
• Consider whether this way of applying the bandpass, using fftconv, is correct.
• Suggest ways of optimising the code.
• Request tests and requirements for merging.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 86.55%. Comparing base (0d14fdc) to head (8f6e244).
Report is 1 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #656      +/-   ##
==========================================
+ Coverage   86.42%   86.55%   +0.12%     
==========================================
  Files          71       71              
  Lines        6358     6418      +60     
==========================================
+ Hits         5495     5555      +60     
  Misses        863      863              

[bruAristimunha]

Small code to test

    from torch import randn

    x = randn(16, 10, 1000)

    layer = FilterBank(sfreq=256, n_chans=10)
    with torch.no_grad():
        out = layer(x)

    print(out.shape)

[bruAristimunha]

The implementation doesn't match the MNE implementation because the direct conversion of the code had low performance and batch execution, but I keep the converted code here

import numpy as np

from mne.filter import _check_zero_phase_length, next_fast_len
from mne.utils import logger


def _smart_pad_torch(x, n_pad, pad="reflect_limited"):
    """Pad vector x."""
    n_pad = np.asarray(n_pad)
    assert n_pad.shape == (2,)
    if (n_pad == 0).all():
        return x
    elif (n_pad < 0).any():
        raise RuntimeError("n_pad must be non-negative")
    if pad == "reflect_limited":
        # need to pad with zeros if len(x) <= npad
        left_zero_pad = torch.zeros(max(n_pad[0] - len(x) + 1, 0))
        right_zero_pad = torch.zeros(max(n_pad[1] - len(x) + 1, 0))
        reflection_left = 2 * x[0]
        reflection_right = 2 * x[-1]
        padded = torch.cat(
            [
                left_zero_pad,
                reflection_left - torch.flip(x[1 : n_pad[0] + 1], dims=[0]),
                x,
                reflection_right - torch.flip(x[-n_pad[1] - 2 + 1 : -1], dims=[0]),
                right_zero_pad,
            ]
        )

        return padded

    else:
        left_pad, right_pad = n_pad.tolist()
        return torch.nn.functional.pad(x, (left_pad, right_pad), mode=pad)


def _1d_overlap_filter_torch(x, n_h, n_edge, phase, h_fft, pad, n_fft):
    """Do one-dimensional overlap-add FFT FIR filtering using PyTorch."""

    # Pad to reduce ringing
    x_ext = _smart_pad_torch(x, (n_edge, n_edge), pad)
    n_x = x_ext.shape[0]
    x_filtered = torch.zeros_like(x_ext)

    n_seg = n_fft - n_h + 1
    n_segments = (n_x + n_seg - 1) // n_seg  # Equivalent to ceil division
    shift = ((n_h - 1) // 2 if phase.startswith("zero") else 0) + n_edge

    # Actual filtering step
    for seg_idx in range(n_segments):
        start = seg_idx * n_seg
        stop = (seg_idx + 1) * n_seg
        seg = x_ext[start:stop]

        # Pad segment to length n_fft
        seg = torch.cat(
            [seg, torch.zeros(n_fft - seg.shape[0], dtype=seg.dtype, device=seg.device)]
        )

        # FFT of the segment
        x_fft = torch.fft.rfft(seg, n=n_fft)
        x_fft *= h_fft
        prod = torch.fft.irfft(x_fft, n=n_fft)

        # Overlap-add operation
        start_filt = max(0, start - shift)
        stop_filt = min(start - shift + n_fft, n_x)
        start_prod = max(0, shift - start)
        stop_prod = start_prod + stop_filt - start_filt
        x_filtered[start_filt:stop_filt] += prod[start_prod:stop_prod]

    # Remove mirrored edges and cast to original dtype
    x_filtered = x_filtered[: n_x - 2 * n_edge].type_as(x)
    return x_filtered


def _overlap_add_filter_torch(
    x,
    h,
    n_fft=None,
    phase="zero",
    pad="reflect_limited",
):
    """Filter the signal x using h with overlap-add FFTs."""
    # set up array for filtering, reshape to 2D, operate on last axis
    orig_shape = x.shape
    nchans = orig_shape[1]
    # reshaping data to 2D
    x = x.view(-1, x.shape[-1])
    # Extend the signal by mirroring the edges to reduce transient filter
    # response
    _check_zero_phase_length(len(h), phase)
    if len(h) == 1:
        return x * h**2 if phase == "zero-double" else x * h
    n_edge = max(min(len(h), x.shape[1]) - 1, 0)
    logger.debug(f"Smart-padding with:  {n_edge} samples on each edge")
    n_x = x.shape[1] + 2 * n_edge

    # Determine FFT length to use
    min_fft = 2 * len(h) - 1
    if n_fft is None:
        max_fft = n_x
        if max_fft >= min_fft:
            # cost function based on number of multiplications
            N = 2 ** np.arange(
                np.ceil(np.log2(min_fft)), np.ceil(np.log2(max_fft)) + 1, dtype=int
            )
            cost = (
                np.ceil(n_x / (N - len(h) + 1).astype(np.float64))
                * N
                * (np.log2(N) + 1)
            )

            # add a heuristic term to prevent too-long FFT's which are slow
            # (not predicted by mult. cost alone, 4e-5 exp. determined)
            cost += 4e-5 * N * n_x

            n_fft = N[np.argmin(cost)]
        else:
            # Use only a single block
            n_fft = next_fast_len(min_fft)
    logger.debug(f"FFT block length:   {n_fft}")
    if n_fft < min_fft:
        raise ValueError(
            f"n_fft is too short, has to be at least 2 * len(h) - 1 ({min_fft}), got "
            f"{n_fft}"
        )

    # Figure out if we should use CUDA
    h_fft_torch = torch.fft.rfft(h, n=n_fft)

    # Process each row separately
    saving = [None] * nchans
    for chan in range(nchans):
        saving[chan] = _1d_overlap_filter_torch(
            x[chan], len(h), n_edge, phase, h_fft=h_fft_torch, pad=pad, n_fft=n_fft
        )

    return torch.stack(saving)

[qbarthelemy]

My two cents for IIR filtering.

filtfilt of torchaudio simply applies lfilter to the signal, then to the returned signal,
see https://github.com/pytorch/audio/blob/main/src/torchaudio/functional/filtering.py#L701

filtfilt of scipy, before applying lfilter to the signal and to its returned version,
starts by estimating the initial inner state z_i
https://github.com/scipy/scipy/blob/v1.14.1/scipy/signal/_signaltools.py#L4202
z_i is then used by lfilter, allowing to avoid the transient at the beginning of the filtered signal (so, at the beginning of each channel).

This initialization of inner state of filter in lfilter of scipy can explain the differences you observe in numerical values.
You can confirm this, adding a comparison on lfilter without z_i.

For FIR filtering, I will look later.

[bruAristimunha]

All good, @qbarthelemy =) Everything matches now, including IIR and FIR. Thank you so much for the revision!

and @sylvchev, indeed, filtering is a kind of dark magic, millions of details.

[qbarthelemy]

Good job @bruAristimunha !

[tomMoral]

Here is a second pass on the review.
Nice for the test.

Once these comments are addressed, this LGTM.

[bruAristimunha]

Thank you so much for all the revision @qbarthelemy, @sylvchev and @tomMoral.

I feel like I'm learning a lot from you! I loved! Thanks again! I particularly enjoyed making this layer, because it delivered good results and there was no other compatible implementation available.

[sylvchev]

You did a very good job and this is a solid foundation that could be reused in many approaches. Kudos Bruno!