Aluode/PerceptionLabPortable
0
1# Authors: The MNE-Python contributors.2# License: BSD-3-Clause3# Copyright the MNE-Python contributors.4 5from collections import defaultdict6from functools import partial7 8import numpy as np9from scipy.optimize import minimize10 11from .._fiff.pick import pick_info, pick_types12from .._fiff.tag import _coil_trans_to_loc, _loc_to_coil_trans13from ..bem import _check_origin14from ..io import BaseRaw15from ..transforms import _find_vector_rotation16from ..utils import (17 _check_fname,18 _check_option,19 _clean_names,20 _ensure_int,21 _pl,22 _reg_pinv,23 _validate_type,24 check_fname,25 logger,26 verbose,27)28from .maxwell import (29 _col_norm_pinv,30 _get_grad_point_coilsets,31 _prep_fine_cal,32 _prep_mf_coils,33 _read_cross_talk,34 _trans_sss_basis,35)36 37 38@verbose39def compute_fine_calibration(40 raw,41 n_imbalance=3,42 t_window=10.0,43 ext_order=2,44 origin=(0.0, 0.0, 0.0),45 cross_talk=None,46 calibration=None,47 *,48 angle_limit=5.0,49 err_limit=5.0,50 verbose=None,51):52 """Compute fine calibration from empty-room data.53 54 Parameters55 ----------56 raw : instance of Raw57 The raw data to use. Should be from an empty-room recording,58 and all channels should be good.59 n_imbalance : int60 Can be 1 or 3 (default), indicating the number of gradiometer61 imbalance components. Only used if gradiometers are present.62 t_window : float63 Time window to use for surface normal rotation in seconds.64 Default is 10.65 %(ext_order_maxwell)s66 Default is 2, which is lower than the default (3) for67 :func:`mne.preprocessing.maxwell_filter` because it tends to yield68 more stable parameter estimates.69 %(origin_maxwell)s70 %(cross_talk_maxwell)s71 calibration : dict | None72 Dictionary with existing calibration. If provided, the magnetometer73 imbalances and adjusted normals will be used and only the gradiometer74 imbalances will be estimated (see step 2 in Notes below).75 angle_limit : float76 The maximum permitted angle in degrees between the original and adjusted77 magnetometer normals. If the angle is exceeded, the segment is treated as78 an outlier and discarded.79 80 .. versionadded:: 1.981 err_limit : float82 The maximum error (in percent) for each channel in order for a segment to83 be used.84 85 .. versionadded:: 1.986 %(verbose)s87 88 Returns89 -------90 calibration : dict91 Fine calibration data.92 count : int93 The number of good segments used to compute the magnetometer94 parameters.95 96 See Also97 --------98 mne.preprocessing.maxwell_filter99 100 Notes101 -----102 This algorithm proceeds in two steps, both optimizing the fit between the103 data and a reconstruction of the data based only on an external multipole104 expansion:105 106 1. Estimate magnetometer normal directions and scale factors. All107 coils (mag and matching grad) are rotated by the adjusted normal108 direction.109 2. Estimate gradiometer imbalance factors. These add point magnetometers110 in just the gradiometer difference direction or in all three directions111 (depending on ``n_imbalance``).112 113 Magnetometer normal and coefficient estimation (1) is typically the most114 time consuming step. Gradiometer imbalance parameters (2) can be115 iteratively reestimated (for example, first using ``n_imbalance=1`` then116 subsequently ``n_imbalance=3``) by passing the previous ``calibration``117 output to the ``calibration`` input in the second call.118 119 MaxFilter processes at most 120 seconds of data, so consider cropping120 your raw instance prior to processing. It also checks to make sure that121 there were some minimal usable ``count`` number of segments (default 5)122 that were included in the estimate.123 124 .. versionadded:: 0.21125 """126 n_imbalance = _ensure_int(n_imbalance, "n_imbalance")127 _check_option("n_imbalance", n_imbalance, (1, 3))128 _validate_type(raw, BaseRaw, "raw")129 ext_order = _ensure_int(ext_order, "ext_order")130 origin = _check_origin(origin, raw.info, "meg", disp=True)131 _check_option("raw.info['bads']", raw.info["bads"], ([],))132 _validate_type(err_limit, "numeric", "err_limit")133 _validate_type(angle_limit, "numeric", "angle_limit")134 for key, val in dict(err_limit=err_limit, angle_limit=angle_limit).items():135 if val < 0:136 raise ValueError(f"{key} must be greater than or equal to 0, got {val}")137 # Fine cal should not include ref channels138 picks = pick_types(raw.info, meg=True, ref_meg=False)139 if raw.info["dev_head_t"] is not None:140 raise ValueError(141 'info["dev_head_t"] is not None, suggesting that the '142 "data are not from an empty-room recording"143 )144 145 info = pick_info(raw.info, picks) # make a copy and pick MEG channels146 mag_picks = pick_types(info, meg="mag", exclude=())147 grad_picks = pick_types(info, meg="grad", exclude=())148 149 # Get cross-talk150 ctc, _ = _read_cross_talk(cross_talk, info["ch_names"])151 152 # Check fine cal153 _validate_type(calibration, (dict, None), "calibration")154 155 #156 # 1. Rotate surface normals using magnetometer information (if present)157 #158 cals = np.ones(len(info["ch_names"]))159 end = len(raw.times) + 1160 time_idxs = np.arange(0, end, int(round(t_window * raw.info["sfreq"])))161 if len(time_idxs) == 1:162 time_idxs = np.concatenate([time_idxs, [end]])163 if time_idxs[-1] != end:164 time_idxs[-1] = end165 count = 0166 locs = np.array([ch["loc"] for ch in info["chs"]])167 zs = locs[mag_picks, -3:].copy()168 if calibration is not None:169 _, calibration, _ = _prep_fine_cal(info, calibration, ignore_ref=True)170 for pi, pick in enumerate(mag_picks):171 idx = calibration["ch_names"].index(info["ch_names"][pick])172 cals[pick] = calibration["imb_cals"][idx].item()173 zs[pi] = calibration["locs"][idx][-3:]174 elif len(mag_picks) > 0:175 cal_list = list()176 z_list = list()177 logger.info(178 f"Adjusting normals for {len(mag_picks)} magnetometers "179 f"(averaging over {len(time_idxs) - 1} time intervals)"180 )181 for start, stop in zip(time_idxs[:-1], time_idxs[1:]):182 logger.info(183 f" Processing interval {start / info['sfreq']:0.3f} - "184 f"{stop / info['sfreq']:0.3f} s"185 )186 data = raw[picks, start:stop][0]187 if ctc is not None:188 data = ctc.dot(data)189 z, cal, good = _adjust_mag_normals(190 info,191 data,192 origin,193 ext_order,194 angle_limit=angle_limit,195 err_limit=err_limit,196 )197 if good:198 z_list.append(z)199 cal_list.append(cal)200 count = len(cal_list)201 if count == 0:202 raise RuntimeError("No usable segments found")203 cals[:] = np.mean(cal_list, axis=0)204 zs[:] = np.mean(z_list, axis=0)205 if len(mag_picks) > 0:206 for ii, new_z in enumerate(zs):207 z_loc = locs[mag_picks[ii]]208 # Find sensors with same NZ and R0 (should be three for VV)209 idxs = _matched_loc_idx(z_loc, locs)210 # Rotate the direction vectors to the plane defined by new normal211 _rotate_locs(locs, idxs, new_z)212 for ci, loc in enumerate(locs):213 info["chs"][ci]["loc"][:] = loc214 del calibration, zs215 216 #217 # 2. Estimate imbalance parameters (always done)218 #219 if len(grad_picks) > 0:220 extra = "X direction" if n_imbalance == 1 else ("XYZ directions")221 logger.info(f"Computing imbalance for {len(grad_picks)} gradimeters ({extra})")222 imb_list = list()223 for start, stop in zip(time_idxs[:-1], time_idxs[1:]):224 logger.info(225 f" Processing interval {start / info['sfreq']:0.3f} - "226 f"{stop / info['sfreq']:0.3f} s"227 )228 data = raw[picks, start:stop][0]229 if ctc is not None:230 data = ctc.dot(data)231 out = _estimate_imbalance(info, data, cals, n_imbalance, origin, ext_order)232 imb_list.append(out)233 imb = np.mean(imb_list, axis=0)234 else:235 imb = np.zeros((len(info["ch_names"]), n_imbalance))236 237 #238 # Put in output structure239 #240 assert len(np.intersect1d(mag_picks, grad_picks)) == 0241 imb_cals = [242 cals[ii : ii + 1] if ii in mag_picks else imb[ii]243 for ii in range(len(info["ch_names"]))244 ]245 ch_names = _clean_names(info["ch_names"], remove_whitespace=True)246 calibration = dict(ch_names=ch_names, locs=locs, imb_cals=imb_cals)247 return calibration, count248 249 250def _matched_loc_idx(mag_loc, all_loc):251 return np.where(252 [253 np.allclose(mag_loc[-3:], loc[-3:]) and np.allclose(mag_loc[:3], loc[:3])254 for loc in all_loc255 ]256 )[0]257 258 259def _rotate_locs(locs, idxs, new_z):260 new_z = new_z / np.linalg.norm(new_z)261 old_z = locs[idxs[0]][-3:]262 old_z = old_z / np.linalg.norm(old_z)263 rot = _find_vector_rotation(old_z, new_z)264 for ci in idxs:265 this_trans = _loc_to_coil_trans(locs[ci])266 this_trans[:3, :3] = np.dot(rot, this_trans[:3, :3])267 locs[ci][:] = _coil_trans_to_loc(this_trans)268 np.testing.assert_allclose(locs[ci][-3:], new_z, atol=1e-4)269 270 271def _vector_angle(x, y):272 """Get the angle between two vectors in degrees."""273 return np.abs(274 np.arccos(275 np.clip(276 (x * y).sum(axis=-1)277 / (np.linalg.norm(x, axis=-1) * np.linalg.norm(y, axis=-1)),278 -1,279 1.0,280 )281 )282 )283 284 285def _adjust_mag_normals(info, data, origin, ext_order, *, angle_limit, err_limit):286 """Adjust coil normals using magnetometers and empty-room data."""287 # in principle we could allow using just mag or mag+grad, but MF uses288 # just mag so let's follow suit289 mag_scale = 100.0290 picks_use = pick_types(info, meg="mag", exclude="bads")291 picks_meg = pick_types(info, meg=True, exclude=())292 picks_mag_orig = pick_types(info, meg="mag", exclude="bads")293 info = pick_info(info, picks_use) # copy294 data = data[picks_use]295 cals = np.ones((len(data), 1))296 angles = np.zeros(len(cals))297 picks_mag = pick_types(info, meg="mag")298 data[picks_mag] *= mag_scale299 # Transform variables so we're only dealing with good mags300 exp = dict(int_order=0, ext_order=ext_order, origin=origin)301 all_coils = _prep_mf_coils(info, ignore_ref=True)302 S_tot = _trans_sss_basis(exp, all_coils, coil_scale=mag_scale)303 first_err = _data_err(data, S_tot, cals)304 count = 0305 # two passes: first do the worst, then do all in order306 zs = np.array([ch["loc"][-3:] for ch in info["chs"]])307 zs /= np.linalg.norm(zs, axis=-1, keepdims=True)308 orig_zs = zs.copy()309 match_idx = dict()310 locs = np.array([ch["loc"] for ch in info["chs"]])311 for pick in picks_mag:312 match_idx[pick] = _matched_loc_idx(locs[pick], locs)313 counts = defaultdict(lambda: 0)314 for ki, kind in enumerate(("worst first", "in order")):315 logger.info(f" Magnetometer normal adjustment ({kind}) ...")316 S_tot = _trans_sss_basis(exp, all_coils, coil_scale=mag_scale)317 for pick in picks_mag:318 err = _data_err(data, S_tot, cals, axis=1)319 320 # First pass: do worst; second pass: do all in order (up to 3x/sen)321 if ki == 0:322 order = list(np.argsort(err[picks_mag]))323 cal_idx = 0324 while len(order) > 0:325 cal_idx = picks_mag[order.pop(-1)]326 if counts[cal_idx] < 3:327 break328 if err[cal_idx] < 2.5:329 break # move on to second loop330 else:331 cal_idx = pick332 counts[cal_idx] += 1333 assert cal_idx in picks_mag334 count += 1335 old_z = zs[cal_idx].copy()336 objective = partial(337 _cal_sss_target,338 old_z=old_z,339 all_coils=all_coils,340 cal_idx=cal_idx,341 data=data,342 cals=cals,343 match_idx=match_idx,344 S_tot=S_tot,345 origin=origin,346 ext_order=ext_order,347 )348 349 # Figure out the additive term for z-component350 zs[cal_idx] = minimize(351 objective,352 old_z,353 bounds=[(-2, 2)] * 3,354 # BFGS is the default for minimize but COBYLA converges faster355 method="COBYLA",356 # Start with a small relative step because nominal geometry information357 # should be fairly accurate to begin with358 options=dict(rhobeg=1e-1),359 ).x360 361 # Do in-place adjustment to all_coils362 cals[cal_idx] = 1.0 / np.linalg.norm(zs[cal_idx])363 zs[cal_idx] *= cals[cal_idx]364 for idx in match_idx[cal_idx]:365 _rotate_coil(zs[cal_idx], old_z, all_coils, idx, inplace=True)366 367 # Recalculate S_tot, taking into account rotations368 S_tot = _trans_sss_basis(exp, all_coils)369 370 # Reprt results371 old_err = err[cal_idx]372 new_err = _data_err(data, S_tot, cals, idx=cal_idx)373 angles[cal_idx] = np.abs(374 np.rad2deg(_vector_angle(zs[cal_idx], orig_zs[cal_idx]))375 )376 ch_name = info["ch_names"][cal_idx]377 logger.debug(378 f" Optimization step {count:3d} ｜ "379 f"{ch_name} ({counts[cal_idx]}) ｜ "380 f"res {old_err:5.2f}→{new_err:5.2f}% ｜ "381 f"×{cals[cal_idx, 0]:0.3f} ｜ {angles[cal_idx]:0.2f}°"382 )383 last_err = _data_err(data, S_tot, cals)384 # Chunk is usable if all angles and errors are both small385 reason = list()386 max_angle = np.max(angles)387 if max_angle >= angle_limit:388 reason.append(f"max angle {max_angle:0.2f} >= {angle_limit:0.1f}°")389 each_err = _data_err(data, S_tot, cals, axis=-1)[picks_mag]390 n_bad = (each_err > err_limit).sum()391 if n_bad:392 bad_max = np.argmax(each_err)393 reason.append(394 f"{n_bad} residual{_pl(n_bad)} > {err_limit:0.1f}% "395 f"(max: {each_err[bad_max]:0.2f}% @ "396 f"{info['ch_names'][picks_mag[bad_max]]})"397 )398 reason = ", ".join(reason)399 if reason:400 reason = f" ({reason})"401 good = not bool(reason)402 assert np.allclose(np.linalg.norm(zs, axis=1), 1.0)403 logger.info(f" Fit mismatch {first_err:0.2f}→{last_err:0.2f}%")404 logger.info(f" Data segment {'' if good else 'un'}usable{reason}")405 # Reformat zs and cals to be the n_mags (including bads)406 assert zs.shape == (len(data), 3)407 assert cals.shape == (len(data), 1)408 imb_cals = np.ones(len(picks_meg))409 imb_cals[picks_mag_orig] = cals[:, 0]410 return zs, imb_cals, good411 412 413def _data_err(data, S_tot, cals, idx=None, axis=None):414 if idx is None:415 idx = slice(None)416 S_tot = S_tot / cals417 data_model = np.dot(np.dot(S_tot[idx], _col_norm_pinv(S_tot.copy())[0]), data)418 err = 100 * (419 np.linalg.norm(data_model - data[idx], axis=axis)420 / np.linalg.norm(data[idx], axis=axis)421 )422 return err423 424 425def _rotate_coil(new_z, old_z, all_coils, idx, inplace=False):426 """Adjust coils."""427 # Turn NX and NY to the plane determined by NZ428 old_z = old_z / np.linalg.norm(old_z)429 new_z = new_z / np.linalg.norm(new_z)430 rot = _find_vector_rotation(old_z, new_z) # additional coil rotation431 this_sl = all_coils[5][idx]432 this_rmag = np.dot(rot, all_coils[0][this_sl].T).T433 this_cosmag = np.dot(rot, all_coils[1][this_sl].T).T434 if inplace:435 all_coils[0][this_sl] = this_rmag436 all_coils[1][this_sl] = this_cosmag437 subset = (438 this_rmag,439 this_cosmag,440 np.zeros(this_rmag.shape[0], int),441 1,442 all_coils[4][[idx]],443 {0: this_sl},444 )445 return subset446 447 448def _cal_sss_target(449 new_z, old_z, all_coils, cal_idx, data, cals, S_tot, origin, ext_order, match_idx450):451 """Evaluate objective function for SSS-based magnetometer calibration."""452 cals[cal_idx] = 1.0 / np.linalg.norm(new_z)453 exp = dict(int_order=0, ext_order=ext_order, origin=origin)454 S_tot = S_tot.copy()455 # Rotate necessary coils properly and adjust correct element in c456 for idx in match_idx[cal_idx]:457 this_coil = _rotate_coil(new_z, old_z, all_coils, idx)458 # Replace correct row of S_tot with new value459 S_tot[idx] = _trans_sss_basis(exp, this_coil)460 # Get the GOF461 return _data_err(data, S_tot, cals, idx=cal_idx)462 463 464def _estimate_imbalance(info, data, cals, n_imbalance, origin, ext_order):465 """Estimate gradiometer imbalance parameters."""466 mag_scale = 100.0467 n_iterations = 3468 mag_picks = pick_types(info, meg="mag", exclude=())469 grad_picks = pick_types(info, meg="grad", exclude=())470 data = data.copy()471 data[mag_picks, :] *= mag_scale472 del mag_picks473 474 grad_imb = np.zeros((len(grad_picks), n_imbalance))475 exp = dict(origin=origin, int_order=0, ext_order=ext_order)476 all_coils = _prep_mf_coils(info, ignore_ref=True)477 grad_point_coils = _get_grad_point_coilsets(info, n_imbalance, ignore_ref=True)478 S_orig = _trans_sss_basis(exp, all_coils, coil_scale=mag_scale)479 S_orig /= cals[:, np.newaxis]480 # Compute point gradiometers for each grad channel481 this_cs = np.array([mag_scale], float)482 S_pt = np.array(483 [_trans_sss_basis(exp, coils, None, this_cs) for coils in grad_point_coils]484 )485 for k in range(n_iterations):486 S_tot = S_orig.copy()487 # In theory we could zero out the homogeneous components with:488 # S_tot[grad_picks, :3] = 0489 # But in practice it doesn't seem to matter490 S_recon = S_tot[grad_picks]491 492 # Add influence of point magnetometers493 S_tot[grad_picks, :] += np.einsum("ij,ijk->jk", grad_imb.T, S_pt)494 495 # Compute multipolar moments496 mm = np.dot(_col_norm_pinv(S_tot.copy())[0], data)497 498 # Use good channels to recalculate499 prev_imb = grad_imb.copy()500 data_recon = np.dot(S_recon, mm)501 assert S_pt.shape == (n_imbalance, len(grad_picks), S_tot.shape[1])502 khi_pts = (S_pt @ mm).transpose(1, 2, 0)503 assert khi_pts.shape == (len(grad_picks), data.shape[1], n_imbalance)504 residual = data[grad_picks] - data_recon505 assert residual.shape == (len(grad_picks), data.shape[1])506 d = (residual[:, np.newaxis, :] @ khi_pts)[:, 0]507 assert d.shape == (len(grad_picks), n_imbalance)508 dinv, _, _ = _reg_pinv(khi_pts.swapaxes(-1, -2) @ khi_pts, rcond=1e-6)509 assert dinv.shape == (len(grad_picks), n_imbalance, n_imbalance)510 grad_imb[:] = (d[:, np.newaxis] @ dinv)[:, 0]511 # This code is equivalent but hits a np.linalg.pinv bug on old NumPy:512 # grad_imb[:] = np.sum( # dot product across the time dim513 # np.linalg.pinv(khi_pts) * residual[:, np.newaxis], axis=-1)514 deltas = np.linalg.norm(grad_imb - prev_imb) / max(515 np.linalg.norm(grad_imb), np.linalg.norm(prev_imb)516 )517 logger.debug(518 f" Iteration {k + 1}/{n_iterations}: "519 f"max ∆ = {100 * deltas.max():7.3f}%"520 )521 imb = np.zeros((len(data), n_imbalance))522 imb[grad_picks] = grad_imb523 return imb524 525 526def read_fine_calibration(fname):527 """Read fine calibration information from a ``.dat`` file.528 529 The fine calibration typically includes improved sensor locations,530 calibration coefficients, and gradiometer imbalance information.531 532 Parameters533 ----------534 fname : path-like535 The filename.536 537 Returns538 -------539 calibration : dict540 Fine calibration information. Key-value pairs are:541 542 - ``ch_names``543 List of str of the channel names.544 - ``locs``545 Coil location and orientation parameters.546 - ``imb_cals``547 For magnetometers, the calibration coefficients.548 For gradiometers, one or three imbalance parameters.549 """550 # Read new sensor locations551 fname = _check_fname(fname, overwrite="read", must_exist=True)552 check_fname(fname, "cal", (".dat",))553 ch_names, locs, imb_cals = list(), list(), list()554 with open(fname) as fid:555 for line in fid:556 if line[0] in "#\n":557 continue558 vals = line.strip().split()559 if len(vals) not in [14, 16]:560 raise RuntimeError(561 "Error parsing fine calibration file, "562 "should have 14 or 16 entries per line "563 f"but found {len(vals)} on line:\n{line}"564 )565 # `vals` contains channel number566 ch_name = vals[0]567 if len(ch_name) in (3, 4): # heuristic for Neuromag fix568 try:569 ch_name = int(ch_name)570 except ValueError: # something other than e.g. 113 or 2642571 pass572 else:573 ch_name = f"MEG{int(ch_name):04}"574 # (x, y, z), x-norm 3-vec, y-norm 3-vec, z-norm 3-vec575 # and 1 or 3 imbalance terms576 ch_names.append(ch_name)577 locs.append(np.array(vals[1:13], float))578 imb_cals.append(np.array(vals[13:], float))579 locs = np.array(locs)580 return dict(ch_names=ch_names, locs=locs, imb_cals=imb_cals)581 582 583def write_fine_calibration(fname, calibration):584 """Write fine calibration information to a ``.dat`` file.585 586 Parameters587 ----------588 fname : path-like589 The filename to write out.590 calibration : dict591 Fine calibration information.592 """593 fname = _check_fname(fname, overwrite=True)594 check_fname(fname, "cal", (".dat",))595 keys = ("ch_names", "locs", "imb_cals")596 with open(fname, "wb") as cal_file:597 for ch_name, loc, imb_cal in zip(*(calibration[key] for key in keys)):598 cal_line = np.concatenate([loc, imb_cal]).round(6)599 cal_line = " ".join(f"{c:0.6f}" for c in cal_line)600 cal_file.write(f"{ch_name} {cal_line}\n".encode("ASCII"))601