generate_epochs.py

import numpy as np
from scipy import signal

import mne
from mne import (read_forward_solution, read_cov, read_label,
                 pick_types_evoked, pick_types_forward, read_evokeds)
from mne.io import Raw
from mne.datasets import sample
from mne.time_frequency import iir_filter_raw, morlet
from mne.simulation import generate_sparse_stc

###############################################################################
# Load real data as templates
data_path = sample.data_path()

raw = Raw(data_path + '/MEG/sample/sample_audvis_raw.fif')
# proj = read_proj(data_path + '/MEG/sample/sample_audvis_ecg_proj.fif')
# raw.info['projs'] += proj
raw.info['bads'] = ['MEG 2443', 'EEG 053']  # mark bad channels

fwd_fname = data_path + '/MEG/sample/sample_audvis-meg-eeg-oct-6-fwd.fif'
ave_fname = data_path + '/MEG/sample/sample_audvis-no-filter-ave.fif'
cov_fname = data_path + '/MEG/sample/sample_audvis-cov.fif'

fwd = read_forward_solution(fwd_fname, force_fixed=True, surf_ori=True)
fwd = pick_types_forward(fwd, meg='grad', eeg=False, exclude=[])

cov = read_cov(cov_fname)

condition = 'Left Auditory'
evoked_template = read_evokeds(ave_fname, condition=condition, baseline=None)
evoked_template = pick_types_evoked(evoked_template, meg=True, eeg=True,
                                    exclude=raw.info['bads'])

label_names = ['Aud-lh', 'Aud-rh']
labels = [read_label(data_path + '/MEG/sample/labels/%s.label' % ln)
          for ln in label_names]

###############################################################################
# Generate source time courses and the correspond evoked data
snr = 60  # dB
tmin = -0.1
sfreq = 1000.  # Hz
tstep = 1. / sfreq
n_samples = 600
times = np.linspace(tmin, tmin + n_samples * tstep, n_samples)

# Generate times series from 2 Morlet wavelets
stc_data = np.zeros((len(labels), len(times)))
# Ws = morlet(sfreq, [3, 15], n_cycles=[1, 1.5])
Ws = morlet(sfreq, [3, 10], n_cycles=[1, 1.5])
stc_data[0][:len(Ws[0])] = np.real(Ws[0])
stc_data[1][:len(Ws[1])] = np.real(Ws[1])
stc_data *= 100 * 1e-9  # use nAm as unit

# time translation
stc_data[1] = np.roll(stc_data[1], 80)
stc = generate_sparse_stc(fwd['src'], labels, stc_data, tmin, tstep,
                          random_state=0)

n_channels, n_times = 204, len(stc.times)
picks = mne.pick_types(raw.info, meg='grad', exclude=[])
info = mne.pick_info(raw.info, picks)
info['bads'] = []
info['sfreq'] = 1000.
tmin = stc.times[0]
evoked_template = mne.evoked.EvokedArray(np.random.randn(n_channels, n_times), info, tmin=tmin)
evoked = mne.forward.apply_forward(fwd, stc, evoked_template)

n_epochs = 30
events = np.zeros((n_epochs, 3))
epochs = mne.epochs.EpochsArray(np.tile(evoked.data, (n_epochs, 1, 1)), info, events, tmin)

# add noise
snr = -5.
rng = np.random.RandomState(42)

noise_model = 'white'
noise_model = 'color'

if noise_model == 'white':
    noise = rng.standard_normal(epochs.get_data().shape)
elif noise_model == 'color':
    noise_cov = mne.cov.pick_channels_cov(cov, include=epochs.ch_names)
    n_channels, n_samples = evoked.data.shape
    cov_data = np.eye(n_channels) * np.mean(np.diag(noise_cov.data))
    noise = np.empty_like(epochs._data)
    iir_filter = iir_filter_raw(raw, order=5, picks=picks, tmin=60, tmax=180)
    for k in range(len(epochs)):
        noise[k] = rng.multivariate_normal(np.zeros(n_channels), cov_data, n_samples).T
        noise[k] = signal.lfilter([1], iir_filter, noise[k], axis=-1)
else:
    1/0

tmp = np.mean((epochs.get_data() ** 2).ravel()) / np.mean((noise ** 2).ravel())
tmp = 10 * np.log10(tmp)
noise.data = 10 ** ((tmp - float(snr)) / 20.) * noise
epochs._data += noise

import matplotlib.pyplot as plt
plt.close('all')
evoked = epochs.average()
evoked.plot()
plt.show()

epochs.save('simu_%s-epo.fif' % noise_model)