add weighted regression

Author: nbara

README.md

@@ -1,6 +1,7 @@
 # MEEGkit
 
 [![Build Status](https://travis-ci.org/nbara/python-meegkit.svg?branch=master)](https://travis-ci.org/nbara/python-meegkit)
+[![codecov](https://codecov.io/gh/nbara/python-meegkit/branch/master/graph/badge.svg)](https://codecov.io/gh/nbara/python-meegkit)
 [![Binder](https://mybinder.org/badge.svg)](https://mybinder.org/v2/gh/nbara/python-meegkit/master)
 
 Denoising tools for M/EEG processing in Python.

examples/example_detrend.py

@@ -0,0 +1,79 @@
+"""Robust detrending examples.
+
+Some toy examples to showcase usage for ``meegkit.detrend`` module.
+
+Robust referencing is adapted from [1]_.
+
+References
+----------
+.. [1] de Cheveigné, A., & Arzounian, D. (2018). Robust detrending,
+   rereferencing, outlier detection, and inpainting for multichannel data.
+   NeuroImage, 172, 903-912.
+
+"""
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib.gridspec import GridSpec
+
+from meegkit.detrend import regress
+
+np.random.seed(9)
+
+# Simple regression example, no weights
+# -----------------------------------------------------------------------------
+# fit random walk
+y = np.cumsum(np.random.randn(1000, 1), axis=0)
+x = np.arange(1000)[:, None]
+x = np.hstack([x, x ** 2, x ** 3])
+[b, z] = regress(y, x)
+
+plt.figure(1)
+plt.plot(y, label='data')
+plt.plot(z, label='fit')
+plt.title('No weights')
+plt.legend()
+
+# Simple regression example, with weights
+# -----------------------------------------------------------------------------
+y = np.cumsum(np.random.randn(1000, 1), axis=0)
+w = np.random.rand(*y.shape)
+[b, z] = regress(y, x, w)
+
+plt.figure(2)
+plt.plot(y, label='data')
+plt.plot(z, label='fit')
+plt.title('Weighted regression')
+plt.legend()
+
+# Downweight 1st half of the data
+# -----------------------------------------------------------------------------
+y = np.cumsum(np.random.randn(1000, 1), axis=0) + 1000
+w = np.ones(y.shape[0])
+w[:500] = 0
+[b, z] = regress(y, x, w)
+
+f = plt.figure(3, constrained_layout=True)
+gs = GridSpec(3, 1, figure=f)
+ax1 = f.add_subplot(gs[:2, 0])
+ax1.plot(y, label='data')
+ax1.plot(z, label='fit')
+ax1.set_xticklabels('')
+ax1.set_title('Split-wise regression')
+ax1.legend()
+ax2 = f.add_subplot(gs[2, 0])
+l, = ax2.plot(np.arange(1000), np.zeros(1000))
+ax2.stackplot(np.arange(1000), w, labels=['weights'], color=l.get_color())
+ax2.legend(loc=2)
+
+# Multichannel regression
+# -----------------------------------------------------------------------------
+y = np.cumsum(np.random.randn(1000, 2), axis=0)
+w = np.ones(y.shape[0])
+[b, z] = regress(y, x, w)
+
+plt.figure(4)
+plt.plot(y, label='data')
+plt.plot(z, ls=':', label='fit')
+plt.title('Channel-wise regression')
+plt.legend()
+plt.show()

meegkit/detrend.py

@@ -0,0 +1,99 @@
+"""Robust detrending."""
+import numpy as np
+
+from .utils import demean, pca, unfold
+from .utils.matrix import _check_weights
+
+
+def regress(y, x, w=None, threshold=1e-7, return_mean=False):
+    """Weighted regression.
+
+    Parameters
+    ----------
+    y : array, shape=(n_times, n_chans)
+        Data.
+    x : array, shape=(n_times, n_chans)
+        Regressor.
+    w :
+        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
+        PCA threshold (default=1e-7).
+    return_mean : bool
+        If True, also return the signal mean prior to regression.
+
+    Returns
+    -------
+    b : array, shape=(n_chans, n_chans)
+        Regression matrix (apply to x to approximate y).
+    z : array, shape=(n_times, n_chans)
+        Regression (x @ b).
+
+    """
+    # check/fix sizes
+    w = _check_weights(w, y)
+    n_times = y.shape[0]
+    n_chans = y.shape[1]
+    x = unfold(x)
+    y = unfold(y)
+    if x.shape[0] != y.shape[0]:
+        raise ValueError('x and y have incompatible shapes!')
+
+    # save weighted mean
+    mn = y - demean(y, w)
+
+    if not w.any():  # simple regression
+        xx = demean(x)
+        yy = demean(y)
+
+        # PCA
+        V, _ = pca(xx.T.dot(xx), thresh=threshold)
+        xxx = xx.dot(V)
+        b = yy.T.dot(xxx) / xxx.T.dot(xxx)
+        b = b.T
+        z = np.dot(demean(x, w).dot(V), b)
+        z = z + mn
+
+    else:  # weighted regression
+        if w.shape[0] != n_times:
+            raise ValueError('!')
+
+        if w.shape[1] == 1:  # same weight for all channels
+            if sum(w.flatten()) == 0:
+                print('weights all zero')
+                b = 0
+            else:
+                yy = demean(y, w) * w
+                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)
+
+            z = demean(x, w).dot(V).dot(b.T)
+            z = z + mn
+
+        else:  # each channel has own weight
+            if w.shape[1] != y.shape[1]:
+                raise ValueError('!')
+            z = np.zeros(y.shape)
+            b = np.zeros((n_chans, n_chans))
+            for i in range(n_chans):
+                if sum(w[:, i]) == 0:
+                    print('weights all zero for channel {}'.format(i))
+                    c = np.zeros(y.shape[1], 1)
+                else:
+                    wc = w[:, i][:, None]  # channel-specific weight
+                    yy = demean(y[:, i], wc) * wc
+                    # remove channel-specific-weighted mean from regressor
+                    x = demean(x, wc)
+                    xx = x * wc
+                    V, _ = pca(xx.T.dot(xx), thresh=threshold)
+                    xx = xx.dot(V)
+                    c = 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

meegkit/utils/matrix.py

@@ -436,6 +436,9 @@ def unfold(X):
     """Unfold 3D X into 2D (concatenate trials)."""
     n_samples, n_chans, n_trials = theshapeof(X)
 
+    if X.shape == (n_samples,):
+        X = X[:, None]
+
     if n_trials > 1:
         return np.reshape(
             np.transpose(X, (0, 2, 1)),
@@ -603,7 +606,8 @@ def _check_weights(weights, X):
             warnings.warn('weights should be a list or a numpy array.')
         weights = np.array([])
 
-    if len(weights) > 0:
+    weights = np.asanyarray(weights)
+    if weights.size > 0:
         dtype = np.complex128 if np.any(np.iscomplex(weights)) else np.float64
         weights = np.asanyarray(weights, dtype=dtype)
         if weights.ndim > 3:
@@ -618,8 +622,9 @@ def _check_weights(weights, X):
         if X.ndim == 3 and weights.ndim == 1:
             weights = weights[:, np.newaxis, np.newaxis]
 
-        if weights.shape[1] > 1:
-            raise ValueError("Weights array should have a single column.")
+        if weights.ndim > 1:
+            if weights.shape[1] > 1 and weights.shape[1] != X.shape[1]:
+                raise ValueError("Weights array should have a single column.")
 
     return weights
 

tests/test_detrend.py

@@ -0,0 +1,45 @@
+"""Test robust detrending."""
+import numpy as np
+
+from meegkit.detrend import regress
+
+
+def test_regress():
+    """Test regression."""
+    # Simple regression example, no weights
+    # fit random walk
+    y = np.cumsum(np.random.randn(1000, 1), axis=0)
+    x = np.arange(1000)[:, None]
+    x = np.hstack([x, x ** 2, x ** 3])
+    [b, z] = regress(y, x)
+
+    # Simple regression example, with weights
+    y = np.cumsum(np.random.randn(1000, 1), axis=0)
+    w = np.random.rand(*y.shape)
+    [b, z] = regress(y, x, w)
+
+    # Downweight 1st half of the data
+    y = np.cumsum(np.random.randn(1000, 1), axis=0) + 1000
+    w = np.ones(y.shape[0])
+    w[:500] = 0
+    [b, z] = regress(y, x, w)
+
+    # # Multichannel regression
+    y = np.cumsum(np.random.randn(1000, 2), axis=0)
+    w = np.ones(y.shape[0])
+    [b, z] = regress(y, x, w)
+    assert z.shape == (1000, 2)
+    assert b.shape == (2, 1)
+
+    # Multichannel regression
+    y = np.cumsum(np.random.randn(1000, 2), axis=0)
+    w = np.ones(y.shape)
+    w[:, 1] == .8
+    [b, z] = regress(y, x, w)
+    assert z.shape == (1000, 2)
+    assert b.shape == (2, 1)
+
+
+if __name__ == '__main__':
+    import pytest
+    pytest.main([__file__])