XFit-style dipole fitting GUI

doc/api/inverse.rst

@@ -85,6 +85,7 @@ Inverse Solutions
    Dipole
    DipoleFixed
    fit_dipole
+   gui.dipolefit
 
 :py:mod:`mne.dipole`:
 

doc/changes/dev/13074.newfeature.rst

@@ -0,0 +1 @@
+Add a GUI for interactive guided dipole fitting (:func:`mne.gui.dipolefit`), by `Marijn van Vliet`_

doc/conf.py

@@ -405,6 +405,7 @@
     "default",
     # unlinkable
     "CoregistrationUI",
+    "DipoleFitUI",
     "mne_qt_browser.figure.MNEQtBrowser",
     # pooch, since its website is unreliable and users will rarely need the links
     "pooch.Unzip",

mne/commands/mne_dipolefit.py

@@ -0,0 +1,161 @@
+# Authors: The MNE-Python contributors.
+# License: BSD-3-Clause
+# Copyright the MNE-Python contributors.
+
+"""Open the dipole fitting GUI.
+
+Examples
+--------
+.. code-block:: console
+
+    $ mne dipolefit
+
+"""
+
+import os.path as op
+
+import mne
+
+
+def run():
+    """Run command."""
+    from mne.commands.utils import _add_verbose_flag, get_optparser
+
+    parser = get_optparser(__file__)
+
+    parser.add_option(
+        "-e",
+        "--evoked",
+        default=None,
+        metavar="EVOKED_FILE",
+        help='The evoked file ("-ave.fif") containing the data to fit dipoles to.',
+    )
+    parser.add_option(
+        "--condition",
+        default=0,
+        help="The condition to use.",
+    )
+    parser.add_option(
+        "--baseline-from",
+        default=None,
+        type=float,
+        metavar="TIME",
+        help="The earliest timepoint to use as baseline.",
+    )
+    parser.add_option(
+        "--baseline-to",
+        default=None,
+        type=float,
+        metavar="TIME",
+        help="The latest timepoint to use as baseline.",
+    )
+    parser.add_option(
+        "-c",
+        "--cov",
+        default=None,
+        metavar="COV_FILE",
+        help='The noise covariance ("-cov.fif") to use.',
+    )
+    parser.add_option(
+        "-b",
+        "--bem",
+        default=None,
+        metavar="BEM_FILE",
+        help='The BEM model ("-bem-sol.fif") to use.',
+    )
+    parser.add_option(
+        "-t",
+        "--initial-time",
+        default=None,
+        type=float,
+        metavar="TIME",
+        help="The initial time to show",
+    )
+    parser.add_option(
+        "--trans",
+        default=None,
+        metavar="TRANS_FILE",
+        help='Head<->MRI transform FIF file ("-trans.fif")',
+    )
+    parser.add_option(
+        "--stc",
+        default=None,
+        metavar="STC_FILE",
+        help="An optional distributed source estimate to show during dipole fitting.",
+    )
+    parser.add_option(
+        "-s", "--subject", dest="subject", default=None, help="Subject name"
+    )
+    parser.add_option(
+        "-d",
+        "--subjects-dir",
+        default=None,
+        help="Subjects directory",
+    )
+    parser.add_option(
+        "--hide-density",
+        action="store_true",
+        default=False,
+        help="Prevent showing the magnetic field density as blobs of color.",
+    )
+    parser.add_option(
+        "--channel-type",
+        default=None,
+        help=(
+            'Restrict channel types to either "meg" or "eeg". By default both are used '
+            "if present."
+        ),
+    )
+    parser.add_option(
+        "-j", "--cpus", default=-1, type=int, help="Number of CPUs to use."
+    )
+    _add_verbose_flag(parser)
+
+    options, args = parser.parse_args()
+
+    # expanduser allows ~ for paths
+    subjects_dir = options.subjects_dir
+    if subjects_dir is not None:
+        subjects_dir = op.expanduser(subjects_dir)
+    bem = options.bem
+    if bem is not None:
+        bem = op.expanduser(bem)
+    trans = options.trans
+    if trans is not None:
+        trans = op.expanduser(trans)
+    stc = options.stc
+    if stc is not None:
+        stc = op.expanduser(stc)
+    import faulthandler
+
+    # Condition can be specified as integer index or string comment.
+    if options.condition is not None:
+        try:
+            condition = int(options.condition)
+        except ValueError:
+            condition = options.condition
+    else:
+        condition = None
+
+    faulthandler.enable()
+    mne.gui.dipolefit(
+        evoked=options.evoked,
+        condition=condition,
+        baseline=(options.baseline_from, options.baseline_to),
+        cov=options.cov,
+        bem=bem,
+        subject=options.subject,
+        subjects_dir=subjects_dir,
+        stc=stc,
+        ch_type=options.channel_type,
+        initial_time=options.initial_time,
+        trans=trans,
+        n_jobs=options.cpus,
+        show_density=not options.hide_density,
+        show=True,
+        block=True,
+        verbose=options.verbose,
+    )
+
+
+mne.utils.run_command_if_main()

mne/dipole.py

@@ -30,7 +30,12 @@
 )
 from .parallel import parallel_func
 from .source_space._source_space import SourceSpaces, _make_volume_source_space
-from .surface import _compute_nearest, _points_outside_surface, transform_surface_to
+from .surface import (
+    _CheckInside,
+    _compute_nearest,
+    _DistanceQuery,
+    transform_surface_to,
+)
 from .transforms import _coord_frame_name, _print_coord_trans, apply_trans
 from .utils import (
     ExtendedTimeMixin,
@@ -911,7 +916,7 @@ def _write_dipole_bdip(fname, dip):
             fid.write(np.array(has_errors, ">i4").tobytes())  # has_errors
             fid.write(np.zeros(1, ">f4").tobytes())  # noise level
             for key in _BDIP_ERROR_KEYS:
-                val = dip.conf[key][ti] if key in dip.conf else 0.0
+                val = dip.conf[key][ti] if key in dip.conf else np.array(0.0)
                 assert val.shape == ()
                 fid.write(np.array(val, ">f4").tobytes())
             fid.write(np.zeros(25, ">f4").tobytes())
@@ -1050,7 +1055,7 @@ def _fit_Q(*, sensors, fwd_data, whitener, B, B2, B_orig, rd, ori=None):
 
 def _fit_dipoles(
     fun,
-    min_dist_to_inner_skull,
+    constraint,
     data,
     times,
     guess_rrs,
@@ -1069,7 +1074,7 @@ def _fit_dipoles(
     # parallel over time points
     res = parallel(
         p_fun(
-            min_dist_to_inner_skull,
+            constraint,
             B,
             t,
             guess_rrs,
@@ -1267,26 +1272,41 @@ def _fit_confidence(*, rd, Q, ori, whitener, fwd_data, sensors):
     return conf
 
 
-def _surface_constraint(rd, surf, min_dist_to_inner_skull):
-    """Surface fitting constraint."""
-    dist = _compute_nearest(surf["rr"], rd[np.newaxis, :], return_dists=True)[1][0]
-    if _points_outside_surface(rd[np.newaxis, :], surf, 1)[0]:
-        dist *= -1.0
-    # Once we know the dipole is below the inner skull,
-    # let's check if its distance to the inner skull is at least
-    # min_dist_to_inner_skull. This can be enforced by adding a
-    # constrain proportional to its distance.
-    dist -= min_dist_to_inner_skull
-    return dist
+def _surface_constraint(surf, min_dist_to_inner_skull):
+    """Create a surface fitting constraint function."""
+    distance_checker = _DistanceQuery(surf["rr"], method="BallTree")
+    try:
+        import pyvista  # noqa F401
+
+        inside_checker = _CheckInside(surf, mode="pyvista")
+    except ImportError:
+        inside_checker = _CheckInside(surf, mode="old")
+
+    def constraint(rd):
+        dist = distance_checker.query(rd[np.newaxis, :])[0][0]
+        if not inside_checker(rd[np.newaxis, :])[0]:
+            dist *= -1.0
+        # Once we know the dipole is below the inner skull,
+        # let's check if its distance to the inner skull is at least
+        # min_dist_to_inner_skull. This can be enforced by adding a
+        # constrain proportional to its distance.
+        dist -= min_dist_to_inner_skull
+        return dist
+
+    return constraint
+
+
+def _sphere_constraint(r0, R_adj):
+    """Create a sphere fitting constraint function."""
 
+    def constraint(rd):
+        return R_adj - np.sqrt(np.sum((rd - r0) ** 2))
 
-def _sphere_constraint(rd, r0, R_adj):
-    """Sphere fitting constraint."""
-    return R_adj - np.sqrt(np.sum((rd - r0) ** 2))
+    return constraint
 
 
 def _fit_dipole(
-    min_dist_to_inner_skull,
+    constraint,
     B_orig,
     t,
     guess_rrs,
@@ -1303,22 +1323,6 @@ def _fit_dipole(
     """Fit a single bit of data."""
     B = np.dot(whitener, B_orig)
 
-    # make constraint function to keep the solver within the inner skull
-    if "rr" in fwd_data["inner_skull"]:  # bem
-        surf = fwd_data["inner_skull"]
-        constraint = partial(
-            _surface_constraint,
-            surf=surf,
-            min_dist_to_inner_skull=min_dist_to_inner_skull,
-        )
-    else:  # sphere
-        surf = None
-        constraint = partial(
-            _sphere_constraint,
-            r0=fwd_data["inner_skull"]["r0"],
-            R_adj=fwd_data["inner_skull"].radius - min_dist_to_inner_skull,
-        )
-
     # Find a good starting point (find_best_guess in C)
     B2 = np.dot(B, B)
     if B2 == 0:
@@ -1388,9 +1392,9 @@ def _fit_dipole(
     )
 
     msg = "---- Fitted : %7.1f ms" % (1000.0 * t)
-    if surf is not None:
+    if "rr" in fwd_data["inner_skull"]:  # bem
         dist_to_inner_skull = _compute_nearest(
-            surf["rr"], rd_final[np.newaxis, :], return_dists=True
+            fwd_data["inner_skull"]["rr"], rd_final[np.newaxis, :], return_dists=True
         )[1][0]
         msg += ", distance to inner skull : %2.4f mm" % (dist_to_inner_skull * 1000.0)
 
@@ -1399,7 +1403,7 @@ def _fit_dipole(
 
 
 def _fit_dipole_fixed(
-    min_dist_to_inner_skull,
+    constraint,
     B_orig,
     t,
     guess_rrs,
@@ -1703,17 +1707,23 @@ def fit_dipole(
         logger.info("Go through all guess source locations...")
 
     # inner_skull goes from mri to head frame
-    if "rr" in inner_skull:
+    if not bem["is_sphere"]:
         transform_surface_to(inner_skull, "head", mri_head_t)
+
+    # make constraint function to keep the solver within the inner skull
+    if bem["is_sphere"]:
+        constraint = _sphere_constraint(
+            r0=inner_skull["r0"],
+            R_adj=inner_skull.radius - min_dist_to_inner_skull,
+        )
+    else:
+        constraint = _surface_constraint(
+            surf=inner_skull,
+            min_dist_to_inner_skull=min_dist_to_inner_skull,
+        )
+
     if fixed_position:
-        if "rr" in inner_skull:
-            check = _surface_constraint(pos, inner_skull, min_dist_to_inner_skull)
-        else:
-            check = _sphere_constraint(
-                pos,
-                inner_skull["r0"],
-                R_adj=inner_skull.radius - min_dist_to_inner_skull,
-            )
+        check = constraint(pos)
         if check <= 0:
             raise ValueError(
                 f"fixed position is {-1000 * check:0.1f}mm outside the inner skull "
@@ -1753,7 +1763,7 @@ def fit_dipole(
     fun = _fit_dipole_fixed if fixed_position else _fit_dipole
     out = _fit_dipoles(
         fun,
-        min_dist_to_inner_skull,
+        constraint,
         data,
         times,
         guess_src["rr"],

mne/gui/__init__.pyi

@@ -1,2 +1,2 @@
-__all__ = ["_GUIScraper", "coregistration"]
-from ._gui import _GUIScraper, coregistration
+__all__ = ["_GUIScraper", "coregistration", "dipolefit"]
+from ._gui import _GUIScraper, coregistration, dipolefit