ENH: Speed up forward computations for iterative fitting

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`_.

mne/dipole.py

@@ -1720,7 +1720,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

@@ -26,7 +26,7 @@
     _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,
@@ -38,7 +38,11 @@
     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(
@@ -533,35 +537,38 @@ 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"))
+            _filter_source_spaces(check_inside_brain, n_jobs=n_jobs, **kwargs)
             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_surface = "outermost sphere shell"
+            check_inside_brain = _CheckInsideSphere(bem)
+            _filter_source_spaces(check_inside_brain, **kwargs)
             if len(bem["layers"]) == 0:
 
-                def check_inside(x):
+                def check_inside_head(x):
                     return np.zeros(len(x), bool)
 
             else:
 
-                def check_inside(x):
+                def check_inside_head(x):
+                    # MEG sensors are in MRI coords at this point, so need to
+                    # transform BEM center to MRI coords, too
                     r0 = apply_trans(invert_transform(mri_head_t), bem["r0"])
-                    return np.linalg.norm(x - r0, axis=1) < bem["layers"][-1]["rad"]
+                    return np.linalg.norm(x - r0, axis=1) < bem.radius
 
         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()
+            n_inside = check_inside_head(meg_loc).sum()
             if n_inside:
                 raise RuntimeError(
                     f"Found {n_inside} MEG sensor{_pl(n_inside)} inside the "
@@ -934,3 +941,82 @@ def use_coil_def(fname):
         yield
     finally:
         _extra_coil_def_fname = None
+
+
+# Class optimized for incremental fitting using the same sensors and BEM at different
+# locations (e.g., for Xfit-like iterative forward solutions)
+class _ForwardModeler:
+    @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
+        # TODO: Make `src` optional in _prepare_for_forward
+        import mne
+
+        src = mne.setup_volume_source_space(
+            "", pos=dict(rr=np.zeros((1, 3)), nn=np.array([[0, 0, 1]])), verbose="error"
+        )
+        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()
+        for s in src:
+            transform_surface_to(s, "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

mne/forward/tests/test_make_forward.py

@@ -602,6 +602,29 @@ def test_make_forward_solution_sphere(tmp_path, fname_src_small):
     assert fwd["mri_head_t"]["trans"][0, 3] == -0.05
 
 
+@testing.requires_testing_data
+def test_make_forward_sphere_exclude():
+    """Test that points are excluded that are outside BEM sphere inner layer."""
+    bem = make_sphere_model(r0=(0.0, 0.0, 0.04), head_radius=0.1)
+    src = setup_volume_source_space(
+        "sample",
+        pos=10.0,
+        sphere=(0.0, 0.0, 0.04, 0.1),
+        subjects_dir=subjects_dir,
+        exclude=10,
+        verbose=True,
+    )
+    assert src[0]["nuse"] == 3694
+    trans = Transform("mri", "head")  # identity for simplicity
+    raw = read_raw_fif(fname_raw)
+    raw.pick(raw.ch_names[:1])
+    fwd = make_forward_solution(raw.info, trans, src, bem, verbose=True)
+    assert fwd["nsource"] == 3694
+    bem_small = make_sphere_model(r0=(0.0, 0.0, 0.04), head_radius=0.09)
+    fwd_small = make_forward_solution(raw.info, trans, src, bem_small, verbose=True)
+    assert fwd_small["nsource"] < 3694
+
+
 @pytest.mark.slowtest
 @testing.requires_testing_data
 def test_forward_mixed_source_space(tmp_path):

mne/source_space/_source_space.py

@@ -45,6 +45,7 @@
 from ..parallel import parallel_func
 from ..surface import (
     _CheckInside,
+    _CheckInsideSphere,
     _compute_nearest,
     _create_surf_spacing,
     _get_ico_surface,
@@ -1907,7 +1908,9 @@ def setup_volume_source_space(
             surf["rr"] *= 1e-3  # must be converted to meters
         else:  # Load an icosahedron and use that as the surface
             logger.info("Setting up the sphere...")
-            surf = dict(R=sphere[3], r0=sphere[:3])
+            surf = ConductorModel(
+                layers=[dict(rad=sphere[3])], r0=sphere[:3], is_sphere=True
+            )
         # Make the grid of sources in MRI space
         sp = _make_volume_source_space(
             surf,
@@ -2063,10 +2066,10 @@ def _make_volume_source_space(
         cm = np.mean(surf["rr"], axis=0)  # center of mass
         maxdist = np.linalg.norm(surf["rr"] - cm, axis=1).max()
     else:
-        mins = surf["r0"] - surf["R"]
-        maxs = surf["r0"] + surf["R"]
+        mins = surf["r0"] - surf.radius
+        maxs = surf["r0"] + surf.radius
         cm = surf["r0"].copy()
-        maxdist = surf["R"]
+        maxdist = surf.radius
 
     # Define the sphere which fits the surface
     logger.info(
@@ -2136,18 +2139,13 @@ def _make_volume_source_space(
         1000 * exclude,
         1000 * maxdist,
     )
+    kwargs = dict(limit=mindist, mri_head_t=None, src=[sp])
+    assert sp["coord_frame"] == FIFF.FIFFV_COORD_MRI
     if "rr" in surf:
-        _filter_source_spaces(surf, mindist, None, [sp], n_jobs)
-    else:  # sphere
-        vertno = np.where(sp["inuse"])[0]
-        bads = (
-            np.linalg.norm(sp["rr"][vertno] - surf["r0"], axis=-1)
-            >= surf["R"] - mindist / 1000.0
-        )
-        sp["nuse"] -= bads.sum()
-        sp["inuse"][vertno[bads]] = False
-        sp["vertno"] = np.where(sp["inuse"])[0]
-        del vertno
+        assert surf["coord_frame"] == FIFF.FIFFV_COORD_MRI
+    else:
+        assert surf["is_sphere"]
+    _filter_source_spaces(surf, n_jobs=n_jobs, **kwargs)
     del surf
     logger.info(
         "%d sources remaining after excluding the sources outside "
@@ -2534,7 +2532,9 @@ def _grid_interp_jit(from_shape, to_shape, trans, order, inuse):
 
 
 @verbose
-def _filter_source_spaces(surf, limit, mri_head_t, src, n_jobs=None, verbose=None):
+def _filter_source_spaces(
+    surf_or_check_inside, *, limit, mri_head_t, src, n_jobs=None, verbose=None
+):
     """Remove all source space points closer than a given limit (in mm)."""
     if src[0]["coord_frame"] == FIFF.FIFFV_COORD_HEAD and mri_head_t is None:
         raise RuntimeError(
@@ -2558,7 +2558,14 @@ def _filter_source_spaces(surf, limit, mri_head_t, src, n_jobs=None, verbose=Non
     logger.info(out_str + " (will take a few...)")
 
     # fit a sphere to a surf quickly
-    check_inside = _CheckInside(surf)
+    if isinstance(surf_or_check_inside, _CheckInside | _CheckInsideSphere):
+        check_inside = surf_or_check_inside
+    else:
+        if "rr" in surf_or_check_inside:
+            check_inside = _CheckInside(surf_or_check_inside)
+        else:
+            check_inside = _CheckInsideSphere(surf_or_check_inside)
+    del surf_or_check_inside
 
     # Check that the source is inside surface (often the inner skull)
     for s in src:
@@ -2568,7 +2575,7 @@ def _filter_source_spaces(surf, limit, mri_head_t, src, n_jobs=None, verbose=Non
         if s["coord_frame"] == FIFF.FIFFV_COORD_HEAD:
             r1s = apply_trans(inv_trans["trans"], r1s)
 
-        inside = check_inside(r1s, n_jobs)
+        inside = check_inside(r1s, n_jobs=n_jobs)
         omit_outside = (~inside).sum()
 
         # vectorized nearest using BallTree (or cdist)
@@ -2578,16 +2585,14 @@ def _filter_source_spaces(surf, limit, mri_head_t, src, n_jobs=None, verbose=Non
             idx = np.where(inside)[0]
             check_r1s = r1s[idx]
             if check_inside.inner_r is not None:
-                # ... and those that are at least inner_sphere + limit away
+                # ... and those that are at least inner_sphere - limit away
                 mask = (
                     np.linalg.norm(check_r1s - check_inside.cm, axis=-1)
                     >= check_inside.inner_r - limit / 1000.0
                 )
                 idx = idx[mask]
                 check_r1s = check_r1s[mask]
-            dists = _compute_nearest(
-                surf["rr"], check_r1s, return_dists=True, method="KDTree"
-            )[1]
+            dists = check_inside.query(check_r1s)[0]
             close = dists < limit / 1000.0
             omit_limit = np.sum(close)
             inside[idx[close]] = False

mne/surface.py

@@ -779,7 +779,7 @@ def _init_pyvista(self):
             self.pdata = _surface_to_polydata(self.surf).clean()
 
     @verbose
-    def __call__(self, rr, n_jobs=None, verbose=None):
+    def __call__(self, rr, *, n_jobs=None, verbose=None):
         n_orig = len(rr)
         logger.info(
             f"Checking surface interior status for {n_orig} point{_pl(n_orig, ' ')}..."
@@ -794,6 +794,14 @@ def __call__(self, rr, n_jobs=None, verbose=None):
         logger.info(f"Interior check completed in {(time.time() - t0) * 1000:0.1f} ms")
         return inside
 
+    def query(self, rr):
+        """Get the distance to the nearest point."""
+        if not hasattr(self, "_tree"):  # compute on the fly only when needed
+            from scipy.spatial import KDTree
+
+            self._tree = KDTree(self.surf["rr"])
+        return self._tree.query(rr)
+
     def _call_pyvista(self, rr):
         pdata = _surface_to_polydata(dict(rr=rr))
         out = pdata.select_enclosed_points(self.pdata, check_surface=False)
@@ -855,6 +863,35 @@ def _call_old(self, rr, n_jobs):
         return inside
 
 
+class _CheckInsideSphere:
+    def __init__(self, sphere):
+        from .bem import ConductorModel
+
+        assert isinstance(sphere, ConductorModel) and sphere["is_sphere"]
+        self.cm = sphere["r0"]
+        self.inner_r = sphere.radius  # float or None
+
+    # No need for verbose dec here because no MNE code is called that would log
+    def __call__(self, rr, *, n_jobs=None, verbose=None):
+        assert isinstance(rr, np.ndarray), type(rr)
+        assert rr.ndim == 2 and rr.shape[1] == 3
+        if self.inner_r is None:
+            return np.ones(rr.shape[0], bool)
+        else:
+            return np.linalg.norm(rr - self.cm, axis=-1) <= self.inner_r
+
+    def query(self, rr):
+        """Return the distance to the sphere surface for each point."""
+        assert isinstance(rr, np.ndarray), type(rr)
+        assert rr.ndim == 2 and rr.shape[1] == 3, rr.shape
+        idx = np.zeros(rr.shape[0], int)
+        if self.inner_r is None:
+            dists = np.full(rr.shape[0], np.inf)
+        else:
+            dists = np.abs(np.linalg.norm(rr - self.cm, axis=-1) - self.inner_r)
+        return dists, idx
+
+
 ###############################################################################
 # Handle freesurfer