Live updating evoked plot when rejecting epochs

Support & Discussions
,
4

Thank you very much for your helpful example. For now, I have come up with a working example of how I want it to be. It gives the epochs plot as well as the initial evoked plot and subsequently updates it when I (un)toggle bad epochs. The update process of the evoked plot isn’t as smooth as I wished it would be (ideally without closing and reopening the plot), but I feel like this is the best I could do. Do you have any recommendations?

import mne
import numpy as np
from mne.viz import plot_evoked_topo
from PySide6.QtWidgets import QApplication
from PySide6.QtCore import QTimer
import sys
import matplotlib.pyplot as plt

def load_data():
    path = mne.datasets.sample.data_path()
    raw = mne.io.read_raw_fif(path / "MEG/sample/sample_audvis_raw.fif")
    events = mne.find_events(raw)
    return mne.Epochs(raw, events, tmin=-0.2, tmax=0.5)

def update_evoked_plot(epochs, fig):
    # Get indices of good epochs (excluding bad epochs)
    good_epochs = np.setdiff1d(np.arange(len(epochs)), fig.mne.bad_epochs)
    evoked = epochs[good_epochs].average()
    
    # Close previous evoked topo plot if it exists
    if hasattr(update_evoked_plot, 'topo_fig'):
        plt.close(update_evoked_plot.topo_fig)
    
    # Create and display new evoked topo plot
    update_evoked_plot.topo_fig = plot_evoked_topo(evoked, show=False)
    update_evoked_plot.topo_fig.show()

def main():
    # Load the epoched data
    epochs = load_data()
    
    # Create Qt application
    app = QApplication(sys.argv)
    
    # Plot epochs and get the figure
    fig = epochs.plot(block=False)
    
    # Store the original toggle_bad_epoch function
    orig_toggle_bad_epoch = fig._toggle_bad_epoch
    
    def toggle_bad_epoch(*params):
        result = orig_toggle_bad_epoch(*params)
        update_evoked_plot(epochs, fig)
        return result
    
    # Replace the original toggle_bad_epoch function with our wrapper
    fig._toggle_bad_epoch = toggle_bad_epoch
    
    # Initial update of evoked plot after a short delay. Without this there seems to be a conflict with the epochs.plot() call.
    QTimer.singleShot(1000, lambda: update_evoked_plot(epochs, fig))
    
    # Start the Qt event loop
    sys.exit(app.exec())

if __name__ == "__main__":
    main()