accusleepy 0.6.0__py3-none-any.whl

accusleepy/gui/mplwidget.py

@@ -0,0 +1,386 @@
+# Widget with a matplotlib FigureCanvas for manual scoring
+from collections.abc import Callable
+
+import matplotlib.ticker as mticker
+import numpy as np
+from matplotlib.backend_bases import MouseButton
+from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg as FigureCanvas
+from matplotlib.figure import Figure
+from matplotlib.gridspec import GridSpec
+from matplotlib.patches import Rectangle
+from matplotlib.widgets import RectangleSelector
+
+# from PySide6.QtWidgets import *
+from PySide6 import QtWidgets
+
+from accusleepy.brain_state_set import BrainStateSet
+
+# upper limit of spectrogram y-axis, in Hz
+SPEC_UPPER_F = 30
+# interval of spectrogram y-axis ticks, in Hz
+SPEC_Y_TICK_INTERVAL = 10
+
+# margins around subplots in the figure
+SUBPLOT_TOP_MARGIN = 0.98
+SUBPLOT_BOTTOM_MARGIN = 0.02
+SUBPLOT_LEFT_MARGIN = 0.07
+SUBPLOT_RIGHT_MARGIN = 0.95
+
+# maximum number of x-axis ticks to show on the lower plot
+MAX_LOWER_X_TICK_N = 7
+
+
+class MplWidget(QtWidgets.QWidget):
+    """Widget that displays a matplotlib FigureCanvas"""
+
+    def __init__(self, parent=None):
+        QtWidgets.QWidget.__init__(self, parent)
+        # set up the canvas and store a reference to its axes
+        self.canvas = FigureCanvas(Figure())
+        self.canvas.axes = None
+
+        # set the widget layout and remove the margins
+        vertical_layout = QtWidgets.QVBoxLayout()
+        vertical_layout.addWidget(self.canvas)
+        vertical_layout.setContentsMargins(0, 0, 0, 0)
+        self.setLayout(vertical_layout)
+
+        # given during the setup process
+        self.epoch_length = None
+
+        # upper plot references
+        self.upper_marker = None
+        self.label_img_ref = None
+        self.spec_ref = None
+        self.roi = None
+        self.editing_patch = None
+        self.roi_patch = None
+
+        # lower plot references
+        self.eeg_line = None
+        self.emg_line = None
+        self.top_marker = None
+        self.bottom_marker = None
+
+    def setup_upper_figure(
+        self,
+        n_epochs: int,
+        label_img: np.array,
+        confidence_scores: np.array,
+        confidence_img: np.array,
+        spec: np.array,
+        f: np.array,
+        emg: np.array,
+        epochs_to_show: int,
+        label_display_options: np.array,
+        brain_state_set: BrainStateSet,
+        roi_function: Callable,
+    ):
+        """Initialize upper FigureCanvas for the manual scoring GUI
+
+        :param n_epochs: number of epochs in the recording
+        :param label_img: brain state labels, as an image
+        :param confidence_scores: confidence scores
+        :param confidence_img: confidence scores, as an image
+        :param spec: EEG spectrogram image
+        :param f: EEG spectrogram frequency axis
+        :param emg: EMG RMS per epoch
+        :param epochs_to_show: number of epochs to show in the lower plot
+        :param label_display_options: valid brain state y-axis locations
+        :param brain_state_set: set of brain states options
+        :param roi_function: callback for ROI selection
+        """
+        # references to parts of the epoch marker
+        self.upper_marker = list()
+
+        # subplot layout
+        height_ratios = [2, 8, 2, 12, 13]
+        gs1 = GridSpec(5, 1, hspace=0.02, height_ratios=height_ratios)
+        gs2 = GridSpec(5, 1, hspace=0.4, height_ratios=height_ratios)
+        axes = list()
+        axes.append(self.canvas.figure.add_subplot(gs1[0]))
+        axes.append(self.canvas.figure.add_subplot(gs1[1]))
+        axes.append(self.canvas.figure.add_subplot(gs1[2]))
+        axes.append(self.canvas.figure.add_subplot(gs1[3]))
+        axes.append(self.canvas.figure.add_subplot(gs2[4]))
+
+        # subplots have different axes limits
+        for i in [0, 1, 2, 4]:
+            axes[i].set_xlim((-0.5, n_epochs - 0.5))
+        axes[3].set_xlim(0, n_epochs)
+
+        # confidence score subplot
+        if confidence_scores is None:
+            axes[0].set_visible(False)
+        else:
+            axes[0].set_ylim([-0.5, 0.5])
+            axes[0].set_xticks([])
+            axes[0].set_yticks([0])
+            axes[0].set_yticklabels(["Conf."])
+            axes[0].tick_params(axis="y", color="white")
+            axes[0].imshow(
+                confidence_img, aspect="auto", origin="lower", interpolation="None"
+            )
+            confidence_x = (
+                np.repeat(list(range(len(confidence_scores) + 1)), 2)[1:-1] - 0.5
+            )
+            confidence_y = np.repeat(confidence_scores, 2) - 0.5
+            axes[0].plot(confidence_x, confidence_y, "k", linewidth=0.5)
+            for side in ["left", "right", "bottom", "top"]:
+                axes[0].spines[side].set_visible(False)
+
+        # brain state subplot
+        axes[1].set_ylim(
+            [-0.5, np.max(label_display_options) - np.min(label_display_options) + 0.5]
+        )
+        axes[1].set_xticks([])
+        axes[1].set_yticks(
+            label_display_options - np.min(label_display_options),
+        )
+        axes[1].set_yticklabels([b.name for b in brain_state_set.brain_states])
+        ax2 = axes[1].secondary_yaxis("right")
+        ax2.set_yticks(
+            label_display_options - np.min(label_display_options),
+        )
+        ax2.set_yticklabels([b.digit for b in brain_state_set.brain_states])
+        self.label_img_ref = axes[1].imshow(
+            label_img, aspect="auto", origin="lower", interpolation="None"
+        )
+        # add patch to dim the display when creating an ROI
+        self.editing_patch = axes[1].add_patch(
+            Rectangle(
+                xy=(-0.5, -0.5),
+                width=n_epochs,
+                height=np.max(label_display_options)
+                - np.min(label_display_options)
+                + 1,
+                color="white",
+                edgecolor=None,
+                alpha=0.4,
+                visible=False,
+            )
+        )
+        # add the ROI selection widget, but disable it until it's needed
+        self.roi = RectangleSelector(
+            ax=axes[1],
+            onselect=roi_function,
+            interactive=False,
+            button=MouseButton(1),
+        )
+        self.roi.set_active(False)
+        # keep a reference to the ROI patch so we can change its color later
+        # index 0 is the "editing_patch" created earlier
+        self.roi_patch = [c for c in axes[1].get_children() if type(c) is Rectangle][1]
+
+        # epoch marker subplot
+        axes[2].set_ylim((0, 1))
+        axes[2].axis("off")
+        self.upper_marker.append(
+            axes[2].plot([-0.5, epochs_to_show - 0.5], [0.5, 0.5], "r")[0]
+        )
+        self.upper_marker.append(axes[2].plot([0], [0.5], "rD")[0])
+
+        # EEG spectrogram subplot
+        # select subset of frequencies to show
+        f = f[f <= SPEC_UPPER_F]
+        spec = spec[0 : len(f), :]
+        axes[3].set_ylabel("EEG", rotation="horizontal", ha="right")
+        axes[3].set_yticks(
+            np.linspace(
+                0,
+                len(f),
+                1 + round(SPEC_UPPER_F / SPEC_Y_TICK_INTERVAL),
+            ),
+        )
+        axes[3].set_yticklabels(
+            np.arange(0, SPEC_UPPER_F + SPEC_Y_TICK_INTERVAL, SPEC_Y_TICK_INTERVAL)
+        )
+        axes[3].tick_params(axis="both", which="major", labelsize=8)
+        axes[3].xaxis.set_major_formatter(mticker.FuncFormatter(self.time_formatter))
+        self.spec_ref = axes[3].imshow(
+            spec,
+            vmin=np.percentile(spec, 2),
+            vmax=np.percentile(spec, 98),
+            aspect="auto",
+            origin="lower",
+            interpolation="None",
+            extent=(
+                0,
+                n_epochs,
+                -0.5,
+                len(f) + 0.5,
+            ),
+        )
+
+        # EMG subplot
+        axes[4].set_xticks([])
+        axes[4].set_yticks([])
+        axes[4].set_ylabel("EMG", rotation="horizontal", ha="right")
+        axes[4].plot(
+            emg,
+            "k",
+            linewidth=0.5,
+        )
+
+        self.canvas.figure.subplots_adjust(
+            left=SUBPLOT_LEFT_MARGIN,
+            right=SUBPLOT_RIGHT_MARGIN,
+            top=SUBPLOT_TOP_MARGIN,
+            bottom=SUBPLOT_BOTTOM_MARGIN,
+        )
+
+        self.canvas.axes = axes
+
+    def setup_lower_figure(
+        self,
+        label_img: np.array,
+        sampling_rate: int | float,
+        epochs_to_show: int,
+        brain_state_set: BrainStateSet,
+        label_display_options: np.array,
+    ):
+        """Initialize lower FigureCanvas for the manual scoring GUI
+
+        :param label_img: brain state labels, as an image
+        :param sampling_rate: EEG/EMG sampling rate, in Hz
+        :param epochs_to_show: number of epochs to show in the lower plot
+        :param brain_state_set: set of brain states options
+        :param label_display_options: valid brain state y-axis locations
+        """
+        # number of samples in one epoch
+        # we can expect this to be very close to an integer
+        samples_per_epoch = round(sampling_rate * self.epoch_length)
+        # number of EEG/EMG samples to plot
+        samples_shown = round(samples_per_epoch * epochs_to_show)
+
+        # references to parts of the epoch markers
+        self.top_marker = list()
+        self.bottom_marker = list()
+        # epoch marker display parameters
+        marker_dy = 0.25
+        marker_y_offset_top = 0.02
+        marker_y_offset_bottom = 0.01
+
+        # subplot layout
+        gs1 = GridSpec(3, 1, hspace=0)
+        gs2 = GridSpec(3, 1, hspace=0.5)
+        axes = list()
+        axes.append(self.canvas.figure.add_subplot(gs1[0]))
+        axes.append(self.canvas.figure.add_subplot(gs1[1]))
+        axes.append(self.canvas.figure.add_subplot(gs2[2]))
+
+        # EEG subplot
+        axes[0].set_xticks([])
+        axes[0].set_yticks([])
+        axes[0].set_xlim((0, samples_shown))
+        axes[0].set_ylim((-1, 1))
+        axes[0].set_ylabel("EEG", rotation="horizontal", ha="right")
+        self.eeg_line = axes[0].plot(
+            np.zeros(samples_shown),
+            "k",
+            linewidth=0.5,
+        )[0]
+        # top epoch marker
+        marker_x = [
+            [0, 0],
+            [0, samples_per_epoch],
+            [samples_per_epoch, samples_per_epoch],
+        ]
+        marker_y = np.array(
+            [
+                [1 - marker_dy, 1],
+                [1, 1],
+                [1 - marker_dy, 1],
+            ]
+        )
+        for x, y in zip(marker_x, marker_y):
+            self.top_marker.append(axes[0].plot(x, y - marker_y_offset_top, "r")[0])
+
+        # EMG subplot
+        axes[1].set_xticks(
+            resample_x_ticks(
+                np.arange(
+                    0,
+                    samples_shown,
+                    samples_per_epoch,
+                )
+            )
+        )
+        axes[1].tick_params(axis="x", which="major", labelsize=8)
+        axes[1].set_yticks([])
+        axes[1].set_xlim((0, samples_shown))
+        axes[1].set_ylim((-1, 1))
+        axes[1].set_ylabel("EMG", rotation="horizontal", ha="right")
+        self.emg_line = axes[1].plot(
+            np.zeros(samples_shown),
+            "k",
+            linewidth=0.5,
+        )[0]
+
+        for x, y in zip(marker_x, marker_y):
+            self.bottom_marker.append(
+                axes[1].plot(x, -1 * (y - marker_y_offset_bottom), "r")[0]
+            )
+
+        # brain state subplot
+        axes[2].set_xticks([])
+        axes[2].set_yticks(
+            label_display_options - np.min(label_display_options),
+        )
+        axes[2].set_yticklabels([b.name for b in brain_state_set.brain_states])
+        ax2 = axes[2].secondary_yaxis("right")
+        ax2.set_yticks(
+            label_display_options - np.min(label_display_options),
+        )
+        ax2.set_yticklabels([b.digit for b in brain_state_set.brain_states])
+        axes[2].set_xlim((-0.5, epochs_to_show - 0.5))
+        axes[2].set_ylim(
+            [-0.5, np.max(label_display_options) - np.min(label_display_options) + 0.5]
+        )
+        self.label_img_ref = axes[2].imshow(
+            label_img[:, 0:epochs_to_show, :],
+            aspect="auto",
+            origin="lower",
+            interpolation="None",
+        )
+
+        self.canvas.figure.subplots_adjust(
+            left=SUBPLOT_LEFT_MARGIN,
+            right=SUBPLOT_RIGHT_MARGIN,
+            top=SUBPLOT_TOP_MARGIN,
+            bottom=SUBPLOT_BOTTOM_MARGIN,
+        )
+
+        self.canvas.axes = axes
+
+    def time_formatter(self, x, pos):
+        x = x * self.epoch_length
+        return "{:02d}:{:02d}:{:05.2f}".format(
+            int(x // 3600), int(x // 60) % 60, (x % 60)
+        )
+
+
+def resample_x_ticks(x_ticks: np.array) -> np.array:
+    """Choose a subset of x_ticks to display
+
+    The x-axis can get crowded if there are too many timestamps shown.
+    This function finds a subset of evenly spaced x-axis ticks that
+    includes the one at the beginning of the central epoch.
+
+    :param x_ticks: full set of x_ticks
+    :return: smaller subset of x_ticks
+    """
+    if len(x_ticks) <= MAX_LOWER_X_TICK_N:
+        return x_ticks
+
+    # number of ticks to the left of the central epoch
+    # this will always be an integer
+    nl = round((len(x_ticks) - 1) / 2)
+
+    # search for even tick spacings that include the central epoch
+    # if necessary, skip the leftmost tick
+    for offset in [0, 1]:
+        if (nl - offset) % 3 == 0:
+            return x_ticks[offset :: round((nl - offset) / 3)]
+        elif (nl - offset) % 2 == 0:
+            return x_ticks[offset :: round((nl - offset) / 2)]