Proposal for an event API for MNE-Python

events.py

"""
Proposal for an event API for MNE-Python.
-----------------------------------------

We want figures to be able to communicate with each other, such that a change
in one figure can trigger a change in another figure. For example, moving the
time cursor in a Brain plot can update the current time in an evoked plot.
Another scenario is two drawing routines drawing into the same window. For
example, one function is drawing a source estimate, while another is drawing
magnetic field lines. The event system will allow both drawing routines to
communicate and update in-sync.

The API is designed to be minimally invasive: to require the least amount of
modifications to the existing figure routines.


Event channels and linking them
===============================
The general idea is to have an event channel where multiple drawing routines
can subscribe to and publish on. All events sent on the channel will be
received by all drawing routines. All drawing routines can send messages on
the channel to be received by all subscribers.

The simplest design would be to have one global event channel, so all figures
can talk to each other. However, with one global event channel, all currently
open figures are always "linked". This is a bit too rigid, as a user may have
multiple figures open that concern very different things. Changing the time on
one figure should not necesarrily change the time on all other figures, unless
we explicitly want it to be so. Hence, event channels are tied to figures. Each
figure has its own event channel.

To broadcast events across figures, we allow message channels to be "linked".
When an event is published on a channel, it is also published on all linked
channels. This way, we have control over which figures talk to each-other and
which don't. Implementation wise, we could have a single event channel object
that is shared between figures. However, that makes it difficult to ever
"unlink" channels. Unlinking is something that needs to happen for example when
the user closes one of the figures. So instead, we explicitly keep track of a
list of linked channels and whenever we publish on one channel, we check the
list for any linked channels on which the event must also be published.


Event objects and their parameters
==================================
The events are modeled after matplotlib's event system. An event has two
important properties: a name and a value. For example, the "time_change" event
should have the new time as a value. Values should be any python object and
some events have more than one value. It is important that events are
consistent. The "time_change" event emitted by one drawing routine should have
the same values with the same variable names. Ideally, this is enforced in the
code. Hence, events are objects with predefined fields. Python's dataclasses
should serve this purpose nicely. API-wise, when publishing an event, we can
create a new instance of the event's class. When subscribing to an event,
having to dig up the correct class is a bit of a hassle. Following matplotlib's
example, we use the string name of the event to refer to it in this case.


Callbacks
=========
When subscribing to an event on a channel, a callback function is passed that
will be called with the event object.
"""
from weakref import WeakKeyDictionary, WeakSet
from dataclasses import dataclass
import matplotlib
import matplotlib.pyplot as plt
import mne


# Global dict {fig: channel} containing all currently active event channels.
# Our figure objects are quite big, so perhaps it's a good idea to use weak
# references here, so we don't accidentally keep references to them when
# something goes wrong in the event system.
event_channels = WeakKeyDictionary()

# The event channels of figures can be linked together. This dict keeps track
# of these links. Links are bi-directional, so if fig1 -> fig2 exists in this
# dictionary, then so must fig2 -> fig1 exist in this dict.
event_channel_links = WeakKeyDictionary()


# Different events
@dataclass
class Event:
    source: matplotlib.figure.Figure | mne.viz.Brain
    name: str

@dataclass
class TimeChange(Event):
    name: str = 'time_change'
    time: float = 0.0


def get_event_channel(fig):
    """Get the event channel associated with a figure.

    If the event channel doesn't exist yet, it gets created and added to the
    global ``event_channels`` dict.

    Parameters
    ----------
    fig : matplotlib.figure.Figure | mne.viz.Brain
        The figure to get the event channel for.

    Returns
    -------
    channel : dict[event -> list]
        The event channel. An event channel is a list mapping string event
        names to a list of callback representing all subscribers to the
        channel.
    """
    # Create the event channel if it doesn't exist yet
    if fig not in event_channels:
        # The channel itself is a dict mapping string event names to a list of
        # subscribers. No subscribers yet for this new event channel.
        event_channels[fig] = dict()

        # When the figure is closed, its associated event channel should be
        # deleted. This is a good time to set this up.
        def delete_event_channel(event=None):
            """Delete the event channel (callback function)."""
            publish(fig, event='close')  # Notify subscribers of imminent close
            unlink(fig)  # Remove channel from the event_channel_links dict
            del event_channels[fig]

        # Hook up the above callback function to the close event of the figure
        # window. How this is done exactly depends on the various figure types
        # MNE-Python has.
        if isinstance(fig, matplotlib.figure.Figure):
            fig.canvas.mpl_connect('close_event', delete_event_channel)
        elif isinstance(fig, mne.viz.Brain):
            fig._renderer._window.signal_close.connect(delete_event_channel)
        else:
            raise NotImplementedError('This figure type is not support yet.')

    # Now the event channel exists for sure.
    return event_channels[fig]


def publish(fig, event):
    """Publish an event to all subscribers of the figure's channel.

    The figure's event channel and all linked event channels are searched for
    subscribers to the given event. Each subscriber had provided a callback
    function when subscribing, so we call that.

    Parameters
    ----------
    fig : matplotlib.figure.Figure | mne.viz.Brain
        The figure that publishes the event.
    event : Event
        Event to publish.
    """
    # Compile a list of all event channels that the event should be published
    # on.
    channels = [get_event_channel(fig)]
    if fig in event_channel_links:
        linked_channels = [get_event_channel(linked_fig)
                           for linked_fig in event_channel_links[fig]]
        channels.extend(linked_channels)

    # Publish the event by calling the registered callback functions.
    for channel in channels:
        if event.name not in channel:
            channel[event.name] = set()
        for callback in channel[event.name]:
            callback(event=event)


def subscribe(fig, event_name, callback):
    """Subscribe to an event on a figure's event channel.

    Parameters
    ----------
    fig : matplotlib.figure.Figure | mne.viz.Brain
        The figure of which event channel to subscribe.
    event_name : str
        The name of the event to listen for.
    callback : func
        The function that should be called whenever the event is published.
    """
    channel = get_event_channel(fig)
    if event_name not in channel:
        channel[event_name] = set()
    channel[event_name].add(callback)


def link(fig1, fig2):
    """Link the event channels of two figures together.

    When event channels are linked, any events that are published on one
    channel are simultaneously published on the other channel.

    Parameters
    ----------
    fig1 : matplotlib.figure.Figure | mne.viz.Brain
        The first figure whose event channel will be linked to the second.
    fig2 : matplotlib.figure.Figure | mne.viz.Brain
        The second figure whose event channel will be linked to the first.
    """
    if fig1 not in event_channel_links:
        event_channel_links[fig1] = WeakSet([fig2])
    else:
        event_channel_links[fig1].add(fig2)
    if fig2 not in event_channel_links:
        event_channel_links[fig2] = WeakSet([fig1])
    else:
        event_channel_links[fig2].add(fig1)


def unlink(fig):
    """Remove all links involving the event channel of the given figure.

    Parameters
    ----------
    fig : matplotlib.figure.Figure | mne.viz.Brain
        The figure whose event channel should be unlinked from all other event
        channels.
    """
    linked_figs = event_channel_links.get(fig)
    if linked_figs is not None:
        for linked_fig in linked_figs:
            event_channel_links[linked_fig].remove(fig)
    del event_channel_links[fig]


# -----------------------------------------------------------------------------
# Example of two-way communication between an evoked plot and a source estimate
# plot: sharing the same current time
# -----------------------------------------------------------------------------

# Load some data. Both evoked and source estimate
data_path = mne.datasets.sample.data_path()
evoked = mne.read_evokeds(f'{data_path}/MEG/sample/sample_audvis-ave.fif',
                          condition = 'Left Auditory')
evoked.apply_baseline()
evoked.resample(100).crop(0, 0.24)
stc = mne.read_source_estimate(f'{data_path}/MEG/sample/sample_audvis-meg-eeg')
assert evoked.data.shape[1] == stc.data.shape[1]

# Figure 1 will be an evoked plot
fig1 = evoked.plot()

# Add some interactivity to the plot: a time slider in each subplot. Whenever
# the mouse moves, the time is updated. The updated time will be published as a
# "time_change" event. The publish() function is clever enough to notice the
# figure does not have an event channel yet and create one for us.
fig1_time_sliders = [ax.axvline(0.1) for ax in fig1.axes[:4]]
plt.connect(
    'motion_notify_event',
    lambda event: publish(fig1, TimeChange(source=fig1, time=event.xdata))
)

# We want to update the time cursor both when the mouse is moving across the
# evoked figure, or when it receives a "time change" event from another
# figure. Hence we don't implement the cursor movement in the matplotlib event
# handler, but rather in the event handler for our new MNE-Python event API.
def fig1_time_change(event):
    """Called whenever a "time_change" event is published."""
    if event.time is not None:  # Can be None when cursor is out of bounds
        # Update the time sliders to reflect the new time
        for slider in fig1_time_sliders:
            slider.set_xdata([event.time])
        fig1.canvas.draw()

# To receive "time_change" events, we subscribe to our figure's event channel
# and attach the callback function. As a rule, a figure only ever subscribes to
# events on its own event channel. Communication between figures happens by
# linking event channels.
subscribe(fig1, 'time_change', fig1_time_change)

# At this point, the time slider should work for the evoked figure. It also
# emits "time_change" events that other figures could listen to. Let's create a
# second figure to do just that. To make it interesting, let's make it not a
# matplotlib figure, but a Brain figure.
fig2 = stc.plot('sample', subjects_dir=f'{data_path}/subjects')

# Also in this figure, we listen for mouse movements and publish "time_change"
# events. In the future, this should be setup in the constructor of `Brain`.
fig2_time_slider = fig2.mpl_canvas.fig.axes[0].lines[1]  # Grab existing slider
fig2.mpl_canvas.canvas.mpl_connect(
    'motion_notify_event',
    lambda event: publish(fig2, TimeChange(source=fig2, time=event.xdata))
)

def fig2_time_change(event):
    """Called whenever a "time_change" event is published."""
    if event.time is not None:
        # Update the time for this Brain figure.
        fig2.set_time(event.time)
        # set_time() does not seem to update the slider, so do that manually
        fig2_time_slider.set_xdata([event.time])
        fig2.mpl_canvas.canvas.draw()

# Start listening to "time_change" events. At the moment just on the event
# channel of the Brain figure.
subscribe(fig2, 'time_change', fig2_time_change)

# We can achieve interaction between the two figures by linking their
# corresponding event channels. Now, "time_change" events are published across
# both of them, regardless of which figure they originated from. Since both
# figures are listening to their own event channels, both figures will get the
# message.
link(fig1, fig2)