Aluode/PerceptionLabPortable
0
1# Authors: The MNE-Python contributors.2# License: BSD-3-Clause3# Copyright the MNE-Python contributors.4 5import numpy as np6from numpy.polynomial.legendre import legval7from scipy.interpolate import RectBivariateSpline8from scipy.linalg import pinv9from scipy.spatial.distance import pdist, squareform10 11from .._fiff.meas_info import _simplify_info12from .._fiff.pick import pick_channels, pick_info, pick_types13from ..surface import _normalize_vectors14from ..utils import _validate_type, logger, verbose, warn15 16 17def _calc_h(cosang, stiffness=4, n_legendre_terms=50):18 """Calculate spherical spline h function between points on a sphere.19 20 Parameters21 ----------22 cosang : array-like | float23 cosine of angles between pairs of points on a spherical surface. This24 is equivalent to the dot product of unit vectors.25 stiffness : float26 stiffnes of the spline. Also referred to as ``m``.27 n_legendre_terms : int28 number of Legendre terms to evaluate.29 """30 factors = [31 (2 * n + 1) / (n ** (stiffness - 1) * (n + 1) ** (stiffness - 1) * 4 * np.pi)32 for n in range(1, n_legendre_terms + 1)33 ]34 return legval(cosang, [0] + factors)35 36 37def _calc_g(cosang, stiffness=4, n_legendre_terms=50):38 """Calculate spherical spline g function between points on a sphere.39 40 Parameters41 ----------42 cosang : array-like of float, shape(n_channels, n_channels)43 cosine of angles between pairs of points on a spherical surface. This44 is equivalent to the dot product of unit vectors.45 stiffness : float46 stiffness of the spline.47 n_legendre_terms : int48 number of Legendre terms to evaluate.49 50 Returns51 -------52 G : np.ndrarray of float, shape(n_channels, n_channels)53 The G matrix.54 """55 factors = [56 (2 * n + 1) / (n**stiffness * (n + 1) ** stiffness * 4 * np.pi)57 for n in range(1, n_legendre_terms + 1)58 ]59 return legval(cosang, [0] + factors)60 61 62def _make_interpolation_matrix(pos_from, pos_to, alpha=1e-5):63 """Compute interpolation matrix based on spherical splines.64 65 Implementation based on [1]66 67 Parameters68 ----------69 pos_from : np.ndarray of float, shape(n_good_sensors, 3)70 The positions to interpolate from.71 pos_to : np.ndarray of float, shape(n_bad_sensors, 3)72 The positions to interpolate.73 alpha : float74 Regularization parameter. Defaults to 1e-5.75 76 Returns77 -------78 interpolation : np.ndarray of float, shape(len(pos_from), len(pos_to))79 The interpolation matrix that maps good signals to the location80 of bad signals.81 82 References83 ----------84 [1] Perrin, F., Pernier, J., Bertrand, O. and Echallier, JF. (1989).85 Spherical splines for scalp potential and current density mapping.86 Electroencephalography Clinical Neurophysiology, Feb; 72(2):184-7.87 """88 pos_from = pos_from.copy()89 pos_to = pos_to.copy()90 n_from = pos_from.shape[0]91 n_to = pos_to.shape[0]92 93 # normalize sensor positions to sphere94 _normalize_vectors(pos_from)95 _normalize_vectors(pos_to)96 97 # cosine angles between source positions98 cosang_from = pos_from.dot(pos_from.T)99 cosang_to_from = pos_to.dot(pos_from.T)100 G_from = _calc_g(cosang_from)101 G_to_from = _calc_g(cosang_to_from)102 assert G_from.shape == (n_from, n_from)103 assert G_to_from.shape == (n_to, n_from)104 105 if alpha is not None:106 G_from.flat[:: len(G_from) + 1] += alpha107 108 C = np.vstack(109 [110 np.hstack([G_from, np.ones((n_from, 1))]),111 np.hstack([np.ones((1, n_from)), [[0]]]),112 ]113 )114 C_inv = pinv(C)115 116 interpolation = np.hstack([G_to_from, np.ones((n_to, 1))]) @ C_inv[:, :-1]117 assert interpolation.shape == (n_to, n_from)118 return interpolation119 120 121def _do_interp_dots(inst, interpolation, goods_idx, bads_idx):122 """Dot product of channel mapping matrix to channel data."""123 from ..epochs import BaseEpochs124 from ..evoked import Evoked125 from ..io import BaseRaw126 127 _validate_type(inst, (BaseRaw, BaseEpochs, Evoked), "inst")128 inst._data[..., bads_idx, :] = np.matmul(129 interpolation, inst._data[..., goods_idx, :]130 )131 132 133@verbose134def _interpolate_bads_eeg(inst, origin, exclude=None, ecog=False, verbose=None):135 if exclude is None:136 exclude = list()137 bads_idx = np.zeros(len(inst.ch_names), dtype=bool)138 goods_idx = np.zeros(len(inst.ch_names), dtype=bool)139 140 picks = pick_types(inst.info, meg=False, eeg=not ecog, ecog=ecog, exclude=exclude)141 inst.info._check_consistency()142 bads_idx[picks] = [inst.ch_names[ch] in inst.info["bads"] for ch in picks]143 144 if len(picks) == 0 or bads_idx.sum() == 0:145 return146 147 goods_idx[picks] = True148 goods_idx[bads_idx] = False149 150 pos = inst._get_channel_positions(picks)151 152 # Make sure only EEG are used153 bads_idx_pos = bads_idx[picks]154 goods_idx_pos = goods_idx[picks]155 156 # test spherical fit157 distance = np.linalg.norm(pos - origin, axis=-1)158 distance = np.mean(distance / np.mean(distance))159 if np.abs(1.0 - distance) > 0.1:160 warn(161 "Your spherical fit is poor, interpolation results are "162 "likely to be inaccurate."163 )164 165 pos_good = pos[goods_idx_pos] - origin166 pos_bad = pos[bads_idx_pos] - origin167 logger.info(f"Computing interpolation matrix from {len(pos_good)} sensor positions")168 interpolation = _make_interpolation_matrix(pos_good, pos_bad)169 170 logger.info(f"Interpolating {len(pos_bad)} sensors")171 _do_interp_dots(inst, interpolation, goods_idx, bads_idx)172 173 174@verbose175def _interpolate_bads_ecog(inst, *, origin, exclude=None, verbose=None):176 _interpolate_bads_eeg(inst, origin, exclude=exclude, ecog=True, verbose=verbose)177 178 179def _interpolate_bads_meg(180 inst, mode="accurate", *, origin, verbose=None, ref_meg=False181):182 return _interpolate_bads_meeg(183 inst, mode, ref_meg=ref_meg, eeg=False, origin=origin, verbose=verbose184 )185 186 187@verbose188def _interpolate_bads_nan(189 inst,190 *,191 ch_type,192 ref_meg=False,193 exclude=(),194 verbose=None,195):196 info = _simplify_info(inst.info)197 picks_type = pick_types(info, ref_meg=ref_meg, exclude=exclude, **{ch_type: True})198 use_ch_names = [inst.info["ch_names"][p] for p in picks_type]199 bads_type = [ch for ch in inst.info["bads"] if ch in use_ch_names]200 if len(bads_type) == 0 or len(picks_type) == 0:201 return202 # select the bad channels to be interpolated203 picks_bad = pick_channels(inst.info["ch_names"], bads_type, exclude=[])204 inst._data[..., picks_bad, :] = np.nan205 206 207@verbose208def _interpolate_bads_meeg(209 inst,210 mode="accurate",211 *,212 meg=True,213 eeg=True,214 ref_meg=False,215 exclude=(),216 origin,217 method=None,218 verbose=None,219):220 from ..forward import _map_meg_or_eeg_channels221 222 if method is None:223 method = {"meg": "MNE", "eeg": "MNE"}224 bools = dict(meg=meg, eeg=eeg)225 info = _simplify_info(inst.info)226 for ch_type, do in bools.items():227 if not do:228 continue229 kw = dict(meg=False, eeg=False)230 kw[ch_type] = True231 picks_type = pick_types(info, ref_meg=ref_meg, exclude=exclude, **kw)232 picks_good = pick_types(info, ref_meg=ref_meg, exclude="bads", **kw)233 use_ch_names = [inst.info["ch_names"][p] for p in picks_type]234 bads_type = [ch for ch in inst.info["bads"] if ch in use_ch_names]235 if len(bads_type) == 0 or len(picks_type) == 0:236 continue237 # select the bad channels to be interpolated238 picks_bad = pick_channels(inst.info["ch_names"], bads_type, exclude=[])239 240 # do MNE based interpolation241 if ch_type == "eeg":242 picks_to = picks_type243 bad_sel = np.isin(picks_type, picks_bad)244 else:245 picks_to = picks_bad246 bad_sel = slice(None)247 info_from = pick_info(inst.info, picks_good)248 info_to = pick_info(inst.info, picks_to)249 mapping = _map_meg_or_eeg_channels(info_from, info_to, mode=mode, origin=origin)250 mapping = mapping[bad_sel]251 _do_interp_dots(inst, mapping, picks_good, picks_bad)252 253 254@verbose255def _interpolate_bads_nirs(inst, exclude=(), verbose=None):256 from mne.preprocessing.nirs import _validate_nirs_info257 258 if len(pick_types(inst.info, fnirs=True, exclude=())) == 0:259 return260 261 # Returns pick of all nirs and ensures channels are correctly ordered262 picks_nirs = _validate_nirs_info(inst.info)263 nirs_ch_names = [inst.info["ch_names"][p] for p in picks_nirs]264 nirs_ch_names = [ch for ch in nirs_ch_names if ch not in exclude]265 bads_nirs = [ch for ch in inst.info["bads"] if ch in nirs_ch_names]266 if len(bads_nirs) == 0:267 return268 picks_bad = pick_channels(inst.info["ch_names"], bads_nirs, exclude=[])269 bads_mask = [p in picks_bad for p in picks_nirs]270 271 chs = [inst.info["chs"][i] for i in picks_nirs]272 locs3d = np.array([ch["loc"][:3] for ch in chs])273 274 dist = pdist(locs3d)275 dist = squareform(dist)276 277 for bad in picks_bad:278 dists_to_bad = dist[bad]279 # Ignore distances to self280 dists_to_bad[dists_to_bad == 0] = np.inf281 # Ignore distances to other bad channels282 dists_to_bad[bads_mask] = np.inf283 # Find closest remaining channels for same frequency284 closest_idx = np.argmin(dists_to_bad) + (bad % 2)285 inst._data[bad] = inst._data[closest_idx]286 287 # TODO: this seems like a bug because it does not respect reset_bads288 inst.info["bads"] = [ch for ch in inst.info["bads"] if ch in exclude]289 290 return inst291 292 293def _find_seeg_electrode_shaft(pos, tol_shaft=0.002, tol_spacing=1):294 # 1) find nearest neighbor to define the electrode shaft line295 # 2) find all contacts on the same line296 # 3) remove contacts with large distances297 298 dist = squareform(pdist(pos))299 np.fill_diagonal(dist, np.inf)300 301 shafts = list()302 shaft_ts = list()303 for i, n1 in enumerate(pos):304 if any([i in shaft for shaft in shafts]):305 continue306 n2 = pos[np.argmin(dist[i])] # 1307 # https://mathworld.wolfram.com/Point-LineDistance3-Dimensional.html308 shaft_dists = np.linalg.norm(309 np.cross((pos - n1), (pos - n2)), axis=1310 ) / np.linalg.norm(n2 - n1)311 shaft = np.where(shaft_dists < tol_shaft)[0] # 2312 shaft_prev = None313 for _ in range(10): # avoid potential cycles314 if np.array_equal(shaft, shaft_prev):315 break316 shaft_prev = shaft317 # compute median shaft line318 v = np.median(319 [320 pos[i] - pos[j]321 for idx, i in enumerate(shaft)322 for j in shaft[idx + 1 :]323 ],324 axis=0,325 )326 c = np.median(pos[shaft], axis=0)327 # recompute distances328 shaft_dists = np.linalg.norm(329 np.cross((pos - c), (pos - c + v)), axis=1330 ) / np.linalg.norm(v)331 shaft = np.where(shaft_dists < tol_shaft)[0]332 ts = np.array([np.dot(c - n0, v) / np.linalg.norm(v) ** 2 for n0 in pos[shaft]])333 shaft_order = np.argsort(ts)334 shaft = shaft[shaft_order]335 ts = ts[shaft_order]336 337 # only include the largest group with spacing with the error tolerance338 # avoid interpolating across spans between contacts339 t_diffs = np.diff(ts)340 t_diff_med = np.median(t_diffs)341 spacing_errors = (t_diffs - t_diff_med) / t_diff_med342 groups = list()343 group = [shaft[0]]344 for j in range(len(shaft) - 1):345 if spacing_errors[j] > tol_spacing:346 groups.append(group)347 group = [shaft[j + 1]]348 else:349 group.append(shaft[j + 1])350 groups.append(group)351 group = [group for group in groups if i in group][0]352 ts = ts[np.isin(shaft, group)]353 shaft = np.array(group, dtype=int)354 355 shafts.append(shaft)356 shaft_ts.append(ts)357 return shafts, shaft_ts358 359 360@verbose361def _interpolate_bads_seeg(362 inst, exclude=None, tol_shaft=0.002, tol_spacing=1, verbose=None363):364 if exclude is None:365 exclude = list()366 picks = pick_types(inst.info, meg=False, seeg=True, exclude=exclude)367 inst.info._check_consistency()368 bads_idx = np.isin(np.array(inst.ch_names)[picks], inst.info["bads"])369 370 if len(picks) == 0 or bads_idx.sum() == 0:371 return372 373 pos = inst._get_channel_positions(picks)374 375 # Make sure only sEEG are used376 bads_idx_pos = bads_idx[picks]377 378 shafts, shaft_ts = _find_seeg_electrode_shaft(379 pos, tol_shaft=tol_shaft, tol_spacing=tol_spacing380 )381 382 # interpolate the bad contacts383 picks_bad = list(np.where(bads_idx_pos)[0])384 for shaft, ts in zip(shafts, shaft_ts):385 bads_shaft = np.array([idx for idx in picks_bad if idx in shaft])386 if bads_shaft.size == 0:387 continue388 goods_shaft = shaft[np.isin(shaft, bads_shaft, invert=True)]389 if goods_shaft.size < 4: # cubic spline requires 3 channels390 msg = "No shaft" if shaft.size < 4 else "Not enough good channels"391 no_shaft_chs = " and ".join(np.array(inst.ch_names)[bads_shaft])392 raise RuntimeError(393 f"{msg} found in a line with {no_shaft_chs} "394 "at least 3 good channels on the same line "395 f"are required for interpolation, {goods_shaft.size} found. "396 f"Dropping {no_shaft_chs} is recommended."397 )398 logger.debug(399 f"Interpolating {np.array(inst.ch_names)[bads_shaft]} using "400 f"data from {np.array(inst.ch_names)[goods_shaft]}"401 )402 bads_shaft_idx = np.where(np.isin(shaft, bads_shaft))[0]403 goods_shaft_idx = np.where(~np.isin(shaft, bads_shaft))[0]404 405 z = inst._data[..., goods_shaft, :]406 is_epochs = z.ndim == 3407 if is_epochs:408 z = z.swapaxes(0, 1)409 z = z.reshape(z.shape[0], -1)410 y = np.arange(z.shape[-1])411 out = RectBivariateSpline(x=ts[goods_shaft_idx], y=y, z=z)(412 x=ts[bads_shaft_idx], y=y413 )414 if is_epochs:415 out = out.reshape(bads_shaft.size, inst._data.shape[0], -1)416 out = out.swapaxes(0, 1)417 inst._data[..., bads_shaft, :] = out418