Filters

def bandpass_filter(data: np.ndarray,
                    N: int,
                    fcl: float,
                    fch: float,
                    fs: float) -> np.ndarray:
    """
    Exact Python replica of Ton van den Bogert's MATLAB `bandpassfilter`.
    @credit Ton van den Bogert, https://biomch-l.isbweb.org/archive/index.php/t-26625.html

    -------------------------------------------------------------------
    Parameters
    ----------
    data : 1-D ndarray
        Signal to be filtered.
    N : int
        *Total* filter order requested by the user (must be even).
        Internally the Butterworth prototype is designed with order N/2 and then
        applied forwards & backwards (filtfilt) which doubles the effective order.
    fcl, fch : float
        Lower and higher cut-off frequencies [Hz].
    fs : float
        Sampling frequency [Hz].
        Must satisfy fs > 2*fch/beta with beta = (sqrt(2)-1)**(1/N), i.e. the
        normalised cut-offs handed to butter() must stay below 1, otherwise a
        ValueError is raised (see the guard below). MATLAB errors on the same
        inputs; earlier versions of this function clipped instead and returned
        a finite but wrong result.

    Returns
    -------
    fdata : ndarray
        Zero-phase band-pass filtered signal (same length as `data`).

    Raises
    ------
    ValueError
        If N is not an even integer >= 2, if the cut-offs do not satisfy
        0 < fcl < fch, or if `fs` is too small for the requested cut-offs.
    """

    # ------------- 1. argument checks  ----------
    if N < 2 or N % 2:
        raise ValueError("N must be an even integer ≥ 2 (bi-directional filtering).")
    if fs <= 0 or fcl <= 0 or fch <= 0 or fcl >= fch:
        raise ValueError("Cut-off frequencies must satisfy 0 < fcl < fch < fs/2.")
    # ------------------------------------------------------------------------

    # ------------- 2. translate MATLAB design rule --------------------------
    # In Ton's routine, N is halved BEFORE the pre-warp constant is computed
    # (`N = N/2;` precedes the Wn line in the original MATLAB), so the
    # exponent below uses halfN, not the original N. Using the original N
    # here previously shifted both cutoffs (e.g. for N=2: beta=0.8022,
    # cutoffs off by ~20%, instead of the correct beta=0.6436) -- confirmed
    # against real MATLAB output on gait EMG data.
    halfN = N // 2  # order actually given to butter()
    beta = (np.sqrt(2) - 1) ** (1 / (2 * halfN))  # pre-warping constant (uses HALVED order)
    Wn = (2.0 * np.asarray([fcl, fch])) / (fs * beta)  # normalised (0–1)
    # ------------------------------------------------------------------------

    # ------------- 2b. sampling-frequency guard -----------------------------
    # MATLAB errors out when the sampling frequency is too small for the
    # requested cut-off. Its own guard,
    #     fs < 2 * (-1 + sqrt(2))^(-(1/2)/N) * fc
    # is evaluated with the UN-halved N (it precedes `N = N/2;`), which makes
    # it weaker than the filter actually needs: for N=2, fc=6 Hz it passes
    # from fs=14.96 Hz on, while Wn stays below 1 only from fs=18.65 Hz on --
    # in that gap MATLAB passes its own guard and then butter() itself errors.
    # Testing Wn < 1 directly therefore rejects exactly the inputs MATLAB
    # rejects, in one check instead of two.
    # Wn used to be clipped into (1e-6, 0.999) here, which silently returned a
    # finite but WRONG result for an under-sampled input instead of raising.
    if Wn[0] <= 0:
        raise ValueError(
            f"fcl={fcl} Hz underflows to a normalised cut-off of {Wn[0]} at "
            f"fs={fs} Hz; the filter is not representable."
        )
    if Wn[0] >= 1.0:
        raise ValueError(
            f"Your sampling frequency is too small for the specified lower "
            f"cut-off frequency: fs={fs} Hz, fcl={fcl} Hz and N={N} give a "
            f"normalised cut-off of {Wn[0]:.6g} >= 1. Need fs > "
            f"{2.0 * fcl / beta:.6g} Hz."
        )
    if Wn[1] >= 1.0:
        raise ValueError(
            f"Your sampling frequency is too small for the specified higher "
            f"cut-off frequency: fs={fs} Hz, fch={fch} Hz and N={N} give a "
            f"normalised cut-off of {Wn[1]:.6g} >= 1. Need fs > "
            f"{2.0 * fch / beta:.6g} Hz."
        )
    # ------------------------------------------------------------------------

    # ------------- 3. design filter & apply zero-phase ----------------------
    # Classical [b,a] + filtfilt with MATLAB's default padding
    # (padtype='odd', padlen=3*(max(len(a),len(b))-1)) rather than
    # scipy's sosfiltfilt defaults, which use a different pad length and
    # measurably diverge from MATLAB on short signals (e.g. epoch-windowed
    # EMG). Confirmed to match real MATLAB output to ~1e-13 relative error.
    b, a = butter(halfN, Wn, btype='bandpass')
    padlen = 3 * (max(len(a), len(b)) - 1)
    fdata = filtfilt(b, a, data, axis=-1, padtype='odd', padlen=padlen)
    return fdata

def bandpass_filter_exact_corners(data, N, fcl, fch, fs):
    """
    BANDPASS_FILTER_EXACT_CORNERS band-passes each channel with -3 dB points
    landing exactly on fcl and fch

      emg = bandpass_filter_exact_corners(emg, 2, 15.0, 450.0, 2048.0)

      INPUT
        data ... signal over time, numeric vector; or a matrix whose rows are
                  the channels and whose columns are the samples over time,
                  then each row is filtered on its own [xV]
        N    ... filter order, number. Must be EVEN and >= 2 []
        fcl  ... lower cutoff frequency, the realised -3 dB point [Hz]
        fch  ... higher cutoff frequency, the realised -3 dB point [Hz]
        fs   ... sampling frequency [Hz]

      OUTPUT
        out  ... filtered signal, same shape as data [xV]

      RAISES
        ValueError ... if N is not an even integer >= 2, or if the cutoffs do
                        not satisfy 0 < fcl < fch < fs/2

      WHY
        bandpass_filter is an exact replica of Ton van den Bogert's MATLAB
        original, whose beta pre-warp is derived for a SINGLE corner. A
        band-pass Butterworth has no such closed form, so dividing BOTH
        corners by beta < 1 pushes both UP by 1/beta = 1.554: at fs = 2048 Hz
        and N = 2 a nominal 15-450 Hz band actually lands at 35.3-574.9 Hz,
        and 30-450 Hz at 68.8-582.3 Hz. This function instead solves
        numerically for the corners to hand to butter(), same goal, no closed
        form.

      *INFO* ... NOT a MATLAB match, use bandpass_filter for that. This is an
                  addition of the Python implementation
      *INFO* ... the solved upper corner sits ABOVE fch for a two-pass
                  filter, so it can still run into the Nyquist clip even when
                  the requested fch does not. Keep fch well below fs/2
    """
    # Check the arguments
    if N < 2 or N % 2:
        raise ValueError("N must be an even integer >= 2 (bi-directional filtering).")
    if fs <= 0 or fcl <= 0 or fch <= 0 or fcl >= fch or fch >= fs / 2:
        raise ValueError("Cut-off frequencies must satisfy 0 < fcl < fch < fs/2.")

    halfN = N // 2
    target = 1 / np.sqrt(2)

    # Solve for the butter() corners whose TWO-PASS response is -3 dB at fcl/fch
    def residual(corners):
        lo = np.clip(corners[0], 1e-3, fs / 2 - 1e-3)
        hi = np.clip(corners[1], lo + 1e-3, fs / 2 - 1e-3)
        Wn = np.clip((2.0 * np.array([lo, hi])) / fs, 1e-6, 0.999)
        b, a = butter(halfN, Wn, btype='bandpass')
        _, h = freqz(b, a, worN=2 * np.pi * np.array([fcl, fch]) / fs)
        return np.abs(h) ** 2 - target

    lo, hi = fsolve(residual, x0=[fcl, fch], xtol=1e-12)
    Wn = np.clip((2.0 * np.array([lo, hi])) / fs, 1e-6, 0.999)

    # Design and apply zero-phase, MATLAB's padding
    b, a = butter(halfN, Wn, btype='bandpass')
    padlen = 3 * (max(len(a), len(b)) - 1)
    return filtfilt(b, a, data, axis=-1, padtype='odd', padlen=padlen)

def lowpass_filter(data, N, fc, fs):
    """
    LOWPASS_FILTER applies a zero-phase Butterworth lowpass to each channel

      env = lowpass_filter(rectified, 2, 15.0, 2048.0)

      INPUT
        data ... signal over time, numeric vector; or a matrix whose rows are
                  the channels and whose columns are the samples over time,
                  then each row is filtered on its own [xV]
        N    ... filter order, number. Must be EVEN and >= 2: it is halved
                  internally and doubled again by the bi-directional
                  (forwards and backwards) filtering []
        fc   ... cutoff frequency, the realised -3 dB point of the whole
                  two-pass filter [Hz]
        fs   ... sampling frequency [Hz]

      OUTPUT
        out  ... filtered signal, same shape as data [xV]

      RAISES
        ValueError ... if N is not an even integer >= 2, if fs or fc are not
                        positive, or if fs is too small for fc

      *INFO* ... MATLAB's own guard uses the UN-halved order and is therefore
                  weaker than the filter needs; butter() then errors instead.
                  Testing the normalised cutoff against 1 rejects exactly the
                  same inputs in one check
      *INFO* ... padding follows MATLAB (padtype 'odd', padlen
                  3*(max(len(a),len(b))-1)), NOT scipy's defaults. For the
                  edge bias that still leaves, see filters.padding
    """
    # Check the arguments
    if N < 2 or N % 2:
        raise ValueError("N must be an even integer >= 2 (bi-directional filtering).")
    if fs <= 0 or fc <= 0:
        raise ValueError("fs and fc must be positive.")

    # Pre-warp with the HALVED order, N is doubled again by filtfilt
    half_n = N // 2
    beta = (np.sqrt(2) - 1) ** (1 / (2 * half_n))
    Wn = float((2.0 * fc) / (fs * beta))

    # Guard the sampling frequency instead of clipping to a wrong-but-finite result
    if Wn <= 0:
        raise ValueError(
            f"fc={fc} Hz underflows to a normalised cut-off of {Wn} at "
            f"fs={fs} Hz; the filter is not representable."
        )
    if Wn >= 1.0:
        raise ValueError(
            f"Your sampling frequency is too small for the specified cut-off "
            f"frequency: fs={fs} Hz, fc={fc} Hz and N={N} give a normalised "
            f"cut-off of {Wn:.6g} >= 1. Need fs > {2.0 * fc / beta:.6g} Hz."
        )

    # Design and apply zero-phase
    b, a = butter(half_n, Wn, btype='low')
    padlen = 3 * (max(len(a), len(b)) - 1)
    return filtfilt(b, a, data, axis=-1, padtype='odd', padlen=padlen)

def moving_average(data, fs, window_s = None, fc = None, kernel = 'bidirectional'):
    """
    MOVING_AVERAGE smooths each channel with a zero-phase moving window

      y = moving_average(squared, fs, fc = 15.0)
      y = moving_average(squared, fs, window_s = 3.0, kernel = 'rectangular')

      INPUT
        data     ... signal over time, numeric vector; or a matrix whose rows
                      are the channels and whose columns are the samples over
                      time, then each row is smoothed on its own [xV]
        fs       ... sampling frequency [Hz]
        window_s ... equivalent rectangular width of the window, give this OR
                      fc, not both [s]
        fc       ... wanted -3 dB cutoff, give this OR window_s [Hz]
        kernel   ... char, 'bidirectional' (default) or 'rectangular', see
                      the module docstring, TWO KERNELS []

      OUTPUT
        y        ... smoothed signal, same shape as data. Exactly zero-phase,
                      a symmetric input comes back symmetric [xV]

      RAISES
        ValueError ... if neither or both of window_s and fc are given, if
                        either is not positive, or if kernel is unknown

      *INFO* ... the ends are divided by the kernel weight that actually
                  overlapped them, so no taper towards zero is introduced
    """
    x = np.asarray(data, dtype=np.float64)
    h = _kernel(kernel, _resolve_window(fs, window_s, fc, kernel))
    if h.size == 1:
        return x.copy()

    # Convolve the signal and a matching all-ones signal with the same kernel,
    # then divide: at the ends only part of the kernel overlaps and the
    # denominator carries exactly that part
    num = convolve1d(x, h, axis=-1, mode='constant', cval=0.0)
    den = convolve1d(np.ones(x.shape[-1], dtype=np.float64), h, mode='constant', cval=0.0)
    return num / den

def window_length_for_cutoff(fc, fs, kernel = 'bidirectional'):
    """
    WINDOW_LENGTH_FOR_CUTOFF returns the boxcar length whose response is
    -3 dB at fc

      w = window_length_for_cutoff(15.0, 2048.0)

      INPUT
        fc     ... wanted -3 dB cutoff frequency [Hz]
        fs     ... sampling frequency [Hz]
        kernel ... char, 'bidirectional' (default) or 'rectangular' []

      OUTPUT
        w      ... length of ONE boxcar pass, samples []

      RAISES
        ValueError ... if fc is not below the Nyquist frequency, or if no
                        window is long enough to reach fc

      *INFO* ... -3 dB means the response of the WHOLE smoother, both passes
                  included, has fallen to 1/sqrt(2). Give this function and
                  lowpass_filter the same fc and both attenuate a sine at fc
                  equally
      *INFO* ... rule of thumb w ~ 0.319 * fs / fc, so at fs = 2048 Hz and
                  fc = 15 Hz that is 44 samples, a 31.9 ms equivalent window
    """
    if fs <= 0 or fc <= 0:
        raise ValueError("fs and fc must be positive.")
    if fc >= fs / 2:
        raise ValueError(f"fc={fc} Hz must be below the Nyquist frequency fs/2={fs / 2} Hz.")

    passes = _passes(kernel)
    target = 1.0 / np.sqrt(2.0)

    # The response falls monotonically in w until the kernel's first null,
    # which sits at w = fs/fc; bracket below it to stay on that branch
    w_hi = min(fs / fc, 1e6)
    if _response(fc, w_hi, fs, passes) > target:
        raise ValueError(f"No boxcar window reaches -3 dB at fc={fc} Hz for fs={fs} Hz.")

    w = brentq(lambda ww: _response(fc, ww, fs, passes) - target, 1.0 + 1e-9, w_hi, xtol=1e-10)
    return max(1, int(round(w)))

def pad_samples(n_samples, fs, pad_s):
    """
    PAD_SAMPLES converts a padding duration into the sample count both
    reflect_pad and trim_pad will use

      pad = pad_samples(emg.shape[-1], fs, 0.25)

      INPUT
        n_samples ... length of the signal along time []
        fs        ... sampling frequency [Hz]
        pad_s     ... wanted padding duration at each end, 0 disables
                       padding [s]

      OUTPUT
        pad       ... number of samples, clipped to what the signal can
                       supply, so reflect_pad and trim_pad stay consistent []
    """
    if pad_s <= 0:
        return 0
    return int(min(round(pad_s * fs), max(n_samples - 1, 0)))

def reflect_pad(data, pad):
    """
    REFLECT_PAD extends a signal at both ends with an even reflection of itself

      padded = reflect_pad(emg, 512)

      INPUT
        data   ... signal over time, numeric vector; or a matrix whose rows
                    are the channels and whose columns are the samples over
                    time, then each row is padded on its own [xV]
        pad    ... number of samples to add at EACH end. Clipped to what the
                    signal can supply, ignored if <= 0 []

      OUTPUT
        padded ... data with `pad` samples prepended and appended, same
                    number of rows [xV]

      *INFO* ... trim the same number of samples off both ends after
                  filtering, see trim_pad
      *INFO* ... the reflection is EVEN (a mirror), not odd (a point
                  reflection). Odd is what filtfilt does and it does not fix
                  the band-pass edge bias
    """
    x = np.asarray(data, dtype=np.float64)
    if pad <= 0:
        return x

    # Cannot reflect more signal than there is
    pad = min(int(pad), x.shape[-1] - 1)
    if pad <= 0:
        return x

    return np.concatenate([x[..., pad:0:-1], x, x[..., -2:-pad - 2:-1]], axis=-1)

def trim_pad(data, pad):
    """
    TRIM_PAD removes the samples reflect_pad added at both ends

      emg = trim_pad(filtered_padded, 512)

      INPUT
        data ... padded signal over time, vector or channels-by-samples
                  matrix [xV]
        pad  ... number of samples reflect_pad added at EACH end. Must be the
                  value reflect_pad actually used, see pad_samples []

      OUTPUT
        out  ... data with `pad` samples removed from both ends [xV]

      *INFO* ... reflect_pad clips `pad` to the signal length, so ask
                  pad_samples for the number both functions agree on
    """
    x = np.asarray(data, dtype=np.float64)
    if pad <= 0:
        return x
    return x[..., int(pad):x.shape[-1] - int(pad)]