ENH: Add public LayeredMesh and brain.layered_meshes (#13970)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: payamsash

doc/api/visualization.rst

@@ -14,6 +14,7 @@ Visualization
    :toctree: ../generated/
 
    Brain
+   LayeredMesh
    ClickableImage
    EvokedField
    Figure3D

doc/changes/dev/13970.newfeature.rst

@@ -0,0 +1 @@
+Added public :class:`mne.viz.LayeredMesh` class (previously private ``_LayeredMesh``) with a ``smooth_mat`` attribute for efficient real-time source-space data streaming. Added ``Brain.layered_meshes`` attribute to access meshes after plotting, and a ``key`` parameter to :meth:`mne.viz.Brain.add_data` for named overlays, by :newcontrib:`Payam Sadeghi-Shabestari`.

doc/changes/names.inc

@@ -267,6 +267,7 @@
 .. _Paul Roujansky: https://github.com/paulroujansky
 .. _Pavel Navratil: https://github.com/navrpa13
 .. _Pavel Popov: https://github.com/paavalipopov
+.. _Payam Sadeghi-Shabestari: https://github.com/payamsash
 .. _Peter Molfese: https://github.com/pmolfese
 .. _Phillip Alday: https://palday.bitbucket.io
 .. _Pierre Ablin: https://pierreablin.com

doc/conf.py

@@ -426,6 +426,8 @@
     "polars",
     "default",
     # unlinkable
+    "_Renderer",
+    "n_triangles",
     "CoregistrationUI",
     "mne_qt_browser.figure.MNEQtBrowser",
     # pooch, since its website is unreliable and users will rarely need the links

examples/visualization/brain.py

@@ -136,6 +136,48 @@
 fig.suptitle("Dipole Fits Scaled by Amplitude and Colored by GOF")
 
 
+# %%
+# Update overlays via ``Brain.layered_meshes``
+# --------------------------------------------
+#
+# After :meth:`~mne.viz.Brain.add_data` is called, each hemisphere's surface
+# is backed by a :class:`~mne.viz.LayeredMesh` stored in
+# ``brain.layered_meshes``.  Calling
+# :meth:`~mne.viz.LayeredMesh.update_overlay` pushes new scalar data without
+# rebuilding the full rendering pipeline — the key operation for packages that
+# stream source-space data onto the brain in real time, such as
+# `MNE-RT <https://payamsash.github.io/mne-rt/>`_.
+#
+# Here we add an initial bottom-to-top gradient as the first data frame.
+
+brain = mne.viz.Brain("sample", subjects_dir=subjects_dir, hemi="lh", **brain_kwargs)
+coords = brain.geo["lh"].coords
+data_t0 = coords[:, 2]
+data_t0 = (data_t0 - data_t0.min()) / (data_t0.max() - data_t0.min())
+brain.add_data(
+    data_t0, hemi="lh", fmin=0, fmax=1, colormap="viridis", smoothing_steps=5
+)
+brain.show_view(azimuth=190, elevation=70, distance=350, focalpoint=(0, 0, 20))
+
+# %%
+# Simulate a new data frame arriving: call
+# :meth:`~mne.viz.LayeredMesh.update_overlay` on the **same** brain to replace
+# the scalars in-place — no new mesh, no new actor, no pipeline rebuild.
+
+# we create a new brain here for comparison purposes
+brain_update = mne.viz.Brain(
+    "sample", subjects_dir=subjects_dir, hemi="lh", **brain_kwargs
+)
+brain_update.add_data(
+    data_t0, hemi="lh", fmin=0, fmax=1, colormap="viridis", smoothing_steps=5
+)
+brain_update.show_view(azimuth=190, elevation=70, distance=350, focalpoint=(0, 0, 20))
+data_t1 = coords[:, 1]
+data_t1 = (data_t1 - data_t1.min()) / (data_t1.max() - data_t1.min())
+mesh = brain_update.layered_meshes["lh"]
+mesh.update_overlay(name="data", scalars=data_t1)
+mesh.update()
+
 # %%
 # Use per-vertex opacity for distributed data
 # --------------------------------------------

mne/viz/_3d_overlay.py

@@ -74,7 +74,45 @@ def _norm(self, rng):
         return (self._scalars - rng[0]) / factor
 
 
-class _LayeredMesh:
+class LayeredMesh:
+    """A mesh with support for layered RGBA overlays and optional smoothing.
+
+    This class manages a single brain-surface mesh and composites multiple
+    named overlays (e.g., curvature, data, labels) on top of each other using
+    alpha blending.  It is the object stored in
+    ``Brain.layered_meshes``.
+
+    .. warning::
+        This class is not meant to be instantiated directly, and the
+        initialization API could change at any time! Use
+        :meth:`mne.viz.Brain.add_data` to create instances.
+
+    Parameters
+    ----------
+    renderer : instance of _Renderer
+        The renderer used to create the underlying mesh polydata.
+    vertices : array, shape (n_vertices, 3)
+        Vertex coordinates in metres.
+    triangles : array, shape (n_triangles, 3)
+        Triangle indices into ``vertices``.
+    normals : array, shape (n_vertices, 3)
+        Vertex normals.
+
+    Attributes
+    ----------
+    smooth_mat : scipy.sparse.csr_array or None
+        Optional spatial smoother / upsampler applied to scalar data before
+        rendering.  When set (e.g., by
+        :meth:`~mne.viz.Brain.set_data_smoothing`), every call to
+        :meth:`update_overlay` will multiply the incoming scalars by this
+        matrix before storing them.  Shape must be
+        ``(n_surface_vertices, n_source_vertices)``.
+
+    Notes
+    -----
+    .. versionadded:: 1.13
+    """
+
     def __init__(self, renderer, vertices, triangles, normals):
         self._renderer = renderer
         self._vertices = vertices
@@ -91,8 +129,9 @@ def __init__(self, renderer, vertices, triangles, normals):
 
         self._default_scalars = np.ones(vertices.shape)
         self._default_scalars_name = "Data"
+        self.smooth_mat = None
 
-    def map(self):
+    def _map(self):
         kwargs = {
             "color": None,
             "pickable": True,
@@ -135,7 +174,30 @@ def _compose_overlays(self):
                 B = self._compute_over(cache, A)
         return B, cache
 
-    def add_overlay(self, scalars, colormap, rng, opacity, name):
+    def add_overlay(self, scalars, colormap, rng, opacity, name, smooth=False):
+        """Add a named overlay to the mesh.
+
+        Parameters
+        ----------
+        scalars : array, shape (n_vertices,)
+            Scalar values to display. If ``smooth=True`` and
+            ``smooth_mat`` is set, shape must be ``(n_src_vertices,)``.
+        colormap : array, shape (n_colors, 4)
+            RGBA colormap table (values in ``[0, 255]``).
+        rng : array-like, shape (2,)
+            ``[min, max]`` range for colormap mapping.
+        opacity : float | None
+            Overlay opacity in ``[0, 1]``. ``None`` keeps the existing value.
+        name : str
+            Unique key identifying this overlay.
+        smooth : bool
+            If ``True`` and ``smooth_mat`` is set, multiply ``scalars``
+            by ``smooth_mat`` before rendering. Use ``True`` only for
+            source-space data; surface-space overlays (curvature, labels,
+            annotations) should use the default ``False``.
+        """
+        if smooth and self.smooth_mat is not None:
+            scalars = self.smooth_mat.dot(scalars)
         overlay = _Overlay(
             scalars=scalars,
             colormap=colormap,
@@ -156,6 +218,13 @@ def add_overlay(self, scalars, colormap, rng, opacity, name):
         self._apply()
 
     def remove_overlay(self, names):
+        """Remove one or more overlays by name.
+
+        Parameters
+        ----------
+        names : str | list of str
+            Name(s) of the overlay(s) to remove.
+        """
         to_update = False
         if not isinstance(names, list):
             names = [names]
@@ -172,6 +241,14 @@ def _apply(self):
         self._polydata[self._default_scalars_name] = self._current_colors
 
     def update(self, colors=None):
+        """Recompose all overlays and refresh the mesh texture.
+
+        Parameters
+        ----------
+        colors : array-like of shape (n_triangles, 4) | None
+            Pre-composited RGBA colors to blend over the cached layer stack.
+            If ``None``, all overlays are recomposed from scratch.
+        """
         if colors is not None and self._cached_colors is not None:
             self._current_colors = self._compute_over(self._cached_colors, colors)
         else:
@@ -187,10 +264,30 @@ def _clean(self):
         self._renderer = None
 
     def update_overlay(self, name, scalars=None, colormap=None, opacity=None, rng=None):
+        """Update an existing overlay in-place.
+
+        Parameters
+        ----------
+        name : str
+            Key of the overlay to update (must already exist).
+        scalars : array, shape (n_vertices,) | None
+            New scalar values. If ``None``, scalars are unchanged.
+            If ``smooth_mat`` is set, smoothing is applied automatically
+            (matching the behaviour of the original :meth:`~mne.viz.Brain.add_data`
+            call that created the overlay).
+        colormap : array, shape (n_colors, 4) | None
+            New RGBA colormap table. If ``None``, colormap is unchanged.
+        opacity : float | None
+            New opacity in ``[0, 1]``. If ``None``, opacity is unchanged.
+        rng : array-like, shape (2,) | None
+            New ``[min, max]`` colormap range. If ``None``, range is unchanged.
+        """
         overlay = self._overlays.get(name, None)
         if overlay is None:
             return
         if scalars is not None:
+            if self.smooth_mat is not None:
+                scalars = self.smooth_mat.dot(scalars)
             overlay._scalars = scalars
         if colormap is not None:
             overlay._colormap = colormap

mne/viz/__init__.pyi

@@ -1,5 +1,6 @@
 __all__ = [
     "Brain",
+    "LayeredMesh",
     "ClickableImage",
     "EvokedField",
     "Figure3D",
@@ -102,7 +103,7 @@ from ._3d import (
     set_3d_options,
     snapshot_brain_montage,
 )
-from ._brain import Brain
+from ._brain import Brain, LayeredMesh
 from ._figure import get_browser_backend, set_browser_backend, use_browser_backend
 from ._proj import plot_projs_joint
 from .backends._abstract import Figure3D

mne/viz/_brain/__init__.py

@@ -4,7 +4,7 @@
 # License: BSD-3-Clause
 # Copyright the MNE-Python contributors.
 
-from ._brain import Brain, _LayeredMesh
+from ._brain import Brain, LayeredMesh
 from ._scraper import _BrainScraper
 from ._linkviewer import _LinkViewer