Aluode/PerceptionLabPortable
0
1# Authors: The MNE-Python contributors.2# License: BSD-3-Clause3# Copyright the MNE-Python contributors.4 5# The computations in this code were primarily derived from Matti Hämäläinen's6# C code.7#8# Many of the idealized equations behind these calculations can be found in:9# 1) Realistic conductivity geometry model of the human head for interpretation10# of neuromagnetic data. Hämäläinen and Sarvas, 1989. Specific to MNE11# 2) EEG and MEG: forward solutions for inverse methods. Mosher, Leahy, and12# Lewis, 1999. Generalized discussion of forward solutions.13 14from copy import deepcopy15 16import numpy as np17 18from .._fiff.constants import FIFF19from ..bem import _import_openmeeg, _make_openmeeg_geometry20from ..fixes import bincount, jit21from ..parallel import parallel_func22from ..surface import _jit_cross, _project_onto_surface23from ..transforms import apply_trans, invert_transform24from ..utils import _check_option, _pl, fill_doc, logger, verbose, warn25 26# #############################################################################27# COIL SPECIFICATION AND FIELD COMPUTATION MATRIX28 29 30def _dup_coil_set(coils, coord_frame, t):31 """Make a duplicate."""32 if t is not None and coord_frame != t["from"]:33 raise RuntimeError("transformation frame does not match the coil set")34 coils = deepcopy(coils)35 if t is not None:36 coord_frame = t["to"]37 for coil in coils:38 assert isinstance(coil, dict), f"Coil must be a dict, got {type(coil)}"39 for key in ("ex", "ey", "ez"):40 if key in coil:41 coil[key] = apply_trans(t["trans"], coil[key], False)42 coil["r0"] = apply_trans(t["trans"], coil["r0"])43 coil["rmag"] = apply_trans(t["trans"], coil["rmag"])44 coil["cosmag"] = apply_trans(t["trans"], coil["cosmag"], False)45 coil["coord_frame"] = t["to"]46 return coils, coord_frame47 48 49def _check_coil_frame(coils, coord_frame, bem):50 """Check to make sure the coils are in the correct coordinate frame."""51 if coord_frame != FIFF.FIFFV_COORD_MRI:52 if coord_frame == FIFF.FIFFV_COORD_HEAD:53 # Make a transformed duplicate54 coils, coord_frame = _dup_coil_set(coils, coord_frame, bem["head_mri_t"])55 else:56 raise RuntimeError(f"Bad coil coordinate frame {coord_frame}")57 return coils, coord_frame58 59 60@fill_doc61def _lin_field_coeff(surf, mult, rmags, cosmags, ws, bins, n_jobs):62 """Parallel wrapper for _do_lin_field_coeff to compute linear coefficients.63 64 Parameters65 ----------66 surf : dict67 Dict containing information for one surface of the BEM68 mult : float69 Multiplier for particular BEM surface (Iso Skull Approach discussed in70 Mosher et al., 1999 and Hämäläinen and Sarvas, 1989 Section III?)71 rmag : ndarray, shape (n_integration_pts, 3)72 3D positions of MEG coil integration points (from coil['rmag'])73 cosmag : ndarray, shape (n_integration_pts, 3)74 Direction of the MEG coil integration points (from coil['cosmag'])75 ws : ndarray, shape (n_integration_pts,)76 Weights for MEG coil integration points77 bins : ndarray, shape (n_integration_points,)78 The sensor assignments for each rmag/cosmag/w.79 %(n_jobs)s80 81 Returns82 -------83 coeff : list84 Linear coefficients with lead fields for each BEM vertex on each sensor85 (?)86 """87 parallel, p_fun, n_jobs = parallel_func(88 _do_lin_field_coeff, n_jobs, max_jobs=len(surf["tris"])89 )90 nas = np.array_split91 coeffs = parallel(92 p_fun(surf["rr"], t, tn, ta, rmags, cosmags, ws, bins)93 for t, tn, ta in zip(94 nas(surf["tris"], n_jobs),95 nas(surf["tri_nn"], n_jobs),96 nas(surf["tri_area"], n_jobs),97 )98 )99 return mult * np.sum(coeffs, axis=0)100 101 102@jit()103def _do_lin_field_coeff(bem_rr, tris, tn, ta, rmags, cosmags, ws, bins):104 """Compute field coefficients (parallel-friendly).105 106 See section IV of Mosher et al., 1999 (specifically equation 35).107 108 Parameters109 ----------110 bem_rr : ndarray, shape (n_BEM_vertices, 3)111 Positions on one BEM surface in 3-space. 2562 BEM vertices for BEM with112 5120 triangles (ico-4)113 tris : ndarray, shape (n_BEM_vertices, 3)114 Vertex indices for each triangle (referring to bem_rr)115 tn : ndarray, shape (n_BEM_vertices, 3)116 Triangle unit normal vectors117 ta : ndarray, shape (n_BEM_vertices,)118 Triangle areas119 rmag : ndarray, shape (n_sensor_pts, 3)120 3D positions of MEG coil integration points (from coil['rmag'])121 cosmag : ndarray, shape (n_sensor_pts, 3)122 Direction of the MEG coil integration points (from coil['cosmag'])123 ws : ndarray, shape (n_sensor_pts,)124 Weights for MEG coil integration points125 bins : ndarray, shape (n_sensor_pts,)126 The sensor assignments for each rmag/cosmag/w.127 128 Returns129 -------130 coeff : ndarray, shape (n_MEG_sensors, n_BEM_vertices)131 Linear coefficients with effect of each BEM vertex on each sensor (?)132 """133 coeff = np.zeros((bins[-1] + 1, len(bem_rr)))134 w_cosmags = ws.reshape(-1, 1) * cosmags135 diff = rmags.reshape(rmags.shape[0], 1, rmags.shape[1]) - bem_rr136 den = np.sum(diff * diff, axis=-1)137 den *= np.sqrt(den)138 den *= 3139 for ti in range(len(tris)):140 tri, tri_nn, tri_area = tris[ti], tn[ti], ta[ti]141 # Accumulate the coefficients for each triangle node and add to the142 # corresponding coefficient matrix143 144 # Simple version (bem_lin_field_coeffs_simple)145 # The following is equivalent to:146 # tri_rr = bem_rr[tri]147 # for j, coil in enumerate(coils['coils']):148 # x = func(coil['rmag'], coil['cosmag'],149 # tri_rr, tri_nn, tri_area)150 # res = np.sum(coil['w'][np.newaxis, :] * x, axis=1)151 # coeff[j][tri + off] += mult * res152 153 c = np.empty((diff.shape[0], tri.shape[0], diff.shape[2]))154 _jit_cross(c, diff[:, tri], tri_nn)155 c *= w_cosmags.reshape(w_cosmags.shape[0], 1, w_cosmags.shape[1])156 for ti in range(3):157 x = np.sum(c[:, ti], axis=-1)158 x /= den[:, tri[ti]] / tri_area159 coeff[:, tri[ti]] += bincount(bins, weights=x, minlength=bins[-1] + 1)160 return coeff161 162 163def _concatenate_coils(coils):164 """Concatenate MEG coil parameters."""165 rmags = np.concatenate([coil["rmag"] for coil in coils])166 cosmags = np.concatenate([coil["cosmag"] for coil in coils])167 ws = np.concatenate([coil["w"] for coil in coils])168 n_int = np.array([len(coil["rmag"]) for coil in coils])169 if n_int[-1] == 0:170 # We assume each sensor has at least one integration point,171 # which should be a safe assumption. But let's check it here, since172 # our code elsewhere relies on bins[-1] + 1 being the number of sensors173 raise RuntimeError("not supported")174 bins = np.repeat(np.arange(len(n_int)), n_int)175 return rmags, cosmags, ws, bins176 177 178@fill_doc179def _bem_specify_coils(bem, coils, coord_frame, mults, n_jobs):180 """Set up for computing the solution at a set of MEG coils.181 182 Parameters183 ----------184 bem : instance of ConductorModel185 BEM information186 coils : list of dict, len(n_MEG_sensors)187 MEG sensor information dicts188 coord_frame : int189 Class constant identifying coordinate frame190 mults : ndarray, shape (1, n_BEM_vertices)191 Multiplier for every vertex in BEM192 %(n_jobs)s193 194 Returns195 -------196 sol: ndarray, shape (n_MEG_sensors, n_BEM_vertices)197 MEG solution198 """199 # Make sure MEG coils are in MRI coordinate frame to match BEM coords200 coils, coord_frame = _check_coil_frame(coils, coord_frame, bem)201 202 # leaving this in in case we want to easily add in the future203 # if method != 'simple': # in ['ferguson', 'urankar']:204 # raise NotImplementedError205 206 # Compute the weighting factors to obtain the magnetic field in the linear207 # potential approximation208 209 # Process each of the surfaces210 rmags, cosmags, ws, bins = _triage_coils(coils)211 del coils212 lens = np.cumsum(np.r_[0, [len(s["rr"]) for s in bem["surfs"]]])213 sol = np.zeros((bins[-1] + 1, bem["solution"].shape[1]))214 215 lims = np.concatenate([np.arange(0, sol.shape[0], 100), [sol.shape[0]]])216 # Put through the bem (in channel-based chunks to save memory)217 for start, stop in zip(lims[:-1], lims[1:]):218 mask = np.logical_and(bins >= start, bins < stop)219 r, c, w, b = rmags[mask], cosmags[mask], ws[mask], bins[mask] - start220 # Compute coeffs for each surface, one at a time221 for o1, o2, surf, mult in zip(222 lens[:-1], lens[1:], bem["surfs"], bem["field_mult"]223 ):224 coeff = _lin_field_coeff(surf, mult, r, c, w, b, n_jobs)225 sol[start:stop] += np.dot(coeff, bem["solution"][o1:o2])226 sol *= mults227 return sol228 229 230def _bem_specify_els(bem, els, mults):231 """Set up for computing the solution at a set of EEG electrodes.232 233 Parameters234 ----------235 bem : instance of ConductorModel236 BEM information237 els : list of dict, len(n_EEG_sensors)238 List of EEG sensor information dicts239 mults: ndarray, shape (1, n_BEM_vertices)240 Multiplier for every vertex in BEM241 242 Returns243 -------244 sol : ndarray, shape (n_EEG_sensors, n_BEM_vertices)245 EEG solution246 """247 sol = np.zeros((len(els), bem["solution"].shape[1]))248 scalp = bem["surfs"][0]249 250 # Operate on all integration points for all electrodes (in MRI coords)251 rrs = np.concatenate(252 [apply_trans(bem["head_mri_t"]["trans"], el["rmag"]) for el in els], axis=0253 )254 ws = np.concatenate([el["w"] for el in els])255 tri_weights, tri_idx = _project_onto_surface(rrs, scalp)256 tri_weights *= ws[:, np.newaxis]257 weights = np.matmul(258 tri_weights[:, np.newaxis], bem["solution"][scalp["tris"][tri_idx]]259 )[:, 0]260 # there are way more vertices than electrodes generally, so let's iterate261 # over the electrodes262 edges = np.concatenate([[0], np.cumsum([len(el["w"]) for el in els])])263 for ii, (start, stop) in enumerate(zip(edges[:-1], edges[1:])):264 sol[ii] = weights[start:stop].sum(0)265 sol *= mults266 return sol267 268 269# #############################################################################270# BEM COMPUTATION271 272_MAG_FACTOR = 1e-7 # μ_0 / (4π)273 274# def _bem_inf_pot(rd, Q, rp):275# """The infinite medium potential in one direction. See Eq. (8) in276# Mosher, 1999"""277# NOTE: the (μ_0 / (4π) factor has been moved to _prep_field_communication278# diff = rp - rd # (Observation point position) - (Source position)279# diff2 = np.sum(diff * diff, axis=1) # Squared magnitude of diff280# # (Dipole moment) dot (diff) / (magnitude ^ 3)281# return np.sum(Q * diff, axis=1) / (diff2 * np.sqrt(diff2))282 283 284@jit()285def _bem_inf_pots(mri_rr, bem_rr, mri_Q=None):286 """Compute the infinite medium potential in all 3 directions.287 288 Parameters289 ----------290 mri_rr : ndarray, shape (n_dipole_vertices, 3)291 Chunk of 3D dipole positions in MRI coordinates292 bem_rr: ndarray, shape (n_BEM_vertices, 3)293 3D vertex positions for one BEM surface294 mri_Q : ndarray, shape (3, 3)295 3x3 head -> MRI transform. I.e., head_mri_t.dot(np.eye(3))296 297 Returns298 -------299 ndarray : shape(n_dipole_vertices, 3, n_BEM_vertices)300 """301 # NOTE: the (μ_0 / (4π) factor has been moved to _prep_field_communication302 # Get position difference vector between BEM vertex and dipole303 diff = np.empty((len(mri_rr), 3, len(bem_rr)))304 for ri in range(mri_rr.shape[0]):305 rr = mri_rr[ri]306 this_diff = bem_rr - rr307 diff_norm = np.sum(this_diff * this_diff, axis=1)308 diff_norm *= np.sqrt(diff_norm)309 diff_norm[diff_norm == 0] = 1.0310 if mri_Q is not None:311 this_diff = np.dot(this_diff, mri_Q.T)312 this_diff /= diff_norm.reshape(-1, 1)313 diff[ri] = this_diff.T314 315 return diff316 317 318# This function has been refactored to process all points simultaneously319# def _bem_inf_field(rd, Q, rp, d):320# """Infinite-medium magnetic field. See (7) in Mosher, 1999"""321# # Get vector from source to sensor integration point322# diff = rp - rd323# diff2 = np.sum(diff * diff, axis=1) # Get magnitude of diff324#325# # Compute cross product between diff and dipole to get magnetic field at326# # integration point327# x = fast_cross_3d(Q[np.newaxis, :], diff)328#329# # Take magnetic field dotted by integration point normal to get magnetic330# # field threading the current loop. Divide by R^3 (equivalently, R^2 * R)331# return np.sum(x * d, axis=1) / (diff2 * np.sqrt(diff2))332 333 334@jit()335def _bem_inf_fields(rr, rmag, cosmag):336 """Compute infinite-medium magnetic field at one MEG sensor.337 338 This operates on all dipoles in all 3 basis directions.339 340 Parameters341 ----------342 rr : ndarray, shape (n_source_points, 3)343 3D dipole source positions344 rmag : ndarray, shape (n_sensor points, 3)345 3D positions of 1 MEG coil's integration points (from coil['rmag'])346 cosmag : ndarray, shape (n_sensor_points, 3)347 Direction of 1 MEG coil's integration points (from coil['cosmag'])348 349 Returns350 -------351 ndarray, shape (n_dipoles, 3, n_integration_pts)352 Magnetic field from all dipoles at each MEG sensor integration point353 """354 # rr, rmag refactored according to Equation (19) in Mosher, 1999355 # Knowing that we're doing all directions, refactor above function:356 357 # rr, 3, rmag358 diff = rmag.T.reshape(1, 3, rmag.shape[0]) - rr.reshape(rr.shape[0], 3, 1)359 diff_norm = np.sum(diff * diff, axis=1) # rr, rmag360 diff_norm *= np.sqrt(diff_norm) # Get magnitude of distance cubed361 diff_norm_ = diff_norm.reshape(-1)362 diff_norm_[diff_norm_ == 0] = 1 # avoid nans363 364 # This is the result of cross-prod calcs with basis vectors,365 # as if we had taken (Q=np.eye(3)), then multiplied by cosmags366 # factor, and then summed across directions367 x = np.empty((rr.shape[0], 3, rmag.shape[0]))368 x[:, 0] = diff[:, 1] * cosmag[:, 2] - diff[:, 2] * cosmag[:, 1]369 x[:, 1] = diff[:, 2] * cosmag[:, 0] - diff[:, 0] * cosmag[:, 2]370 x[:, 2] = diff[:, 0] * cosmag[:, 1] - diff[:, 1] * cosmag[:, 0]371 diff_norm = diff_norm_.reshape((rr.shape[0], 1, rmag.shape[0]))372 x /= diff_norm373 # x.shape == (rr.shape[0], 3, rmag.shape[0])374 return x375 376 377@fill_doc378def _bem_pot_or_field(rr, mri_rr, mri_Q, coils, solution, bem_rr, n_jobs, coil_type):379 """Calculate the magnetic field or electric potential forward solution.380 381 The code is very similar between EEG and MEG potentials, so combine them.382 This does the work of "fwd_comp_field" (which wraps to "fwd_bem_field")383 and "fwd_bem_pot_els" in MNE-C.384 385 Parameters386 ----------387 rr : ndarray, shape (n_dipoles, 3)388 3D dipole source positions389 mri_rr : ndarray, shape (n_dipoles, 3)390 3D source positions in MRI coordinates391 mri_Q :392 3x3 head -> MRI transform. I.e., head_mri_t.dot(np.eye(3))393 coils : list of dict, len(sensors)394 List of sensors where each element contains sensor specific information395 solution : ndarray, shape (n_sensors, n_BEM_rr)396 Comes from _bem_specify_coils397 bem_rr : ndarray, shape (n_BEM_vertices, 3)398 3D vertex positions for all surfaces in the BEM399 %(n_jobs)s400 coil_type : str401 'meg' or 'eeg'402 403 Returns404 -------405 B : ndarray, shape (n_dipoles * 3, n_sensors)406 Forward solution for a set of sensors407 """408 # Both MEG and EEG have the inifinite-medium potentials409 # This could be just vectorized, but eats too much memory, so instead we410 # reduce memory by chunking within _do_inf_pots and parallelize, too:411 parallel, p_fun, n_jobs = parallel_func(_do_inf_pots, n_jobs, max_jobs=len(rr))412 nas = np.array_split413 B = np.sum(414 parallel(415 p_fun(416 mri_rr, sr.copy(), np.ascontiguousarray(mri_Q), np.array(sol)417 ) # copy and contig418 for sr, sol in zip(nas(bem_rr, n_jobs), nas(solution.T, n_jobs))419 ),420 axis=0,421 )422 # The copy()s above should make it so the whole objects don't need to be423 # pickled...424 425 # Only MEG coils are sensitive to the primary current distribution.426 if coil_type == "meg":427 # Primary current contribution (can be calc. in coil/dipole coords)428 parallel, p_fun, n_jobs = parallel_func(_do_prim_curr, n_jobs)429 pcc = np.concatenate(parallel(p_fun(r, coils) for r in nas(rr, n_jobs)), axis=0)430 B += pcc431 B *= _MAG_FACTOR432 return B433 434 435def _do_prim_curr(rr, coils):436 """Calculate primary currents in a set of MEG coils.437 438 See Mosher et al., 1999 Section II for discussion of primary vs. volume439 currents.440 441 Parameters442 ----------443 rr : ndarray, shape (n_dipoles, 3)444 3D dipole source positions in head coordinates445 coils : list of dict446 List of MEG coils where each element contains coil specific information447 448 Returns449 -------450 pc : ndarray, shape (n_sources, n_MEG_sensors)451 Primary current for set of MEG coils due to all sources452 """453 rmags, cosmags, ws, bins = _triage_coils(coils)454 n_coils = bins[-1] + 1455 del coils456 pc = np.empty((len(rr) * 3, n_coils))457 for start, stop in _rr_bounds(rr, chunk=1):458 pp = _bem_inf_fields(rr[start:stop], rmags, cosmags)459 pp *= ws460 pp.shape = (3 * (stop - start), -1)461 pc[3 * start : 3 * stop] = [462 bincount(bins, this_pp, bins[-1] + 1) for this_pp in pp463 ]464 return pc465 466 467def _rr_bounds(rr, chunk=200):468 # chunk data nicely469 bounds = np.concatenate([np.arange(0, len(rr), chunk), [len(rr)]])470 return zip(bounds[:-1], bounds[1:])471 472 473def _do_inf_pots(mri_rr, bem_rr, mri_Q, sol):474 """Calculate infinite potentials for MEG or EEG sensors using chunks.475 476 Parameters477 ----------478 mri_rr : ndarray, shape (n_dipoles, 3)479 3D dipole source positions in MRI coordinates480 bem_rr : ndarray, shape (n_BEM_vertices, 3)481 3D vertex positions for all surfaces in the BEM482 mri_Q :483 3x3 head -> MRI transform. I.e., head_mri_t.dot(np.eye(3))484 sol : ndarray, shape (n_sensors_subset, n_BEM_vertices_subset)485 Comes from _bem_specify_coils486 487 Returns488 -------489 B : ndarray, (n_dipoles * 3, n_sensors)490 Forward solution for sensors due to volume currents491 """492 # Doing work of 'fwd_bem_pot_calc' in MNE-C493 # The following code is equivalent to this, but saves memory494 # v0s = _bem_inf_pots(rr, bem_rr, Q) # n_rr x 3 x n_bem_rr495 # v0s.shape = (len(rr) * 3, v0s.shape[2])496 # B = np.dot(v0s, sol)497 498 # We chunk the source mri_rr's in order to save memory499 B = np.empty((len(mri_rr) * 3, sol.shape[1]))500 for start, stop in _rr_bounds(mri_rr):501 # v0 in Hämäläinen et al., 1989 == v_inf in Mosher, et al., 1999502 v0s = _bem_inf_pots(mri_rr[start:stop], bem_rr, mri_Q)503 v0s = v0s.reshape(-1, v0s.shape[2])504 B[3 * start : 3 * stop] = np.dot(v0s, sol)505 return B506 507 508# #############################################################################509# SPHERE COMPUTATION510 511 512def _sphere_pot_or_field(rr, mri_rr, mri_Q, coils, solution, bem_rr, n_jobs, coil_type):513 """Do potential or field for spherical model."""514 fun = _eeg_spherepot_coil if coil_type == "eeg" else _sphere_field515 parallel, p_fun, n_jobs = parallel_func(fun, n_jobs, max_jobs=len(rr))516 B = np.concatenate(517 parallel(p_fun(r, coils, sphere=solution) for r in np.array_split(rr, n_jobs))518 )519 return B520 521 522def _sphere_field(rrs, coils, sphere):523 """Compute field for spherical model using Jukka Sarvas' field computation.524 525 Jukka Sarvas, "Basic mathematical and electromagnetic concepts of the526 biomagnetic inverse problem", Phys. Med. Biol. 1987, Vol. 32, 1, 11-22.527 528 The formulas have been manipulated for efficient computation529 by Matti Hämäläinen, February 1990530 """531 rmags, cosmags, ws, bins = _triage_coils(coils)532 return _do_sphere_field(rrs, rmags, cosmags, ws, bins, sphere["r0"])533 534 535@jit()536def _do_sphere_field(rrs, rmags, cosmags, ws, bins, r0):537 n_coils = bins[-1] + 1538 # Shift to the sphere model coordinates539 rrs = rrs - r0540 B = np.zeros((3 * len(rrs), n_coils))541 for ri in range(len(rrs)):542 rr = rrs[ri]543 # Check for a dipole at the origin544 if np.sqrt(np.dot(rr, rr)) <= 1e-10:545 continue546 this_poss = rmags - r0547 548 # Vector from dipole to the field point549 a_vec = this_poss - rr550 a = np.sqrt(np.sum(a_vec * a_vec, axis=1))551 r = np.sqrt(np.sum(this_poss * this_poss, axis=1))552 rr0 = np.sum(this_poss * rr, axis=1)553 ar = (r * r) - rr0554 ar0 = ar / a555 F = a * (r * a + ar)556 gr = (a * a) / r + ar0 + 2.0 * (a + r)557 g0 = a + 2 * r + ar0558 # Compute the dot products needed559 re = np.sum(this_poss * cosmags, axis=1)560 r0e = np.sum(rr * cosmags, axis=1)561 g = (g0 * r0e - gr * re) / (F * F)562 good = (a > 0) | (r > 0) | ((a * r) + 1 > 1e-5)563 rr_ = rr.reshape(1, 3)564 v1 = np.empty((cosmags.shape[0], 3))565 _jit_cross(v1, rr_, cosmags)566 v2 = np.empty((cosmags.shape[0], 3))567 _jit_cross(v2, rr_, this_poss)568 xx = (good * ws).reshape(-1, 1) * (569 v1 / F.reshape(-1, 1) + v2 * g.reshape(-1, 1)570 )571 for jj in range(3):572 zz = bincount(bins, xx[:, jj], n_coils)573 B[3 * ri + jj, :] = zz574 B *= _MAG_FACTOR575 return B576 577 578def _eeg_spherepot_coil(rrs, coils, sphere):579 """Calculate the EEG in the sphere model."""580 rmags, cosmags, ws, bins = _triage_coils(coils)581 n_coils = bins[-1] + 1582 del coils583 584 # Shift to the sphere model coordinates585 rrs = rrs - sphere["r0"]586 587 B = np.zeros((3 * len(rrs), n_coils))588 for ri, rr in enumerate(rrs):589 # Only process dipoles inside the innermost sphere590 if np.sqrt(np.dot(rr, rr)) >= sphere["layers"][0]["rad"]:591 continue592 # fwd_eeg_spherepot_vec593 vval_one = np.zeros((len(rmags), 3))594 595 # Make a weighted sum over the equivalence parameters596 for eq in range(sphere["nfit"]):597 # Scale the dipole position598 rd = sphere["mu"][eq] * rr599 rd2 = np.sum(rd * rd)600 rd2_inv = 1.0 / rd2601 # Go over all electrodes602 this_pos = rmags - sphere["r0"]603 604 # Scale location onto the surface of the sphere (not used)605 # if sphere['scale_pos']:606 # pos_len = (sphere['layers'][-1]['rad'] /607 # np.sqrt(np.sum(this_pos * this_pos, axis=1)))608 # this_pos *= pos_len609 610 # Vector from dipole to the field point611 a_vec = this_pos - rd612 613 # Compute the dot products needed614 a = np.sqrt(np.sum(a_vec * a_vec, axis=1))615 a3 = 2.0 / (a * a * a)616 r2 = np.sum(this_pos * this_pos, axis=1)617 r = np.sqrt(r2)618 rrd = np.sum(this_pos * rd, axis=1)619 ra = r2 - rrd620 rda = rrd - rd2621 622 # The main ingredients623 F = a * (r * a + ra)624 c1 = a3 * rda + 1.0 / a - 1.0 / r625 c2 = a3 + (a + r) / (r * F)626 627 # Mix them together and scale by lambda/(rd*rd)628 m1 = c1 - c2 * rrd629 m2 = c2 * rd2630 631 vval_one += (632 sphere["lambda"][eq]633 * rd2_inv634 * (m1[:, np.newaxis] * rd + m2[:, np.newaxis] * this_pos)635 )636 637 # compute total result638 xx = vval_one * ws[:, np.newaxis]639 zz = np.array([bincount(bins, x, bins[-1] + 1) for x in xx.T])640 B[3 * ri : 3 * ri + 3, :] = zz641 # finishing by scaling by 1/(4*M_PI)642 B *= 0.25 / np.pi643 return B644 645 646def _triage_coils(coils):647 return coils if isinstance(coils, tuple) else _concatenate_coils(coils)648 649 650# #############################################################################651# MAGNETIC DIPOLE (e.g. CHPI)652 653_MIN_DIST_LIMIT = 1e-5654 655 656def _magnetic_dipole_field_vec(rrs, coils, too_close="raise"):657 rmags, cosmags, ws, bins = _triage_coils(coils)658 fwd, min_dist = _compute_mdfv(rrs, rmags, cosmags, ws, bins, too_close)659 if min_dist < _MIN_DIST_LIMIT:660 msg = f"Coil too close (dist = {min_dist * 1000:g} mm)"661 if too_close == "raise":662 raise RuntimeError(msg)663 func = warn if too_close == "warning" else logger.info664 func(msg)665 return fwd666 667 668@jit()669def _compute_mdfv(rrs, rmags, cosmags, ws, bins, too_close):670 """Compute an MEG forward solution for a set of magnetic dipoles."""671 # The code below is a more efficient version (~30x) of this:672 # for ri, rr in enumerate(rrs):673 # for k in range(len(coils)):674 # this_coil = coils[k]675 # # Go through all points676 # diff = this_coil['rmag'] - rr677 # dist2 = np.sum(diff * diff, axis=1)[:, np.newaxis]678 # dist = np.sqrt(dist2)679 # if (dist < 1e-5).any():680 # raise RuntimeError('Coil too close')681 # dist5 = dist2 * dist2 * dist682 # sum_ = (3 * diff * np.sum(diff * this_coil['cosmag'],683 # axis=1)[:, np.newaxis] -684 # dist2 * this_coil['cosmag']) / dist5685 # fwd[3*ri:3*ri+3, k] = 1e-7 * np.dot(this_coil['w'], sum_)686 fwd = np.zeros((3 * len(rrs), bins[-1] + 1))687 min_dist = np.inf688 ws2 = ws.reshape(-1, 1)689 for ri in range(len(rrs)):690 rr = rrs[ri]691 diff = rmags - rr692 dist2_ = np.sum(diff * diff, axis=1)693 dist2 = dist2_.reshape(-1, 1)694 dist = np.sqrt(dist2)695 min_dist = min(dist.min(), min_dist)696 if min_dist < _MIN_DIST_LIMIT and too_close == "raise":697 break698 t_ = np.sum(diff * cosmags, axis=1)699 t = t_.reshape(-1, 1)700 sum_ = ws2 * (3 * diff * t - dist2 * cosmags) / (dist2 * dist2 * dist)701 for ii in range(3):702 fwd[3 * ri + ii] = bincount(bins, sum_[:, ii], bins[-1] + 1)703 fwd *= _MAG_FACTOR704 return fwd, min_dist705 706 707# #############################################################################708# MAIN TRIAGING FUNCTION709 710 711@verbose712def _prep_field_computation(*, sensors, bem, n_jobs, verbose=None):713 """Precompute and store some things that are used for both MEG and EEG.714 715 Calculation includes multiplication factors, coordinate transforms,716 compensations, and forward solutions. All are stored in modified fwd_data.717 718 Parameters719 ----------720 rr : ndarray, shape (n_dipoles, 3)721 3D dipole source positions in head coordinates722 bem : instance of ConductorModel723 Boundary Element Model information724 fwd_data : dict725 Dict containing sensor information in the head coordinate frame.726 Gets updated here with BEM and sensor information for later forward727 calculations.728 %(n_jobs)s729 %(verbose)s730 """731 bem_rr = mults = mri_Q = head_mri_t = None732 if not bem["is_sphere"]:733 if bem["bem_method"] != FIFF.FIFFV_BEM_APPROX_LINEAR:734 raise RuntimeError("only linear collocation supported")735 # Store (and apply soon) μ_0/(4π) factor before source computations736 mults = np.repeat(737 bem["source_mult"] / (4.0 * np.pi), [len(s["rr"]) for s in bem["surfs"]]738 )[np.newaxis, :]739 # Get positions of BEM points for every surface740 bem_rr = np.concatenate([s["rr"] for s in bem["surfs"]])741 742 # The dipole location and orientation must be transformed743 head_mri_t = bem["head_mri_t"]744 mri_Q = bem["head_mri_t"]["trans"][:3, :3].T745 746 solutions = dict()747 for coil_type in sensors:748 coils = sensors[coil_type]["defs"]749 if not bem["is_sphere"]:750 if coil_type == "meg":751 # MEG field computation matrices for BEM752 start = "Composing the field computation matrix"753 logger.info("\n" + start + "...")754 cf = FIFF.FIFFV_COORD_HEAD755 # multiply solution by "mults" here for simplicity756 solution = _bem_specify_coils(bem, coils, cf, mults, n_jobs)757 else:758 # Compute solution for EEG sensor759 logger.info("Setting up for EEG...")760 solution = _bem_specify_els(bem, coils, mults)761 else:762 solution = bem763 if coil_type == "eeg":764 logger.info(765 "Using the equivalent source approach in the "766 "homogeneous sphere for EEG"767 )768 sensors[coil_type]["defs"] = _triage_coils(coils)769 solutions[coil_type] = solution770 771 # Get appropriate forward physics function depending on sphere or BEM model772 fun = _sphere_pot_or_field if bem["is_sphere"] else _bem_pot_or_field773 774 # Update fwd_data with775 # bem_rr (3D BEM vertex positions)776 # mri_Q (3x3 Head->MRI coord transformation applied to identity matrix)777 # head_mri_t (head->MRI coord transform dict)778 # fun (_bem_pot_or_field if not 'sphere'; otherwise _sph_pot_or_field)779 # solutions (len 2 list; [ndarray, shape (n_MEG_sens, n BEM vertices),780 # ndarray, shape (n_EEG_sens, n BEM vertices)]781 fwd_data = dict(782 bem_rr=bem_rr, mri_Q=mri_Q, head_mri_t=head_mri_t, fun=fun, solutions=solutions783 )784 return fwd_data785 786 787@fill_doc788def _compute_forwards_meeg(rr, *, sensors, fwd_data, n_jobs, silent=False):789 """Compute MEG and EEG forward solutions for all sensor types."""790 Bs = dict()791 # The dipole location and orientation must be transformed to mri coords792 mri_rr = None793 if fwd_data["head_mri_t"] is not None:794 mri_rr = np.ascontiguousarray(apply_trans(fwd_data["head_mri_t"]["trans"], rr))795 mri_Q, bem_rr, fun = fwd_data["mri_Q"], fwd_data["bem_rr"], fwd_data["fun"]796 solutions = fwd_data["solutions"]797 del fwd_data798 rr = np.ascontiguousarray(rr) # usually true but not guaranteed, e.g. in dipole.py799 for coil_type, sens in sensors.items():800 coils = sens["defs"]801 compensator = sens.get("compensator", None)802 post_picks = sens.get("post_picks", None)803 solution = solutions.get(coil_type, None)804 805 # Do the actual forward calculation for a list MEG/EEG sensors806 if not silent:807 logger.info(808 f"Computing {coil_type.upper()} at {len(rr)} source location{_pl(rr)} "809 "(free orientations)..."810 )811 # Calculate forward solution using spherical or BEM model812 B = fun(813 rr,814 mri_rr,815 mri_Q,816 coils=coils,817 solution=solution,818 bem_rr=bem_rr,819 n_jobs=n_jobs,820 coil_type=coil_type,821 )822 823 # Compensate if needed (only done for MEG systems w/compensation)824 if compensator is not None:825 B = B @ compensator.T826 if post_picks is not None:827 B = B[:, post_picks]828 Bs[coil_type] = B829 return Bs830 831 832@verbose833def _compute_forwards(rr, *, bem, sensors, n_jobs, verbose=None):834 """Compute the MEG and EEG forward solutions."""835 # Split calculation into two steps to save (potentially) a lot of time836 # when e.g. dipole fitting837 solver = bem.get("solver", "mne")838 _check_option("solver", solver, ("mne", "openmeeg"))839 if bem["is_sphere"] or solver == "mne":840 # This modifies "sensors" in place, so let's copy it in case the calling841 # function needs to reuse it (e.g., in simulate_raw.py)842 sensors = deepcopy(sensors)843 fwd_data = _prep_field_computation(sensors=sensors, bem=bem, n_jobs=n_jobs)844 Bs = _compute_forwards_meeg(845 rr, sensors=sensors, fwd_data=fwd_data, n_jobs=n_jobs846 )847 else:848 Bs = _compute_forwards_openmeeg(rr, bem=bem, sensors=sensors)849 n_sensors_want = sum(len(s["ch_names"]) for s in sensors.values())850 n_sensors = sum(B.shape[1] for B in Bs.values())851 n_sources = list(Bs.values())[0].shape[0]852 assert (n_sources, n_sensors) == (len(rr) * 3, n_sensors_want)853 return Bs854 855 856def _compute_forwards_openmeeg(rr, *, bem, sensors):857 """Compute the MEG and EEG forward solutions for OpenMEEG."""858 if len(bem["surfs"]) != 3:859 raise RuntimeError("Only 3-layer BEM is supported for OpenMEEG.")860 om = _import_openmeeg("compute a forward solution using OpenMEEG")861 hminv = om.SymMatrix(bem["solution"])862 geom = _make_openmeeg_geometry(bem, invert_transform(bem["head_mri_t"]))863 864 # Make dipoles for all XYZ orientations865 dipoles = np.c_[866 np.kron(rr.T, np.ones(3)[None, :]).T,867 np.kron(np.ones(len(rr))[:, None], np.eye(3)),868 ]869 dipoles = np.asfortranarray(dipoles)870 dipoles = om.Matrix(dipoles)871 dsm = om.DipSourceMat(geom, dipoles, "Brain")872 Bs = dict()873 if "eeg" in sensors:874 rmags, _, ws, bins = _concatenate_coils(sensors["eeg"]["defs"])875 rmags = np.asfortranarray(rmags.astype(np.float64))876 eeg_sensors = om.Sensors(om.Matrix(np.asfortranarray(rmags)), geom)877 h2em = om.Head2EEGMat(geom, eeg_sensors)878 eeg_fwd_full = om.GainEEG(hminv, dsm, h2em).array()879 Bs["eeg"] = np.array(880 [bincount(bins, ws * x, bins[-1] + 1) for x in eeg_fwd_full.T], float881 )882 if "meg" in sensors:883 rmags, cosmags, ws, bins = _concatenate_coils(sensors["meg"]["defs"])884 rmags = np.asfortranarray(rmags.astype(np.float64))885 cosmags = np.asfortranarray(cosmags.astype(np.float64))886 labels = [str(ii) for ii in range(len(rmags))]887 weights = radii = np.ones(len(labels))888 meg_sensors = om.Sensors(labels, rmags, cosmags, weights, radii)889 h2mm = om.Head2MEGMat(geom, meg_sensors)890 ds2mm = om.DipSource2MEGMat(dipoles, meg_sensors)891 meg_fwd_full = om.GainMEG(hminv, dsm, h2mm, ds2mm).array()892 B = np.array(893 [bincount(bins, ws * x, bins[-1] + 1) for x in meg_fwd_full.T], float894 )895 compensator = sensors["meg"].get("compensator", None)896 post_picks = sensors["meg"].get("post_picks", None)897 if compensator is not None:898 B = B @ compensator.T899 if post_picks is not None:900 B = B[:, post_picks]901 Bs["meg"] = B902 return Bs903