# Authors: Mark Wronkiewicz # Yousra Bekhti # Eric Larson # # License: BSD (3-clause) import os.path as op from copy import deepcopy import numpy as np from numpy.testing import assert_allclose, assert_array_equal import pytest from mne import (read_source_spaces, pick_types, read_trans, read_cov, make_sphere_model, create_info, setup_volume_source_space, find_events, Epochs, fit_dipole, transform_surface_to, make_ad_hoc_cov, SourceEstimate, setup_source_space, read_bem_solution, make_forward_solution, convert_forward_solution) from mne.chpi import _calculate_chpi_positions, read_head_pos, _get_hpi_info from mne.tests.test_chpi import _assert_quats from mne.datasets import testing from mne.simulation import simulate_sparse_stc, simulate_raw from mne.source_space import _compare_source_spaces from mne.io import read_raw_fif, RawArray from mne.time_frequency import psd_welch from mne.utils import _TempDir, run_tests_if_main data_path = testing.data_path(download=False) raw_fname = op.join(data_path, 'MEG', 'sample', 'sample_audvis_trunc_raw.fif') cov_fname = op.join(data_path, 'MEG', 'sample', 'sample_audvis_trunc-cov.fif') trans_fname = op.join(data_path, 'MEG', 'sample', 'sample_audvis_trunc-trans.fif') subjects_dir = op.join(data_path, 'subjects') bem_path = op.join(subjects_dir, 'sample', 'bem') src_fname = op.join(bem_path, 'sample-oct-2-src.fif') bem_fname = op.join(bem_path, 'sample-320-320-320-bem-sol.fif') bem_1_fname = op.join(bem_path, 'sample-320-bem-sol.fif') raw_chpi_fname = op.join(data_path, 'SSS', 'test_move_anon_raw.fif') pos_fname = op.join(data_path, 'SSS', 'test_move_anon_raw_subsampled.pos') def _make_stc(raw, src): """Make a STC.""" seed = 42 sfreq = raw.info['sfreq'] # Hz tstep = 1. / sfreq n_samples = len(raw.times) // 10 times = np.arange(0, n_samples) * tstep stc = simulate_sparse_stc(src, 10, times, random_state=seed) return stc def _get_data(): """Get some starting data.""" # raw with ECG channel raw = read_raw_fif(raw_fname).crop(0., 5.0).load_data() data_picks = pick_types(raw.info, meg=True, eeg=True) other_picks = pick_types(raw.info, meg=False, stim=True, eog=True) picks = np.sort(np.concatenate((data_picks[::16], other_picks))) raw = raw.pick_channels([raw.ch_names[p] for p in picks]) raw.info.normalize_proj() ecg = RawArray(np.zeros((1, len(raw.times))), create_info(['ECG 063'], raw.info['sfreq'], 'ecg')) for key in ('dev_head_t', 'highpass', 'lowpass', 'dig'): ecg.info[key] = raw.info[key] raw.add_channels([ecg]) src = read_source_spaces(src_fname) trans = read_trans(trans_fname) sphere = make_sphere_model('auto', 'auto', raw.info) stc = _make_stc(raw, src) return raw, src, stc, trans, sphere @testing.requires_testing_data def test_simulate_raw_sphere(): """Test simulation of raw data with sphere model.""" seed = 42 raw, src, stc, trans, sphere = _get_data() assert len(pick_types(raw.info, meg=False, ecg=True)) == 1 # head pos head_pos_sim = dict() # these will be at 1., 2., ... sec shifts = [[0.001, 0., -0.001], [-0.001, 0.001, 0.]] for time_key, shift in enumerate(shifts): # Create 4x4 matrix transform and normalize temp_trans = deepcopy(raw.info['dev_head_t']) temp_trans['trans'][:3, 3] += shift head_pos_sim[time_key + 1.] = temp_trans['trans'] # # Test raw simulation with basic parameters # raw_sim = simulate_raw(raw, stc, trans, src, sphere, read_cov(cov_fname), head_pos=head_pos_sim, blink=True, ecg=True, random_state=seed, use_cps=True) raw_sim_2 = simulate_raw(raw, stc, trans_fname, src_fname, sphere, cov_fname, head_pos=head_pos_sim, blink=True, ecg=True, random_state=seed, use_cps=True) assert_array_equal(raw_sim_2[:][0], raw_sim[:][0]) std = dict(grad=2e-13, mag=10e-15, eeg=0.1e-6) raw_sim = simulate_raw(raw, stc, trans, src, sphere, make_ad_hoc_cov(raw.info, std=std), head_pos=head_pos_sim, blink=True, ecg=True, random_state=seed, use_cps=True) raw_sim_2 = simulate_raw(raw, stc, trans_fname, src_fname, sphere, cov=std, head_pos=head_pos_sim, blink=True, ecg=True, random_state=seed, use_cps=True) assert_array_equal(raw_sim_2[:][0], raw_sim[:][0]) sphere_norad = make_sphere_model('auto', None, raw.info) raw_meg = raw.copy().pick_types() raw_sim = simulate_raw(raw_meg, stc, trans, src, sphere_norad, make_ad_hoc_cov(raw.info, std=None), head_pos=head_pos_sim, blink=True, ecg=True, random_state=seed, use_cps=True) raw_sim_2 = simulate_raw(raw_meg, stc, trans_fname, src_fname, sphere_norad, cov='simple', head_pos=head_pos_sim, blink=True, ecg=True, random_state=seed, use_cps=True) assert_array_equal(raw_sim_2[:][0], raw_sim[:][0]) # Test IO on processed data tempdir = _TempDir() test_outname = op.join(tempdir, 'sim_test_raw.fif') raw_sim.save(test_outname) raw_sim_loaded = read_raw_fif(test_outname, preload=True) assert_allclose(raw_sim_loaded[:][0], raw_sim[:][0], rtol=1e-6, atol=1e-20) del raw_sim, raw_sim_2 # with no cov (no noise) but with artifacts, most time periods should match # but the EOG/ECG channels should not for ecg, eog in ((True, False), (False, True), (True, True)): raw_sim_3 = simulate_raw(raw, stc, trans, src, sphere, cov=None, head_pos=head_pos_sim, blink=eog, ecg=ecg, random_state=seed, use_cps=True) raw_sim_4 = simulate_raw(raw, stc, trans, src, sphere, cov=None, head_pos=head_pos_sim, blink=False, ecg=False, random_state=seed, use_cps=True) picks = np.arange(len(raw.ch_names)) diff_picks = pick_types(raw.info, meg=False, ecg=ecg, eog=eog) these_picks = np.setdiff1d(picks, diff_picks) close = np.isclose(raw_sim_3[these_picks][0], raw_sim_4[these_picks][0], atol=1e-20) assert np.mean(close) > 0.7 far = ~np.isclose(raw_sim_3[diff_picks][0], raw_sim_4[diff_picks][0], atol=1e-20) assert np.mean(far) > 0.99 del raw_sim_3, raw_sim_4 # make sure it works with EEG-only and MEG-only raw_sim_meg = simulate_raw(raw.copy().pick_types(meg=True, eeg=False), stc, trans, src, sphere, cov=None, ecg=True, blink=True, random_state=seed, use_cps=True) raw_sim_eeg = simulate_raw(raw.copy().pick_types(meg=False, eeg=True), stc, trans, src, sphere, cov=None, ecg=True, blink=True, random_state=seed, use_cps=True) raw_sim_meeg = simulate_raw(raw.copy().pick_types(meg=True, eeg=True), stc, trans, src, sphere, cov=None, ecg=True, blink=True, random_state=seed, use_cps=True) assert_allclose(np.concatenate((raw_sim_meg[:][0], raw_sim_eeg[:][0])), raw_sim_meeg[:][0], rtol=1e-7, atol=1e-20) del raw_sim_meg, raw_sim_eeg, raw_sim_meeg # check that different interpolations are similar given small movements raw_sim = simulate_raw(raw, stc, trans, src, sphere, cov=None, head_pos=head_pos_sim, interp='linear', use_cps=True) raw_sim_hann = simulate_raw(raw, stc, trans, src, sphere, cov=None, head_pos=head_pos_sim, interp='hann', use_cps=True) assert_allclose(raw_sim[:][0], raw_sim_hann[:][0], rtol=1e-1, atol=1e-14) del raw_sim, raw_sim_hann # Make impossible transform (translate up into helmet) and ensure failure head_pos_sim_err = deepcopy(head_pos_sim) head_pos_sim_err[1.][2, 3] -= 0.1 # z trans upward 10cm pytest.raises(RuntimeError, simulate_raw, raw, stc, trans, src, sphere, ecg=False, blink=False, head_pos=head_pos_sim_err, use_cps=True) pytest.raises(RuntimeError, simulate_raw, raw, stc, trans, src, bem_fname, ecg=False, blink=False, head_pos=head_pos_sim_err, use_cps=True) # other degenerate conditions pytest.raises(TypeError, simulate_raw, 'foo', stc, trans, src, sphere, use_cps=True) pytest.raises(TypeError, simulate_raw, raw, 'foo', trans, src, sphere, use_cps=True) pytest.raises(ValueError, simulate_raw, raw, stc.copy().crop(0, 0), trans, src, sphere, use_cps=True) stc_bad = stc.copy() stc_bad.tstep += 0.1 pytest.raises(ValueError, simulate_raw, raw, stc_bad, trans, src, sphere, use_cps=True) pytest.raises(TypeError, simulate_raw, raw, stc, trans, src, sphere, cov=0, use_cps=True) # wrong covariance type pytest.raises(RuntimeError, simulate_raw, raw, stc, trans, src, sphere, chpi=True, use_cps=True) # no cHPI info pytest.raises(ValueError, simulate_raw, raw, stc, trans, src, sphere, interp='foo', use_cps=True) pytest.raises(TypeError, simulate_raw, raw, stc, trans, src, sphere, head_pos=1., use_cps=True) pytest.raises(RuntimeError, simulate_raw, raw, stc, trans, src, sphere, head_pos=pos_fname, use_cps=True) # ends up with t>t_end head_pos_sim_err = deepcopy(head_pos_sim) head_pos_sim_err[-1.] = head_pos_sim_err[1.] # negative time pytest.raises(RuntimeError, simulate_raw, raw, stc, trans, src, sphere, head_pos=head_pos_sim_err, use_cps=True) raw_bad = raw.copy() raw_bad.info['dig'] = None pytest.raises(RuntimeError, simulate_raw, raw_bad, stc, trans, src, sphere, blink=True, use_cps=True) @pytest.mark.slowtest @testing.requires_testing_data def test_simulate_raw_bem(): """Test simulation of raw data with BEM.""" raw, src, stc, trans, sphere = _get_data() src = setup_source_space('sample', 'oct1', subjects_dir=subjects_dir) for s in src: s['nuse'] = 3 s['vertno'] = src[1]['vertno'][:3] s['inuse'].fill(0) s['inuse'][s['vertno']] = 1 # use different / more complete STC here vertices = [s['vertno'] for s in src] stc = SourceEstimate(np.eye(sum(len(v) for v in vertices)), vertices, 0, 1. / raw.info['sfreq']) raw_sim_sph = simulate_raw(raw, stc, trans, src, sphere, cov=None, use_cps=True) raw_sim_bem = simulate_raw(raw, stc, trans, src, bem_fname, cov=None, n_jobs=2, use_cps=True) # some components (especially radial) might not match that well, # so just make sure that most components have high correlation assert_array_equal(raw_sim_sph.ch_names, raw_sim_bem.ch_names) picks = pick_types(raw.info, meg=True, eeg=True) n_ch = len(picks) corr = np.corrcoef(raw_sim_sph[picks][0], raw_sim_bem[picks][0]) assert_array_equal(corr.shape, (2 * n_ch, 2 * n_ch)) med_corr = np.median(np.diag(corr[:n_ch, -n_ch:])) assert med_corr > 0.65 # do some round-trip localization for s in src: transform_surface_to(s, 'head', trans) locs = np.concatenate([s['rr'][s['vertno']] for s in src]) tmax = (len(locs) - 1) / raw.info['sfreq'] cov = make_ad_hoc_cov(raw.info) # The tolerance for the BEM is surprisingly high (28) but I get the same # result when using MNE-C and Xfit, even when using a proper 5120 BEM :( for use_raw, bem, tol in ((raw_sim_sph, sphere, 2), (raw_sim_bem, bem_fname, 31)): events = find_events(use_raw, 'STI 014') assert len(locs) == 6 evoked = Epochs(use_raw, events, 1, 0, tmax, baseline=None).average() assert len(evoked.times) == len(locs) fits = fit_dipole(evoked, cov, bem, trans, min_dist=1.)[0].pos diffs = np.sqrt(np.sum((locs - fits) ** 2, axis=-1)) * 1000 med_diff = np.median(diffs) assert med_diff < tol, '%s: %s' % (bem, med_diff) @testing.requires_testing_data def test_simulate_round_trip(): """Test simulate_raw round trip calculations.""" # Check a diagonal round-trip raw, src, stc, trans, sphere = _get_data() raw.pick_types(meg=True, stim=True) bem = read_bem_solution(bem_1_fname) old_bem = bem.copy() old_src = src.copy() old_trans = trans.copy() fwd = make_forward_solution(raw.info, trans, src, bem) # no omissions assert (sum(len(s['vertno']) for s in src) == sum(len(s['vertno']) for s in fwd['src']) == 36) # make sure things were not modified assert (old_bem['surfs'][0]['coord_frame'] == bem['surfs'][0]['coord_frame']) assert trans == old_trans _compare_source_spaces(src, old_src) data = np.eye(fwd['nsource']) raw.crop(0, (len(data) - 1) / raw.info['sfreq']) stc = SourceEstimate(data, [s['vertno'] for s in fwd['src']], 0, 1. / raw.info['sfreq']) for use_fwd in (None, fwd): if use_fwd is None: use_trans, use_src, use_bem = trans, src, bem else: use_trans = use_src = use_bem = None for use_cps in (False, True): this_raw = simulate_raw(raw, stc, use_trans, use_src, use_bem, cov=None, use_cps=use_cps, forward=use_fwd) this_raw.pick_types(meg=True, eeg=True) assert (old_bem['surfs'][0]['coord_frame'] == bem['surfs'][0]['coord_frame']) assert trans == old_trans _compare_source_spaces(src, old_src) this_fwd = convert_forward_solution(fwd, force_fixed=True, use_cps=use_cps) assert_allclose(this_raw[:][0], this_fwd['sol']['data'], atol=1e-12, rtol=1e-6) with pytest.raises(ValueError, match='If forward is not None then'): simulate_raw(raw, stc, trans, src, bem, forward=fwd) fwd['info']['dev_head_t']['trans'][0, 0] = 1. with pytest.raises(ValueError, match='dev_head_t.*does not match'): simulate_raw(raw, stc, None, None, None, forward=fwd) @pytest.mark.slowtest @testing.requires_testing_data def test_simulate_raw_chpi(): """Test simulation of raw data with cHPI.""" raw = read_raw_fif(raw_chpi_fname, allow_maxshield='yes') picks = np.arange(len(raw.ch_names)) picks = np.setdiff1d(picks, pick_types(raw.info, meg=True, eeg=True)[::4]) raw.load_data().pick_channels([raw.ch_names[pick] for pick in picks]) raw.info.normalize_proj() sphere = make_sphere_model('auto', 'auto', raw.info) # make sparse spherical source space sphere_vol = tuple(sphere['r0'] * 1000.) + (sphere.radius * 1000.,) src = setup_volume_source_space('sample', sphere=sphere_vol, pos=70.) stc = _make_stc(raw, src) # simulate data with cHPI on raw_sim = simulate_raw(raw, stc, None, src, sphere, cov=None, chpi=False, interp='zero', use_cps=True) # need to trim extra samples off this one raw_chpi = simulate_raw(raw, stc, None, src, sphere, cov=None, chpi=True, head_pos=pos_fname, interp='zero', use_cps=True) # test cHPI indication hpi_freqs, hpi_pick, hpi_ons = _get_hpi_info(raw.info) assert_allclose(raw_sim[hpi_pick][0], 0.) assert_allclose(raw_chpi[hpi_pick][0], hpi_ons.sum()) # test that the cHPI signals make some reasonable values picks_meg = pick_types(raw.info, meg=True, eeg=False) picks_eeg = pick_types(raw.info, meg=False, eeg=True) for picks in [picks_meg[:3], picks_eeg[:3]]: psd_sim, freqs_sim = psd_welch(raw_sim, picks=picks) psd_chpi, freqs_chpi = psd_welch(raw_chpi, picks=picks) assert_array_equal(freqs_sim, freqs_chpi) freq_idx = np.sort([np.argmin(np.abs(freqs_sim - f)) for f in hpi_freqs]) if picks is picks_meg: assert (psd_chpi[:, freq_idx] > 100 * psd_sim[:, freq_idx]).all() else: assert_allclose(psd_sim, psd_chpi, atol=1e-20) # test localization based on cHPI information quats_sim = _calculate_chpi_positions(raw_chpi, t_step_min=10.) quats = read_head_pos(pos_fname) _assert_quats(quats, quats_sim, dist_tol=5e-3, angle_tol=3.5) run_tests_if_main()