MRG, ENH: Speed up forward computations with Numba (mne-tools#7133)

* ENH: Speed up forwards with Numba

* FIX: Order

* FIX: Missing mult

* FIX: Refactor and add env option not to use numba (at import)

* DOC: whats_new.rst

* ENH: Faster cHPI

* FIX: Missed one

Author: larsoner

doc/changes/latest.inc

@@ -67,6 +67,8 @@ Changelog
 
 - Speed up :func:`mne.beamformer.make_lcmv` and :func:`mne.beamformer.make_dics` calculations by vectorizing linear algebra calls by `Dmitrii Altukhov`_ and `Eric Larson`_
 
+- Speed up :func:`mne.make_forward_solution` using Numba, by `Eric Larson`_
+
 - For KIT systems without built-in layout, :func:`mne.channels.find_layout` now falls back on an automatically generated layout, by `Christian Brodbeck`_
 
 - :meth:`mne.Epochs.plot` now takes a ``epochs_colors`` parameter to color specific epoch segments by `Mainak Jas`_

mne/chpi.py

@@ -45,11 +45,11 @@
 from .forward import (_magnetic_dipole_field_vec, _create_meg_coils,
                       _concatenate_coils)
 from .cov import make_ad_hoc_cov, compute_whitener
+from .fixes import jit
 from .transforms import (apply_trans, invert_transform, _angle_between_quats,
                          quat_to_rot, rot_to_quat)
 from .utils import (verbose, logger, use_log_level, _check_fname, warn,
-                    _check_option, _svd_lwork, _repeated_svd,
-                    ddot, dgemm, dgemv)
+                    _check_option)
 
 # Eventually we should add:
 #   hpicons
@@ -380,8 +380,7 @@ def _get_hpi_initial_fit(info, adjust=False, verbose=None):
     return hpi_rrs
 
 
-def _magnetic_dipole_objective(x, B, B2, coils, scale, method, too_close,
-                               lwork):
+def _magnetic_dipole_objective(x, B, B2, coils, scale, method, too_close):
     """Project data onto right eigenvectors of whitened forward."""
     if method == 'forward':
         fwd = _magnetic_dipole_field_vec(x[np.newaxis, :], coils, too_close)
@@ -390,27 +389,32 @@ def _magnetic_dipole_objective(x, B, B2, coils, scale, method, too_close,
         # Eventually we can try incorporating external bases here, which
         # is why the :3 is on the SVD below
         fwd = _sss_basis(dict(origin=x, int_order=1, ext_order=0), coils).T
+    return _magnetic_dipole_delta(fwd, scale, B, B2)
+
+
+@jit()
+def _magnetic_dipole_delta(fwd, scale, B, B2):
     # Here we use .T to get scale to Fortran order, which speeds things up
-    fwd = dgemm(alpha=1., a=fwd, b=scale.T)  # np.dot(fwd, scale.T)
-    one = _repeated_svd(fwd, lwork, overwrite_a=True)[2]
-    one = dgemv(alpha=1, a=one, x=B)
-    Bm2 = ddot(one, one)
+    fwd = np.dot(fwd, scale.T)
+    one = np.linalg.svd(fwd, full_matrices=False)[2]
+    one = np.dot(one, B)
+    Bm2 = np.dot(one, one)
     return B2 - Bm2
 
 
 def _fit_magnetic_dipole(B_orig, x0, coils, scale, method, too_close):
     """Fit a single bit of data (x0 = pos)."""
     from scipy.optimize import fmin_cobyla
-    B = dgemv(alpha=1, a=scale, x=B_orig)  # np.dot(scale, B_orig)
-    B2 = ddot(B, B)  # np.dot(B, B)
-    lwork = _svd_lwork((3, B_orig.shape[0]))
+    B = np.dot(scale, B_orig)
+    B2 = np.dot(B, B)
     objective = partial(_magnetic_dipole_objective, B=B, B2=B2,
                         coils=coils, scale=scale, method=method,
-                        too_close=too_close, lwork=lwork)
-    x = fmin_cobyla(objective, x0, (), rhobeg=1e-4, rhoend=1e-5, disp=False)
+                        too_close=too_close)
+    x = fmin_cobyla(objective, x0, (), rhobeg=1e-3, rhoend=1e-5, disp=False)
     return x, 1. - objective(x) / B2
 
 
+@jit()
 def _chpi_objective(x, coil_dev_rrs, coil_head_rrs):
     """Compute objective function."""
     d = np.dot(coil_dev_rrs, quat_to_rot(x[:3]).T)
@@ -428,7 +432,8 @@ def _unit_quat_constraint(x):
 def _fit_chpi_quat(coil_dev_rrs, coil_head_rrs, x0):
     """Fit rotation and translation (quaternion) parameters for cHPI coils."""
     from scipy.optimize import fmin_cobyla
-    denom = np.sum((coil_head_rrs - np.mean(coil_head_rrs, axis=0)) ** 2)
+    denom = np.linalg.norm(coil_head_rrs - np.mean(coil_head_rrs, axis=0))
+    denom *= denom
     objective = partial(_chpi_objective, coil_dev_rrs=coil_dev_rrs,
                         coil_head_rrs=coil_head_rrs)
     x0 = x0.copy()

mne/fixes.py

@@ -15,6 +15,7 @@
 import inspect
 from distutils.version import LooseVersion
 from math import log
+import os
 from pathlib import Path
 import warnings
 
@@ -1195,14 +1196,25 @@ def jit(nopython=True, nogil=True, fastmath=True, cache=True,
         return numba.jit(nopython=nopython, nogil=nogil, fastmath=fastmath,
                          cache=cache, **kwargs)
 except ImportError:
+    has_numba = False
+else:
+    has_numba = (os.getenv('MNE_USE_NUMBA', 'true').lower() == 'true')
+
+
+if not has_numba:
     def jit(**kwargs):  # noqa
         def _jit(func):
             return func
         return _jit
     prange = range
-    has_numba = False
+    bincount = np.bincount
 else:
-    has_numba = True
+    @jit()
+    def bincount(x, weights, minlength):  # noqa: D103
+        out = np.zeros(minlength)
+        for idx, w in zip(x, weights):
+            out[idx] += w
+        return out
 
 
 ###############################################################################