ENH: Speed up forward computations for iterative fitting (#13407)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel McCloy <[email protected]>

Author: 3 people

doc/changes/dev/13407.bugfix.rst

@@ -0,0 +1 @@
+Fix bug with :func:`mne.make_forward_solution` where sources were not checked to make sure they're inside the inner skull for spherical BEMs, by `Eric Larson`_.

examples/simulation/simulate_raw_data.py

@@ -31,8 +31,6 @@
     simulate_sparse_stc,
 )
 
-print(__doc__)
-
 data_path = sample.data_path()
 meg_path = data_path / "MEG" / "sample"
 raw_fname = meg_path / "sample_audvis_raw.fif"

examples/simulation/simulated_raw_data_using_subject_anatomy.py

@@ -26,8 +26,6 @@
 import mne
 from mne.datasets import sample
 
-print(__doc__)
-
 # %%
 # In this example, raw data will be simulated for the sample subject, so its
 # information needs to be loaded. This step will download the data if it not

mne/_ola.py

@@ -309,7 +309,7 @@ def __init__(
         # Create our window boundaries
         window_name = window if isinstance(window, str) else "custom"
         self._window = get_window(
-            window, self._n_samples, fftbins=(self._n_samples - 1) % 2
+            window, self._n_samples, fftbins=bool((self._n_samples - 1) % 2)
         )
         self._window /= _check_cola(
             self._window, self._n_samples, self._step, window_name, tol=tol

mne/chpi.py

@@ -38,6 +38,7 @@
     pick_types,
 )
 from ._fiff.proj import Projection, setup_proj
+from .bem import ConductorModel
 from .channels.channels import _get_meg_system
 from .cov import compute_whitener, make_ad_hoc_cov
 from .dipole import _make_guesses
@@ -1343,9 +1344,8 @@ def compute_chpi_locs(
 
     # Make some location guesses (1 cm grid)
     R = np.linalg.norm(meg_coils[0], axis=1).min()
-    guesses = _make_guesses(
-        dict(R=R, r0=np.zeros(3)), 0.01, 0.0, 0.005, verbose=safe_false
-    )[0]["rr"]
+    sphere = ConductorModel(layers=[dict(rad=R)], r0=np.zeros(3), is_sphere=True)
+    guesses = _make_guesses(sphere, 0.01, 0.0, 0.005, verbose=safe_false)[0]["rr"]
     logger.info(
         f"Computing {len(guesses)} HPI location guesses "
         f"(1 cm grid in a {R * 100:.1f} cm sphere)"

mne/conftest.py

@@ -729,14 +729,16 @@ def _evoked_cov_sphere(_evoked):
     evoked.pick(evoked.ch_names[::4])
     assert len(evoked.ch_names) == 77
     cov = mne.read_cov(fname_cov)
-    sphere = mne.make_sphere_model("auto", "auto", evoked.info)
+    sphere = mne.make_sphere_model(
+        (0.0, 0.0, 0.04), 0.1, relative_radii=(0.995, 0.997, 0.998, 1.0)
+    )
     return evoked, cov, sphere
 
 
 @pytest.fixture(scope="session")
 def _fwd_surf(_evoked_cov_sphere):
     """Compute a forward for a surface source space."""
-    evoked, cov, sphere = _evoked_cov_sphere
+    evoked, _, sphere = _evoked_cov_sphere
     src_surf = mne.read_source_spaces(fname_src)
     return mne.make_forward_solution(
         evoked.info, fname_trans, src_surf, sphere, mindist=5.0
@@ -747,7 +749,7 @@ def _fwd_surf(_evoked_cov_sphere):
 def _fwd_subvolume(_evoked_cov_sphere):
     """Compute a forward for a surface source space."""
     pytest.importorskip("nibabel")
-    evoked, cov, sphere = _evoked_cov_sphere
+    evoked, _, sphere = _evoked_cov_sphere
     volume_labels = ["Left-Cerebellum-Cortex", "right-Cerebellum-Cortex"]
     with pytest.raises(ValueError, match=r"Did you mean one of \['Right-Cere"):
         mne.setup_volume_source_space(
@@ -761,9 +763,12 @@ def _fwd_subvolume(_evoked_cov_sphere):
         subjects_dir=subjects_dir,
         add_interpolator=False,
     )
-    return mne.make_forward_solution(
-        evoked.info, fname_trans, src_vol, sphere, mindist=5.0
+    fwd = mne.make_forward_solution(
+        evoked.info, fname_trans, src_vol, sphere, mindist=1.0
     )
+    nsrc = sum(s["nuse"] for s in src_vol)
+    assert fwd["nsource"] == nsrc
+    return fwd
 
 
 @pytest.fixture

mne/dipole.py

@@ -17,7 +17,7 @@
 from ._fiff.pick import pick_types
 from ._fiff.proj import _needs_eeg_average_ref_proj, make_projector
 from ._freesurfer import _get_aseg, head_to_mni, head_to_mri, read_freesurfer_lut
-from .bem import _bem_find_surface, _bem_surf_name, _fit_sphere
+from .bem import ConductorModel, _bem_find_surface, _bem_surf_name, _fit_sphere
 from .cov import _ensure_cov, compute_whitener
 from .evoked import _aspect_rev, _read_evoked, _write_evokeds
 from .fixes import _safe_svd
@@ -960,9 +960,8 @@ def _make_guesses(surf, grid, exclude, mindist, n_jobs=None, verbose=None):
         )
     else:
         logger.info(
-            "Making a spherical guess space with radius {:7.1f} mm...".format(
-                1000 * surf["R"]
-            )
+            f"Making a spherical guess space with radius {1000 * surf.radius:7.1f} "
+            "mm..."
         )
     logger.info("Filtering (grid = %6.f mm)..." % (1000 * grid))
     src = _make_volume_source_space(
@@ -1317,7 +1316,7 @@ def _fit_dipole(
         constraint = partial(
             _sphere_constraint,
             r0=fwd_data["inner_skull"]["r0"],
-            R_adj=fwd_data["inner_skull"]["R"] - min_dist_to_inner_skull,
+            R_adj=fwd_data["inner_skull"].radius - min_dist_to_inner_skull,
         )
 
     # Find a good starting point (find_best_guess in C)
@@ -1600,13 +1599,14 @@ def fit_dipole(
                 raise RuntimeError(
                     "No MEG channels found, but MEG-only sphere model used"
                 )
-            R = np.min(np.sqrt(np.sum(R * R, axis=1)))  # use dist to sensors
-            kind = "max_rad"
+            R = np.min(np.linalg.norm(R, axis=1))
+            kind = "min_rad"
         logger.info(
             f"Sphere model      : origin at ({1000 * r0[0]: 7.2f} {1000 * r0[1]: 7.2f} "
             f"{1000 * r0[2]: 7.2f}) mm, {kind} = {R:6.1f} mm"
         )
-        inner_skull = dict(R=R, r0=r0)  # NB sphere model defined in head frame
+        # NB sphere model defined in head frame
+        inner_skull = ConductorModel(layers=[dict(rad=R)], r0=r0, is_sphere=True)
         del R, r0
 
     # Deal with DipoleFixed cases here
@@ -1710,7 +1710,9 @@ def fit_dipole(
             check = _surface_constraint(pos, inner_skull, min_dist_to_inner_skull)
         else:
             check = _sphere_constraint(
-                pos, inner_skull["r0"], R_adj=inner_skull["R"] - min_dist_to_inner_skull
+                pos,
+                inner_skull["r0"],
+                R_adj=inner_skull.radius - min_dist_to_inner_skull,
             )
         if check <= 0:
             raise ValueError(
@@ -1720,7 +1722,7 @@ def fit_dipole(
 
     # C code computes guesses w/sphere model for speed, don't bother here
     fwd_data = _prep_field_computation(
-        guess_src["rr"], sensors=sensors, bem=bem, n_jobs=n_jobs, verbose=safe_false
+        sensors=sensors, bem=bem, n_jobs=n_jobs, verbose=safe_false
     )
     fwd_data["inner_skull"] = inner_skull
     guess_fwd, guess_fwd_orig, guess_fwd_scales = _dipole_forwards(

mne/forward/_compute_forward.py

@@ -709,7 +709,7 @@ def _compute_mdfv(rrs, rmags, cosmags, ws, bins, too_close):
 
 
 @verbose
-def _prep_field_computation(rr, *, sensors, bem, n_jobs, verbose=None):
+def _prep_field_computation(*, sensors, bem, n_jobs, verbose=None):
     """Precompute and store some things that are used for both MEG and EEG.
 
     Calculation includes multiplication factors, coordinate transforms,
@@ -840,7 +840,7 @@ def _compute_forwards(rr, *, bem, sensors, n_jobs, verbose=None):
         # This modifies "sensors" in place, so let's copy it in case the calling
         # function needs to reuse it (e.g., in simulate_raw.py)
         sensors = deepcopy(sensors)
-        fwd_data = _prep_field_computation(rr, sensors=sensors, bem=bem, n_jobs=n_jobs)
+        fwd_data = _prep_field_computation(sensors=sensors, bem=bem, n_jobs=n_jobs)
         Bs = _compute_forwards_meeg(
             rr, sensors=sensors, fwd_data=fwd_data, n_jobs=n_jobs
         )

mne/forward/_make_forward.py

@@ -21,12 +21,13 @@
 from ..bem import ConductorModel, _bem_find_surface, read_bem_solution
 from ..source_estimate import VolSourceEstimate
 from ..source_space._source_space import (
+    SourceSpaces,
     _complete_vol_src,
     _ensure_src,
     _filter_source_spaces,
     _make_discrete_source_space,
 )
-from ..surface import _CheckInside, _normalize_vectors
+from ..surface import _CheckInside, _CheckInsideSphere, _normalize_vectors
 from ..transforms import (
     Transform,
     _coord_frame_name,
@@ -35,10 +36,13 @@
     _print_coord_trans,
     apply_trans,
     invert_transform,
-    transform_surface_to,
 )
 from ..utils import _check_fname, _pl, _validate_type, logger, verbose, warn
-from ._compute_forward import _compute_forwards
+from ._compute_forward import (
+    _compute_forwards,
+    _compute_forwards_meeg,
+    _prep_field_computation,
+)
 from .forward import _FWD_ORDER, Forward, _merge_fwds, convert_forward_solution
 
 _accuracy_dict = dict(
@@ -459,7 +463,7 @@ def _prepare_for_forward(
     # let's make a copy in case we modify something
     src = _ensure_src(src).copy()
     nsource = sum(s["nuse"] for s in src)
-    if nsource == 0:
+    if len(src) and nsource == 0:
         raise RuntimeError(
             "No sources are active in these source spaces. "
             '"do_all" option should be used.'
@@ -517,11 +521,12 @@ def _prepare_for_forward(
 
     # Transform the source spaces into the appropriate coordinates
     # (will either be HEAD or MRI)
-    for s in src:
-        transform_surface_to(s, "head", mri_head_t)
-    logger.info(
-        f"Source spaces are now in {_coord_frame_name(s['coord_frame'])} coordinates."
-    )
+    src._transform_to("head", mri_head_t)
+    if len(src):
+        logger.info(
+            f"Source spaces are now in {_coord_frame_name(src[0]['coord_frame'])} "
+            "coordinates."
+        )
 
     # Prepare the BEM model
     eegnames = sensors.get("eeg", dict()).get("ch_names", [])
@@ -533,48 +538,50 @@ def _prepare_for_forward(
     # Circumvent numerical problems by excluding points too close to the skull,
     # and check that sensors are not inside any BEM surface
     if bem is not None:
+        kwargs = dict(limit=mindist, mri_head_t=mri_head_t, src=src)
         if not bem["is_sphere"]:
             check_surface = "inner skull surface"
-            inner_skull = _bem_find_surface(bem, "inner_skull")
-            check_inside = _filter_source_spaces(
-                inner_skull, mindist, mri_head_t, src, n_jobs
-            )
+            check_inside_brain = _CheckInside(_bem_find_surface(bem, "inner_skull"))
             logger.info("")
             if len(bem["surfs"]) == 3:
                 check_surface = "scalp surface"
-                check_inside = _CheckInside(_bem_find_surface(bem, "head"))
+                check_inside_head = _CheckInside(_bem_find_surface(bem, "head"))
+            else:
+                check_inside_head = check_inside_brain
         else:
             check_surface = "outermost sphere shell"
-            if len(bem["layers"]) == 0:
+            check_inside_brain = _CheckInsideSphere(bem)
+            if bem.radius is not None:
+                check_inside_head = _CheckInsideSphere(bem, check="outer")
+            else:
 
-                def check_inside(x):
+                def check_inside_head(x):
                     return np.zeros(len(x), bool)
 
-            else:
-
-                def check_inside(x):
-                    r0 = apply_trans(invert_transform(mri_head_t), bem["r0"])
-                    return np.linalg.norm(x - r0, axis=1) < bem["layers"][-1]["rad"]
+        if len(src):
+            _filter_source_spaces(check_inside_brain, **kwargs)
 
         if "meg" in sensors:
-            meg_loc = apply_trans(
-                invert_transform(mri_head_t),
-                np.array([coil["r0"] for coil in sensors["meg"]["defs"]]),
-            )
-            n_inside = check_inside(meg_loc).sum()
+            meg_loc = np.array([coil["r0"] for coil in sensors["meg"]["defs"]])
+            if not bem["is_sphere"]:
+                meg_loc = apply_trans(invert_transform(mri_head_t), meg_loc)
+            n_inside = check_inside_head(meg_loc).sum()
             if n_inside:
                 raise RuntimeError(
                     f"Found {n_inside} MEG sensor{_pl(n_inside)} inside the "
                     f"{check_surface}, perhaps coordinate frames and/or "
                     "coregistration must be incorrect"
                 )
 
-    rr = np.concatenate([s["rr"][s["vertno"]] for s in src])
-    if len(rr) < 1:
-        raise RuntimeError(
-            "No points left in source space after excluding "
-            "points close to inner skull."
-        )
+    if len(src):
+        rr = np.concatenate([s["rr"][s["vertno"]] for s in src])
+        if len(rr) < 1:
+            raise RuntimeError(
+                "No points left in source space after excluding "
+                "points close to inner skull."
+            )
+    else:
+        rr = np.zeros((0, 3))
 
     # deal with free orientations:
     source_nn = np.tile(np.eye(3), (len(rr), 1))
@@ -934,3 +941,75 @@ def use_coil_def(fname):
         yield
     finally:
         _extra_coil_def_fname = None
+
+
+class _ForwardModeler:
+    """Optimized incremental fitting using the same sensors and BEM."""
+
+    @verbose
+    def __init__(
+        self,
+        info,
+        trans,
+        bem,
+        *,
+        mindist=0.0,
+        n_jobs=1,
+        verbose=None,
+    ):
+        self.mri_head_t, _ = _get_trans(trans)
+        self.mindist = mindist
+        self.n_jobs = n_jobs
+        src = SourceSpaces([])
+        self.sensors, _, _, _, self.bem = _prepare_for_forward(
+            src,
+            self.mri_head_t,
+            info,
+            bem,
+            mindist,
+            n_jobs,
+            bem_extra="",
+            trans="",
+            info_extra="",
+            meg=True,
+            eeg=True,
+            ignore_ref=False,
+        )
+        self.fwd_data = _prep_field_computation(
+            sensors=self.sensors,
+            bem=self.bem,
+            n_jobs=self.n_jobs,
+        )
+        if self.bem["is_sphere"]:
+            self.check_inside = _CheckInsideSphere(self.bem)
+        else:
+            self.check_inside = _CheckInside(_bem_find_surface(self.bem, "inner_skull"))
+
+    def compute(self, src):
+        src = _ensure_src(src).copy()
+        src._transform_to("head", self.mri_head_t)
+        kwargs = dict(limit=self.mindist, mri_head_t=self.mri_head_t, src=src)
+        _filter_source_spaces(self.check_inside, n_jobs=self.n_jobs, **kwargs)
+        rr = np.concatenate([s["rr"][s["vertno"]] for s in src])
+        if len(rr) < 1:
+            raise RuntimeError(
+                "No points left in source space after excluding "
+                "points close to inner skull."
+            )
+
+        sensors = deepcopy(self.sensors)
+        fwd_data = deepcopy(self.fwd_data)
+        fwds = _compute_forwards_meeg(
+            rr,
+            sensors=sensors,
+            fwd_data=fwd_data,
+            n_jobs=self.n_jobs,
+        )
+        fwds = {
+            key: _to_forward_dict(fwds[key], sensors[key]["ch_names"])
+            for key in _FWD_ORDER
+            if key in fwds
+        }
+        fwd = _merge_fwds(fwds, verbose=False)
+        del fwds
+        return fwd