See how we can speed up dipole fitting

fwd_opt.py

import numpy as np
import mne
import time

data_path = mne.datasets.sample.data_path()
subjects_dir = data_path / 'subjects'
n_dips = 10
radius = 0.05
rr = np.random.default_rng(0).normal(size=(n_dips, 3))
rr /= np.linalg.norm(rr, axis=1)[:, np.newaxis] / radius * 0.9
nn = np.cross([1, 0, 0], rr)
nn /= np.linalg.norm(nn, axis=1)[:, np.newaxis]
assert nn.shape == rr.shape
bem = mne.read_bem_solution(data_path / "subjects" / "sample" / "bem" / "sample-5120-5120-5120-bem-sol.fif")
assert bem["surfs"][-1]["id"] == mne.io.constants.FIFF.FIFFV_BEM_SURF_ID_BRAIN
rr += bem["surfs"][-1]["rr"].mean(axis=0)  # center the points around the BEM center
evoked = mne.read_evokeds(data_path / "MEG" / "sample" / "sample_audvis-ave.fif", condition=0)
evoked.crop(tmin=0.1, tmax=0.1)

# %%
# _CheckInside
t0 = time.time()
ci = mne.surface._CheckInside(bem["surfs"][-1])
assert all(ci(rr))
print(f"Single _CheckInside: {time.time() - t0:0.1f}s")
t0 = time.time()
for r in rr:
    ci = mne.surface._CheckInside(bem["surfs"][-1])
    assert all(ci(r[np.newaxis]))
print(f"Loop _CheckInside:   {time.time() - t0:0.1}s")
print()

# %%
# make_forward_solution
trans = mne.read_trans(data_path / "MEG" / "sample" / "sample_audvis_raw-trans.fif")
t0 = time.time()
src = mne.setup_volume_source_space("sample", pos=dict(rr=rr, nn=nn))
fwd = mne.make_forward_solution(evoked.info, trans, src, bem)
print(f"Single make_forward_solution: {time.time() - t0:0.1f}s")
t0 = time.time()
fwds = list()
for ri in range(n_dips):
    this_src = mne.setup_volume_source_space("sample", pos=dict(rr=rr[[ri]], nn=nn[[ri]]))
    fwds.append(mne.make_forward_solution(evoked.info, trans, this_src, bem))
print(f"Loop make_forward_solution:   {time.time() - t0:0.1f}s")
print()
check = np.concatenate([f["sol"]["data"] for f in fwds], axis=1)
np.testing.assert_allclose(check, fwd["sol"]["data"])

# %%
# make_forward_dipole
t0 = time.time()
mri_head_t = mne.transforms.invert_transform(trans)
rr_head = mne.transforms.apply_trans(mri_head_t, rr)
nn_head = mne.transforms.apply_trans(mri_head_t, nn, move=False)
assert mri_head_t["from"] == mne.io.constants.FIFF.FIFFV_COORD_MRI
dip = mne.Dipole(np.arange(n_dips), rr_head, np.ones(n_dips), nn_head, np.ones(n_dips))
fwd_dip, _ = mne.make_forward_dipole(dip, bem, evoked.info, trans)
print(f"Single make_forward_dipole: {time.time() - t0:0.1f}s")
fwd_fixed = mne.convert_forward_solution(fwd, force_fixed=True)
np.testing.assert_allclose(fwd_dip["source_rr"], fwd_fixed["source_rr"])
np.testing.assert_allclose(fwd_dip["source_nn"], fwd_fixed["source_nn"])
np.testing.assert_allclose(fwd_dip["sol"]["data"], fwd_fixed["sol"]["data"], rtol=1e-6)
t0 = time.time()
fwd_dips = list()
for ri in range(n_dips):
    this_dip = mne.Dipole(np.arange(ri, ri + 1), rr_head[[ri]], np.ones(1), nn_head[[ri]], np.ones(1))
    fwd_dips.append(mne.make_forward_dipole(this_dip, bem, evoked.info, trans)[0])
print(f"Loop make_forward_dipole:   {time.time() - t0:0.1f}s")
check = np.concatenate([f["sol"]["data"] for f in fwd_dips], axis=1)
np.testing.assert_allclose(check, fwd_fixed["sol"]["data"], rtol=1e-6)