cov test

covtest.py

# -*- coding: utf-8 -*-
"""
Test covariance generation
"""
from __future__ import print_function

from os import path as op
import warnings
import numpy as np

import mne

use_sample = False

if use_sample:
    data_dir = mne.datasets.sample.data_path()
else:
    data_dir = op.join('.')
subjects = ['test']
tmin, tmax = -0.2, 0

data_types = ['real', 'noise', 'rank-deficient']
n_epochs_use = [10, 20, 40]  # , 40, 80]
lowpasses = [None]  # XXX fix 2, 10, 50, 250, 
rng = np.random.RandomState(0)
scalings = dict(grad=1e13, mag=1e15, eeg=1e6)

methods = None
logliks = np.empty((len(subjects), len(data_types), len(lowpasses),
                    len(n_epochs_use), 4))
for si, subj in enumerate(subjects):
    print('Subject #%s/%s' % (si + 1, len(subjects)))
    for di, data_type in enumerate(data_types):
        print('  Data type: %s' % data_type)
        if not use_sample:
            fname_raw = op.join(data_dir, subj + '_raw.fif')
        else:
            fname_raw = op.join(data_dir, 'MEG', 'sample', 'sample_audvis_raw.fif')
        with warnings.catch_warnings(record=True):
            raw = mne.io.Raw(fname_raw, allow_maxshield=True)
        raw.crop(0, max(n_epochs_use) * 0.2 + 1)
        raw.preload_data()
        if data_type in ('noise', 'rank-deficient'):
            data = rng.randn(len(raw.ch_names), raw.n_times)
            for meg, eeg, key in zip(['grad', 'mag', 'false'],
                                     [False, False, True],
                                     ['grad', 'mag', 'eeg']):
                picks = mne.pick_types(raw.info, meg=meg, eeg=eeg, exclude=[])
                data[picks] /= scalings[key]
                if data_type == 'rank-deficient' and not eeg:
                    data[picks[-len(picks)//2:]] = data[picks[:len(picks)//2]]
            raw = mne.io.RawArray(data, raw.info)
            raw.info['projs'] = []
        else:
            assert data_type == 'real'
        events = mne.make_fixed_length_events(raw, 1, duration=tmax-tmin)

        for li, lowpass in enumerate(lowpasses):
            raw_lp = raw.copy()
            if lowpass is not None:
                print('    Low-passing at %s Hz' % lowpass)
                raw_lp.filter(None, lowpass, filter_length='2s', n_jobs='cuda')
            else:
                print('    Not filtering')
            epochs_all = mne.Epochs(raw_lp, events, 1, tmin, tmax,
                                    preload=True)
            for ni, n_epochs in enumerate(n_epochs_use):
                print('      Using %s epochs... ' % n_epochs)
                cov = mne.compute_covariance(epochs_all[:n_epochs],
                                             method='auto',
                                             return_estimators=True)
                if methods is None:
                    methods = np.sort([c['method'] for c in cov])
                these_methods = np.array([c['method'] for c in cov])
                order = np.argsort(these_methods)
                assert np.array_equal(methods, these_methods[order])
                logliks[si, di, li, ni, order] = [c['loglik'] for c in cov]
assert lowpasses[-1] is None
lowpasses[-1] = int(round(raw.info['sfreq'] / 2))
n_samp_eff = ((tmax - tmin) * raw.info['sfreq'] * np.array(n_epochs_use)
              * np.array(lowpasses)[:, np.newaxis] / float(lowpasses[-1]))
assert n_samp_eff.shape == (len(lowpasses), len(n_epochs_use))
for di, data_type in enumerate(data_types):
    these_liks = logliks[:, di, ...]
    print('inf:  %s%% for %s:' % (100*np.isinf(these_liks).mean(), data_type))
    for mi, method in enumerate(methods):
        print('    (%s%% for %s)' % (100*np.isinf(these_liks[..., mi]).mean(),
                                     method))

# plot the result
import matplotlib.pyplot as plt
plt.ion()
fig, axs = plt.subplots(1, len(data_types))
for ai, (ax, title) in enumerate(zip(axs, data_types)):
    data = np.mean(logliks[:, ai, -1], axis=0)
    lines = ax.plot(data.T)
    ax.set_title(title)
plt.legend(lines, methods)