ENH: add a vtk.js backend for MNE's 3D renderer (JupyterLite split 3/5)

doc/changes/dev/14144.newfeature.rst

@@ -0,0 +1 @@
+Add a vtk.js drawing backend for MNE's 3D renderer, selected with ``mne.viz.set_3d_backend("jupyterlite_notebook")`` and used by the JupyterLite documentation where VTK cannot load, by `Natneal B`_.

doc/sphinxext/jupyterlite_lite_renderer.py

@@ -0,0 +1,23 @@
+"""Turn on MNE's pyvista-js 3D renderer inside the JupyterLite kernel.
+
+VTK has no WebAssembly build, so the browser draws with pyvista-js instead. The
+renderer itself is ordinary library code in ``mne/viz/backends/_lite.py``; this
+module only exposes the few lines of notebook code that switch MNE over to it.
+``LITE_RENDERER_CELL`` is appended to ``LITE_SETUP_CELL`` in
+``jupyterlite_setup_cell.py``, which the docs build prepends to each notebook.
+"""
+
+# Authors: The MNE-Python contributors.
+# License: BSD-3-Clause
+# Copyright the MNE-Python contributors.
+
+LITE_RENDERER_CELL = """
+# Using pyvista-js (vtk.js) to draw MNE's 3D rendering in JupyterLite.
+# See mne/viz/backends/_lite.py for more details.
+try:
+    import mne.viz
+
+    mne.viz.set_3d_backend("jupyterlite_notebook")
+except Exception as _e:
+    print("[JupyterLite] could not select the pyvista-js renderer: " + repr(_e))
+"""

mne/conftest.py

@@ -740,7 +740,12 @@ def browser_backend(request, garbage_collect, monkeypatch):
             mne_qt_browser._browser_instances.clear()
 
 
-@pytest.fixture(params=[pytest.param("pyvistaqt", marks=pytest.mark.pvtest)])
+@pytest.fixture(
+    params=[
+        pytest.param("pyvistaqt", marks=pytest.mark.pvtest),
+        pytest.param("jupyterlite_notebook", marks=pytest.mark.pvtest),
+    ]
+)
 def renderer(request, options_3d, garbage_collect):
     """Yield the 3D backends."""
     with _use_backend(request.param, interactive=False) as renderer:
@@ -761,6 +766,13 @@ def renderer_notebook(request, options_3d):
         yield renderer
 
 
+@pytest.fixture(params=[pytest.param("jupyterlite_notebook", marks=pytest.mark.pvtest)])
+def renderer_lite(request, options_3d):
+    """Yield the JupyterLite (vtk.js) renderer alone, for its own tests."""
+    with _use_backend(request.param, interactive=False) as renderer:
+        yield renderer
+
+
 @pytest.fixture(params=[pytest.param("pyvistaqt", marks=pytest.mark.pvtest)])
 def renderer_interactive_pyvistaqt(request, options_3d, qt_windows_closed):
     """Yield the interactive PyVista backend."""
@@ -788,22 +800,32 @@ def _use_backend(backend_name, interactive):
     # figure-count test (in other modules) fails. Restore it on teardown.
     mpl_backend = matplotlib.get_backend()
     _check_skip_backend(backend_name)
+    from mne.viz.backends import renderer
+
+    # use_3d_backend only puts a backend back if one was already selected, so
+    # the first renderer test of a session would otherwise decide the backend
+    # every later test inherits; the JupyterLite one draws for a browser, so
+    # that must never be it
+    was = (renderer.MNE_3D_BACKEND, renderer.backend)
     try:
         with _use_test_3d_backend(backend_name, interactive=interactive):
-            from mne.viz.backends import renderer
-
             try:
                 yield renderer
             finally:
                 renderer.backend._close_all()
     finally:
+        renderer.MNE_3D_BACKEND, renderer.backend = was
         if matplotlib.get_backend() != mpl_backend:
             matplotlib.use(mpl_backend, force=True)
 
 
 def _check_skip_backend(name):
     from mne.viz.backends._utils import _notebook_vtk_works
 
+    if name == "jupyterlite_notebook":
+        # draws with vtk.js in a browser: no VTK, no Qt, no ffmpeg
+        pytest.importorskip("pyvista_js")
+        return
     pytest.importorskip("pyvista")
     pytest.importorskip("imageio_ffmpeg")
     if name == "pyvistaqt":

mne/report/tests/test_report.py

@@ -401,7 +401,7 @@ def test_report_raw_psd_and_date(tmp_path):
 
 @pytest.mark.slowtest  # slow on Azure
 @testing.requires_testing_data
-def test_render_add_sections(renderer, tmp_path):
+def test_render_add_sections(renderer_pyvistaqt, tmp_path):
     """Test adding figures/images to section."""
     pytest.importorskip("nibabel")
     try:
@@ -451,7 +451,7 @@ def test_render_add_sections(renderer, tmp_path):
 
 @pytest.mark.slowtest
 @testing.requires_testing_data
-def test_render_mri(renderer, tmp_path):
+def test_render_mri(renderer_pyvistaqt, tmp_path):
     """Test rendering MRI for mne report."""
     pytest.importorskip("nibabel")
     trans_fname_new = tmp_path / "temp-trans.fif"
@@ -1247,7 +1247,7 @@ def test_report_backward_compat(tmp_path):
 
 @pytest.mark.slowtest  # 30 s on Azure
 @testing.requires_testing_data
-def test_manual_report_3d(tmp_path, renderer):
+def test_manual_report_3d(tmp_path, renderer_pyvistaqt):
     """Simulate adding 3D sections."""
     pytest.importorskip("nibabel")
     r = Report(title="My Report")

mne/tests/test_transforms.py

@@ -386,6 +386,19 @@ def test_vector_rotation():
     quat_1 = rot_to_quat(rot)
     quat_2 = rot_to_quat(np.eye(3))
     assert_allclose(_angle_between_quats(quat_1, quat_2), np.pi / 2.0)
+    # many at once, including the parallel, antiparallel and nearly antiparallel
+    # cases, the first two of which have no rotation axis of their own
+    b = np.random.default_rng(0).normal(size=(20, 3))
+    b = np.concatenate([b, [x, -x, [-1, 1e-8, 0]]])
+    b /= np.linalg.norm(b, axis=1, keepdims=True)
+    rots = _find_vector_rotation(x, b)
+    assert_allclose(rots @ x, b, atol=1e-12)
+    eye = np.broadcast_to(np.eye(3), rots.shape)
+    assert_allclose(rots @ rots.transpose(0, 2, 1), eye, atol=1e-12)
+    assert_allclose(rots[-2], np.diag([-1, -1, 1]), atol=1e-12)
+    for rot, this_b in zip(rots[:3], b):  # each is the minimal rotation
+        angle = _angle_between_quats(rot_to_quat(rot), np.zeros(3))
+        assert_allclose(angle, np.arccos(x @ this_b))
 
 
 def test_average_quats():

mne/transforms.py

@@ -1413,26 +1413,68 @@ def _quat_mult(one, two):
 
 
 def _skew_symmetric_cross(a):
-    """Compute the skew-symmetric cross product of a vector."""
-    return np.array([[0.0, -a[2], a[1]], [a[2], 0.0, -a[0]], [-a[1], a[0], 0.0]])
+    """Compute the skew-symmetric cross product matrix of (..., 3) vector(s)."""
+    a = np.asarray(a, float)
+    ax = np.zeros(a.shape + (3,))
+    ax[..., 0, 1], ax[..., 0, 2] = -a[..., 2], a[..., 1]
+    ax[..., 1, 0], ax[..., 1, 2] = a[..., 2], -a[..., 0]
+    ax[..., 2, 0], ax[..., 2, 1] = -a[..., 1], a[..., 0]
+    return ax
 
 
 def _find_vector_rotation(a, b):
-    """Find the rotation matrix that maps unit vector a to b."""
+    """Find the rotation matrix that maps unit vector a to unit vector(s) b.
+
+    Parameters
+    ----------
+    a : array, shape (3,)
+        The unit vector to rotate.
+    b : array, shape (3,) | shape (..., 3)
+        The unit vector(s) to rotate ``a`` onto.
+
+    Returns
+    -------
+    R : array, shape (3, 3) | shape (..., 3, 3)
+        The rotation(s) about ``a x b`` by the angle between them, so that
+        ``R @ a`` is ``b``. Antiparallel vectors, where that axis vanishes,
+        get a half turn about an arbitrary axis perpendicular to ``a``.
+
+    Notes
+    -----
+    Mapping one vector onto another leaves a free parameter: the roll about
+    ``b``. Any rotation about ``b`` composed with the result maps ``a`` onto
+    ``b`` just as well, and this function settles it by taking the minimal
+    rotation, about ``a x b``. So it is right for things that look the same
+    however they are rolled about their axis, like arrows, tubes and the EEG
+    electrode cylinders, and wrong for a flat MEG coil, whose orientation
+    needs the full rotation from ``_loc_to_coil_trans``.
+    """
     # Rodrigues' rotation formula:
     #   https://en.wikipedia.org/wiki/Rodrigues%27_rotation_formula
     #   http://math.stackexchange.com/a/476311
+    a = np.asarray(a, float)
+    b = np.asarray(b, float)
+    assert a.shape == (3,), a.shape
     assert np.isclose(np.linalg.norm(a), 1.0), np.linalg.norm(a)
-    assert np.isclose(np.linalg.norm(b), 1.0), np.linalg.norm(b)
-    R = np.eye(3)
-    v = np.cross(a, b)
-    if np.allclose(v, 0.0):  # identical
-        return R
-    s = np.dot(v, v)  # sine of the angle between them
-    c = np.dot(a, b)  # cosine of the angle between them
+    assert b.shape[-1:] == (3,), b.shape
+    assert np.allclose(np.linalg.norm(b, axis=-1), 1.0), np.linalg.norm(b, axis=-1)
+    v = np.cross(a, b)  # rotation axis, with the sine of the angle as its length
+    s = (v * v).sum(-1)  # sine squared
+    c = b @ a  # cosine
     vx = _skew_symmetric_cross(v)
-    R += vx + np.dot(vx, vx) * (1 - c) / s
-    # Now we have: np.allclose(R @ a, b)
+    # (1 - c) / s is 1 / (1 + c), but written this way it stays accurate as b
+    # approaches -a, where 1 + c cancels and s does not. Only an s that has
+    # vanished outright needs special handling below.
+    degenerate = s < np.finfo(float).tiny
+    factor = (1.0 - c) / np.where(degenerate, 1.0, s)
+    R = np.eye(3) + vx + vx @ vx * factor[..., np.newaxis, np.newaxis]
+    if degenerate.any():
+        # parallel is the identity (vx is zero); antiparallel is a half turn
+        # about a unit vector k perpendicular to a, for which the coordinate
+        # axis least aligned with a serves, since it is never parallel to it
+        k = np.cross(a, np.eye(3)[np.argmin(np.abs(a))])
+        k /= np.linalg.norm(k)
+        R[degenerate & (c < 0)] = 2 * np.outer(k, k) - np.eye(3)
     return R
 
 

mne/utils/config.py

@@ -942,6 +942,7 @@ def sys_info(
             "nbclient",
             "nbformat",
             "nitime",
+            "pyvista-js",
             "imageio",
             "imageio-ffmpeg",
             "snirf",

mne/utils/tests/test_config.py

@@ -167,7 +167,7 @@ def test_sys_info_complete():
     ]
     missing = []
     for dep in deps:
-        dep = dep.split("[")[0].split(">")[0].strip()
+        dep = dep.split(";")[0].split("[")[0].split(">")[0].strip()
         if f" {dep}" not in out:
             missing.append(dep)
     if missing:

mne/viz/_3d.py

@@ -841,7 +841,7 @@ def plot_alignment(
 
     # initialize figure
     renderer = _get_renderer(
-        fig,
+        fig=fig,
         name=f"Sensor alignment: {subject}",
         bgcolor=(0.5, 0.5, 0.5),
         size=(800, 800),
@@ -1442,8 +1442,7 @@ def _plot_glyphs(
         )
         x_axis = np.array([1.0, 0.0, 0.0])
         nn = vectors / np.linalg.norm(vectors, axis=1, keepdims=True)
-        rots = np.array([_find_vector_rotation(x_axis, this_nn) for this_nn in nn])
-        quats = rot_to_quat(rots)
+        quats = rot_to_quat(_find_vector_rotation(x_axis, nn))
     rr, tris = renderer._glyph_template(kind, **template_kw)
     actor, cloud = renderer.instanced_mesh(
         rr=rr,
@@ -3953,7 +3952,7 @@ def snapshot_brain_montage(fig, montage, hide_sensors=True):
         )
 
     # initialize figure
-    renderer = _get_renderer(fig, show=True)
+    renderer = _get_renderer(fig=fig, show=True)
 
     xyz = np.vstack(xyz)
     proj = renderer.project(xyz=xyz, ch_names=ch_names)

mne/viz/_brain/tests/test_brain.py

@@ -223,9 +223,9 @@ def test_brain_data_gc(renderer_interactive_pyvistaqt, brain_gc):
 
 
 @testing.requires_testing_data
-def test_brain_routines(renderer, brain_gc):
+def test_brain_routines(renderer_pyvistaqt, brain_gc):
     """Test backend agnostic Brain routines."""
-    brain_klass = renderer.get_brain_class()
+    brain_klass = renderer_pyvistaqt.get_brain_class()
     from mne.viz._brain import Brain
 
     assert brain_klass == Brain
@@ -996,7 +996,7 @@ def test_single_hemi(hemi, renderer_interactive_pyvistaqt, brain_gc):
 @testing.requires_testing_data
 @pytest.mark.slowtest
 @pytest.mark.parametrize("interactive_state", (False, True))
-def test_brain_save_movie(tmp_path, renderer, brain_gc, interactive_state):
+def test_brain_save_movie(tmp_path, renderer_pyvistaqt, brain_gc, interactive_state):
     """Test saving a movie of a Brain instance."""
     pytest.importorskip("imageio")
     imageio_ffmpeg = pytest.importorskip("imageio_ffmpeg")