# Author: Denis A. Engemann # Victoria Peterson # License: BSD-3-Clause import numpy as np import pytest from numpy.testing import (assert_array_almost_equal, assert_array_equal) from mne import io from mne.time_frequency import psd_array_welch from mne.decoding.ssd import SSD from mne.utils import requires_sklearn from mne.filter import filter_data from mne import create_info from mne.decoding import CSP freqs_sig = 9, 12 freqs_noise = 8, 13 def simulate_data(freqs_sig=[9, 12], n_trials=100, n_channels=20, n_samples=500, samples_per_second=250, n_components=5, SNR=0.05, random_state=42): """Simulate data according to an instantaneous mixin model. Data are simulated in the statistical source space, where n=n_components sources contain the peak of interest. """ rng = np.random.RandomState(random_state) filt_params_signal = dict(l_freq=freqs_sig[0], h_freq=freqs_sig[1], l_trans_bandwidth=1, h_trans_bandwidth=1, fir_design='firwin') # generate an orthogonal mixin matrix mixing_mat = np.linalg.svd(rng.randn(n_channels, n_channels))[0] # define sources S_s = rng.randn(n_trials * n_samples, n_components) # filter source in the specific freq. band of interest S_s = filter_data(S_s.T, samples_per_second, **filt_params_signal).T S_n = rng.randn(n_trials * n_samples, n_channels - n_components) S = np.hstack((S_s, S_n)) # mix data X_s = np.dot(mixing_mat[:, :n_components], S_s.T).T X_n = np.dot(mixing_mat[:, n_components:], S_n.T).T # add noise X_s = X_s / np.linalg.norm(X_s, 'fro') X_n = X_n / np.linalg.norm(X_n, 'fro') X = SNR * X_s + (1 - SNR) * X_n X = X.T S = S.T return X, mixing_mat, S @pytest.mark.slowtest def test_ssd(): """Test Common Spatial Patterns algorithm on raw data.""" X, A, S = simulate_data() sf = 250 n_channels = X.shape[0] info = create_info(ch_names=n_channels, sfreq=sf, ch_types='eeg') n_components_true = 5 # Init filt_params_signal = dict(l_freq=freqs_sig[0], h_freq=freqs_sig[1], l_trans_bandwidth=1, h_trans_bandwidth=1) filt_params_noise = dict(l_freq=freqs_noise[0], h_freq=freqs_noise[1], l_trans_bandwidth=1, h_trans_bandwidth=1) ssd = SSD(info, filt_params_signal, filt_params_noise) # freq no int freq = 'foo' filt_params_signal = dict(l_freq=freq, h_freq=freqs_sig[1], l_trans_bandwidth=1, h_trans_bandwidth=1) filt_params_noise = dict(l_freq=freqs_noise[0], h_freq=freqs_noise[1], l_trans_bandwidth=1, h_trans_bandwidth=1) with pytest.raises(TypeError, match='must be an instance '): ssd = SSD(info, filt_params_signal, filt_params_noise) # Wrongly specified noise band freq = 2 filt_params_signal = dict(l_freq=freq, h_freq=freqs_sig[1], l_trans_bandwidth=1, h_trans_bandwidth=1) filt_params_noise = dict(l_freq=freqs_noise[0], h_freq=freqs_noise[1], l_trans_bandwidth=1, h_trans_bandwidth=1) with pytest.raises(ValueError, match='Wrongly specified '): ssd = SSD(info, filt_params_signal, filt_params_noise) # filt param no dict filt_params_signal = freqs_sig filt_params_noise = freqs_noise with pytest.raises(ValueError, match='must be defined'): ssd = SSD(info, filt_params_signal, filt_params_noise) # Data type filt_params_signal = dict(l_freq=freqs_sig[0], h_freq=freqs_sig[1], l_trans_bandwidth=1, h_trans_bandwidth=1) filt_params_noise = dict(l_freq=freqs_noise[0], h_freq=freqs_noise[1], l_trans_bandwidth=1, h_trans_bandwidth=1) ssd = SSD(info, filt_params_signal, filt_params_noise) raw = io.RawArray(X, info) pytest.raises(TypeError, ssd.fit, raw) # check non-boolean return_filtered with pytest.raises(ValueError, match='return_filtered'): ssd = SSD(info, filt_params_signal, filt_params_noise, return_filtered=0) # check non-boolean sort_by_spectral_ratio with pytest.raises(ValueError, match='sort_by_spectral_ratio'): ssd = SSD(info, filt_params_signal, filt_params_noise, sort_by_spectral_ratio=0) # More than 1 channel type ch_types = np.reshape([['mag'] * 10, ['eeg'] * 10], n_channels) info_2 = create_info(ch_names=n_channels, sfreq=sf, ch_types=ch_types) with pytest.raises(ValueError, match='At this point SSD'): ssd = SSD(info_2, filt_params_signal, filt_params_noise) # Number of channels info_3 = create_info(ch_names=n_channels + 1, sfreq=sf, ch_types='eeg') ssd = SSD(info_3, filt_params_signal, filt_params_noise) pytest.raises(ValueError, ssd.fit, X) # Fit n_components = 10 ssd = SSD(info, filt_params_signal, filt_params_noise, n_components=n_components) # Call transform before fit pytest.raises(AttributeError, ssd.transform, X) # Check outputs ssd.fit(X) assert (ssd.filters_.shape == (n_channels, n_channels)) assert (ssd.patterns_.shape == (n_channels, n_channels)) # Transform X_ssd = ssd.fit_transform(X) assert (X_ssd.shape[0] == n_components) # back and forward ssd = SSD(info, filt_params_signal, filt_params_noise, n_components=None, sort_by_spectral_ratio=False) ssd.fit(X) X_denoised = ssd.apply(X) assert_array_almost_equal(X_denoised, X) # denoised by low-rank-factorization ssd = SSD(info, filt_params_signal, filt_params_noise, n_components=n_components, sort_by_spectral_ratio=True) ssd.fit(X) X_denoised = ssd.apply(X) assert (np.linalg.matrix_rank(X_denoised) == n_components) # Power ratio ordering ssd = SSD(info, filt_params_signal, filt_params_noise, n_components=None, sort_by_spectral_ratio=False) ssd.fit(X) spec_ratio, sorter_spec = ssd.get_spectral_ratio(ssd.transform(X)) # since we now that the number of true components is 5, the relative # difference should be low for the first 5 components and then increases index_diff = np.argmax(-np.diff(spec_ratio)) assert index_diff == n_components_true - 1 # Check detected peaks # fit ssd n_components = n_components_true filt_params_signal = dict(l_freq=freqs_sig[0], h_freq=freqs_sig[1], l_trans_bandwidth=1, h_trans_bandwidth=1) filt_params_noise = dict(l_freq=freqs_noise[0], h_freq=freqs_noise[1], l_trans_bandwidth=1, h_trans_bandwidth=1) ssd = SSD(info, filt_params_signal, filt_params_noise, n_components=n_components, sort_by_spectral_ratio=False) ssd.fit(X) out = ssd.transform(X) psd_out, _ = psd_array_welch(out[0], sfreq=250, n_fft=250) psd_S, _ = psd_array_welch(S[0], sfreq=250, n_fft=250) corr = np.abs(np.corrcoef((psd_out, psd_S))[0, 1]) assert np.abs(corr) > 0.95 # Check pattern estimation # Since there is no exact ordering of the recovered patterns # a pair-wise greedy search will be done error = list() for ii in range(n_channels): corr = np.abs(np.corrcoef(ssd.patterns_[ii, :].T, A[:, 0])[0, 1]) error.append(1 - corr) min_err = np.min(error) assert min_err < 0.3 # threshold taken from SSD original paper def test_ssd_epoched_data(): """Test Common Spatial Patterns algorithm on epoched data. Compare the outputs when raw data is used. """ X, A, S = simulate_data(n_trials=100, n_channels=20, n_samples=500) sf = 250 n_channels = X.shape[0] info = create_info(ch_names=n_channels, sfreq=sf, ch_types='eeg') n_components_true = 5 # Build epochs as sliding windows over the continuous raw file # Epoch length is 1 second X_e = np.reshape(X, (100, 20, 500)) # Fit filt_params_signal = dict(l_freq=freqs_sig[0], h_freq=freqs_sig[1], l_trans_bandwidth=4, h_trans_bandwidth=4) filt_params_noise = dict(l_freq=freqs_noise[0], h_freq=freqs_noise[1], l_trans_bandwidth=4, h_trans_bandwidth=4) # ssd on epochs ssd_e = SSD(info, filt_params_signal, filt_params_noise) ssd_e.fit(X_e) # ssd on raw ssd = SSD(info, filt_params_signal, filt_params_noise) ssd.fit(X) # Check if the 5 first 5 components are the same for both _, sorter_spec_e = ssd_e.get_spectral_ratio(ssd_e.transform(X_e)) _, sorter_spec = ssd.get_spectral_ratio(ssd.transform(X)) assert_array_equal(sorter_spec_e[:n_components_true], sorter_spec[:n_components_true]) @requires_sklearn def test_ssd_pipeline(): """Test if SSD works in a pipeline.""" from sklearn.pipeline import Pipeline sf = 250 X, A, S = simulate_data(n_trials=100, n_channels=20, n_samples=500) X_e = np.reshape(X, (100, 20, 500)) # define bynary random output y = np.random.randint(2, size=100) info = create_info(ch_names=20, sfreq=sf, ch_types='eeg') filt_params_signal = dict(l_freq=freqs_sig[0], h_freq=freqs_sig[1], l_trans_bandwidth=4, h_trans_bandwidth=4) filt_params_noise = dict(l_freq=freqs_noise[0], h_freq=freqs_noise[1], l_trans_bandwidth=4, h_trans_bandwidth=4) ssd = SSD(info, filt_params_signal, filt_params_noise) csp = CSP() pipe = Pipeline([('SSD', ssd), ('CSP', csp)]) pipe.set_params(SSD__n_components=5) pipe.set_params(CSP__n_components=2) out = pipe.fit_transform(X_e, y) assert (out.shape == (100, 2)) assert (pipe.get_params()['SSD__n_components'] == 5) def test_sorting(): """Test sorting learning during training.""" X, _, _ = simulate_data(n_trials=100, n_channels=20, n_samples=500) # Epoch length is 1 second X = np.reshape(X, (100, 20, 500)) # split data Xtr, Xte = X[:80], X[80:] sf = 250 n_channels = Xtr.shape[1] info = create_info(ch_names=n_channels, sfreq=sf, ch_types='eeg') filt_params_signal = dict(l_freq=freqs_sig[0], h_freq=freqs_sig[1], l_trans_bandwidth=4, h_trans_bandwidth=4) filt_params_noise = dict(l_freq=freqs_noise[0], h_freq=freqs_noise[1], l_trans_bandwidth=4, h_trans_bandwidth=4) # check sort_by_spectral_ratio set to False ssd = SSD(info, filt_params_signal, filt_params_noise, n_components=None, sort_by_spectral_ratio=False) ssd.fit(Xtr) _, sorter_tr = ssd.get_spectral_ratio(ssd.transform(Xtr)) _, sorter_te = ssd.get_spectral_ratio(ssd.transform(Xte)) assert any(sorter_tr != sorter_te) # check sort_by_spectral_ratio set to True ssd = SSD(info, filt_params_signal, filt_params_noise, n_components=None, sort_by_spectral_ratio=True) ssd.fit(Xtr) # check sorters sorter_in = ssd.sorter_spec ssd = SSD(info, filt_params_signal, filt_params_noise, n_components=None, sort_by_spectral_ratio=False) ssd.fit(Xtr) _, sorter_out = ssd.get_spectral_ratio(ssd.transform(Xtr)) assert all(sorter_in == sorter_out) def test_return_filtered(): """Test return filtered option.""" # Check return_filtered # Simulated more noise data and with broader frequency than the desired X, _, _ = simulate_data(SNR=0.9, freqs_sig=[4, 13]) sf = 250 n_channels = X.shape[0] info = create_info(ch_names=n_channels, sfreq=sf, ch_types='eeg') filt_params_signal = dict(l_freq=freqs_sig[0], h_freq=freqs_sig[1], l_trans_bandwidth=1, h_trans_bandwidth=1) filt_params_noise = dict(l_freq=freqs_noise[0], h_freq=freqs_noise[1], l_trans_bandwidth=1, h_trans_bandwidth=1) # return filtered to true ssd = SSD(info, filt_params_signal, filt_params_noise, sort_by_spectral_ratio=False, return_filtered=True) ssd.fit(X) out = ssd.transform(X) psd_out, freqs = psd_array_welch(out[0], sfreq=250, n_fft=250) freqs_up = int(freqs[psd_out > 0.5][0]), int(freqs[psd_out > 0.5][-1]) assert (freqs_up == freqs_sig) # return filtered to false ssd = SSD(info, filt_params_signal, filt_params_noise, sort_by_spectral_ratio=False, return_filtered=False) ssd.fit(X) out = ssd.transform(X) psd_out, freqs = psd_array_welch(out[0], sfreq=250, n_fft=250) freqs_up = int(freqs[psd_out > 0.5][0]), int(freqs[psd_out > 0.5][-1]) assert (freqs_up != freqs_sig)