add robust detrending

Author: nbara

examples/example_detrend.py

@@ -15,10 +15,13 @@
 import numpy as np
 from matplotlib.gridspec import GridSpec
 
-from meegkit.detrend import regress
+from meegkit.detrend import regress, detrend
 
 np.random.seed(9)
 
+# Regression
+# =============================================================================
+
 # Simple regression example, no weights
 # -----------------------------------------------------------------------------
 # fit random walk
@@ -72,8 +75,55 @@
 [b, z] = regress(y, x, w)
 
 plt.figure(4)
-plt.plot(y, label='data')
-plt.plot(z, ls=':', label='fit')
+plt.plot(y, label='data', color='C0')
+plt.plot(z, ls=':', label='fit', color='C1')
 plt.title('Channel-wise regression')
 plt.legend()
+
+
+# Detrending
+# =============================================================================
+
+# Basic example with a linear trend
+# -----------------------------------------------------------------------------
+x = np.arange(100)[:, None]
+x = x + np.random.randn(*x.shape)
+y, _, _ = detrend(x, 1)
+
+plt.figure(5)
+plt.plot(x, label='original')
+plt.plot(y, label='detrended')
+plt.legend()
+
+# Detrend biased random walk
+# -----------------------------------------------------------------------------
+x = np.cumsum(np.random.randn(1000, 1) + 0.1)
+y, _, _ = detrend(x, 3)
+
+plt.figure(6)
+plt.plot(x, label='original')
+plt.plot(y, label='detrended')
+plt.legend()
+
+# Detrend with weights
+# -----------------------------------------------------------------------------
+x = np.linspace(0, 100, 1000)[:, None]
+x = x + 3 * np.random.randn(*x.shape)
+
+# introduce some strong artifact on the first 100 samples
+x[:100, :] = 100
+
+# Detrend
+y, _, _ = detrend(x, 3, None)
+
+# Same process but this time downweight artifactual window
+w = np.ones(x.shape)
+w[:100, :] = 0
+yy, _, _ = detrend(x, 3, w)
+
+plt.figure(7)
+plt.plot(x, label='original')
+plt.plot(y, label='detrended - no weights')
+plt.plot(yy, label='detrended - weights')
+plt.legend()
 plt.show()

meegkit/detrend.py

@@ -1,10 +1,126 @@
 """Robust detrending."""
 import numpy as np
 
+from scipy.linalg import pinv, lstsq, solve
+
 from .utils import demean, pca, unfold
 from .utils.matrix import _check_weights
 
 
+def detrend(x, order, w=None, basis='polynomials', threshold=3, n_iter=4):
+    """Robustly remove trend.
+
+    The data are fit to the basis using weighted least squares. The weight is
+    updated by setting samples for which the residual is greater than 'thresh'
+    times its std to zero, and the fit is repeated at most 'niter'-1 times.
+
+    The choice of order (and basis) determines what complexity of the trend
+    that can be removed.  It may be useful to first detrend with a low order
+    to avoid fitting outliers, and then increase the order.
+
+    Parameters
+    ----------
+    x : array, shape=(n_times, n_channels[, n_trials])
+        Raw data
+    order : int
+        Order of polynomial or number of sin/cosine pairs
+    w: weights
+    basis: {'polynomials', 'sinusoids'} | ndarray
+        Basis for regression.
+    threshold : int
+        Threshold for outliers, in number of standard deviations (default=3).
+    niter : int
+        Number of iterations (default=5).
+
+    Returns
+    -------
+    y: detrended data
+    w: updated weights
+    r: basis matrix used
+
+    Examples
+    --------
+    Fit linear trend, ignoring samples > 3*sd from it, and remove:
+    >> y = detrend(x, 1)
+
+    Fit/remove polynomial order=5 with initial weighting w, threshold = 4*sd:
+    >> y = detrend(x, 5, w, [],4 )
+
+    Fit/remove linear then 3rd order polynomial:
+    >> [y, w]= detrend(x, 1)
+    >> [yy, ww] = detrend(y, 3)
+
+    """
+    if threshold == 0:
+        raise ValueError('thresh=0 is not what you want...')
+
+    # check/fix sizes
+    dims = x.shape
+    w = _check_weights(w, x)
+    x = unfold(x)
+    w = unfold(w)
+    n_times, n_chans = x.shape
+
+    # regressors
+    if isinstance(basis, np.ndarray):
+        r = basis
+    else:
+        lin = np.linspace(-1, 1, n_times)
+        if basis == 'polynomials' or basis is None:
+            r = np.zeros((n_times, order))
+            for i, o in enumerate(range(1, order + 1)):
+                r[:, i] = lin ** o
+        elif basis == 'sinusoids':
+            r = np.zeros((n_times, order * 2))
+            for i, o in enumerate(range(1, order + 1)):
+                r[:, 2 * i] = np.sin[2 * np.pi * o * lin / 2]
+                r[:, 2 * i + 1] = np.cos[2 * np.pi * o * lin / 2]
+        else:
+            raise ValueError('!')
+
+    # iteratively remove trends
+    # the tricky bit is to ensure that weighted means are removed before
+    # calculating the regression (see regress()).
+    for iIter in range(n_iter):
+        # weighted regression on basis
+        _, y = regress(x, r, w)
+
+        # find outliers
+        d = x - y
+        if w.any():
+            d = d * w
+        ww = np.ones_like(x)
+        ww[(abs(d) > threshold * np.std(d))] = 0
+
+        # update weights
+        if not w.any():
+            w = ww
+        else:
+            w = np.amin((w, ww), axis=0)
+        del ww
+
+    y = x - y
+    y = np.reshape(y, dims)
+    w = np.reshape(w, dims)
+
+    # if show: # don't return, just plot
+    #     figure(1)
+    #     subplot 411
+    #     plot(x)
+    #     title('raw')
+    #     subplot 412
+    #     plot(y)
+    #     title('detrended')
+    #     subplot 413
+    #     plot(x-y)
+    #     title('trend')
+    #     subplot 414
+    #     nt_imagescc(w')
+    #     title('weight')
+
+    return y, w, r
+
+
 def regress(y, x, w=None, threshold=1e-7, return_mean=False):
     """Weighted regression.
 
@@ -14,7 +130,7 @@ def regress(y, x, w=None, threshold=1e-7, return_mean=False):
         Data.
     x : array, shape=(n_times, n_chans)
         Regressor.
-    w :
+    w : array, shape=(n_times, n_chans)
         Weight to apply to `y`. `w` is either a matrix of same size as `y`, or
         a column vector to be applied to each column of `y`.
     threshold : float
@@ -49,7 +165,7 @@ def regress(y, x, w=None, threshold=1e-7, return_mean=False):
         # PCA
         V, _ = pca(xx.T.dot(xx), thresh=threshold)
         xxx = xx.dot(V)
-        b = yy.T.dot(xxx) / xxx.T.dot(xxx)
+        b = mrdivide(yy.T.dot(xxx), xxx.T.dot(xxx))
         b = b.T
         z = np.dot(demean(x, w).dot(V), b)
         z = z + mn
@@ -67,7 +183,7 @@ def regress(y, x, w=None, threshold=1e-7, return_mean=False):
                 xx = demean(x, w) * w
                 V, _ = pca(xx.T.dot(xx), thresh=threshold)
                 xxx = xx.dot(V)
-                b = yy.T.dot(xxx) / xxx.T.dot(xxx)
+                b = mrdivide(yy.T.dot(xxx), xxx.T.dot(xxx))
 
             z = demean(x, w).dot(V).dot(b.T)
             z = z + mn
@@ -89,11 +205,51 @@ def regress(y, x, w=None, threshold=1e-7, return_mean=False):
                     xx = x * wc
                     V, _ = pca(xx.T.dot(xx), thresh=threshold)
                     xx = xx.dot(V)
-                    c = yy.T.dot(xx) / xx.T.dot(xx)
+                    c = mrdivide(yy.T.dot(xx), xx.T.dot(xx))
 
                 z[:, i] = x.dot(V.dot(c.T)).flatten()
                 z[:, i] += mn[:, i]
                 b[i] = c
             b = b[:, :V.shape[1]]
 
     return b, z
+
+
+def mrdivide(A, B):
+    r"""Matrix right-division (A/B).
+
+    Solves the linear system XB = A for X. We can write equivalently:
+
+    1) XB = A
+    2) (XB).T = A.T
+    3) B.T X.T = A.T
+
+    Therefore A/B amounts to solving B.T X.T = A.T for X.T:
+
+    >> mldivide(B.T, A.T).T
+
+    References
+    ----------
+    .. [1] https://docs.scipy.org/doc/numpy/user/numpy-for-matlab-users.html
+
+    """
+    return mldivide(B.T, A.T).T
+
+
+def mldivide(A, B):
+    r"""Matrix left-division (A\B).
+
+    Solves the AX = B for X. In other words, X minimizes norm(A*X - B), the
+    length of the vector AX - B:
+    - linalg.solve(A, B) if A is square
+    - linalg.lstsq(A, B) otherwise
+
+    References
+    ----------
+    .. [1] https://docs.scipy.org/doc/numpy/user/numpy-for-matlab-users.html
+
+    """
+    if A.shape[0] == A.shape[1]:
+        return solve(A, B)
+    else:
+        return lstsq(A, B)

meegkit/utils/matrix.py

@@ -435,6 +435,8 @@ def fold(X, epoch_size):
 def unfold(X):
     """Unfold 3D X into 2D (concatenate trials)."""
     n_samples, n_chans, n_trials = theshapeof(X)
+    if X.size == 0:
+        return X
 
     if X.shape == (n_samples,):
         X = X[:, None]

tests/test_detrend.py

@@ -1,7 +1,9 @@
 """Test robust detrending."""
 import numpy as np
 
-from meegkit.detrend import regress
+from numpy.testing import assert_almost_equal
+
+from meegkit.detrend import regress, detrend
 
 
 def test_regress():
@@ -40,6 +42,37 @@ def test_regress():
     assert b.shape == (2, 1)
 
 
+def test_detrend():
+    """Test detrending."""
+    # basic
+    # x = np.arange(100)[:, None]
+    # x = x + np.random.randn(*x.shape)
+    # y, _, _ = detrend(x, 1)
+
+    # assert y.shape == x.shape
+
+    # detrend biased random walk
+    x = np.cumsum(np.random.randn(1000, 1) + 0.1)
+    y, _, _ = detrend(x, 3)
+
+    assert y.shape == x.shape
+
+    # weights
+    trend = np.linspace(0, 100, 1000)[:, None]
+    data = 3 * np.random.randn(*trend.shape)
+    x = trend + data
+    x[:100, :] = 100
+    w = np.ones(x.shape)
+    w[:100, :] = 0
+    y, _, _ = detrend(x, 3, None)
+    yy, _, _ = detrend(x, 3, w)
+
+    assert y.shape == x.shape
+    assert yy.shape == x.shape
+
+    assert_almost_equal(yy[100:], data[100:], decimal=.3)
+
 if __name__ == '__main__':
-    import pytest
-    pytest.main([__file__])
+    # import pytest
+    # pytest.main([__file__])
+    test_detrend()