ENH: Implement TimeChange behaviour for evoked_topo and recursive option in link (#13968)

Co-authored-by: Marijn van Vliet <w.m.vanvliet@gmail.com>

Author: Gnefil

doc/changes/dev/13968.apichange.rst

@@ -0,0 +1 @@
+In :func:`mne.viz.iter_topography`, the ``on_pick`` callable for sub-figures now has a new parameter ``orig_fig`` that refers to the main figure. To preserve backward compatibility, it is made optional. By `Lifeng Qiu Lin`_.

doc/changes/dev/13968.newfeature.rst

@@ -0,0 +1,2 @@
+Add ``TimeChange`` behaviour for :func:`mne.viz.plot_evoked_topo`, where time is marked with solid vertical line after clicking on subfigure axis. 
+Add ``recursive`` parameter to :func:`mne.viz.ui_events.link` that enables linking to a target figure, as well as all figures that the target figure links to, by `Lifeng Qiu Lin`_.

mne/evoked.py

@@ -614,12 +614,6 @@ def plot_topo(
         select=False,
         show=True,
     ):
-        """.
-
-        Notes
-        -----
-        .. versionadded:: 0.10.0
-        """
         return plot_evoked_topo(
             self,
             layout=layout,

mne/viz/evoked.py

@@ -1243,6 +1243,14 @@ def plot_evoked_topo(
     -------
     fig : instance of matplotlib.figure.Figure
         Images of evoked responses at sensor locations.
+
+    Notes
+    -----
+    The figure will publish and subscribe to the following UI events:
+
+    * :class:`~mne.viz.ui_events.TimeChange`
+
+    .. versionadded:: 1.13.0
     """
     if type(evoked) not in (tuple, list):
         evoked = [evoked]

mne/viz/tests/test_topo.py

@@ -25,6 +25,7 @@
     mne_analyze_colormap,
     plot_evoked_topo,
     plot_topo_image_epochs,
+    ui_events,
 )
 from mne.viz.evoked import _line_plot_onselect
 from mne.viz.topo import _imshow_tfr, _plot_update_evoked_topo_proj, iter_topography
@@ -274,6 +275,16 @@ def test_plot_topo_single_ch():
     assert isinstance(ax._cursorline, matplotlib.lines.Line2D)
     _fake_click(fig, ax, (1.5, 1.5), kind="motion")  # cursor should disappear
     assert ax._cursorline is None
+    # test select bar
+    _fake_click(fig, ax, (1.5, 1.5), kind="press")
+    assert ax._selectline is None
+    _fake_click(fig, ax, (0.5, 0.5), kind="press")
+    assert isinstance(ax._selectline, matplotlib.lines.Line2D)
+    init_time = ax._selectline.get_xdata()[0]
+    assert init_time == 0.0
+    _fake_click(fig, ax, (0.6, 0.5), kind="press")
+    changed_time = ax._selectline.get_xdata()[0]
+    assert changed_time != init_time
     plt.close("all")
 
 
@@ -334,6 +345,32 @@ def test_plot_topo_select():
     assert fig.lasso.selection == ["MEG 0111", "MEG 0132", "MEG 0133", "MEG 0131"]
 
 
+def test_plot_topo_timechange():
+    """Test that time change events are properly handled in plot_evoked_topo."""
+    evoked = _get_epochs().average()
+    fig = plot_evoked_topo(evoked)
+    _fake_click(fig, fig.axes[0], (0.08, 0.65))  # open single channel
+    subfig = plt.gcf()
+    ax = plt.gca()
+    _fake_click(subfig, ax, (0.5, 0.5), kind="press")
+    init_time = ax._selectline.get_xdata()[0]
+    assert init_time == 0.0
+
+    # test existence of _current_time in main figure
+    assert hasattr(fig, "_current_time")
+    assert fig._current_time == 0.0
+
+    # test time change event from subfig and main fig
+    ui_events.publish(subfig, ui_events.TimeChange(time=0.1))
+    assert ax._selectline.get_xdata()[0] == 0.1
+    assert fig._current_time == 0.1
+    ui_events.publish(subfig, ui_events.TimeChange(time=0.2))
+    assert ax._selectline.get_xdata()[0] == 0.2
+    assert fig._current_time == 0.2
+
+    plt.close("all")
+
+
 def test_plot_tfr_topo():
     """Test plotting of TFR data."""
     epochs = _get_epochs()

mne/viz/tests/test_ui_events.py

@@ -234,9 +234,19 @@ def callback(event):
     ui_events.publish(fig2, ui_events.TimeChange(time=10.2))
     assert len(callback_calls) == 2  # Only called for both figures once
 
+    # Test recursive linking.
+    fig3 = plt.figure()
+    ui_events.subscribe(fig3, "time_change", callback)
+    ui_events.link(fig2, fig3, recursive=True)  # fig1 and fig2 are already linked
+    callback_calls.clear()
+    ui_events.publish(fig3, ui_events.TimeChange(time=10.2))
+    ui_events.publish(fig1, ui_events.TimeChange(time=10.2))
+    assert len(callback_calls) == 6  # Called for all three figures twice
+
     # Test cleanup
     fig1.canvas.callbacks.process("close_event", None)
     fig2.canvas.callbacks.process("close_event", None)
+    fig3.canvas.callbacks.process("close_event", None)
     assert len(event_channels) == 0
     assert len(event_channel_links) == 0
 

mne/viz/topo.py

@@ -6,14 +6,21 @@
 
 from copy import deepcopy
 from functools import partial
+from inspect import signature
 
 import numpy as np
 from scipy import ndimage
 
 from .._fiff.pick import _FNIRS_CH_TYPES_SPLIT, _picks_to_idx, channel_type, pick_types
 from ..defaults import _handle_default
 from ..utils import Bunch, _check_option, _clean_names, _is_numeric, _to_rgb, fill_doc
-from .ui_events import ChannelsSelect, publish, subscribe
+from .ui_events import (
+    ChannelsSelect,
+    TimeChange,
+    link,
+    publish,
+    subscribe,
+)
 from .utils import (
     DraggableColorbar,
     SelectFromCollection,
@@ -57,7 +64,7 @@ def iter_topography(
     on_pick : callable | None
         The callback function to be invoked on clicking one
         of the axes. Is supposed to instantiate the following
-        API: ``function(axis, channel_index)``.
+        API: ``function(axis, channel_index, orig_fig)``.
     fig : matplotlib.figure.Figure | None
         The figure object to be considered. If None, a new
         figure will be created.
@@ -413,14 +420,18 @@ def _plot_topo_onpick(event, show_func):
             return
         ch_idx = orig_ax._mne_ch_idx
         face_color = orig_ax._mne_ax_face_color
-        fig, ax = plt.subplots(1)
+        subfig, ax = plt.subplots(1)
 
         plt.title(orig_ax._mne_ch_name)
         ax.set_facecolor(face_color)
 
         # allow custom function to override parameters
-        show_func(ax, ch_idx)
-        plt_show(fig=fig)
+        if "orig_fig" in signature(show_func).parameters:
+            show_func(ax, ch_idx, orig_fig=fig)
+            plt_show(fig=subfig)
+        else:
+            show_func(ax, ch_idx)
+            plt_show(fig=fig)
 
     except Exception as err:
         # matplotlib silently ignores exceptions in event handlers,
@@ -568,6 +579,7 @@ def _plot_timeseries(
     hline=None,
     hvline_color="w",
     labels=None,
+    orig_fig=None,
 ):
     """Show time series on topo split across multiple axes."""
     import matplotlib.pyplot as plt
@@ -627,7 +639,7 @@ def _cursor_vline(event):
             return
         if ax._cursorline is not None:
             ax._cursorline.remove()
-        ax._cursorline = ax.axvline(event.xdata, color=ax._cursorcolor)
+        ax._cursorline = ax.axvline(event.xdata, color=ax._cursorcolor, alpha=0.2)
         ax.figure.canvas.draw()
 
     def _rm_cursor(event):
@@ -637,14 +649,35 @@ def _rm_cursor(event):
             ax._cursorline = None
         ax.figure.canvas.draw()
 
+    def _on_click(event):
+        if event.inaxes == ax:
+            publish(ax.figure, TimeChange(time=event.xdata))
+
+    def _on_time_change_sub(event):
+        _update_selectline(event.time)
+
+    def _update_selectline(time):
+        if ax._selectline is not None:
+            ax._selectline.remove()
+        ax._selectline = ax.axvline(time, color=ax._selectcolor, alpha=1)
+        ax.figure.canvas.draw()
+
     ax._cursorline = None
     # choose cursor color based on perceived brightness of background
     facecol = _to_rgb(ax.get_facecolor())
     face_brightness = np.dot(facecol, [299, 587, 114])
     ax._cursorcolor = "white" if face_brightness < 150 else "black"
 
+    ax._selectline = None
+    ax._selectcolor = "white" if face_brightness < 150 else "black"
+
     plt.connect("motion_notify_event", _cursor_vline)
     plt.connect("axes_leave_event", _rm_cursor)
+    plt.connect("button_press_event", _on_click)
+
+    subscribe(ax.figure, "time_change", _on_time_change_sub)
+
+    link(orig_fig, ax.figure, recursive=True)
 
     ymin, ymax = ax.get_ylim()
     # don't pass vline or hline here (this fxn doesn't do hvline_color):
@@ -1139,6 +1172,14 @@ def _plot_evoked_topo(
         select=select,
     )
 
+    setattr(fig, "_current_time", None)
+
+    def on_time_change(event, fig):
+        """Respond to a time change UI event."""
+        fig._current_time = event.time
+
+    subscribe(fig, "time_change", partial(on_time_change, fig=fig))
+
     add_background_image(fig, fig_background)
 
     if legend is not False:

mne/viz/ui_events.py

@@ -411,7 +411,9 @@ def unsubscribe(fig, event_names, callback=None, *, verbose=None):
 
 
 @verbose
-def link(*figs, include_events=None, exclude_events=None, verbose=None):
+def link(
+    *figs, include_events=None, exclude_events=None, recursive=False, verbose=None
+):
     """Link the event channels of two figures together.
 
     When event channels are linked, any events that are published on one
@@ -430,6 +432,9 @@ def link(*figs, include_events=None, exclude_events=None, verbose=None):
     exclude_events : list of str | None
         Select which events not to publish across figures. By default (``None``),
         no events are excluded.
+    recursive : bool
+        If ``True``, also link the existing link-groups that figs already belong
+        to, so all members are mutually linked.
     %(verbose)s
     """
     if include_events is not None:
@@ -443,7 +448,17 @@ def link(*figs, include_events=None, exclude_events=None, verbose=None):
         if fig not in _event_channel_links:
             _event_channel_links[fig] = weakref.WeakKeyDictionary()
 
-    # Link the event channels
+    if recursive:
+        # Also link to anything that is linked to any of the figures in `figs`.
+        figs_set = weakref.WeakSet()
+        for fig in figs:
+            figs_set.add(fig)
+            for linked_fig in _event_channel_links[fig].keys():
+                figs_set.add(linked_fig)
+        # Deliberately let current include and exclude events dominate over past ones.
+        figs = figs_set
+
+    # Do the actual linking.
     for fig1 in figs:
         for fig2 in figs:
             if fig1 is not fig2:

tutorials/evoked/20_visualize_evoked.py

@@ -269,7 +269,10 @@ def custom_func(x):
 #
 #     In interactive sessions, both approaches to topographical plotting allow
 #     you to click one of the sensor subplots to open a larger version of the
-#     evoked plot at that sensor.
+#     evoked plot at that sensor. In this view, you can select a time by clicking
+#     somewhere in the timecourse. The selected time is marked with a vertical line on
+#     all channels, also in the original figure and any other single-channel figures
+#     that are open.
 #
 #
 # 3D Field Maps