"""Utility functions for autoreject."""
# Authors: Mainak Jas <mainak.jas@telecom-paristech.fr>
# Denis A. Engemann <denis.engemann@gmail.com>
from collections import defaultdict
import contextlib
import inspect
import warnings
import numpy as np
import mne
from mne import pick_types, pick_info, channel_type
from mne._fiff.pick import _picks_to_idx
from mne.channels.interpolation import _do_interp_dots
from mne.forward._field_interpolation import _setup_dots
def _check_ch_locs(chs):
"""Check if channel locations exist.
Parameters
----------
chs : dict
The channels from info['chs']
"""
locs3d = np.array([ch['loc'][:3] for ch in chs])
return not ((locs3d == 0).all() or
(~np.isfinite(locs3d)).any() or
np.allclose(locs3d, 0.))
def _check_data(epochs, picks, ch_constraint='data_channels', check_loc=True,
verbose=True):
BaseEpochs = _get_epochs_type()
if not isinstance(epochs, BaseEpochs):
raise ValueError('Only accepts MNE epochs objects.')
if epochs.preload is False:
raise ValueError('Data must be preloaded.')
n_bads = len(epochs.info['bads'])
picked_info = pick_info(epochs.info, picks)
ch_types_picked = {
channel_type(picked_info, idx)
for idx in range(len(picks))}
if check_loc and not _check_ch_locs(picked_info['chs']):
raise RuntimeError('Valid channel positions are needed '
'for autoreject to work')
# XXX : ch_constraint -> allow_many_types=True | False
supported_channels = ['mag', 'grad', 'eeg', 'hbo', 'hbr', 'ecog', 'seeg']
channel_text = ', '.join(x for x in supported_channels)
message = f'AutoReject only supports {channel_text} currently.'
if ch_constraint == 'data_channels':
if not all(ch in supported_channels for ch in ch_types_picked):
raise ValueError(message)
elif ch_constraint == 'single_channel_type':
if sum(ch in ch_types_picked for ch in supported_channels) > 1:
raise ValueError(message) # XXX: to check
else:
raise ValueError('bad value for ch_constraint.')
if n_bads > 0:
if verbose is not False:
warnings.warn(
'%i channels are marked as bad. These will be ignored. '
'If you want them to be considered by autoreject please '
'remove them from epochs.info["bads"].' % n_bads)
def _handle_picks(info, picks):
"""Pick the data channls or return picks."""
if picks is None:
out = mne.pick_types(info, meg=True, eeg=True, ref_meg=False,
fnirs=True, ecog=True, seeg=True, exclude='bads')
else:
out = _picks_to_idx(info, picks, exclude='bads')
return out
def _get_picks_by_type(info, picks):
"""Get the picks grouped by channel type."""
# do magic here
sub_picks_ = defaultdict(list)
keys = list()
for pp in picks:
key = channel_type(info=info, idx=pp)
sub_picks_[key].append(pp)
if key not in keys:
keys.append(key)
picks_by_type = [(kk, sub_picks_[kk]) for kk in keys]
return picks_by_type
[docs]
def set_matplotlib_defaults(plt, style='ggplot'):
"""Set publication quality defaults for matplotlib.
Parameters
----------
plt : instance of matplotlib.pyplot
The plt instance.
"""
import matplotlib
matplotlib.style.use(style)
fontsize = 17
params = {'axes.labelsize': fontsize + 2,
'legend.fontsize': fontsize,
'xtick.labelsize': fontsize,
'ytick.labelsize': fontsize,
'axes.titlesize': fontsize + 2}
plt.rcParams.update(params)
def _pprint(params, offset=0, printer=repr):
"""Pretty print the dictionary 'params' (copied from sklearn)
Parameters
----------
params : dict
The dictionary to pretty print
offset : int
The offset in characters to add at the begin of each line.
printer : callable
The function to convert entries to strings, typically
the builtin str or repr
Returns
-------
lines : str
The pretty print of the dictionary as a string.
"""
# Do a multi-line justified repr:
options = np.get_printoptions()
np.set_printoptions(precision=5, threshold=64, edgeitems=2)
params_list = list()
this_line_length = offset
line_sep = ',\n' + (1 + offset // 2) * ' '
for i, (k, v) in enumerate(sorted(params.items())):
if type(v) is float:
# use str for representing floating point numbers
# this way we get consistent representation across
# architectures and versions.
this_repr = '%s=%s' % (k, str(v))
else:
# use repr of the rest
this_repr = '%s=%s' % (k, printer(v))
if len(this_repr) > 500:
this_repr = this_repr[:300] + '...' + this_repr[-100:]
if i > 0:
if (this_line_length + len(this_repr) >= 75 or '\n' in this_repr):
params_list.append(line_sep)
this_line_length = len(line_sep)
else:
params_list.append(', ')
this_line_length += 2
params_list.append(this_repr)
this_line_length += len(this_repr)
np.set_printoptions(**options)
lines = ''.join(params_list)
# Strip trailing space to avoid nightmare in doctests
lines = '\n'.join(li.rstrip(' ') for li in lines.split('\n'))
return lines
def _pbar(iterable, desc, verbose=True, **kwargs):
raise_error = False
if isinstance(verbose, str):
if verbose not in {"tqdm", "tqdm_notebook", "progressbar"}:
raise_error = True
verbose = bool(verbose)
elif isinstance(verbose, (int, bool)):
verbose = bool(verbose) # this can happen with pickling
else:
raise_error = True
warnings.warn(
(f"verbose flag only supports boolean inputs. Option {verbose} "
f"coerced into type {bool(verbose)}"), DeprecationWarning)
verbose = bool(verbose)
if raise_error:
raise ValueError(
f"verbose must be a boolean value. Got {repr(verbose)}")
if verbose:
from mne.utils.progressbar import ProgressBar
pbar = ProgressBar(iterable, mesg=desc, **kwargs)
else:
pbar = iterable
return pbar
def _get_epochs_type():
return mne.epochs.BaseEpochs
def clean_by_interp(inst, picks=None, verbose=True):
"""Clean epochs/evoked by LOOCV.
Parameters
----------
inst : mne.Evoked | mne.Epochs
The evoked or epochs object.
picks : str | list | slice | None
Channels to include. Slices and lists of integers will be interpreted
as channel indices. In lists, channel *type* strings (e.g.,
``['meg', 'eeg']``) will pick channels of those types, channel *name*
strings (e.g., ``['MEG0111', 'MEG2623']`` will pick the given channels.
Can also be the string values ``'all'`` to pick all channels, or
``'data'`` to pick data channels. None (default) will pick data
channels {'meg', 'eeg'}. Note that channels in ``info['bads']`` *will
be included* if their names or indices are explicitly provided.
verbose : bool
The verbosity of progress messages.
If False, suppress all output messages.
Returns
-------
inst_clean : instance of mne.Evoked or mne.Epochs
Instance after interpolation of bad channels.
"""
return _clean_by_interp(inst, picks=picks, verbose=verbose)
def _clean_by_interp(inst, picks=None, dots=None, verbose=True):
if not isinstance(inst, (mne.Evoked, _get_epochs_type())):
raise ValueError('Unrecognized type for inst')
inst_interp = inst.copy()
picks = _handle_picks(info=inst_interp.info, picks=picks)
dots = _get_interp_data(inst.info, picks, dots, mode='fast')
# interpolate each channel from all the other picks, in one go
ch_names = [inst.ch_names[p] for p in picks]
loo = np.concatenate(list(_interp_matrices(
inst.info, picks, dots, [[ch] for ch in ch_names], mode='fast',
desc='Creating augmented epochs', verbose=verbose)))
inst_interp._data[..., picks, :] = np.matmul(
loo, inst._data[..., picks, :])
return inst_interp
def _interp_matrices(info, picks, dots, bad_sets, mode, desc, verbose,
**kwargs):
"""Get the matrix that interpolates each set of bads from the picks.
Interpolating an identity matrix gives the interpolation matrix itself,
so this works on n_picks x n_picks data rather than on all of the data.
Yields, for each set of bad channel names, an array of shape
(n_bads, n_picks) to apply to the data of the picks. Rows of bads that
are not interpolated (e.g., not MEG or EEG) leave the data as is.
"""
picks = np.asarray(picks)
eye = np.eye(len(picks))
evoked = mne.EvokedArray(eye.copy(), pick_info(info, picks),
verbose=False)
if not _has_interp_data() and dots is not None:
# the dots are for the MEG channels of info, so pick those of evoked
meg_picks = _meg_dots_picks(info)
pos = np.searchsorted(meg_picks, picks[np.isin(picks, meg_picks)])
dots = _pick_dots(dots, pos, pos)
idx = {name: ii for ii, name in enumerate(evoked.ch_names)}
# the Info was checked when creating evoked, and only its bads change, so
# do not let MNE-Python (1.13+) check it again for every set of bads
skip_checks = getattr(evoked.info, '_skip_checks', contextlib.nullcontext)
with skip_checks():
for bads in _pbar(bad_sets, desc=desc, verbose=verbose, **kwargs):
bads_idx = [idx[ch] for ch in bads]
evoked.info['bads'] = list(bads)
interpolate_bads(evoked, picks=np.arange(len(picks)), dots=dots,
reset_bads=True, mode=mode)
yield evoked.data[bads_idx] # a copy, as it is fancy indexing
evoked.data[bads_idx] = eye[bads_idx]
# Origins (head frame, in meters) used for EEG (spherical spline) and MEG
# (field mapping) interpolation. The EEG origin of zero matches the
# _make_interpolation_matrix call in _interpolate_bads_eeg below.
_INTERP_ORIGINS = dict(eeg=(0., 0., 0.), meg=(0., 0., 0.04))
def _has_interp_data():
"""Check if MNE-Python can precompute interpolation data (1.14+)."""
return hasattr(mne.channels, 'compute_interpolation_data')
def _precompute_interp(info, picks, mode='fast', n_jobs=1):
"""Precompute the interpolation data for the channel types in picks."""
if _has_interp_data():
return _get_interp_data(info, picks, mode=mode, n_jobs=n_jobs)
ch_types = {channel_type(info, pick) for pick in picks}
if {'mag', 'grad'} & ch_types:
this_info = pick_info(info, _meg_dots_picks(info), copy=True)
return _compute_dots(this_info, mode=mode, templates=None,
n_jobs=n_jobs)
return None
def _get_interp_data(info, picks, dots=None, mode='fast', n_jobs=1):
"""Get MNE's precomputed interpolation data, computing what is missing.
The data is computed for the EEG and MEG channels in picks only, so other
channels (e.g., those in info['bads']) cannot change the results, even
by rounding. It can be reused for any bads and subset of picks. With
MNE-Python < 1.14, ``dots`` is returned unchanged.
"""
if not _has_interp_data():
return dots
from mne.channels import compute_interpolation_data, InterpolationData
# n_jobs was added after compute_interpolation_data
kwargs = dict(n_jobs=n_jobs) if 'n_jobs' in inspect.signature(
compute_interpolation_data).parameters else dict()
# anything else (e.g., the dot products of files written with
# MNE-Python < 1.14) is recomputed
dots = {key: val for key, val in (
dots.items() if isinstance(dots, dict) else ())
if isinstance(val, InterpolationData)}
for kind, origin in _INTERP_ORIGINS.items():
type_picks = pick_types(info, meg=kind == 'meg', eeg=kind == 'eeg',
ref_meg=False, exclude=())
type_picks = type_picks[np.isin(type_picks, picks)]
if len(type_picks) == 0:
continue
names = {info['ch_names'][pick] for pick in type_picks}
if kind not in dots or not names <= set(dots[kind]['ch_names']):
dots.update(compute_interpolation_data(
pick_info(info, type_picks), origin=origin, mode=mode,
verbose='error', **kwargs))
return dots
def _restore_interp_data(dots):
"""Restore the types of interpolation data read from HDF5."""
if _has_interp_data() and isinstance(dots, dict):
from mne.channels import InterpolationData
dots = {key: InterpolationData(val) for key, val in dots.items()}
return dots
def _interpolate_bads_mne(inst, picks, dots, reset_bads, mode):
"""Interpolate bad MEG and EEG channels using MNE's interpolate_bads."""
picks = set(picks)
eeg_meg = pick_types(inst.info, meg=True, eeg=True, ref_meg=False,
exclude=())
names = [inst.ch_names[p] for p in eeg_meg if p in picks]
# channels not in picks are neither interpolated nor interpolated from
exclude = [inst.ch_names[p] for p in eeg_meg if p not in picks]
bads = [ch for ch in inst.info['bads'] if ch in names]
if len(bads) > 0:
dots = _get_interp_data(inst.info, sorted(picks), dots, mode=mode)
orig_bads = inst.info['bads']
inst.info['bads'] = bads + exclude # keeps excluded MEG out of goods
try:
inst.interpolate_bads(
reset_bads=False, origin=_INTERP_ORIGINS['eeg'],
exclude=exclude, on_bad_position='ignore',
interpolation=dots, verbose='error')
finally:
inst.info['bads'] = orig_bads
if reset_bads is True:
inst.info['bads'] = []
return inst
def interpolate_bads(inst, picks, dots=None, reset_bads=True, mode='accurate'):
"""Interpolate bad MEG and EEG channels."""
import mne
if _has_interp_data():
return _interpolate_bads_mne(inst, picks, dots=dots,
reset_bads=reset_bads, mode=mode)
if isinstance(dots, dict): # written with MNE-Python 1.14+
dots = None
eeg_picks = set(pick_types(inst.info, meg=False, eeg=True, exclude=[]))
eeg_picks_interp = [p for p in picks if p in eeg_picks]
meg_picks = set(pick_types(inst.info, meg=True, eeg=False, exclude=[]))
meg_picks_interp = [p for p in picks if p in meg_picks]
# to prevent cobyla printf error
with mne.utils.use_log_level('error'):
if len(eeg_picks_interp) > 0:
_interpolate_bads_eeg(inst, picks=eeg_picks_interp)
if len(meg_picks_interp) > 0:
_interpolate_bads_meg_fast(inst, picks=meg_picks_interp,
dots=dots, mode=mode)
if reset_bads is True:
inst.info['bads'] = []
return inst
def _interpolate_bads_eeg(inst, picks=None):
""" Interpolate bad EEG channels.
Operates in place.
Parameters
----------
inst : mne.io.Raw, mne.Epochs or mne.Evoked
The data to interpolate. Must be preloaded.
picks : str | list | slice | None
Channels to include for interpolation. Slices and lists of integers
will be interpreted as channel indices. In lists, channel *name*
strings (e.g., ``['EEG 01', 'EEG 02']``) will pick the given channels.
None (default) will pick all EEG channels. Note that channels in
``info['bads']`` *will be included* if their names or indices are
explicitly provided.
"""
from mne.bem import _fit_sphere
from mne.utils import logger, warn
from mne.channels.interpolation import _do_interp_dots
from mne.channels.interpolation import _make_interpolation_matrix
import numpy as np
inst.info._check_consistency()
if picks is None:
picks = pick_types(inst.info, meg=False, eeg=True, exclude=[])
else:
picks = _handle_picks(inst.info, picks)
bads_idx = np.zeros(len(inst.ch_names), dtype=bool)
goods_idx = np.zeros(len(inst.ch_names), dtype=bool)
bads_idx[picks] = [inst.ch_names[ch] in inst.info['bads'] for ch in picks]
if len(picks) == 0 or bads_idx.sum() == 0:
return
goods_idx[picks] = True
goods_idx[bads_idx] = False
pos = inst._get_channel_positions(picks)
# Make sure only good EEG are used
bads_idx_pos = bads_idx[picks]
goods_idx_pos = goods_idx[picks]
pos_good = pos[goods_idx_pos]
pos_bad = pos[bads_idx_pos]
# test spherical fit
radius, center = _fit_sphere(pos_good)
distance = np.sqrt(np.sum((pos_good - center) ** 2, 1))
distance = np.mean(distance / radius)
if np.abs(1. - distance) > 0.1:
warn('Your spherical fit is poor, interpolation results are '
'likely to be inaccurate.')
logger.info('Computing interpolation matrix from {0} sensor '
'positions'.format(len(pos_good)))
interpolation = _make_interpolation_matrix(pos_good, pos_bad)
logger.info('Interpolating {0} sensors'.format(len(pos_bad)))
_do_interp_dots(inst, interpolation, goods_idx, bads_idx)
def _interpolate_bads_meg_fast(inst, picks, mode='accurate',
dots=None):
"""Interpolate bad channels from data in good channels."""
# We can have pre-picked instances or not.
# And we need to handle it.
inst_picked = True
if len(inst.ch_names) > len(picks):
picked_info = pick_info(inst.info, picks)
# the dots are those of the MEG channels of inst, in their order
pos = np.searchsorted(_meg_dots_picks(inst.info), picks)
dots = _pick_dots(dots, pos, pos)
inst_picked = False
else:
picked_info = inst.info.copy()
def get_picks_bad_good(info, picks_meg):
picks_good = [p for p in picks_meg
if info['ch_names'][p] not in info['bads']]
# select the bad meg channel to be interpolated
if len(info['bads']) == 0:
picks_bad = []
else:
picks_bad = [p for p in picks_meg
if info['ch_names'][p] in info['bads']]
return picks_meg, picks_good, picks_bad
picks_meg, picks_good, picks_bad = get_picks_bad_good(
picked_info, range(picked_info['nchan']))
# return without doing anything if there are no meg channels
if len(picks_meg) == 0 or len(picks_bad) == 0:
return
# we need to make sure that only meg channels are passed here
# as the MNE interpolation code is not fogriving.
# This is why we picked the info.
mapping = _fast_map_meg_channels(
picked_info, pick_from=picks_good, pick_to=picks_bad,
dots=dots, mode=mode)
# the downside is that the mapping matrix now does not match
# the unpicked info of the data.
# Since we may have picked the info, we need to double map
# the indices.
_, picks_good_, picks_bad_orig = get_picks_bad_good(
inst.info, picks)
ch_names_a = [picked_info['ch_names'][pp] for pp in picks_bad]
ch_names_b = [inst.info['ch_names'][pp] for pp in picks_bad_orig]
assert ch_names_a == ch_names_b
if not inst_picked:
picks_good_ = [pp for pp in picks if pp in picks_good_]
assert len(picks_good_) == len(picks_good)
ch_names_a = [picked_info['ch_names'][pp] for pp in picks_good]
ch_names_b = [inst.info['ch_names'][pp] for pp in picks_good_]
assert ch_names_a == ch_names_b
_do_interp_dots(inst, mapping, picks_good_, picks_bad_orig)
def _meg_dots_picks(info):
"""Get the channels that the dots of _precompute_interp are for."""
return pick_types(info, meg=True, eeg=False, exclude=[])
def _compute_dots(info, mode='fast', *, templates, n_jobs=1):
"""Compute all-to-all dots."""
from mne.forward._lead_dots import _do_self_dots
from mne.forward._make_forward import _create_meg_coils, _read_coil_defs
from mne.bem import _check_origin
if templates is None:
templates = _read_coil_defs(verbose='error')
coils = _create_meg_coils(info['chs'], 'normal', info['dev_head_t'],
templates)
my_origin = _check_origin((0., 0., 0.04), info)
int_rad, noise, lut_fun, n_fact = _setup_dots(mode, info, coils, 'meg')
self_dots = _do_self_dots(int_rad, False, coils, my_origin, 'meg',
lut_fun, n_fact, n_jobs=n_jobs)
# the cross dots of the coils with themselves are the self dots
return self_dots, self_dots
def _pick_dots(dots, pick_from, pick_to):
if dots is None:
return dots
self_dots, cross_dots = dots
self_dots = self_dots[pick_from, :][:, pick_from]
cross_dots = cross_dots[pick_to, :][:, pick_from]
return [self_dots, cross_dots]
def _fast_map_meg_channels(info, pick_from, pick_to,
dots=None, mode='fast'):
from mne.forward._field_interpolation import _compute_mapping_matrix
from mne.bem import _check_origin
miss = 1e-4 # Smoothing criterion for MEG
info_from = pick_info(info, pick_from, copy=True)
my_origin = _check_origin((0., 0., 0.04), info_from)
# the coils are only needed to compute the dots (_setup_dots ignores them)
int_rad, noise, lut_fun, n_fact = _setup_dots(mode, info_from, None,
'meg')
# This function needs a clean input. It hates the presence of other
# channels than MEG channels. Make sure all is picked.
if dots is None:
dots = _compute_dots(info, mode=mode, templates=None)
self_dots, cross_dots = _pick_dots(dots, pick_from, pick_to)
ch_names = [c['ch_name'] for c in info_from['chs']]
fmd = dict(kind='meg', ch_names=ch_names,
origin=my_origin, noise=noise, self_dots=self_dots,
surface_dots=cross_dots, int_rad=int_rad, miss=miss)
# XXX: hack to silence _compute_mapping_matrix
with mne.utils.use_log_level('error'):
fmd['data'] = _compute_mapping_matrix(fmd, info_from)
return fmd['data']
def _get_channel_type(epochs, picks):
"""return whether a set of picks are all meg or all eeg channels."""
picked_info = pick_info(epochs.info, picks)
ch_types_picked = {channel_type(picked_info, idx)
for idx in range(len(picks))}
invalid_ch_types_present = [key for key in ch_types_picked
if key not in ['mag', 'grad', 'eeg'] and
key in epochs]
if len(invalid_ch_types_present) > 0:
raise ValueError('Invalid channel types present in epochs.'
' Expected ONLY `meg` or ONLY `eeg`. Got %s'
% ', '.join(invalid_ch_types_present))
has_meg = any(kk in ch_types_picked for kk in ('mag', 'grad'))
if 'eeg' in ch_types_picked and has_meg:
raise ValueError('Got mixed channel types. Pick either eeg or meg'
' but not both')
if 'eeg' in ch_types_picked:
return 'eeg'
elif has_meg:
return 'meg'
else:
raise ValueError('Oh no! Your channel type is not known.')