Allow subclasses of FigureClass to be passed to plot_raw/plot_epochs (#13979)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: Eric Larson <larson.eric.d@gmail.com>

Author: 3 people

doc/changes/dev/13979.newfeature.rst

@@ -0,0 +1 @@
+Allow subclasses of ``MNEBrowseFigure`` to be passed to plot_raw/plot_epochs, as well as the corresponding ``plot(...)`` methods of the raw and epochs classes, by :newcontrib:`Frankie Robertson`

doc/changes/names.inc

@@ -114,6 +114,7 @@
 .. _Felix Raimundo: https://github.com/gamazeps
 .. _Florian Hofer: https://github.com/hofaflo
 .. _Florin Pop: https://github.com/florin-pop
+.. _Frankie Robertson: https://github.com/frankier
 .. _Frederik Weber: https://github.com/Frederik-D-Weber
 .. _Fu-Te Wong: https://github.com/zuxfoucault
 .. _Gennadiy Belonosov: https://github.com/Genuster

doc/conf.py

@@ -323,6 +323,7 @@
     "and",
     "as",
     "between",
+    "class",
     "data",
     "instance",
     "instances",

examples/visualization/onionskin.py

@@ -0,0 +1,161 @@
+"""
+==============================================================
+Advanced plotting customization by subclassing MNEBrowseFigure
+==============================================================
+
+This example shows how plot_epochs(...) and plot_raw(...) can be customized by
+subclassing MNEBrowseFigure and using the ``figure_class`` argument.
+It plots one EEG trace overlaid ("onion-skinned") on top of another.
+
+This example is "bad code" in a few ways:
+
+* Since the interface for MNEBrowseFigure is not public, it is liable to
+  break between minor and even patch versions of MNE without warning
+* Some functionality is reimplemented from MNEBrowseFigure in a more or
+  less copy-paste style
+* The code is backend-specific, in particular it is limited to the
+  matplotlib backend, and will not work with the qt browser
+* Since there is no way to pass another EEG directly to the MNEBrowseFigure,
+  it is passed through a global variable
+
+Nevertheless, the example shows that the "escape hatch" of using a subclass is
+available when other customization possibilities offered by MNE are not
+sufficient.
+"""
+
+# Authors: The MNE-Python contributors.
+# License: BSD-3-Clause
+# Copyright the MNE-Python contributors.
+
+from mne.datasets import eegbci
+from mne.io import read_raw_edf
+from mne.viz import set_browser_backend
+from mne.viz._mpl_figure import MNEBrowseFigure as MNEBrowseFigureOrig
+
+set_browser_backend("matplotlib")
+
+
+onionskin_eeg = None
+
+
+def _set_onionskin_eeg(eeg):
+    global onionskin_eeg
+    onionskin_eeg = eeg
+
+
+class OnionskinMNEBrowseFigure(MNEBrowseFigureOrig):
+    """
+    Subclass of MNEBrowseFigure adding in onion-skin functionality,
+    i.e. plotting one EEG trace overlaid on top of another.
+    """
+
+    def __init__(self, *args, **kwargs):
+        import numpy as np
+
+        super().__init__(*args, **kwargs)
+        onionskin_kwargs = {
+            **self.mne.trace_kwargs,
+        }
+        self.mne.onionskins = self.mne.ax_main.plot(
+            np.full((1, self.mne.n_channels), np.nan), **onionskin_kwargs
+        )
+
+    def _update_data(self):
+        import numpy as np
+
+        from mne.io.base import BaseRaw
+
+        super()._update_data()
+        if not onionskin_eeg:
+            self.mne.onionskin_data = None
+            return
+        start, stop = self._get_start_stop()
+        if isinstance(onionskin_eeg, BaseRaw):
+            if stop is None:
+                data = onionskin_eeg[:, start:]
+            else:
+                data = onionskin_eeg[:, start:stop]
+            data = data[0]
+        else:
+            ix_start = np.searchsorted(
+                self.mne.boundary_times, self.mne.t_start - self.mne.sampling_period
+            )
+            ix_stop = ix_start + self.mne.n_epochs
+            item = slice(ix_start, ix_stop)
+            print(type(onionskin_eeg))
+            data = np.concatenate(
+                onionskin_eeg.get_data(item=item, copy=False), axis=-1
+            )
+        data = self._process_data(data, start, stop, picks=self.mne.picks)
+        self.mne.onionskin_data = data
+
+    def _draw_traces(self):
+        import numpy as np
+        from matplotlib.colors import to_rgba_array
+        from matplotlib.patches import Rectangle
+
+        super()._draw_traces()
+        if self.mne.onionskin_data is None:
+            return
+        picks = self.mne.picks
+        offset_ixs = (
+            picks
+            if self.mne.butterfly and self.mne.ch_selections is None
+            else slice(None)
+        )
+        offsets = self.mne.trace_offsets[offset_ixs]
+
+        ch_colors = to_rgba_array(self.mne.ch_colors)
+        ch_colors[:, 3] *= 0.5
+
+        decim = np.ones_like(picks)
+        data_picks_mask = np.isin(picks, self.mne.picks_data)
+        decim[data_picks_mask] = self.mne.decim
+        # decim can vary by channel type, so compute different `times` vectors
+        decim_times = {
+            decim_value: self.mne.times[::decim_value] + self.mne.first_time
+            for decim_value in set(decim)
+        }
+
+        time_range = (self.mne.times + self.mne.first_time)[[0, -1]]
+        ylim = self.mne.ax_main.get_ylim()
+        for ii, line in enumerate(self.mne.onionskins):
+            this_offset = offsets[ii]
+            this_times = decim_times[decim[ii]]
+            this_data = (
+                this_offset - self.mne.onionskin_data[ii] * self.mne.scale_factor
+            )
+            this_data = this_data[..., :: decim[ii]]
+            clip = 0.2 if self.mne.butterfly else 0.5
+            bottom = max(this_offset - clip, ylim[1])
+            height = min(2 * clip, ylim[0] - bottom)
+            rect = Rectangle(
+                xy=np.array([time_range[0], bottom]),
+                width=time_range[1] - time_range[0],
+                height=height,
+                transform=self.mne.ax_main.transData,
+            )
+            line.set_clip_path(rect)
+            line.set_xdata(this_times)
+            line.set_ydata(this_data)
+            color = ch_colors[ii]
+            line.set_color(color)
+            line.set_zorder(self.mne.zorder["data"] - 1)
+
+
+subjects = [1]
+runs = [1, 2]
+raw_fnames = eegbci.load_data(subjects, runs)
+first_data = read_raw_edf(raw_fnames[0], preload=True)
+second_data = read_raw_edf(raw_fnames[1], preload=True)
+
+first_data.plot(title="First plot", n_channels=3)
+second_data.plot(title="Second plot", n_channels=3)
+
+_set_onionskin_eeg(first_data)
+second_data.plot(
+    title="Onionskinned plot",
+    n_channels=3,
+    block=True,
+    figure_class=OnionskinMNEBrowseFigure,
+)

mne/epochs.py

@@ -1327,6 +1327,7 @@ def plot(
         overview_mode=None,
         splash=True,
         annotation_colors=None,
+        figure_class=None,
     ):
         return plot_epochs(
             self,
@@ -1354,6 +1355,7 @@ def plot(
             overview_mode=overview_mode,
             splash=splash,
             annotation_colors=annotation_colors,
+            figure_class=figure_class,
         )
 
     @copy_function_doc_to_method_doc(plot_topo_image_epochs)

mne/io/base.py

@@ -2001,6 +2001,7 @@ def plot(
         overview_mode=None,
         splash=True,
         verbose=None,
+        figure_class=None,
     ):
         return plot_raw(
             self,
@@ -2042,6 +2043,7 @@ def plot(
             overview_mode=overview_mode,
             splash=splash,
             verbose=verbose,
+            figure_class=figure_class,
         )
 
     @property

mne/utils/docs.py

@@ -1721,6 +1721,14 @@ def _reflow_param_docstring(docstring, has_first_line=True, width=75):
     A matplotlib-compatible color to use for the figure background. Defaults to black.
 """
 
+docdict["figure_class"] = """
+figure_class : class
+    The backend specific ``MNEBrowseFigure`` class to use. This is typically used
+    to pass a subclass in order to customize the plot. This parameter requires
+    cooperation from the backend, and is currently only supported by the
+    ``matplotlib`` backend.
+"""
+
 docdict["filter_length"] = """
 filter_length : str | int
     Length of the FIR filter to use (if applicable):

mne/viz/_mpl_figure.py

@@ -2638,7 +2638,11 @@ def _init_browser(**kwargs):
     """Instantiate a new MNE browse-style figure."""
     from mne.io import BaseRaw
 
-    fig = _figure(toolbar=False, FigureClass=MNEBrowseFigure, layout=None, **kwargs)
+    figure_class = kwargs.pop("figure_class", None)
+    if figure_class is None:
+        figure_class = MNEBrowseFigure
+
+    fig = _figure(toolbar=False, FigureClass=figure_class, layout=None, **kwargs)
 
     # splash is ignored (maybe we could do it for mpl if we get_backend() and
     # check if it's Qt... but seems overkill)

mne/viz/epochs.py

@@ -765,6 +765,7 @@ def plot_epochs(
     overview_mode=None,
     splash=True,
     annotation_colors=None,
+    figure_class=None,
 ):
     """Visualize epochs.
 
@@ -875,6 +876,9 @@ def plot_epochs(
         will trigger a warning. If ``None`` (default), automatic colors are used.
 
         .. versionadded:: 1.12.1
+    %(figure_class)s
+
+        .. versionadded:: 1.13
 
     Returns
     -------
@@ -1089,6 +1093,7 @@ def plot_epochs(
         theme=theme,
         overview_mode=overview_mode,
         splash=splash,
+        figure_class=figure_class,
     )
 
     fig = _get_browser(show=show, block=block, **params)

mne/viz/raw.py

@@ -75,6 +75,7 @@ def plot_raw(
     overview_mode=None,
     splash=True,
     verbose=None,
+    figure_class=None,
 ):
     """Plot raw data.
 
@@ -227,6 +228,9 @@ def plot_raw(
 
         .. versionadded:: 1.6
     %(verbose)s
+    %(figure_class)s
+
+        .. versionadded:: 1.13
 
     Returns
     -------
@@ -435,6 +439,7 @@ def plot_raw(
         theme=theme,
         overview_mode=overview_mode,
         splash=splash,
+        figure_class=figure_class,
     )
 
     fig = _get_browser(show=show, block=block, **params)