Filtering

.github/workflows/main.yml

@@ -6,6 +6,11 @@ on:
     branches: [ main ]
   pull_request:
     branches: [ main, dev ]
+    types:
+      - opened
+      - reopened
+      - synchronize
+      - ready_for_review
 
 jobs:
   lint:

pynapple/process/__init__.py

@@ -4,6 +4,7 @@
     compute_eventcorrelogram,
 )
 from .decoding import decode_1d, decode_2d
+from .filtering import compute_filtered_signal
 from .perievent import (
     compute_event_trigger_average,
     compute_perievent,

pynapple/process/filtering.py

@@ -0,0 +1,91 @@
+"""Filtering module."""
+
+from numbers import Number
+
+import numpy as np
+from scipy.signal import butter, sosfiltfilt
+
+from .. import core as nap
+
+
+def compute_filtered_signal(
+    data, freq_band, filter_type="bandpass", order=4, sampling_frequency=None
+):
+    """
+    Apply a Butterworth filter to the provided signal.
+
+    This function performs bandpass filtering on time series data
+    using a Butterworth filter. The filter can be configured to be of
+    type "bandpass", "bandstop", "highpass", or "lowpass".
+
+    Parameters
+    ----------
+    data : Tsd, TsdFrame, or TsdTensor
+        The signal to be filtered.
+    freq_band : tuple of (float, float) or float
+        Cutoff frequency(ies) in Hz.
+        - For "bandpass" and "bandstop" filters, provide a tuple specifying the two cutoff frequencies.
+        - For "lowpass" and "highpass" filters, provide a single float specifying the cutoff frequency.
+    filter_type : {'bandpass', 'bandstop', 'highpass', 'lowpass'}, optional
+        The type of frequency filter to apply. Default is "bandpass".
+    order : int, optional
+        The order of the Butterworth filter. Higher values result in sharper frequency cutoffs.
+        Default is 4.
+    sampling_frequency : float, optional
+        The sampling frequency of the signal in Hz. If not provided, it will be inferred from the time axis of the data.
+
+    Returns
+    -------
+    filtered_data : Tsd, TsdFrame, or TsdTensor
+        The filtered signal, with the same data type as the input.
+
+    Raises
+    ------
+    ValueError
+        If `filter_type` is not one of {"bandpass", "bandstop", "highpass", "lowpass"}.
+        If `freq_band` is not a float for "lowpass" and "highpass" filters.
+        If `freq_band` is not a tuple of two floats for "bandpass" and "bandstop" filters.
+
+    Notes
+    -----
+    The cutoff frequency is defined as the frequency at which the amplitude of the signal
+    is reduced by -3 dB (decibels).
+    """
+    if sampling_frequency is None:
+        sampling_frequency = data.rate
+
+    if filter_type not in ["lowpass", "highpass", "bandpass", "bandstop"]:
+        raise ValueError(
+            f"Unrecognized filter type {filter_type}. "
+            "filter_type must be either 'lowpass', 'highpass', 'bandpass',or 'bandstop'."
+        )
+    elif filter_type in ["lowpass", "highpass"] and not isinstance(freq_band, Number):
+        raise ValueError(
+            "Low/high-pass filter specification requires a single frequency. "
+            f"{freq_band} provided instead!"
+        )
+    elif filter_type in ["bandpass", "bandstop"]:
+        try:
+            if len(freq_band) != 2 or not all(
+                isinstance(fq, Number) for fq in freq_band
+            ):
+                raise ValueError
+        except Exception:
+            raise ValueError(
+                "Band-pass/stop filter specification requires two frequencies. "
+                f"{freq_band} provided instead!"
+            )
+
+    sos = butter(
+        order, freq_band, btype=filter_type, fs=sampling_frequency, output="sos"
+    )
+
+    out = np.zeros_like(data.d)
+    for ep in data.time_support:
+        slc = data.get_slice(start=ep.start[0], end=ep.end[0])
+        out[slc] = sosfiltfilt(sos, data.d[slc], axis=0)
+
+    kwargs = dict(t=data.t, d=out, time_support=data.time_support)
+    if isinstance(data, nap.TsdFrame):
+        kwargs["columns"] = data.columns
+    return data.__class__(**kwargs)

tests/test_filtering.py

@@ -0,0 +1,172 @@
+import pytest
+import pynapple as nap
+import numpy as np
+from scipy import signal
+from contextlib import nullcontext as does_not_raise
+
+
+@pytest.fixture
+def sample_data():
+    # Create a sample Tsd data object
+    t = np.linspace(0, 1, 500)
+    d = np.sin(2 * np.pi * 10 * t) + np.random.normal(0, 0.5, t.shape)
+    time_support = nap.IntervalSet(start=[0], end=[1])
+    return nap.Tsd(t=t, d=d, time_support=time_support)
+
+
+@pytest.mark.parametrize("freq", [10, 100])
+@pytest.mark.parametrize("order", [2, 4, 6])
+@pytest.mark.parametrize("btype", ["lowpass", "highpass"])
+@pytest.mark.parametrize("shape", [(5000,), (5000, 2), (5000, 2, 3)])
+@pytest.mark.parametrize(
+    "ep",
+    [
+        nap.IntervalSet(start=[0], end=[1]),
+        nap.IntervalSet(start=[0, 0.5], end=[0.4, 1]),
+        nap.IntervalSet(start=[0, 0.5, 0.95], end=[0.4, 0.9, 1])
+    ]
+)
+def test_filtering_single_freq_match_sci(freq, order, btype, shape: tuple, ep):
+
+    t = np.linspace(0, 1, shape[0])
+    y = np.squeeze(np.cos(np.pi * 2 * 80 * t).reshape(-1, *[1]*(len(shape) - 1)) + np.random.normal(size=shape))
+
+    if len(shape) == 1:
+        tsd = nap.Tsd(t, y, time_support=ep)
+    elif len(shape) == 2:
+        tsd = nap.TsdFrame(t, y, time_support=ep)
+    else:
+        tsd = nap.TsdTensor(t, y, time_support=ep)
+    out = nap.compute_filtered_signal(tsd, freq_band=freq, filter_type=btype, order=order)
+    sos = signal.butter(order, freq, fs=tsd.rate, btype=btype, output="sos")
+    out_sci = []
+    for iset in ep:
+        out_sci.append(signal.sosfiltfilt(sos, tsd.restrict(iset).d, axis=0))
+    out_sci = np.concatenate(out_sci, axis=0)
+    np.testing.assert_array_equal(out.d, out_sci)
+
+
+@pytest.mark.parametrize("freq", [[10, 30], [100,150]])
+@pytest.mark.parametrize("order", [2, 4, 6])
+@pytest.mark.parametrize("btype", ["bandpass", "bandstop"])
+@pytest.mark.parametrize("shape", [(5000,), (5000, 2), (5000, 2, 3)])
+@pytest.mark.parametrize(
+    "ep",
+    [
+        nap.IntervalSet(start=[0], end=[1]),
+        nap.IntervalSet(start=[0, 0.5], end=[0.4, 1]),
+        nap.IntervalSet(start=[0, 0.5, 0.95], end=[0.4, 0.9, 1])
+    ]
+)
+def test_filtering_freq_band_match_sci(freq, order, btype, shape: tuple, ep):
+
+    t = np.linspace(0, 1, shape[0])
+    y = np.squeeze(np.cos(np.pi * 2 * 80 * t).reshape(-1, *[1]*(len(shape) - 1)) + np.random.normal(size=shape))
+
+    if len(shape) == 1:
+        tsd = nap.Tsd(t, y, time_support=ep)
+    elif len(shape) == 2:
+        tsd = nap.TsdFrame(t, y, time_support=ep)
+    else:
+        tsd = nap.TsdTensor(t, y, time_support=ep)
+    out = nap.compute_filtered_signal(tsd, freq_band=freq, filter_type=btype, order=order)
+    sos = signal.butter(order, freq, fs=tsd.rate, btype=btype, output="sos")
+    out_sci = []
+    for iset in ep:
+        out_sci.append(signal.sosfiltfilt(sos, tsd.restrict(iset).d, axis=0))
+    out_sci = np.concatenate(out_sci, axis=0)
+    np.testing.assert_array_equal(out.d, out_sci)
+
+
+@pytest.mark.parametrize("freq", [10, 100])
+@pytest.mark.parametrize("order", [2, 4, 6])
+@pytest.mark.parametrize("btype", ["lowpass", "highpass"])
+@pytest.mark.parametrize("shape", [(5000,), (5000, 2), (5000, 2, 3)])
+@pytest.mark.parametrize(
+    "ep",
+    [
+        nap.IntervalSet(start=[0], end=[1]),
+        nap.IntervalSet(start=[0, 0.5], end=[0.4, 1]),
+        nap.IntervalSet(start=[0, 0.5, 0.95], end=[0.4, 0.9, 1])
+    ]
+)
+def test_filtering_single_freq_dtype(freq, order, btype, shape: tuple, ep):
+    t = np.linspace(0, 1, shape[0])
+    y = np.squeeze(np.cos(np.pi * 2 * 80 * t).reshape(-1, *[1]*(len(shape) - 1)) + np.random.normal(size=shape))
+
+    if len(shape) == 1:
+        tsd = nap.Tsd(t, y, time_support=ep)
+    elif len(shape) == 2:
+        tsd = nap.TsdFrame(t, y, time_support=ep, columns=np.arange(10, 10 + y.shape[1]))
+    else:
+        tsd = nap.TsdTensor(t, y, time_support=ep)
+    out = nap.compute_filtered_signal(tsd, freq_band=freq, filter_type=btype, order=order)
+    assert isinstance(out, type(tsd))
+    assert np.all(out.t == tsd.t)
+    assert np.all(out.time_support == tsd.time_support)
+    if isinstance(tsd, nap.TsdFrame):
+        assert np.all(tsd.columns == out.columns)
+
+
+@pytest.mark.parametrize("freq", [[10, 30], [100, 150]])
+@pytest.mark.parametrize("order", [2, 4, 6])
+@pytest.mark.parametrize("btype", ["bandpass", "bandstop"])
+@pytest.mark.parametrize("shape", [(5000,), (5000, 2), (5000, 2, 3)])
+@pytest.mark.parametrize(
+    "ep",
+    [
+        nap.IntervalSet(start=[0], end=[1]),
+        nap.IntervalSet(start=[0, 0.5], end=[0.4, 1]),
+        nap.IntervalSet(start=[0, 0.5, 0.95], end=[0.4, 0.9, 1])
+    ]
+)
+def test_filtering_freq_band_dtype(freq, order, btype, shape: tuple, ep):
+    t = np.linspace(0, 1, shape[0])
+    y = np.squeeze(np.cos(np.pi * 2 * 80 * t).reshape(-1, *[1]*(len(shape) - 1)) + np.random.normal(size=shape))
+
+    if len(shape) == 1:
+        tsd = nap.Tsd(t, y, time_support=ep)
+    elif len(shape) == 2:
+        tsd = nap.TsdFrame(t, y, time_support=ep, columns=np.arange(10, 10 + y.shape[1]))
+    else:
+        tsd = nap.TsdTensor(t, y, time_support=ep)
+    out = nap.compute_filtered_signal(tsd, freq_band=freq, filter_type=btype, order=order)
+    assert isinstance(out, type(tsd))
+    assert np.all(out.t == tsd.t)
+    assert np.all(out.time_support == tsd.time_support)
+    if isinstance(tsd, nap.TsdFrame):
+        assert np.all(tsd.columns == out.columns)
+
+
+@pytest.mark.parametrize("freq_band, filter_type, order, expected_exception", [
+    ((5, 15), "bandpass", 4, does_not_raise()),
+    ((5, 15), "bandstop", 4, does_not_raise()),
+    (10, "highpass", 4, does_not_raise()),
+    (10, "lowpass", 4, does_not_raise()),
+    ((5, 15), "invalid_filter", 4, pytest.raises(ValueError, match="Unrecognized filter type")),
+    (10, "bandpass", 4, pytest.raises(ValueError,  match="Band-pass/stop filter specification requires two frequencies")),
+    ((5, 15), "highpass", 4, pytest.raises(ValueError, match="Low/high-pass filter specification requires a single frequency")),
+    (None, "bandpass", 4, pytest.raises(ValueError,  match="Band-pass/stop filter specification requires two frequencies")),
+    ((None, 1), "highpass", 4, pytest.raises(ValueError, match="Low/high-pass filter specification requires a single frequency"))
+])
+def test_compute_filtered_signal(sample_data, freq_band, filter_type, order, expected_exception):
+    with expected_exception:
+        filtered_data = nap.filtering.compute_filtered_signal(sample_data, freq_band, filter_type, order)
+        if not expected_exception:
+            assert isinstance(filtered_data, type(sample_data))
+            assert filtered_data.d.shape == sample_data.d.shape
+
+
+# Test with edge-case frequencies close to Nyquist frequency
+@pytest.mark.parametrize("nyquist_fraction", [0.99, 0.999])
+@pytest.mark.parametrize("order", [2, 4])
+def test_filtering_nyquist_edge_case(nyquist_fraction, order, sample_data):
+    nyquist_freq = 0.5 * sample_data.rate
+    freq = nyquist_freq * nyquist_fraction
+
+    out = nap.filtering.compute_filtered_signal(
+        sample_data, freq_band=freq, filter_type="lowpass", order=order
+    )
+    assert isinstance(out, type(sample_data))
+    np.testing.assert_allclose(out.t, sample_data.t)
+    np.testing.assert_allclose(out.time_support, sample_data.time_support)