Speed up evoked time plotting

doc/changes/dev/14249.newfeature.rst

@@ -0,0 +1 @@
+Sped up various evoked plotting functions by taking advantage of blitting, by `Eric Larson`_

mne/gui/tests/test_dipolefit.py

@@ -173,7 +173,7 @@ def process_and_reenter():
     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
+    blit_artists = g._renderer._mplcanvas._blit._artists
     assert blit_artists == [g._time_line, g._time_text]
 
     g.fit_dipole()

mne/report/report.py

@@ -23,7 +23,6 @@
 
 import matplotlib
 import numpy as np
-from matplotlib.animation import AbstractMovieWriter
 
 from .. import __version__ as MNE_VERSION
 from .._fiff.meas_info import Info, read_info
@@ -85,7 +84,7 @@
 from ..viz._brain.view import views_dicts
 from ..viz._scraper import _mne_qt_browser_screenshot
 from ..viz.misc import _get_bem_plotting_surfaces, _plot_mri_contours
-from ..viz.utils import _ndarray_to_fig
+from ..viz.utils import _BlitManager, _ndarray_to_fig
 
 _BEM_VIEWS = ("axial", "sagittal", "coronal")
 
@@ -369,27 +368,6 @@ def _check_tags(tags) -> tuple[str]:
 # PLOTTING FUNCTIONS
 
 
-class _NdArrayCapture(AbstractMovieWriter):
-    def __init__(self, frames: list):
-        super().__init__(fps=1, metadata={}, bitrate=0)
-        self.frames = frames
-
-    def grab_frame(self, **savefig_kwargs):
-        img = _fig_to_img(
-            fig=self.fig, image_format="ndarray", pad_inches=0, **savefig_kwargs
-        )
-        self.frames.append(img)
-
-    def save(self, filename, *args, **kwargs):
-        pass
-
-    def finish(self):
-        pass
-
-    def setup(self, fig, outfile, dpi=None):
-        self.fig = fig
-
-
 def _use_agg(func):
     @functools.wraps(func)
     def wrapper(*args, **kwargs):
@@ -486,7 +464,17 @@ def _fig_to_img(
     logger.debug(
         f"Saving figure with dimension {fig.get_size_inches()} inches with {dpi} dpi"
     )
-    mpl_format = "svg" if image_format == "svg" else "png"
+    if image_format == "ndarray":
+        # Raw RGBA: the caller wants the rendered pixels, so encoding them as PNG
+        # only to decode them again below is pure overhead.
+        mpl_format = "rgba"
+        # Agg truncates the figure size to whole pixels (`RendererAgg.__init__`),
+        # so rounding here would mis-shape the buffer at fractional DPI
+        shape = (int(fig.bbox.size[1]), int(fig.bbox.size[0]), 4)
+    elif image_format == "svg":
+        mpl_format = "svg"
+    else:
+        mpl_format = "png"
     fig.savefig(output, format=mpl_format, dpi=dpi, **mpl_kwargs)
 
     if own_figure:
@@ -512,8 +500,9 @@ def _fig_to_img(
             new.save(output, format=image_format, dpi=(dpi, dpi), **pil_kwargs)
 
     if image_format == "ndarray":
-        output.seek(0)
-        output = plt.imread(output, format="png")
+        # float in [0, 1], like the PNG this used to go through
+        output = np.frombuffer(output.getbuffer(), np.uint8).reshape(shape)
+        output = output.astype(np.float32) / 255
     else:
         output = output.getvalue()
         if image_format == "svg":
@@ -3893,8 +3882,6 @@ def _plot_evoked_topomap_timepoints(
             fig.delaxes(axes[1, 1])
             axes = axes.ravel()[:3]
             axes[0].set_title(ch_type)
-            frames[ch_type] = list()
-            this_writer = _NdArrayCapture(frames[ch_type])
             _, ch_anim = evoked.animate_topomap(
                 times=times,
                 ch_type=ch_type,
@@ -3903,10 +3890,22 @@ def _plot_evoked_topomap_timepoints(
                 show=False,
                 time_format="",  # we impose our own in HTML
                 butterfly=True,
+                blit=False,  # we do our own, `Animation.save` cannot blit at all
                 **topomap_kwargs,
             )
+            _constrain_fig_resolution(fig, max_width=MAX_IMG_WIDTH, max_res=MAX_IMG_RES)
+            fig.canvas.draw()  # the animation does its initial draw here
             ch_anim.pause()
-            ch_anim.save("", writer=this_writer)
+            # Only the topomap image, its contours and the butterfly cursor change
+            # from one frame to the next, so blit those onto a cached picture of the
+            # rest of the figure and read the pixels straight out of the canvas.
+            blit = _BlitManager(fig)
+            frames[ch_type] = list()
+            for frame in range(len(times)):
+                blit.update(ch_anim.mne_frame_func(frame))
+                frames[ch_type].append(
+                    np.asarray(fig.canvas.buffer_rgba(), dtype=np.float32) / 255
+                )
             plt.close(fig)
             del (
                 fig,

mne/report/tests/test_report.py

@@ -33,6 +33,7 @@
 from mne.report.report import (
     _ALLOWED_IMAGE_FORMATS,
     CONTENT_ORDER,
+    _fig_to_img,
 )
 from mne.utils import Bunch, _record_warnings
 from mne.utils._testing import assert_object_equal
@@ -207,6 +208,9 @@ def test_render_report(renderer_pyvistaqt, tmp_path, invisible_fig):
 
     # ndarray support smoke test
     report.add_figure(fig=np.zeros((2, 3, 3)), title="title")
+    # ... and the reverse: a figure whose size is not a whole number of pixels
+    fig = plt.figure(figsize=(2.8, 2.8), dpi=89.6)
+    assert _fig_to_img(fig, image_format="ndarray").shape == (250, 250, 4)
 
     with pytest.raises(TypeError, match="It seems you passed a path"):
         report.add_figure(fig="foo", title="title")

mne/viz/_brain/tests/test_brain.py

@@ -1682,29 +1682,26 @@ def test_brain_time_line_blitting(renderer_interactive_pyvistaqt, brain_gc):
     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()
+    assert brain.time_line in canvas._blit._artists
 
     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
+    brain.set_time(brain._times[-1])  # one redraw caches the background ...
+    assert brain.time_line.get_xdata()[0] == brain._times[-1]
+    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]
+    brain.set_time(brain._times[len(brain._times) // 2])  # ... then it only blits
     assert len(n_draws) == 1
 
-    # adding a trace still redraws in full, and anything can be blitted
+    # a full redraw invalidates the background, 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
+    canvas.update_blit_artists()  # the 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 text not in canvas._blit._artists
     assert len(n_draws) == 3  # restored to the background by a full redraw
     brain.close()
 

mne/viz/backends/_abstract.py

@@ -8,6 +8,7 @@
 from abc import ABC, abstractmethod
 
 from ..ui_events import TimeChange, publish
+from ..utils import _BlitManager
 
 
 class Figure3D(ABC):
@@ -1429,12 +1430,7 @@ 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
+        self._blit = _BlitManager(self.fig, draw=self.update_plot)
 
     def _connect(self):
         for event in ("button_press", "motion_notify") + self._extra_events:
@@ -1458,36 +1454,19 @@ def plot_time_line(self, x, label, update=True, **kwargs):
     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``).
+            The artist to draw separately. Must live in this canvas's axes, and be
+            drawn on top of the curves, as blitting draws it over a cached picture
+            of the rest of the figure.
         """
-        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)
+        self._blit.add(artist)
 
     def remove_blit_artist(self, artist):
         """Stop drawing an artist separately, putting it back in the background.
@@ -1498,36 +1477,15 @@ def remove_blit_artist(self, 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
+        self._blit.remove(artist)
 
     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()
+        self._blit.update()
 
     def update_plot(self):
         """Update the plot."""
@@ -1584,11 +1542,7 @@ 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._blit.close()
         self.axes.clear()
         self.fig.clear()
         self.canvas = None

mne/viz/evoked.py

@@ -49,6 +49,7 @@
 from .ui_events import TimeChange, publish, subscribe
 from .utils import (
     DraggableColorbar,
+    _BlitManager,
     _check_cov,
     _check_delayed_ssp,
     _check_option,
@@ -587,6 +588,10 @@ def _plot_lines(
     sphere = _check_sphere(sphere, info)
     path_effects = [patheffects.withStroke(linewidth=2, foreground="w", alpha=0.75)]
     gfp_path_effects = [patheffects.withStroke(linewidth=5, foreground="w", alpha=0.75)]
+    # The time cursors and the hover label are the only artists that move, so draw
+    # them on top of a cached background rather than redrawing every channel's trace.
+    blit_manager = _BlitManager(fig)
+
     if selectable:
         selectables = np.ones(len(ch_types_used), dtype=bool)
         for type_idx, this_type in enumerate(ch_types_used):
@@ -632,22 +637,29 @@ def _on_hover(event):
                 else:
                     text.set_alpha(0.0)
                     text.set_path_effects([])
+                blit_manager.add(text)
 
             # vertical line to indicate time point
             for ax in axes:
                 line = getattr(ax, "_cursorline", None)
                 if line is None:
-                    ax._cursorline = ax.axvline(event.xdata, color="black", alpha=0.2)
+                    # zorder: blitting draws the cursor over a cached picture of the
+                    # rest of the figure, so it has to be on top of the traces for
+                    # the blitted figure to match a full redraw
+                    line = ax._cursorline = ax.axvline(
+                        event.xdata, color="black", alpha=0.2, zorder=len(ax.lines)
+                    )
+                    blit_manager.add(line)
                 else:
                     line.set_xdata([event.xdata, event.xdata])
-            ax.figure.canvas.draw_idle()
+                    line.set_visible(True)
+            blit_manager.update()
 
         def _rm_cursor(event):
             for ax in axes:
                 if getattr(ax, "_cursorline", None) is not None:
-                    ax._cursorline.remove()
-                    ax._cursorline = None
-            ax.figure.canvas.draw_idle()
+                    ax._cursorline.set_visible(False)
+            blit_manager.update()
 
         def _select_time(event):
             for ax in axes:
@@ -886,10 +898,13 @@ def on_time_change(event):
         for ax in axes:
             line = getattr(ax, "_selectline", None)
             if line is None:
-                ax._selectline = ax.axvline(event.time, color="black", alpha=1)
+                ax._selectline = ax.axvline(
+                    event.time, color="black", alpha=1, zorder=len(ax.lines)
+                )
+                blit_manager.add(ax._selectline)
             else:
                 line.set_xdata([event.time, event.time])
-        ax.figure.canvas.draw()
+        blit_manager.update()
 
     subscribe(fig, "time_change", on_time_change)
 

mne/viz/tests/test_topomap.py

@@ -184,6 +184,17 @@ def test_plot_topomap_animation(capsys, tmp_path):
     assert "extrapolation mode local to mean" in out
     assert fig.axes[0].images[0].get_cmap().name == "viridis"
 
+    # everything drawn on top of the topomap image must be returned by the animation
+    # function, otherwise blitting leaves it frozen at the first frame (gh-14242)
+    items = anim.mne_frame_func(1)  # has to be tested separately on the 'Agg' backend
+    ax = fig.axes[0]
+    zorder = ax.images[0].get_zorder()
+    on_top = [
+        a for a in ax.lines + ax.collections + ax.texts if a.get_zorder() > zorder
+    ]
+    assert len(on_top) > 2  # at least the time label, head outlines and sensors
+    assert set(on_top).issubset(items)
+
     # saving
     PIL = pytest.importorskip("PIL")
     gif_path = tmp_path / "test.gif"
@@ -222,7 +233,7 @@ def test_plot_topomap_animation_csd(capsys):
     _, anim = evoked_csd.animate_topomap(
         ch_type="csd", times=[0, 0.1], butterfly=False, time_unit="s", verbose="debug"
     )
-    anim._func(1)  # _animate has to be tested separately on 'Agg' backend.
+    anim.mne_frame_func(1)  # has to be tested separately on the 'Agg' backend
     out, _ = capsys.readouterr()
     assert "extrapolation mode head to mean" in out
 
@@ -964,7 +975,7 @@ def test_plot_projs_topomap_opm(triaxial_evoked):
 def test_animate_topomap_opm(triaxial_evoked):
     """Test animate_topomap does not crash on colocated OPM channels (gh-13866)."""
     fig, anim = triaxial_evoked.animate_topomap(ch_type="mag", times=[0.0], show=False)
-    anim._func(0)
+    anim.mne_frame_func(0)
     assert len(fig.axes) >= 1