Speed up dipole fit rendering

Author: larsoner

mne/gui/_dipolefit.py

@@ -387,6 +387,11 @@ def _configure_main_display(self, show_sensors=True):
                 time_viewer=False,
                 initial_time=self._current_time,
                 time_label=None,  # the traces plot shows the current time
+                # The source estimate is only a rough guide for where to put
+                # dipoles, so map each surface vertex to its nearest source rather
+                # than smoothing: the upsampling is then a gather instead of a
+                # sparse matrix product, which is cheaper on every time change.
+                smoothing_steps="nearest",
                 brain_kwargs=dict(units="m", show=False),
                 figure=fig_into,
                 # the GUI renders on a white figure, so the Brain (and hence its
@@ -707,7 +712,9 @@ def _on_time_change(self, event):
         if self._time_line is not None:
             self._time_line.set_xdata([new_time])
             self._update_time_text()
-            self._renderer._mplcanvas.update_plot()
+            # only the time line and its label moved, so the traces can be blitted
+            # from the cached background instead of being redrawn
+            self._renderer._mplcanvas.update_blit_artists()
         self._update_arrows()
 
     def _update_time_text(self):
@@ -1364,6 +1371,8 @@ def _setup_mplcanvas(self):
                 fontsize=8,
                 color="black",
             )
+            # the label travels with the time line, so it is drawn along with it
+            canvas.add_blit_artist(self._time_text)
         return self._renderer._mplcanvas
 
     def close(self):

mne/gui/tests/test_dipolefit.py

@@ -171,6 +171,10 @@ def process_and_reenter():
     assert not hasattr(g._fig, "_time_label_actor")
     assert re.fullmatch(r"90 ms · GOF \d+%", g._time_text.get_text())
     assert g._time_text.get_position()[0] == 0.09
+    # both move with the time, so they are drawn on top of a cached background
+    # rather than triggering a full redraw of the traces plot
+    blit_artists = g._renderer._mplcanvas._blit_artists
+    assert blit_artists == [g._time_line, g._time_text]
 
     g.fit_dipole()
     assert len(g._dipoles) == len(g.dipoles) == 2

mne/viz/_3d_overlay.py

@@ -151,17 +151,28 @@ def _map(self):
         self._is_mapped = True
 
     def _compute_over(self, B, A):
+        # Alpha-composite A ("over") on top of B, both RGBA in [0, 1].
+        #
+        # This runs on every time point of an interactive time course, on surfaces
+        # with >1e5 vertices, so it is written to touch the (n_vertices, 4) arrays
+        # as few times as possible and only ever whole: expressing it in terms of
+        # the RGB columns (``C[:, :3] *= ...``) makes every operation strided,
+        # which costs ~4x more than the same work on the full array. The alpha
+        # column is included in the arithmetic and simply overwritten at the end.
         assert A.ndim == B.ndim == 2
         assert A.shape[1] == B.shape[1] == 4
-        A_w = A[:, 3:]  # * 1
-        B_w = B[:, 3:] * (1 - A_w)
-        C = A.copy()
-        C[:, :3] *= A_w
-        C[:, :3] += B[:, :3] * B_w
-        C[:, 3:] += B_w
-        C_alpha_zero = C[:, 3] == 0
-        C[~C_alpha_zero, :3] /= C[~C_alpha_zero, 3:]
-        C[C_alpha_zero, :3] = 0
+        A_w = A[:, 3].copy()  # copy: column slices of a (n, 4) array are strided
+        B_w = B[:, 3].copy()
+        B_w *= 1 - A_w
+        C = A * A_w[:, None]
+        C += B * B_w[:, None]
+        alpha = A_w + B_w
+        # Where the composite is fully transparent the color is undefined: divide
+        # by one there instead, and zero those rows out afterwards.
+        opaque = alpha != 0
+        np.divide(C, np.where(opaque, alpha, 1)[:, None], out=C)
+        C *= opaque[:, None]
+        C[:, 3] = alpha
         return np.clip(C, 0, 1, out=C)
 
     def _compose_overlays(self):

mne/viz/_brain/_brain.py

@@ -1907,7 +1907,9 @@ def plot_time_line(self, update=True):
                 )
             self.time_line.set_xdata([current_time])
             if update:
-                self.mpl_canvas.update_plot()
+                # only the time line moved, so the rest of the figure can be
+                # blitted from the cached background instead of being redrawn
+                self.mpl_canvas.update_blit_artists()
 
     def _configure_help(self):
         pairs = [
@@ -4295,7 +4297,7 @@ def _update_current_time_idx(self, time_idx):
         time_actor = active.get("time_actor", None)
         time_label = active.get("time_label", None)
         for hemi in ["lh", "rh", "vol"]:
-            hemi_needs_recompose = False
+            staged_keys = list()
             for data_key, key_data in self._all_data.items():
                 hemi_data = key_data.get(hemi)
                 if hemi_data is None:
@@ -4353,10 +4355,10 @@ def _update_current_time_idx(self, time_idx):
                         key_data["fmax"],
                     ]
                     if data_key in mesh._overlays:
-                        # Stage without recomposing; a single mesh.update() below
-                        # handles all overlays in O(N) instead of O(N²).
+                        # Stage without recomposing; a single update below handles
+                        # all overlays in O(N) instead of O(N²).
                         mesh.update_overlay(data_key, scalars=act_data, update=False)
-                        hemi_needs_recompose = True
+                        staged_keys.append(data_key)
                     else:
                         mesh.add_overlay(
                             scalars=act_data,
@@ -4371,8 +4373,15 @@ def _update_current_time_idx(self, time_idx):
                 if vectors is not None and data_key == self._active_data_key:
                     self._update_glyphs(hemi, vectors)
 
-            if hemi_needs_recompose and hemi in self.layered_meshes:
-                self.layered_meshes[hemi].update()
+            if staged_keys and hemi in self.layered_meshes:
+                if len(staged_keys) == 1:
+                    # Let update_overlay pick the cached path when the overlay we
+                    # staged is the topmost one: the layers below it (curvature,
+                    # labels, ...) have not changed, so their composite can be
+                    # reused instead of color-mapping them all again.
+                    self.layered_meshes[hemi].update_overlay(staged_keys[0])
+                else:
+                    self.layered_meshes[hemi].update()
 
         active["time_idx"] = time_idx
         self._renderer._update()

mne/viz/_brain/tests/test_brain.py

@@ -190,6 +190,21 @@ def test_layered_mesh(renderer_interactive_pyvistaqt):
             opacity=np.array([0.1, 0.2, 0.3]),
             name="bad-opacity",
         )
+
+    # alpha compositing: transparent top keeps the bottom color, opaque top wins,
+    # and a half-transparent white over opaque black is grey
+    bottom = np.array([[0.0, 0, 0, 1]] * 3)
+    top = np.array([[1.0, 1, 1, 0], [1, 1, 1, 1], [1, 1, 1, 0.5]])
+    assert_allclose(
+        mesh._compute_over(bottom, top),
+        [[0, 0, 0, 1], [1, 1, 1, 1], [0.5, 0.5, 0.5, 1]],
+    )
+    # a fully transparent result is black, and the inputs are left alone
+    bottom, top = np.zeros((1, 4)), np.zeros((1, 4))
+    assert_allclose(mesh._compute_over(bottom, top), [[0, 0, 0, 0]])
+    assert_allclose(bottom, 0)
+    assert_allclose(top, 0)
+
     mesh._clean()
 
 
@@ -1661,6 +1676,39 @@ def row_text(row):
     assert_allclose(peak_line3.get_ydata(), 2.0 * y1)
 
 
+@testing.requires_testing_data
+def test_brain_time_line_blitting(renderer_interactive_pyvistaqt, brain_gc):
+    """Test that moving the time line blits instead of redrawing the traces."""
+    brain = _create_testing_brain(hemi="lh", show_traces=True, initial_time=0)
+    canvas = brain.mpl_canvas
+    assert canvas.canvas.supports_blit
+    assert brain.time_line in canvas._blit_artists
+    assert brain.time_line.get_animated()
+
+    n_draws = list()
+    canvas.canvas.mpl_connect("draw_event", lambda event: n_draws.append(event))
+    canvas.update_plot()  # a full redraw caches the background ...
+    assert canvas._blit_background is not None
+    assert len(n_draws) == 1
+
+    brain.set_time(brain._times[-1])  # ... so moving the time line only blits
+    assert brain.time_line.get_xdata()[0] == brain._times[-1]
+    assert len(n_draws) == 1
+
+    # adding a trace still redraws in full, and anything can be blitted
+    text = canvas.axes.text(0, 0, "hello")
+    canvas.add_blit_artist(text)
+    assert text.get_animated()
+    canvas.update_blit_artists()  # background was dropped, so this redraws
+    assert len(n_draws) == 2
+
+    canvas.remove_blit_artist(text)
+    assert not text.get_animated()
+    assert text not in canvas._blit_artists
+    assert len(n_draws) == 3  # restored to the background by a full redraw
+    brain.close()
+
+
 def _send_mouse_move(widget, point, buttons=None):
     """Deliver a synthetic Qt mouse move (QTest.mouseMove warps the real cursor)."""
     from qtpy.QtCore import QEvent, QPointF, Qt

mne/viz/backends/_abstract.py

@@ -1429,6 +1429,12 @@ def __init__(self, width, height, dpi):
         self.axes = self.fig.add_subplot(111)
         self.axes.set(xlabel="Time (s)", ylabel="Activation (AU)")
         self.manager = None
+        # Artists that are redrawn on their own (see `add_blit_artist`), the
+        # background they are drawn onto, and the draw_event callback id that
+        # keeps that background up to date.
+        self._blit_artists = list()
+        self._blit_background = None
+        self._blit_cid = None
 
     def _connect(self):
         for event in ("button_press", "motion_notify") + self._extra_events:
@@ -1444,10 +1450,85 @@ def plot(self, x, y, label, update=True, **kwargs):
     def plot_time_line(self, x, label, update=True, **kwargs):
         """Plot the vertical line."""
         line = self.axes.axvline(x, label=label, **kwargs)
+        self.add_blit_artist(line)
         if update:
             self.update_plot()
         return line
 
+    def add_blit_artist(self, artist):
+        """Mark an artist as fast-updating, to be drawn by :meth:`update_blit_artists`.
+
+        Such an artist is excluded from the canvas background, so that moving it
+        (e.g. the time line, or a label that travels with it) costs a blit of the
+        cached background rather than a full redraw of the figure.
+
+        Parameters
+        ----------
+        artist : instance of matplotlib.artist.Artist
+            The artist to draw separately. Must live in this canvas's axes: an
+            artist added to the figure itself would be left out of saved images,
+            because Matplotlib only exempts *Axes* children from the rule that
+            animated artists are not drawn (see ``_AxesBase.draw``).
+        """
+        if not self.canvas.supports_blit:  # e.g. ipympl in a notebook
+            return
+        if artist.axes is not self.axes:
+            raise RuntimeError(
+                f"{artist!r} must be an artist of this canvas's axes to be drawn "
+                "separately, got one in " + repr(artist.axes)
+            )
+        if artist in self._blit_artists:
+            return
+        artist.set_animated(True)
+        self._blit_artists.append(artist)
+        # the cached background may already contain this artist, so drop it and
+        # let the next update redraw (and re-cache) the figure without it
+        self._blit_background = None
+        if self._blit_cid is None:
+            # Grab a fresh background after every full redraw, whatever caused it
+            # (update_plot, draw_idle, a resize, a DPI change, ...).
+            self._blit_cid = self.canvas.mpl_connect("draw_event", self._on_draw)
+
+    def remove_blit_artist(self, artist):
+        """Stop drawing an artist separately, putting it back in the background.
+
+        Parameters
+        ----------
+        artist : instance of matplotlib.artist.Artist
+            The artist to stop drawing separately. Artists that were never added
+            are ignored.
+        """
+        if artist not in self._blit_artists:
+            return
+        self._blit_artists.remove(artist)
+        artist.set_animated(False)
+        self.update_plot()  # redraw so the artist becomes part of the background
+
+    def update_blit_artists(self):
+        """Redraw only the artists added with :meth:`add_blit_artist`.
+
+        This is the fast path taken while the time line moves; any other change
+        to the figure needs :meth:`update_plot` instead.
+        """
+        if self._blit_background is None or not self._blit_artists:
+            self.update_plot()  # nothing cached yet (or nothing to draw fast)
+            return
+        self.canvas.restore_region(self._blit_background)
+        self._draw_blit_artists()
+        self.canvas.blit(self.fig.bbox)
+
+    def _draw_blit_artists(self):
+        for artist in self._blit_artists:
+            self.fig.draw_artist(artist)
+
+    def _on_draw(self, event=None):
+        """Cache the background after a full redraw (draw_event callback)."""
+        self._blit_background = self.canvas.copy_from_bbox(self.fig.bbox)
+        if not self.canvas.is_saving():
+            # When saving, Matplotlib draws animated artists itself; drawing them
+            # again here would just double up their antialiasing.
+            self._draw_blit_artists()
+
     def update_plot(self):
         """Update the plot."""
         with warnings.catch_warnings(record=True):
@@ -1503,6 +1584,11 @@ def close(self):
     def clear(self):
         """Clear internal variables."""
         self.close()
+        if self._blit_cid is not None:
+            self.canvas.mpl_disconnect(self._blit_cid)
+            self._blit_cid = None
+        self._blit_artists.clear()  # the artists go away with the figure below
+        self._blit_background = None
         self.axes.clear()
         self.fig.clear()
         self.canvas = None