Add trial-masked detrending (nbara#41)

* add create_maksed_weight

* Formatting docstring

* format public function

* add test + correct dimension of weights

* docstring

Co-authored-by: nicolas barascud <10333715+nbara@users.noreply.github.com>

Author: romquentin

meegkit/detrend.py

@@ -297,3 +297,45 @@ def _plot_detrend(x, y, w):
     ax2.set_ylabel('ch. weights')
     ax2.set_xlabel('samples')
     plt.show()
+
+
+def create_masked_weight(x, events, tmin, tmax, sfreq):
+    """Output a weight matrix for trial-masked detrending.
+
+    Creates a  (n_times, n_channels) weight matrix with masked
+    periods (value of zero) in order to mask the trials of interest during
+    detrending [1]_.
+
+    Parameters
+    ----------
+    x : ndarray, shape=(n_times, n_channels)
+        Raw data matrix.
+    events : ndarray, shape=(n_events)
+        Time samples of the events.
+    tmin : float
+        Start time before event (in seconds).
+    tmax : float
+        End time after event (in seconds).
+    sfreq : float
+        The sampling frequency of the data.
+
+    Returns
+    -------
+    weights : ndarray, shape=(n_times, n_channels)
+        Weight for each channel and each time sample (zero is masked).
+
+    References
+    ----------
+    .. [1] van Driel, J., Olivers, C. N., & Fahrenfort, J. J. (2021). High-pass
+       filtering artifacts in multivariate classification of neural time series
+       data. Journal of Neuroscience Methods, 352, 109080.
+
+    """
+    if x.ndim != 2:
+        raise ValueError('The shape of x must be (n_times, n_channels)')
+
+    weights = np.ones(x.shape)
+    for e in events:
+        weights[int(e + tmin * sfreq): int(e + tmax * sfreq + 1), :] = 0
+
+    return weights

tests/test_detrend.py

@@ -1,7 +1,7 @@
 """Test robust detrending."""
 import numpy as np
 
-from meegkit.detrend import regress, detrend, reduce_ringing
+from meegkit.detrend import regress, detrend, reduce_ringing, create_masked_weight
 
 from scipy.signal import butter, lfilter
 
@@ -89,6 +89,16 @@ def test_detrend(show=False):
     x += 2 * np.sin(2 * np.pi * np.arange(1000) / 200)[:, None]
     y, _, _ = detrend(x, 5, basis='sinusoids', show=True)
 
+    # trial-masked detrending
+    trend = np.linspace(0, 100, 1000)[:, None]
+    data = 3 * np.random.randn(*trend.shape)
+    data[:100, :] = 100
+    x = trend + data
+    events = np.arange(30, 970, 40)
+    tmin, tmax, sfreq = -0.2, 0.3, 20
+    w = create_masked_weight(x, events, tmin, tmax, sfreq)
+    y, _, _ = detrend(x, 1, w, basis='polynomials', show=show)
+
 
 def test_ringing():
     """Test reduce_ringing function."""