Make plot generation asynchronous to keep the UI responsive (#171)

Co-authored-by: Patrick Shriwise <pshriwise@gmail.com>

Author: paulromano

openmc_plotter/main_window.py

@@ -56,6 +56,8 @@ def __init__(self,
         self.model_path = Path(model_path)
         self.threads = threads
         self.default_res = resolution
+        self.model = None
+        self.plot_manager = None
 
     def loadGui(self, use_settings_pkl=True):
 
@@ -114,15 +116,19 @@ def loadGui(self, use_settings_pkl=True):
         self.statusBar().addPermanentWidget(self.coord_label)
         self.coord_label.hide()
 
+        self.plot_manager = self.model.plot_manager
+        self.plot_manager.plot_started.connect(self._on_plot_started)
+        self.plot_manager.plot_queued.connect(self._on_plot_queued)
+        self.plot_manager.plot_finished.connect(self._on_plot_finished)
+        self.plot_manager.plot_error.connect(self._on_plot_error)
+        self.plot_manager.plot_idle.connect(self._on_plot_idle)
+
         # Load Plot
-        self.statusBar().showMessage('Generating Plot...')
         self.geometryPanel.update()
         self.tallyPanel.update()
         self.colorDialog.updateDialogValues()
-        self.statusBar().showMessage('')
 
-        # Timer allows GUI to render before plot finishes loading
-        QtCore.QTimer.singleShot(0, self.showCurrentView)
+        QtCore.QTimer.singleShot(0, self.requestPlotUpdate)
 
         self.plotIm.frozen = False
 
@@ -421,10 +427,17 @@ def updateEditMenu(self):
         changed = self.model.currentView != self.model.defaultView
         self.restoreAction.setDisabled(not changed)
 
+        toggle_actions = (self.maskingAction, self.highlightingAct,
+                          self.outlineAct, self.overlapAct)
+        # Temporarily block signals to avoid triggering plot update
+        for action in toggle_actions:
+            action.blockSignals(True)
         self.maskingAction.setChecked(self.model.currentView.masking)
         self.highlightingAct.setChecked(self.model.currentView.highlighting)
         self.outlineAct.setChecked(self.model.currentView.outlinesCell)
         self.overlapAct.setChecked(self.model.currentView.color_overlaps)
+        for action in toggle_actions:
+            action.blockSignals(False)
 
         num_previous_views = len(self.model.previousViews)
         self.undoAction.setText('&Undo ({})'.format(num_previous_views))
@@ -466,11 +479,14 @@ def saveBatchImage(self, view_file):
         cv = self.model.currentView
         # load the view from file
         self.loadViewFile(view_file)
+        self.waitForPlotIdle()
         self.plotIm.saveImage(view_file.replace('.pltvw', ''))
 
     # Menu and shared methods
     def loadModel(self, reload=False, use_settings_pkl=True):
         if reload:
+            if self.plot_manager is not None:
+                self.plot_manager.wait_for_idle()
             self.resetModels()
         else:
             self.model = PlotModel(use_settings_pkl, self.model_path, self.default_res)
@@ -632,66 +648,48 @@ def plotSourceSites(self):
 
     def applyChanges(self):
         if self.model.activeView != self.model.currentView:
-            self.statusBar().showMessage('Generating Plot...')
-            QApplication.processEvents()
             if self.model.activeView.selectedTally is not None:
                 self.tallyPanel.updateModel()
             self.updateMeshAnnotations()
             self.model.storeCurrent()
             self.model.subsequentViews = []
-            self.plotIm.generatePixmap()
-            self.resetModels()
-            self.showCurrentView()
-            self.statusBar().showMessage('')
+            self.requestPlotUpdate()
         else:
             self.statusBar().showMessage('No changes to apply.', 3000)
 
     def undo(self):
-        self.statusBar().showMessage('Generating Plot...')
-        QApplication.processEvents()
-
+        if not self.model.previousViews:
+            return
         self.model.undo()
-        self.resetModels()
-        self.showCurrentView()
         self.geometryPanel.update()
         self.colorDialog.updateDialogValues()
+        self.requestPlotUpdate()
 
         if not self.model.previousViews:
             self.undoAction.setDisabled(True)
         self.redoAction.setDisabled(False)
-        self.statusBar().showMessage('')
 
     def redo(self):
-        self.statusBar().showMessage('Generating Plot...')
-        QApplication.processEvents()
-
+        if not self.model.subsequentViews:
+            return
         self.model.redo()
-        self.resetModels()
-        self.showCurrentView()
         self.geometryPanel.update()
         self.colorDialog.updateDialogValues()
+        self.requestPlotUpdate()
 
         if not self.model.subsequentViews:
             self.redoAction.setDisabled(True)
         self.undoAction.setDisabled(False)
-        self.statusBar().showMessage('')
 
     def restoreDefault(self):
         if self.model.currentView != self.model.defaultView:
-
-            self.statusBar().showMessage('Generating Plot...')
-            QApplication.processEvents()
-
             self.model.storeCurrent()
             self.model.activeView.adopt_plotbase(self.model.defaultView)
-            self.plotIm.generatePixmap()
-            self.resetModels()
-            self.showCurrentView()
             self.geometryPanel.update()
             self.colorDialog.updateDialogValues()
+            self.requestPlotUpdate()
 
             self.model.subsequentViews = []
-            self.statusBar().showMessage('')
 
     def editBasis(self, basis, apply=False):
         self.model.activeView.basis = basis
@@ -1199,6 +1197,9 @@ def resizeEvent(self, event):
             self.shortcutOverlay.resize(self.width(), self.height())
 
     def closeEvent(self, event):
+        if self.plot_manager is not None:
+            self.plot_manager.wait_for_idle(timeout_ms=250)
+            self.plot_manager.shutdown()
         settings = QtCore.QSettings()
         settings.setValue("mainWindow/Size", self.size())
         settings.setValue("mainWindow/Position", self.pos())
@@ -1213,6 +1214,53 @@ def closeEvent(self, event):
 
         self.saveSettings()
 
+    def requestPlotUpdate(self, view=None):
+        if self.model is None:
+            return
+        if view is None:
+            view = self.model.activeView
+        view_snapshot = copy.deepcopy(view)
+        if self.model.can_reuse_maps(view_snapshot):
+            view_params = self.model.view_params_payload(view_snapshot)
+            self.plot_manager.set_latest_view_params(view_params)
+            self.plot_manager.clear_pending()
+            self.model.makePlot(view_snapshot, self.model.ids_map, self.model.properties)
+            self.resetModels()
+            self.showCurrentView()
+            if not self.plot_manager.is_busy:
+                self._on_plot_idle()
+            return
+        view_params = self.model.view_params_payload(view_snapshot)
+        self.plot_manager.enqueue(view_snapshot, view_params)
+
+    def waitForPlotIdle(self, timeout_ms=None):
+        if self.plot_manager is not None:
+            return self.plot_manager.wait_for_idle(timeout_ms)
+        return True
+
+    def _on_plot_started(self):
+        self.plotIm.showUpdatingOverlay("Generating Plot...")
+
+    def _on_plot_queued(self):
+        self.plotIm.showUpdatingOverlay("Generating Plot... (update queued)")
+
+    def _on_plot_finished(self, view_snapshot, view_params, ids_map, properties):
+        if view_params != self.plot_manager.latest_view_params:
+            return
+        self.model.makePlot(view_snapshot, ids_map, properties)
+        self.resetModels()
+        self.showCurrentView()
+
+    def _on_plot_error(self, error_msg):
+        msg_box = QMessageBox()
+        msg_box.setText(f"Failed to generate plot:\n\n{error_msg}")
+        msg_box.setIcon(QMessageBox.Warning)
+        msg_box.setStandardButtons(QMessageBox.Ok)
+        msg_box.exec()
+
+    def _on_plot_idle(self):
+        self.plotIm.hideUpdatingOverlay()
+
     def saveSettings(self):
         if self.model.statepoint:
             self.model.statepoint.close()

openmc_plotter/plotgui.py

@@ -2,7 +2,7 @@
 
 from PySide6 import QtCore, QtGui
 from PySide6.QtWidgets import (QWidget, QPushButton, QHBoxLayout, QVBoxLayout,
-                               QFormLayout, QComboBox, QSpinBox,
+                               QFormLayout, QComboBox, QSpinBox, QLabel,
                                QDoubleSpinBox, QSizePolicy, QMessageBox,
                                QCheckBox, QRubberBand, QMenu, QDialog,
                                QTabWidget, QTableView, QHeaderView)
@@ -21,6 +21,33 @@
 from .custom_widgets import HorizontalLine
 
 
+class PlotUpdateOverlay(QWidget):
+
+    def __init__(self, parent=None):
+        super().__init__(parent)
+        self.setAttribute(QtCore.Qt.WA_TransparentForMouseEvents, True)
+        self.setAttribute(QtCore.Qt.WA_StyledBackground, True)
+        self.setFocusPolicy(QtCore.Qt.NoFocus)
+        self.setStyleSheet("background-color: rgba(20, 20, 20, 140);")
+
+        layout = QVBoxLayout(self)
+        layout.setContentsMargins(0, 0, 0, 0)
+        layout.setAlignment(QtCore.Qt.AlignCenter)
+
+        self.label = QLabel("Generating Plot...", self)
+        self.label.setAlignment(QtCore.Qt.AlignCenter)
+        font = self.label.font()
+        font.setPointSize(max(12, font.pointSize() + 6))
+        self.label.setFont(font)
+        self.label.setStyleSheet("color: white; background-color: transparent;")
+        layout.addWidget(self.label)
+
+        self.hide()
+
+    def set_message(self, message: str):
+        self.label.setText(message)
+
+
 
 class PlotImage(FigureCanvas):
 
@@ -51,13 +78,15 @@ def __init__(self, model: PlotModel, parent, main_window):
         self.tally_colorbar = None
         self.tally_image = None
         self.image = None
+        self.ax = None
 
         self._property_colorbar_bg = None
         self._tally_colorbar_bg = None
         self._last_tally_indicator_value = None
         self._last_data_indicator_value = None
 
         self.menu = QMenu(self)
+        self.update_overlay = PlotUpdateOverlay(self)
 
     def enterEvent(self, event):
         self.setCursor(QtCore.Qt.CrossCursor)
@@ -79,6 +108,8 @@ def mousePressEvent(self, event):
                                                   QtCore.QSize()))
 
     def getPlotCoords(self, pos):
+        if self.ax is None:
+            return (0.0, 0.0)
         x, y = self.mouseEventCoords(pos)
 
         # get the normalized axis coordinates from the event display units
@@ -119,6 +150,24 @@ def _resize(self):
         self.resize(self.parent.width() * z,
                     self.parent.height() * z)
 
+    def resizeEvent(self, event):
+        super().resizeEvent(event)
+        if self.update_overlay is not None:
+            self.update_overlay.setGeometry(self.rect())
+
+    def showUpdatingOverlay(self, message: str = "Generating Plot..."):
+        if self.update_overlay is None:
+            return
+        self.update_overlay.set_message(message)
+        self.update_overlay.setGeometry(self.rect())
+        self.update_overlay.raise_()
+        self.update_overlay.show()
+
+    def hideUpdatingOverlay(self):
+        if self.update_overlay is None:
+            return
+        self.update_overlay.hide()
+
     def saveImage(self, filename):
         """Save an image of the current view
 
@@ -238,6 +287,8 @@ def mouseDoubleClickEvent(self, event):
 
     def mouseMoveEvent(self, event):
         cv = self.model.currentView
+        if self.ax is None or self.model.image is None:
+            return
         # Show Cursor position relative to plot in status bar
         xPlotPos, yPlotPos = self.getPlotCoords(event.pos())
 
@@ -343,6 +394,11 @@ def mouseReleaseEvent(self, event):
             self.rubber_band.hide()
             self.main_window.applyChanges()
         else:
+            plot_manager = self.main_window.plot_manager
+            if plot_manager.is_busy or plot_manager.has_pending:
+                return
+            if self.main_window.model.activeView != self.main_window.model.currentView:
+                return
             self.main_window.revertDockControls()
 
     def wheelEvent(self, event):
@@ -479,9 +535,7 @@ def generatePixmap(self, update=False):
         if self.frozen:
             return
 
-        self.model.generatePlot()
-        if update:
-            self.updatePixmap()
+        self.main_window.requestPlotUpdate()
 
     def updatePixmap(self):
 
@@ -500,9 +554,8 @@ def updatePixmap(self):
                        cv.origin[self.main_window.yBasis] - cv.height/2.,
                        cv.origin[self.main_window.yBasis] + cv.height/2.]
 
-        # make sure we have a domain image to load
-        if not hasattr(self.model, 'image'):
-            self.model.generatePlot()
+        if not hasattr(self.model, 'image') or self.model.image is None:
+            return
 
         ### DRAW DOMAIN IMAGE ###