diff --git a/doc/changes/dev/14034.newfeature.rst b/doc/changes/dev/14034.newfeature.rst new file mode 100644 index 00000000000..127a64545d6 --- /dev/null +++ b/doc/changes/dev/14034.newfeature.rst @@ -0,0 +1 @@ +Added a ``weighted`` option to :func:`mne.chpi.compute_head_pos` that fits all cHPI coils simultaneously with per-coil weights derived from their goodness of fit and inter-coil distance error, avoiding discontinuous jumps ("teleporting") in the estimated head position caused by coils switching in and out of the best-three subset (:gh:`11330`). The default will change from ``False`` to ``True`` in version 1.14; until then, leaving it unset emits a :class:`FutureWarning`. By `Eric Larson`_. diff --git a/examples/preprocessing/movement_detection.py b/examples/preprocessing/movement_detection.py index dd468feb464..8ffd2da282f 100644 --- a/examples/preprocessing/movement_detection.py +++ b/examples/preprocessing/movement_detection.py @@ -50,7 +50,7 @@ # Get cHPI time series and compute average chpi_locs = mne.chpi.extract_chpi_locs_ctf(raw) -head_pos = mne.chpi.compute_head_pos(raw.info, chpi_locs) +head_pos = mne.chpi.compute_head_pos(raw.info, chpi_locs, weighted=True) original_head_dev_t = mne.transforms.invert_transform(raw.info["dev_head_t"]) average_head_dev_t = mne.transforms.invert_transform( compute_average_dev_head_t(raw, head_pos) diff --git a/mne/chpi.py b/mne/chpi.py index 7ea3d1b407d..21122336c8f 100644 --- a/mne/chpi.py +++ b/mne/chpi.py @@ -4,9 +4,11 @@ 1. Drop coils whose GOF are below ``gof_limit``. If fewer than 3 coils remain, abandon fitting for the chunk. -2. Fit dev_head_t quaternion (using ``_fit_chpi_quat_subset``), - iteratively dropping coils (as long as 3 remain) to find the best GOF - (using ``_fit_chpi_quat``). +2. Fit dev_head_t quaternion. With ``weighted=False`` this uses + ``_fit_chpi_quat_subset``, iteratively dropping coils (as long as 3 remain) + to find the best GOF. With ``weighted=True`` all remaining coils are fit at + once, smoothly down-weighting each coil by its GOF and inter-coil distance error + so that coils do not switch in and out of the fit discontinuously (see :gh:`11330`). 3. If fewer than 3 coils meet the ``dist_limit`` criteria following projection of the fitted device coil locations into the head frame, abandon fitting for the chunk. @@ -522,7 +524,7 @@ def _magnetic_dipole_objective( out, u, s, one = _magnetic_dipole_delta(fwd, whitener, B, B2) if return_moment: one /= s - Q = np.dot(one, u.T) + Q = one @ u.T out = (out, Q) return out @@ -530,24 +532,24 @@ def _magnetic_dipole_objective( @jit() def _magnetic_dipole_delta(fwd, whitener, B, B2): # Here we use .T to get whitener to Fortran order, which speeds things up - fwd = np.dot(fwd, whitener.T) + fwd = fwd @ whitener.T u, s, v = np.linalg.svd(fwd, full_matrices=False) - one = np.dot(v, B) - Bm2 = np.dot(one, one) + one = v @ B + Bm2 = one @ one return B2 - Bm2, u, s, one def _magnetic_dipole_delta_multi(whitened_fwd_svd, B, B2): # Here we use .T to get whitener to Fortran order, which speeds things up - one = np.matmul(whitened_fwd_svd, B) + one = whitened_fwd_svd @ B Bm2 = np.sum(one * one, axis=1) return B2 - Bm2 def _fit_magnetic_dipole(B_orig, x0, too_close, whitener, coils, guesses): """Fit a single bit of data (x0 = pos).""" - B = np.dot(whitener, B_orig) - B2 = np.dot(B, B) + B = whitener @ B_orig + B2 = B @ B objective = partial( _magnetic_dipole_objective, B=B, @@ -570,25 +572,32 @@ def _fit_magnetic_dipole(B_orig, x0, too_close, whitener, coils, guesses): @jit() -def _chpi_objective(x, coil_dev_rrs, coil_head_rrs): +def _chpi_objective(x, coil_dev_rrs, coil_head_rrs, weights): """Compute objective function.""" - d = np.dot(coil_dev_rrs, quat_to_rot(x[:3]).T) + d = coil_dev_rrs @ quat_to_rot(x[:3]).T d += x[3:] d -= coil_head_rrs d *= d - return d.sum() + return d.sum(axis=1) @ weights # sum over coils, weighted -def _fit_chpi_quat(coil_dev_rrs, coil_head_rrs, *, quat=None): - """Fit rotation and translation (quaternion) parameters for cHPI coils.""" - denom = np.linalg.norm(coil_head_rrs - np.mean(coil_head_rrs, axis=0)) - denom *= denom - # We could try to solve it the analytic way: - # TODO someday we could choose to weight these points by their goodness - # of fit somehow, see also https://github.com/mne-tools/mne-python/issues/11330 +def _fit_chpi_quat(coil_dev_rrs, coil_head_rrs, *, weights=None, quat=None): + """Fit rotation and translation (quaternion) parameters for cHPI coils. + + ``weights`` optionally down-weights individual coils (e.g. by goodness of fit + and inter-coil distance error) so that coils can smoothly enter or leave the + fit rather than switching in and out discontinuously; see :gh:`11330`. + """ + if weights is None: + weights = np.ones(len(coil_head_rrs)) + # The GOF ratio is invariant to the overall weight scale, so we can pass the + # (possibly zero-containing) weights through unnormalized to keep the + # unweighted case bit-for-bit identical to before. + mu = (weights @ coil_head_rrs) / weights.sum() + denom = ((coil_head_rrs - mu) ** 2).sum(axis=1) @ weights if quat is None: - quat = _fit_matched_points(coil_dev_rrs, coil_head_rrs)[0] - gof = 1.0 - _chpi_objective(quat, coil_dev_rrs, coil_head_rrs) / denom + quat = _fit_matched_points(coil_dev_rrs, coil_head_rrs, weights=weights)[0] + gof = 1.0 - _chpi_objective(quat, coil_dev_rrs, coil_head_rrs, weights) / denom return quat, gof @@ -954,7 +963,14 @@ def _check_chpi_param(chpi_, name): @verbose def compute_head_pos( - info, chpi_locs, dist_limit=0.005, gof_limit=0.98, adjust_dig=False, verbose=None + info, + chpi_locs, + dist_limit=0.005, + gof_limit=0.98, + adjust_dig=False, + *, + weighted=None, + verbose=None, ): """Compute time-varying head positions. @@ -969,6 +985,15 @@ def compute_head_pos( gof_limit : float Minimum goodness of fit to accept for each coil. %(adjust_dig_chpi)s + weighted : bool + If ``True``, fit all coils that pass the ``gof_limit`` and ``dist_limit`` + criteria simultaneously, weighting each coil by its goodness of fit and its + inter-coil distance error. If ``False``, subselect the three coils that yield + the best fit. Weighting avoids discontinuous jumps in the estimated head + position caused by coils switching in and out of the fit (see :gh:`11330`). + The default (False) will change to True in 1.14. + + .. versionadded:: 1.13 %(verbose)s Returns @@ -990,8 +1015,17 @@ def compute_head_pos( """ _check_chpi_param(chpi_locs, "chpi_locs") _validate_type(info, Info, "info") + if weighted is None: + warn( + "The default for weighted will change from False to True in 1.14, set it " + "explicitly to avoid this warning. Using False.", + FutureWarning, + ) + weighted = False hpi_dig_head_rrs = _get_hpi_initial_fit(info, adjust=adjust_dig, verbose="error") n_coils = len(hpi_dig_head_rrs) + # reference inter-coil distances (rigid, so invariant to the dev_head_t we fit) + hpi_coil_dists = cdist(hpi_dig_head_rrs, hpi_dig_head_rrs) coil_dev_rrs = apply_trans(invert_transform(info["dev_head_t"]), hpi_dig_head_rrs) dev_head_t = info["dev_head_t"]["trans"] pos_0 = dev_head_t[:3, 3] @@ -1021,12 +1055,38 @@ def compute_head_pos( # # 2. Fit the head translation and rotation params (minimize error - # between coil positions and the head coil digitization - # positions) iteratively using different sets of coils. + # between coil positions and the head coil digitization positions). # - this_quat, g, use_idx = _fit_chpi_quat_subset( - this_coil_dev_rrs, hpi_dig_head_rrs, use_idx - ) + if weighted: + # Fit all good coils at once, smoothly down-weighting each coil by its + # GOF and its inter-coil distance error, so that coils enter and leave + # the fit continuously rather than switching in and out (gh-11330). + these_dists = cdist(this_coil_dev_rrs, this_coil_dev_rrs) + dist_err = np.abs(hpi_coil_dists - these_dists).sum(axis=1) / max( + n_coils - 1, 1 + ) + w_gof = np.clip( + (g_coils - gof_limit) / max(1.0 - gof_limit, 1e-12), 0.0, 1.0 + ) + w_dist = np.clip(1.0 - dist_err / dist_limit, 0.0, 1.0) + weights = w_gof * w_dist + use_idx = np.where(weights > 0)[0] + if len(use_idx) < 3: + gofs = ", ".join(f"{g:0.2f}" for g in g_coils) + warn( + f"{_time_prefix(fit_time)}{len(use_idx)}/{n_coils} " + "usable HPI coils, cannot determine the transformation " + f"({gofs} GOF)!" + ) + continue + this_quat, g = _fit_chpi_quat( + this_coil_dev_rrs, hpi_dig_head_rrs, weights=weights + ) + else: + # iteratively drop coils (keeping the best-GOF subset of >= 3) + this_quat, g, use_idx = _fit_chpi_quat_subset( + this_coil_dev_rrs, hpi_dig_head_rrs, use_idx + ) # # 3. Stop if < 3 good @@ -1293,7 +1353,6 @@ def _compute_chpi_amp_or_snr( # note that mean residual is a scalar (same for all HPI freqs) but # is returned as a (tiled) vector (again, because Numba) so that's # why below we take amps_or_snrs[0, 2] instead of [:, 2] - ch_types = raw.get_channel_types() if "mag" in ch_types: sin_fits["mag_snr"][mi] = amps_or_snrs[:, 0] # SNR sin_fits["mag_power"][mi] = amps_or_snrs[:, 1] # mean power @@ -1389,7 +1448,7 @@ def compute_chpi_locs( f"(1 cm grid in a {R * 100:.1f} cm sphere)" ) fwd = _magnetic_dipole_field_vec(guesses, meg_coils, too_close) - fwd = np.dot(fwd, whitener.T) + fwd = fwd @ whitener.T fwd = _reshape_view(fwd, (guesses.shape[0], 3, -1)) fwd = np.linalg.svd(fwd, full_matrices=False)[2] guesses = dict(rr=guesses, whitened_fwd_svd=fwd) @@ -1556,7 +1615,7 @@ def filter_chpi( msg += f" and {len(hpi['line_freqs'])} line harmonic" msg += f" frequencies from {len(meg_picks)} MEG channels" - recon = np.dot(hpi["model"][:, :n_remove], hpi["inv_model"][:n_remove]).T + recon = (hpi["model"][:, :n_remove] @ hpi["inv_model"][:n_remove]).T logger.info(msg) chunks = list() # the chunks to subtract last_endpt = 0 @@ -1572,13 +1631,13 @@ def filter_chpi( else: # first or last window model = hpi["model"][:this_len] inv_model = np.linalg.pinv(model) - this_recon = np.dot(model[:, :n_remove], inv_model[:n_remove]).T + this_recon = (model[:, :n_remove] @ inv_model[:n_remove]).T this_data = raw._data[meg_picks, time_sl] subt_pt = min(midpt + n_step, n_times) if last_endpt != subt_pt: fit_left_edge = left_edge - time_sl.start + hpi["n_window"] // 2 fit_sl = slice(fit_left_edge, fit_left_edge + (subt_pt - last_endpt)) - chunks.append((subt_pt, np.dot(this_data, this_recon[:, fit_sl]))) + chunks.append((subt_pt, this_data @ this_recon[:, fit_sl])) last_endpt = subt_pt # Consume (trailing) chunks that are now safe to remove because @@ -1948,7 +2007,7 @@ def refit_hpi( # entry. At some point we should try recording data where multiple fits are stored # (maybe there actually aren't any...) info["hpi_meas"][-1] = meas - info["hpi_results"][-1] = result + info["hpi_results"][-1] = results return info diff --git a/mne/simulation/tests/test_raw.py b/mne/simulation/tests/test_raw.py index 9babdb7a777..3dcef02610b 100644 --- a/mne/simulation/tests/test_raw.py +++ b/mne/simulation/tests/test_raw.py @@ -572,7 +572,7 @@ def test_simulate_raw_chpi(): # test localization based on cHPI information chpi_amplitudes = compute_chpi_amplitudes(raw, t_step_min=10.0) coil_locs = compute_chpi_locs(raw.info, chpi_amplitudes) - quats_sim = compute_head_pos(raw_chpi.info, coil_locs) + quats_sim = compute_head_pos(raw_chpi.info, coil_locs, weighted=False) quats = read_head_pos(pos_fname) _assert_quats( quats, quats_sim, dist_tol=5e-3, angle_tol=3.5, vel_atol=0.03 diff --git a/mne/tests/test_chpi.py b/mne/tests/test_chpi.py index 925ad14018a..dcb99b2b1f3 100644 --- a/mne/tests/test_chpi.py +++ b/mne/tests/test_chpi.py @@ -15,6 +15,7 @@ from mne.chpi import ( _chpi_locs_to_times_dig, _compute_good_distances, + _fit_chpi_quat, _get_hpi_initial_fit, _setup_ext_proj, compute_chpi_amplitudes, @@ -45,6 +46,7 @@ from mne.transforms import ( _angle_between_quats, angle_distance_between_rigid, + quat_to_rot, rot_to_quat, ) from mne.utils import ( @@ -144,7 +146,7 @@ def test_read_write_head_pos(tmp_path): """Test reading and writing head position quaternion parameters.""" temp_name = tmp_path / "temp.pos" # This isn't a 100% valid quat matrix but it should be okay for tests - head_pos_rand = np.random.RandomState(0).randn(20, 10) + head_pos_rand = np.random.default_rng(0).standard_normal((20, 10)) # This one is valid head_pos_read = read_head_pos(pos_fname) for head_pos_orig in (head_pos_rand, head_pos_read): @@ -292,6 +294,7 @@ def _calculate_chpi_positions( dist_limit=0.005, gof_limit=0.98, ext_order=1, + weighted=False, verbose=None, ): chpi_amplitudes = compute_chpi_amplitudes( @@ -309,7 +312,12 @@ def _calculate_chpi_positions( verbose=verbose, ) head_pos = compute_head_pos( - raw.info, chpi_locs, dist_limit=dist_limit, gof_limit=gof_limit, verbose=verbose + raw.info, + chpi_locs, + dist_limit=dist_limit, + gof_limit=gof_limit, + weighted=weighted, + verbose=verbose, ) return head_pos @@ -472,8 +480,10 @@ def test_initial_fit_redo(): angles = np.rad2deg(np.arccos(np.abs(np.sum(coil_ori * py_ori, axis=1)))) assert_array_less(angles, 20) - # check resulting dev_head_t - head_pos = compute_head_pos(raw.info, chpi_locs) + # check resulting dev_head_t (also the one place we exercise the deprecated + # default value of ``weighted``, see gh-11330) + with pytest.warns(FutureWarning, match="weighted will change"): + head_pos = compute_head_pos(raw.info, chpi_locs) assert head_pos.shape == (1, 10) nm_pos = raw.info["dev_head_t"]["trans"] dist = 1000 * np.linalg.norm(nm_pos[:3, 3] - head_pos[0, 4:7]) @@ -486,6 +496,40 @@ def test_initial_fit_redo(): assert_allclose(gof, 0.9999, atol=1e-4) +def test_fit_chpi_quat_weighted(): + """Test weighted cHPI quaternion fitting (gh-11330).""" + rng = np.random.default_rng(0) + head_rrs = rng.standard_normal((5, 3)) * 0.05 + quat = np.array([0.05, -0.03, 0.02, 0.01, -0.02, 0.03]) + rot = quat_to_rot(quat[:3]) + # device coil positions such that ``rot @ dev + trans == head`` + dev_rrs = (head_rrs - quat[3:]) @ rot + + # uniform weights reproduce the unweighted fit exactly + quat_none, gof_none = _fit_chpi_quat(dev_rrs, head_rrs) + quat_ones, gof_ones = _fit_chpi_quat(dev_rrs, head_rrs, weights=np.ones(5)) + assert_allclose(quat_none, quat_ones, atol=1e-12) + assert_allclose(gof_none, gof_ones, atol=1e-12) + assert_allclose(quat_none, quat, atol=1e-6) + assert_allclose(gof_none, 1.0, atol=1e-10) + + # corrupt one coil: zero-weighting it recovers the exact transform, whereas + # including it unweighted spoils the fit + dev_bad = dev_rrs.copy() + dev_bad[4] += 0.02 + quat_zero = _fit_chpi_quat(dev_bad, head_rrs, weights=np.array([1.0, 1, 1, 1, 0])) + assert_allclose(quat_zero[0], quat, atol=1e-6) + assert np.linalg.norm(_fit_chpi_quat(dev_bad, head_rrs)[0][3:] - quat[3:]) > 1e-3 + + # smoothly down-weighting the bad coil improves the fit + errs = list() + for coil_weight in (1.0, 0.5, 0.1, 0.0): + weights = np.array([1.0, 1.0, 1.0, 1.0, coil_weight]) + quat_fit = _fit_chpi_quat(dev_bad, head_rrs, weights=weights)[0] + errs.append(np.linalg.norm(quat_fit[3:] - quat[3:])) + assert errs[0] > errs[1] > errs[2] > errs[3] + + @testing.requires_testing_data def test_calculate_head_pos_chpi_on_chpi5_in_one_second_steps(): """Comparing estimated cHPI positions with MF results (one second).""" @@ -497,12 +541,18 @@ def test_calculate_head_pos_chpi_on_chpi5_in_one_second_steps(): # maxfilter estimates a wrong head position for interval 16: 41.-42. s raw = _decimate_chpi(raw.crop(0.0, 10.0).load_data(), decim=8) # needs no interpolation, because maxfilter pos files comes with 1 s steps - py_quats = _calculate_chpi_positions( - raw, t_step_min=1.0, t_step_max=1.0, t_window=1.0, verbose="debug" - ) + chpi_amplitudes = compute_chpi_amplitudes(raw, t_step_min=1.0, t_window=1.0) + chpi_locs = compute_chpi_locs(raw.info, chpi_amplitudes, t_step_max=1.0) + py_quats = compute_head_pos(raw.info, chpi_locs, weighted=False, verbose="debug") _assert_quats( py_quats, mf_quats, dist_tol=0.002, angle_tol=1.2, vel_atol=3e-3 ) # 3 mm/s + # the weighted variant (gh-11330) fits all good coils at once rather than + # collapsing to a 3-coil subset, and should still agree with MaxFilter as well + py_quats_w = compute_head_pos(raw.info, chpi_locs, weighted=True) + _assert_quats(py_quats_w, mf_quats, dist_tol=0.002, angle_tol=1.2, vel_atol=3e-3) + # weighting actually changes the fit (uses all coils, not the best-3 subset) + assert not np.array_equal(py_quats, py_quats_w) @pytest.mark.slowtest @@ -672,7 +722,7 @@ def test_calculate_chpi_coil_locs_artemis(): coil_amplitudes["slopes"].fill(np.nan) chpi_locs = compute_chpi_locs(raw.info, coil_amplitudes) assert chpi_locs["rrs"].shape == (0, 3, 3) - pos = compute_head_pos(raw.info, chpi_locs) + pos = compute_head_pos(raw.info, chpi_locs, weighted=False) assert pos.shape == (0, 10) @@ -777,7 +827,7 @@ def test_calculate_head_pos_ctf(tmp_path): """Test extracting of cHPI positions from CTF data.""" raw = read_raw_ctf(ctf_chpi_fname) chpi_locs = extract_chpi_locs_ctf(raw) - quats = compute_head_pos(raw.info, chpi_locs) + quats = compute_head_pos(raw.info, chpi_locs, weighted=False) mc_quats = read_head_pos(ctf_chpi_pos_fname) mc_quats[:, 9] /= 10000 # had old factor in there twice somehow... _assert_quats( @@ -801,11 +851,11 @@ def test_calculate_head_pos_ctf(tmp_path): with pytest.warns(RuntimeWarning, match="is poor"): head_rrs_2 = _get_hpi_initial_fit(raw_read.info, verbose="debug") assert_allclose(head_rrs, head_rrs_2, atol=1e-5) - quats_2 = compute_head_pos(raw_read.info, chpi_locs) + quats_2 = compute_head_pos(raw_read.info, chpi_locs, weighted=False) _assert_quats(quats, quats_2, dist_tol=1e-5, angle_tol=0.1) chpi_locs_2 = extract_chpi_locs_ctf(raw_read) assert_allclose(chpi_locs["rrs"], chpi_locs_2["rrs"], atol=1e-5) - quats_3 = compute_head_pos(raw_read.info, chpi_locs_2) + quats_3 = compute_head_pos(raw_read.info, chpi_locs_2, weighted=False) _assert_quats(quats, quats_3, dist_tol=1e-5, angle_tol=0.1) @@ -818,7 +868,7 @@ def test_calculate_head_pos_kit(): assert chpi_locs["rrs"].shape == (2, 5, 3) assert_array_less(chpi_locs["gofs"], 1.0) assert_array_less(0.98, chpi_locs["gofs"]) - quats = compute_head_pos(raw.info, chpi_locs) + quats = compute_head_pos(raw.info, chpi_locs, weighted=False) assert quats.shape == (2, 10) # plotting works plot_head_positions(quats, info=raw.info) @@ -832,7 +882,7 @@ def test_calculate_head_pos_kit(): with pytest.raises(RuntimeError, match="not find appropriate"): extract_chpi_locs_kit(raw_berlin, "STI 014") with pytest.raises(RuntimeError, match="no initial cHPI"): - compute_head_pos(raw_berlin.info, chpi_locs) + compute_head_pos(raw_berlin.info, chpi_locs, weighted=False) @testing.requires_testing_data diff --git a/tutorials/preprocessing/59_head_positions.py b/tutorials/preprocessing/59_head_positions.py index 4085c4792b6..dc876479b68 100644 --- a/tutorials/preprocessing/59_head_positions.py +++ b/tutorials/preprocessing/59_head_positions.py @@ -63,7 +63,7 @@ # %% # Lastly, compute head positions from the coil locations: -head_pos = mne.chpi.compute_head_pos(raw.info, chpi_locs, verbose=True) +head_pos = mne.chpi.compute_head_pos(raw.info, chpi_locs, weighted=True, verbose=True) # %% # Note that these can then be written to disk or read from disk with