add erds topomap plotting

CHANGELOG.md

@@ -10,6 +10,7 @@
 - Add montage name and location count to infowidget ([#271](https://github.com/cbrnr/mnelab/pull/271) by [Florian Hofer](https://github.com/hofaflo))
 - Add possibility to specify `match_case`, `match_alias`, and `on_missing` to "Set montage..." ([#271](https://github.com/cbrnr/mnelab/pull/271) by [Florian Hofer](https://github.com/hofaflo))
 - Add "Clear montage" to "Edit" menu  ([#271](https://github.com/cbrnr/mnelab/pull/271) by [Florian Hofer](https://github.com/hofaflo))
+- Add support for plotting ERDS topomaps (Plot -> ERDS topomaps...) ([#278](https://github.com/cbrnr/mnelab/pull/278) by [Florian Hofer](https://github.com/hofaflo))
 
 ### Changed
 - Simplify rereferencing workflow ([#258](https://github.com/cbrnr/mnelab/pull/258) by [Florian Hofer](https://github.com/hofaflo))

mnelab/dialogs/__init__.py

@@ -8,7 +8,7 @@
 from .channel_properties import ChannelPropertiesDialog
 from .crop import CropDialog
 from .epoch import EpochDialog
-from .erds import ERDSDialog
+from .erds import ERDSDialog, ERDSTopomapsDialog
 from .error_message import ErrorMessageBox
 from .events import EventsDialog
 from .filter import FilterDialog

mnelab/dialogs/erds.py

@@ -1,16 +1,19 @@
 # Copyright (c) MNELAB developers
 #
 # License: BSD (3-clause)
-
+from PySide6.QtCore import Qt, Slot
 from PySide6.QtWidgets import (
     QDialog,
     QDialogButtonBox,
     QDoubleSpinBox,
     QGridLayout,
     QLabel,
+    QListWidget,
     QVBoxLayout,
 )
 
+from .utils import select_all
+
 
 class ERDSDialog(QDialog):
     def __init__(self, parent, t_range, f_range):
@@ -111,3 +114,120 @@ def b1(self):
     @property
     def b2(self):
         return self._b2.value()
+
+
+class ERDSTopomapsDialog(QDialog):
+    def __init__(self, parent, t_range, f_range, events):
+        super().__init__(parent)
+        self.setWindowTitle("ERDS topomaps")
+        vbox = QVBoxLayout(self)
+        grid = QGridLayout()
+
+        grid.addWidget(QLabel("Frequency range:"), 0, 0)
+        self._f1 = QDoubleSpinBox()
+        self._f1.setRange(*f_range)
+        self._f1.setValue(f_range[0])
+        self._f1.setDecimals(1)
+        self._f1.setSuffix(" Hz")
+        grid.addWidget(self._f1, 0, 1)
+
+        self._f2 = QDoubleSpinBox()
+        self._f2.setRange(*f_range)
+        self._f2.setValue(f_range[1])
+        self._f2.setDecimals(1)
+        self._f2.setSuffix(" Hz")
+        grid.addWidget(self._f2, 0, 2)
+
+        grid.addWidget(QLabel("Step size:"), 1, 0)
+        self._step = QDoubleSpinBox()
+        self._step.setRange(0.1, 5)
+        self._step.setValue(1)
+        self._step.setDecimals(1)
+        self._step.setSingleStep(0.1)
+        self._step.setSuffix(" Hz")
+        grid.addWidget(self._step, 1, 1)
+
+        grid.addWidget(QLabel("Time range:"), 2, 0)
+        self._t1 = QDoubleSpinBox()
+        self._t1.setRange(*t_range)
+        self._t1.setValue(t_range[0])
+        self._t1.setDecimals(1)
+        self._step.setSingleStep(0.1)
+        self._t1.setSuffix(" s")
+        grid.addWidget(self._t1, 2, 1)
+
+        self._t2 = QDoubleSpinBox()
+        self._t2.setRange(*t_range)
+        self._t2.setValue(t_range[1])
+        self._t2.setDecimals(1)
+        self._step.setSingleStep(0.1)
+        self._t2.setSuffix(" s")
+        grid.addWidget(self._t2, 2, 2)
+
+        grid.addWidget(QLabel("Baseline:"), 3, 0)
+        self._b1 = QDoubleSpinBox()
+        self._b1.setRange(*t_range)
+        self._b1.setValue(t_range[0])
+        self._b1.setDecimals(1)
+        self._step.setSingleStep(0.1)
+        self._b1.setSuffix(" s")
+        grid.addWidget(self._b1, 3, 1)
+
+        self._b2 = QDoubleSpinBox()
+        self._b2.setRange(*t_range)
+        self._b2.setValue(0)
+        self._b2.setDecimals(1)
+        self._step.setSingleStep(0.1)
+        self._b2.setSuffix(" s")
+        grid.addWidget(self._b2, 3, 2)
+
+        label = QLabel("Events:")
+        label.setAlignment(Qt.AlignTop)
+        grid.addWidget(label, 4, 0)
+        self.events = QListWidget()
+        self.events.insertItems(0, events)
+        self.events.setSelectionMode(QListWidget.ExtendedSelection)
+        self.events.setMaximumHeight(self.events.sizeHintForRow(0) * 5.5)
+        select_all(self.events)
+        grid.addWidget(self.events, 4, 1, 1, 2)
+        self.events.itemSelectionChanged.connect(self.toggle_ok)
+
+        vbox.addLayout(grid)
+        buttonbox = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel)
+        vbox.addWidget(buttonbox)
+        buttonbox.accepted.connect(self.accept)
+        buttonbox.rejected.connect(self.reject)
+        vbox.setSizeConstraint(QVBoxLayout.SetFixedSize)
+
+    @Slot()
+    def toggle_ok(self):
+        enable = bool(self.events.selectedItems())
+        self.buttonbox.button(QDialogButtonBox.Ok).setEnabled(enable)
+
+    @property
+    def f1(self):
+        return self._f1.value()
+
+    @property
+    def f2(self):
+        return self._f2.value()
+
+    @property
+    def step(self):
+        return self._step.value()
+
+    @property
+    def t1(self):
+        return self._t1.value()
+
+    @property
+    def t2(self):
+        return self._t2.value()
+
+    @property
+    def b1(self):
+        return self._b1.value()
+
+    @property
+    def b2(self):
+        return self._b2.value()

mnelab/mainwindow.py

@@ -40,8 +40,20 @@
 from .io import writers
 from .io.xdf import get_xml, list_chunks
 from .model import InvalidAnnotationsError, LabelsNotFoundError, Model
-from .utils import count_locations, have, image_path, interface_style, natural_sort
-from .viz import plot_erds, plot_evoked, plot_evoked_comparison, plot_evoked_topomaps
+from .utils import (
+    count_locations,
+    have,
+    image_path,
+    interface_style,
+    natural_sort,
+)
+from .viz import (
+    plot_erds,
+    plot_erds_topomaps,
+    plot_evoked,
+    plot_evoked_comparison,
+    plot_evoked_topomaps,
+)
 from .widgets import InfoWidget
 
 MAX_RECENT = 6  # maximum number of recent files
@@ -205,7 +217,12 @@ def __init__(self, model: Model):
         icon = QIcon.fromTheme("plot-locations")
         self.actions["plot_locations"] = plot_menu.addAction(icon, "&Channel locations",
                                                              self.plot_locations)
+        plot_menu.addSeparator()
         self.actions["plot_erds"] = plot_menu.addAction("&ERDS maps...", self.plot_erds)
+        self.actions["plot_erds_topomaps"] = plot_menu.addAction(
+            "ERDS topomaps...",
+            self.plot_erds_topomaps,
+        )
         plot_menu.addSeparator()
         self.actions["plot_evoked"] = plot_menu.addAction(
             "Evoked...",
@@ -423,6 +440,9 @@ def data_changed(self):
             self.actions["plot_erds"].setEnabled(
                 enabled and self.model.current["dtype"] == "epochs"
             )
+            self.actions["plot_erds_topomaps"].setEnabled(
+                enabled and self.model.current["dtype"] == "epochs"
+            )
             self.actions["plot_evoked"].setEnabled(
                 enabled and self.model.current["dtype"] == "epochs"
             )
@@ -772,6 +792,24 @@ def plot_erds(self):
             for fig in figs:
                 fig.show()
 
+    def plot_erds_topomaps(self):
+        """Plot ERDS topomaps."""
+        epochs = self.model.current["data"]
+        t_range = [epochs.tmin, epochs.tmax]
+        f_range = [1, epochs.info["sfreq"] / 2]
+
+        dialog = ERDSTopomapsDialog(self, t_range, f_range, epochs.event_id)
+        if dialog.exec():
+            figs = plot_erds_topomaps(
+                epochs,
+                events=[item.text() for item in dialog.events.selectedItems()],
+                freqs=np.arange(dialog.f1, dialog.f2, dialog.step),
+                baseline=[dialog.b1, dialog.b2],
+                times=[dialog.t1, dialog.t2],
+            )
+            for fig in figs:
+                fig.show()
+
     def plot_evoked(self):
         """Plot evoked potentials for individual channels."""
         epochs = self.model.current["data"]

mnelab/viz.py

@@ -48,6 +48,40 @@ def plot_erds(data, freqs, n_cycles, baseline, times=(None, None)):
     return figs
 
 
+def plot_erds_topomaps(epochs, events, freqs, baseline, times):
+    """
+    Plot ERDS topomaps, one figure per event.
+
+    Parameters
+    ----------
+    epochs : mne.epochs.Epochs
+        Epochs extracted from a Raw instance.
+    events : list[str]
+        Events to include.
+    freqs : np.ndarray
+        Array of frequencies over which the average is taken.
+    baseline : tuple[float, float]
+        Start and end times for baseline correction.
+    times : tuple[float, float]
+        Start and end times between which the average is taken.
+
+    Returns
+    -------
+    list[matplotlib.figure.Figure]
+        A list of the figure(s) generated.
+    """
+    figs = []
+    for event in events:
+        tfr = tfr_multitaper(epochs[event], freqs, freqs, average=True, return_itc=False)
+        tfr.apply_baseline(baseline, mode="percent")
+        tfr.crop(*times)
+        fig = tfr.plot_topomap(title=f"Event: {event}")
+        fig.set_size_inches(4, 3)
+        fig.set_tight_layout(True)
+        figs.append(fig)
+    return figs
+
+
 def plot_evoked(
     epochs,
     picks,