[ENH] Outputs maps in RESS

examples/example_ress.py

@@ -12,9 +12,8 @@
 import matplotlib.pyplot as plt
 import numpy as np
 import scipy.signal as ss
-
 from meegkit import ress
-from meegkit.utils import unfold, rms, fold, snr_spectrum
+from meegkit.utils import fold, matmul3d, rms, snr_spectrum, unfold
 
 # import config
 
@@ -26,11 +25,11 @@
 # Create synthetic data containing a single oscillatory component at 12 hz.
 n_times = 1000
 n_chans = 10
-n_trials = 50
+n_trials = 30
 target = 12
 sfreq = 250
 noise_dim = 8
-SNR = 1.
+SNR = .2
 t0 = 100
 
 # source
@@ -66,13 +65,13 @@
 # -----------------------------------------------------------------------------
 
 # Apply RESS
-out = ress.RESS(data, sfreq=sfreq, peak_freq=target)
+out, maps = ress.RESS(data, sfreq=sfreq, peak_freq=target, return_maps=True)
 
 # Compute PSD
 nfft = 250
 df = sfreq / nfft  # frequency resolution
-bins, psd = ss.welch(out, sfreq, window="hamming", nperseg=nfft,
-                        noverlap=125, axis=0)
+bins, psd = ss.welch(out.squeeze(1), sfreq, window="hamming", nperseg=nfft,
+                     noverlap=125, axis=0)
 psd = psd.mean(axis=1, keepdims=True)  # average over trials
 snr = snr_spectrum(psd, bins, skipbins=2, n_avg=2)
 
@@ -84,4 +83,21 @@
 ax.set_ylabel('SNR (a.u.)')
 ax.set_xlabel('Frequency (Hz)')
 ax.set_xlim([0, 40])
+
+###############################################################################
+# Project components back into sensor space to see the effects of RESS on the
+# average SSVEP.
+
+proj = matmul3d(out, maps.T)
+f, ax = plt.subplots(n_chans, 2, sharey='col')
+for c in range(n_chans):
+    ax[c, 0].plot(data[:, c].mean(-1), lw=.5)
+    ax[c, 1].plot(proj[:, c].mean(-1), lw=.5)
+    ax[c, 0].set_ylabel(f'ch{c}')
+    if c < n_chans:
+        ax[c, 0].set_xticks([])
+        ax[c, 1].set_xticks([])
+
+ax[0, 0].set_title('Trial average (before)')
+ax[0, 1].set_title('Trial average (after)')
 plt.show()

meegkit/ress.py

@@ -6,7 +6,7 @@
 
 
 def RESS(X, sfreq: int, peak_freq: float, neig_freq: float = 1,
-         neig_width: float = 1, n_keep: int = 1,
+         neig_width: float = 1, n_keep: int = 1, return_maps: bool = False,
          show: bool = False):
     """Rhythmic entrainment source separation [1]_.
 
@@ -23,7 +23,26 @@ def RESS(X, sfreq: int, peak_freq: float, neig_freq: float = 1,
     neig_width : float
         FWHM of the neighboring frequencies (default=1).
     n_keep : int
-        Number of components to keep.
+        Number of components to keep (default=1). -1 keeps all components.
+    return_maps : bool
+        If True, also output maps (mixing matrix).
+
+    Returns
+    -------
+    out : array, shape=(n_samples, n_keep, n_trials)
+        RESS time series.
+    maps : array, shape=(n_channels, n_keep)
+        If return_maps is True, also output mixing matrix.
+
+    Notes
+    -----
+    To project the RESS components back into sensor space, one can proceed as
+    follows. First apply RESS:
+    >> out, maps = ress.RESS(data, sfreq, peak_freq, return_maps=True)
+
+    Then multiply each trial by the mixing matrix:
+    >> from meegkit.utils import matmul3d
+    >> proj = matmul3d(out, maps.T)
 
     References
     ----------
@@ -35,6 +54,9 @@ def RESS(X, sfreq: int, peak_freq: float, neig_freq: float = 1,
     n_samples, n_chans, n_trials = theshapeof(X)
     X = demean(X)
 
+    if n_keep == -1:
+        n_keep = n_chans
+
     # Covariance of signal and covariance of noise
     c01, _ = tscov(gaussfilt(X, sfreq, peak_freq + neig_freq,
                              fwhm=neig_width, n_harm=1))
@@ -58,25 +80,21 @@ def RESS(X, sfreq: int, peak_freq: float, neig_freq: float = 1,
     #     d = d[idx]
     #     V = V[:, idx]
 
-    # Keep a fixed number of components
-    # max_comps = np.min((n_keep, V.shape[1]))
-    # V = V[:, np.arange(max_comps)]
-    # d = d[np.arange(max_comps)]
-
     # Normalize components
     V /= np.sqrt(np.sum(V, axis=0) ** 2)
 
-    # extract components and force sign
+    # extract components
     maps = mrdivide(c1 @ V, V.T @ c1 @ V)
-    idx = np.argmax(np.abs(maps[:, 0]))  # find biggest component
-    maps = maps * np.sign(maps[idx, 0])  # force to positive sign
+    maps = maps[:, :n_keep]
+    # idx = np.argmax(np.abs(maps[:, 0]))  # find biggest component
+    # maps = maps * np.sign(maps[idx, 0])  # force to positive sign
 
     # reconstruct RESS component time series
     out = np.zeros((n_samples, n_keep, n_trials))
     for t in range(n_trials):
         out[..., t] = X[:, :, t] @ V[:, np.arange(n_keep)]
 
-    if n_keep == 1:
-        out = out.squeeze(1)
-
-    return out
+    if return_maps:
+        return out, maps
+    else:
+        return out

meegkit/utils/__init__.py

@@ -6,8 +6,9 @@
                           nonlinear_eigenspace, pca, regcov, tscov, tsxcov)
 from .denoise import (demean, find_outlier_samples, find_outlier_trials,
                       mean_over_trials, wpwr)
-from .matrix import (fold, multishift, multismooth, normcol, relshift, shift,
-                     shiftnd, theshapeof, unfold, unsqueeze, widen_mask)
+from .matrix import (fold, matmul3d, multishift, multismooth, normcol,
+                     relshift, shift, shiftnd, theshapeof, unfold, unsqueeze,
+                     widen_mask)
 from .sig import (gaussfilt, hilbert_envelope, slope_sum, smooth,
                   spectral_envelope, teager_kaiser)
 from .stats import (bootstrap_ci, bootstrap_snr, cronbach, rms, robust_mean,

meegkit/utils/matrix.py

@@ -606,6 +606,27 @@ def normcol(X, weights=None, return_norm=False):
         return X_norm
 
 
+def matmul3d(X, mixin):
+    """Apply mixing matrix to each trial of X.
+
+    This is a simple wrapper or np.einsum.
+
+    Parameters
+    ----------
+    X : array, shape=(n_samples, n_chans, n_trials)
+    mixing : array, shape=(n_chans, n_components)
+
+    Returns
+    -------
+    proj : array, shape=(n_samples, n_components, n_trials)
+        Projection.
+
+    """
+    assert X.ndim == 3, 'data must be of shape (n_samples, n_chans, n_trials)'
+    assert mixin.ndim == 2, 'mixing matrix must be 2D'
+    return np.einsum('sct,ck->skt', X, mixin)
+
+
 def _check_shifts(shifts, allow_floats=False):
     """Check shifts."""
     types = (int, np.int_)

tests/test_ress.py

@@ -4,11 +4,11 @@
 import pytest
 import scipy.signal as ss
 from meegkit import ress
-from meegkit.utils import fold, rms, unfold, snr_spectrum
+from meegkit.utils import fold, rms, unfold, snr_spectrum, matmul3d
 
 
 def create_data(n_times, n_chans=10, n_trials=20, freq=12, sfreq=250,
-                noise_dim=8, SNR=1, t0=100, show=False):
+                noise_dim=8, SNR=.8, t0=100, show=False):
     """Create synthetic data.
 
     Returns
@@ -50,7 +50,7 @@ def create_data(n_times, n_chans=10, n_trials=20, freq=12, sfreq=250,
 
 
 @pytest.mark.parametrize('target', [12, 15, 20])
-@pytest.mark.parametrize('n_trials', [10, 20])
+@pytest.mark.parametrize('n_trials', [16, 20])
 def test_ress(target, n_trials, show=False):
     """Test RESS."""
     sfreq = 250
@@ -60,7 +60,7 @@ def test_ress(target, n_trials, show=False):
     out = ress.RESS(data, sfreq=sfreq, peak_freq=target)
 
     nfft = 500
-    bins, psd = ss.welch(out, sfreq, window="boxcar", nperseg=nfft,
+    bins, psd = ss.welch(out.squeeze(1), sfreq, window="boxcar", nperseg=nfft,
                          noverlap=0, axis=0, average='mean')
     # psd = np.abs(np.fft.fft(out, nfft, axis=0))
     # psd = psd[0:psd.shape[0] // 2 + 1]
@@ -85,11 +85,33 @@ def test_ress(target, n_trials, show=False):
         ax[1].set_ylabel('PSD')
         ax[1].set_xlabel('Frequency (Hz)')
         ax[1].set_xlim([0, 40])
-        plt.show()
+        # plt.show()
 
     assert snr[bins == target] > 10
     assert (snr[(bins <= target - 2) | (bins >= target + 2)] < 2).all()
 
+    # test multiple components
+    out, maps = ress.RESS(data, sfreq=sfreq, peak_freq=target, n_keep=1,
+                          return_maps=True)
+    _ = ress.RESS(data, sfreq=sfreq, peak_freq=target, n_keep=2)
+    _ = ress.RESS(data, sfreq=sfreq, peak_freq=target, n_keep=-1)
+
+    proj = matmul3d(out, maps.T)
+    assert proj.shape == data.shape
+
+    if show:
+        f, ax = plt.subplots(data.shape[1], 2, sharey='col')
+        for c in range(data.shape[1]):
+            ax[c, 0].plot(data[:, c].mean(-1), lw=.5, label='data')
+            ax[c, 1].plot(proj[:, c].mean(-1), lw=.5, label='projection')
+            if c < data.shape[1]:
+                ax[c, 0].set_xticks([])
+                ax[c, 1].set_xticks([])
+
+        ax[0, 0].set_title('Before')
+        ax[0, 1].set_title('After')
+        plt.legend()
+        plt.show()
 
 if __name__ == '__main__':
     import pytest