RESS: Add option to return mixing and unmixing matrices

examples/example_ress.py

@@ -65,7 +65,7 @@
 # -----------------------------------------------------------------------------
 
 # Apply RESS
-out, maps = ress.RESS(data, sfreq=sfreq, peak_freq=target, return_maps=True)
+out, maps, _ = ress.RESS(data, sfreq=sfreq, peak_freq=target, return_maps=True)
 
 # Compute PSD
 nfft = 250
@@ -88,7 +88,7 @@
 # Project components back into sensor space to see the effects of RESS on the
 # average SSVEP.
 
-proj = matmul3d(out, maps.T)
+proj = matmul3d(out, maps)
 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)

meegkit/ress.py

@@ -7,7 +7,7 @@
 
 def RESS(X, sfreq: int, peak_freq: float, neig_freq: float = 1,
          peak_width: float = .5, neig_width: float = 1, n_keep: int = 1,
-         return_maps: bool = False, show: bool = False):
+         return_maps: bool = False):
     """Rhythmic entrainment source separation [1]_.
 
     Parameters
@@ -28,24 +28,35 @@ def RESS(X, sfreq: int, peak_freq: float, neig_freq: float = 1,
     n_keep : int
         Number of components to keep (default=1). -1 keeps all components.
     return_maps : bool
-        If True, also output maps (mixing matrix).
+        If True, also output mixing (to_ress) and unmixing matrices
+        (from_ress), used to transform the data into RESS component space and
+        back into sensor space, respectively.
 
     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.
+    from_ress : array, shape=(n_components, n_channels)
+        Unmixing matrix (projects to sensor space).
+    to_ress : array, shape=(n_channels, n_components)
+        Mixing matrix (projects to component space).
 
-    Notes
-    -----
+    Examples
+    --------
     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)
+    follows:
+    >>> # First apply RESS
+    >>> from meegkit.utils import matmul3d  # handles 3D matrix multiplication
+    >>> out, fromRESS, _ = ress.RESS(data, sfreq, peak_freq, return_maps=True)
+    >>> # Then matrix multiply each trial by the unmixing matrix:
+    >>> proj = matmul3d(out, fromRESS)
+
+    To transform a new observation into RESS component space (e.g. in the
+    context of a cross-validation, with separate train/test sets):
+    >>> # Start by applying RESS to the train set:
+    >>> out, _, toRESS = ress.RESS(data, sfreq, peak_freq, return_maps=True)
+    >>> # Then multiply your test data by the toRESS:
+    >>> new_comp = new_data @ toRESS
 
     References
     ----------
@@ -68,36 +79,38 @@ def RESS(X, sfreq: int, peak_freq: float, neig_freq: float = 1,
     c1, _ = tscov(gaussfilt(X, sfreq, peak_freq, fwhm=peak_width, n_harm=1))
 
     # perform generalized eigendecomposition
-    d, V = linalg.eig(c1, (c01 + c02) / 2)
+    d, to_ress = linalg.eig(c1, (c01 + c02) / 2)
     d = d.real
-    V = V.real
+    to_ress = to_ress.real
 
     # Sort eigenvectors by decreasing eigenvalues
     idx = np.argsort(d)[::-1]
     d = d[idx]
-    V = V[:, idx]
+    to_ress = to_ress[:, idx]
 
     # Truncate weak components
     # if thresh is not None:
     #     idx = np.where(d / d.max() > thresh)[0]
     #     d = d[idx]
-    #     V = V[:, idx]
+    #     to_ress = to_ress[:, idx]
+
+    # Normalize components (yields mixing matrix)
+    to_ress /= np.sqrt(np.sum(to_ress, axis=0) ** 2)
+    to_ress = to_ress[:, np.arange(n_keep)]
 
-    # Normalize components
-    V /= np.sqrt(np.sum(V, axis=0) ** 2)
+    # Compute unmixing matrix
+    from_ress = mrdivide(c1 @ to_ress, to_ress.T @ c1 @ to_ress).T
+    from_ress = from_ress[:n_keep, :]
 
-    # extract components
-    maps = mrdivide(c1 @ V, V.T @ c1 @ V)
-    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
+    # idx = np.argmax(np.abs(from_ress[:, 0]))  # find biggest component
+    # from_ress = from_ress * np.sign(from_ress[idx, 0])  # force positive sign
 
-    # reconstruct RESS component time series
+    # Output `n_keep` 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)]
+        out[..., t] = X[:, :, t] @ to_ress
 
     if return_maps:
-        return out, maps
+        return out, from_ress, to_ress
     else:
         return out

tests/test_ress.py

@@ -3,8 +3,9 @@
 import numpy as np
 import pytest
 import scipy.signal as ss
+from scipy.linalg import pinv
 from meegkit import ress
-from meegkit.utils import fold, rms, unfold, snr_spectrum, matmul3d
+from meegkit.utils import fold, matmul3d, rms, snr_spectrum, unfold
 
 
 def create_data(n_times, n_chans=10, n_trials=20, freq=12, sfreq=250,
@@ -50,18 +51,24 @@ 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', [16, 20])
+@pytest.mark.parametrize('n_trials', [16])
 @pytest.mark.parametrize('peak_width', [.5, 1])
-@pytest.mark.parametrize('neig_width', [.5, 1])
-@pytest.mark.parametrize('neig_freq', [.5, 1])
+@pytest.mark.parametrize('neig_width', [1])
+@pytest.mark.parametrize('neig_freq', [1])
 def test_ress(target, n_trials, peak_width, neig_width, neig_freq, show=False):
     """Test RESS."""
     sfreq = 250
-    data, source = create_data(n_times=1000, n_trials=n_trials, freq=target,
-                               sfreq=sfreq, show=False)
-
-    out = ress.RESS(data, sfreq=sfreq, peak_freq=target, neig_freq=neig_freq,
-                    peak_width=peak_width, neig_width=neig_width)
+    n_keep = 1
+    n_chans = 10
+    n_times = 1000
+    data, source = create_data(n_times=n_times, n_trials=n_trials,
+                               n_chans=n_chans, freq=target, sfreq=sfreq,
+                               show=False)
+
+    out = ress.RESS(
+        data, sfreq=sfreq, peak_freq=target, neig_freq=neig_freq,
+        peak_width=peak_width, neig_width=neig_width, n_keep=n_keep
+    )
 
     nfft = 500
     bins, psd = ss.welch(out.squeeze(1), sfreq, window="boxcar",
@@ -96,14 +103,14 @@ def test_ress(target, n_trials, peak_width, neig_width, neig_freq, show=False):
     assert (snr[(bins <= target - 2) | (bins >= target + 2)] < 2).all()
 
     # test multiple components
-    out, maps = ress.RESS(data, sfreq=sfreq, peak_freq=target,
-                          neig_freq=neig_freq, peak_width=peak_width,
-                          neig_width=neig_width, 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)
+    out, fromress, toress = ress.RESS(
+        data, sfreq=sfreq, peak_freq=target, neig_freq=neig_freq,
+        peak_width=peak_width, neig_width=neig_width, n_keep=n_keep,
+        return_maps=True
+    )
 
-    proj = matmul3d(out, maps.T)
-    assert proj.shape == data.shape
+    proj = matmul3d(out, fromress)
+    assert proj.shape == (n_times, n_chans, n_trials)
 
     if show:
         f, ax = plt.subplots(data.shape[1], 2, sharey='col')
@@ -117,9 +124,36 @@ def test_ress(target, n_trials, peak_width, neig_width, neig_freq, show=False):
         ax[0, 0].set_title('Before')
         ax[0, 1].set_title('After')
         plt.legend()
+
+    # 2 comps
+    _ = ress.RESS(
+        data, sfreq=sfreq, peak_freq=target, n_keep=2
+    )
+
+    # All comps
+    out, fromress, toress = ress.RESS(
+        data, sfreq=sfreq, peak_freq=target, n_keep=-1, return_maps=True
+    )
+
+    if show:
+        # Inspect mixing/unmixing matrices
+        combined_data = np.array([toress, fromress, pinv(toress)])
+        _max = np.amax(combined_data)
+
+        f, ax = plt.subplots(3)
+        ax[0].imshow(toress, label='toRESS')
+        ax[0].set_title('toRESS')
+        ax[1].imshow(fromress, label='fromRESS', vmin=-_max, vmax=_max)
+        ax[1].set_title('fromRESS')
+        ax[2].imshow(pinv(toress), vmin=-_max, vmax=_max)
+        ax[2].set_title('toRESS$^{-1}$')
+        plt.tight_layout()
         plt.show()
 
+    print(np.sum(np.abs(pinv(toress) - fromress) >= .1))
+
+
 if __name__ == '__main__':
     import pytest
     pytest.main([__file__])
-    # test_ress(12, 20, show=True)
+    # test_ress(12, 20, 1, 1, 1, show=True)