Mesh annotations (#153)

Co-authored-by: Paul Romano <paul.k.romano@gmail.com>

openmc_plotter/__init__.py

@@ -1 +1 @@
-__version__ = '0.4.0'
+__version__ = '0.4.1'

openmc_plotter/docks.py

@@ -33,6 +33,57 @@ def __init__(self, model, font_metric, parent=None):
         self.setSizePolicy(QSizePolicy.Fixed, QSizePolicy.Expanding)
 
 
+class MeshAnnotationDock(PlotterDock):
+    """Dock for mesh annotation options"""
+
+    def __init__(self, model, font_metric, parent=None):
+        super().__init__(model, font_metric, parent)
+
+        self.treeLayout = QVBoxLayout()
+        self.meshTree = QTreeWidget()
+        self.treeExpander = Expander("Meshes:", layout=self.treeLayout)
+        self.treeExpander.expand()  # start with meshes expanded
+
+        self.meshTree.setColumnCount(1)
+
+        self.mesh_items = []
+        for mesh_id in self.model.cpp_mesh_ids():
+            mesh_item = QTreeWidgetItem(self.meshTree, (f'Mesh {mesh_id}',))
+            mesh_item.setFlags(mesh_item.flags() | QtCore.Qt.ItemIsUserCheckable)
+            mesh_item.setCheckState(0, QtCore.Qt.Unchecked)
+            self.mesh_items.append((mesh_id, mesh_item))
+            self.meshTree.addTopLevelItem(mesh_item)
+
+        self.meshTree.setHeaderHidden(True)
+
+        # Create submit button
+        self.applyButton = QPushButton("Apply Changes")
+        # Mac bug fix
+        self.applyButton.setMinimumHeight(self.font_metric.height() * 1.6)
+        self.applyButton.clicked.connect(self.main_window.applyChanges)
+
+        label = QLabel("Mesh Annotations")
+        self.treeLayout.addWidget(label)
+        self.treeLayout.addWidget(self.meshTree)
+        self.treeLayout.addWidget(HorizontalLine())
+        self.treeLayout.addWidget(self.applyButton)
+
+        self.optionsWidget = QWidget()
+        self.optionsWidget.setLayout(self.treeLayout)
+        self.setWidget(self.optionsWidget)
+
+    def get_checked_meshes(self):
+        return [id for id, item in self.mesh_items if item.checkState(0) == QtCore.Qt.Checked]
+
+    def update(self):
+        pass
+
+    def resizeEvent(self, event):
+        self.main_window.resizeEvent(event)
+
+    hideEvent = showEvent = moveEvent = resizeEvent
+
+
 class DomainDock(PlotterDock):
     """
     Domain options dock

openmc_plotter/main_window.py

@@ -22,7 +22,7 @@
 
 from .plotmodel import PlotModel, DomainTableModel, hash_model
 from .plotgui import PlotImage, ColorDialog
-from .docks import DomainDock, TallyDock
+from .docks import DomainDock, TallyDock, MeshAnnotationDock
 from .overlays import ShortcutsOverlay
 from .tools import ExportDataDialog, SourceSitesDialog
 
@@ -91,6 +91,12 @@ def loadGui(self, use_settings_pkl=True):
         self.tallyDock.setObjectName("Tally Options Dock")
         self.addDockWidget(QtCore.Qt.RightDockWidgetArea, self.tallyDock)
 
+        # Mesh Annotation Dock
+        self.meshAnnotationDock = MeshAnnotationDock(self.model, self.font_metric, self)
+        self.meshAnnotationDock.update()
+        self.meshAnnotationDock.setObjectName("Mesh Annotation Dock")
+        self.addDockWidget(QtCore.Qt.RightDockWidgetArea, self.meshAnnotationDock)
+
         # Color DialogtallyDock
         self.colorDialog = ColorDialog(self.model, self.font_metric, self)
         self.colorDialog.hide()
@@ -386,6 +392,12 @@ def createMenuBar(self):
         self.tallyDockAction.setStatusTip('Toggle tally dock visibility')
         self.tallyDockAction.triggered.connect(self.toggleTallyDockView)
 
+        self.meshAnnotationDockAction = QAction('Mesh &Annotation Dock', self)
+        self.meshAnnotationDockAction.setShortcut("Ctrl+E")
+        self.meshAnnotationDockAction.setToolTip('Toggle mesh annotation dock visibility')
+        self.meshAnnotationDockAction.setStatusTip('Toggle mesh annotation dock visibility')
+        self.meshAnnotationDockAction.triggered.connect(self.toggleMeshAnnotationDockView)
+
         self.zoomAction = QAction('&Zoom...', self)
         self.zoomAction.setShortcut('Alt+Shift+Z')
         self.zoomAction.setToolTip('Edit zoom factor')
@@ -395,6 +407,7 @@ def createMenuBar(self):
         self.viewMenu = self.mainMenu.addMenu('&View')
         self.viewMenu.addAction(self.dockAction)
         self.viewMenu.addAction(self.tallyDockAction)
+        self.viewMenu.addAction(self.meshAnnotationDockAction)
         self.viewMenu.addSeparator()
         self.viewMenu.addAction(self.zoomAction)
         self.viewMenu.aboutToShow.connect(self.updateViewMenu)
@@ -624,6 +637,10 @@ def updateDataMenu(self):
         elif hasattr(self, "closeStatePointAction"):
             self.dataMenu.removeAction(self.closeStatePointAction)
 
+
+    def updateMeshAnnotations(self):
+        self.model.activeView.mesh_annotations = self.meshAnnotationDock.get_checked_meshes()
+
     def plotSourceSites(self):
         self.sourceSitesDialog.show()
         self.sourceSitesDialog.raise_()
@@ -635,6 +652,7 @@ def applyChanges(self):
             QApplication.processEvents()
             if self.model.activeView.selectedTally is not None:
                 self.tallyDock.updateModel()
+            self.updateMeshAnnotations()
             self.model.storeCurrent()
             self.model.subsequentViews = []
             self.plotIm.generatePixmap()
@@ -811,6 +829,18 @@ def toggleTallyDockView(self):
         self.resizePixmap()
         self.showMainWindow()
 
+    def toggleMeshAnnotationDockView(self):
+        if self.meshAnnotationDock.isVisible():
+            self.meshAnnotationDock.hide()
+            if not self.isMaximized() and not self.meshAnnotationDock.isFloating():
+                self.resize(self.width() - self.meshAnnotationDock.width(), self.height())
+        else:
+            self.meshAnnotationDock.setVisible(True)
+            if not self.isMaximized() and not self.meshAnnotationDock.isFloating():
+                self.resize(self.width() + self.meshAnnotationDock.width(), self.height())
+        self.resizePixmap()
+        self.showMainWindow()
+
     def editZoomAct(self):
         percent, ok = QInputDialog.getInt(self, "Edit Zoom", "Zoom Percent:",
                                           self.dock.zoomBox.value(), 25, 2000)

openmc_plotter/overlays.py

@@ -40,6 +40,7 @@ class ShortcutsOverlay(QWidget):
                             ("Horizontal Scroll", "Alt+Scroll")],
                   "Menus": [("Hide/Show Geometry Dock", c_key + "+D"),
                             ("Hide/Show Tally Dock", c_key + "+T"),
+                            ("Hide/Show Mesh Annotation Dock", c_key + "+E"),
                             ("Reload Model", "Shift+" + c_key + "+R"),
                             ("Quit", c_key + "+Q"),
                             ("Display Shortcuts", "?")],

openmc_plotter/plotgui.py

@@ -641,6 +641,10 @@ def updatePixmap(self):
         self.add_outlines()
         self.plotSourceSites()
 
+        # annotate mesh boundaries
+        for mesh_id in cv.mesh_annotations:
+            self.annotate_mesh(mesh_id)
+
         # always make sure the data bounds are set correctly
         self.ax.set_xbound(data_bounds[0], data_bounds[1])
         self.ax.set_ybound(data_bounds[2], data_bounds[3])
@@ -652,6 +656,26 @@ def updatePixmap(self):
         self.draw()
         return "Done"
 
+    def current_view_data_bounds(self):
+        cv = self.model.currentView
+        return [cv.origin[self.main_window.xBasis] - cv.width/2.,
+                cv.origin[self.main_window.xBasis] + cv.width/2.,
+                cv.origin[self.main_window.yBasis] - cv.height/2.,
+                cv.origin[self.main_window.yBasis] + cv.height/2.]
+
+    def annotate_mesh(self, mesh_id):
+        mesh_bins = self.model.mesh_plot_bins(mesh_id)
+
+        data_bounds = self.current_view_data_bounds()
+        self.mesh_contours = self.ax.contour(
+            mesh_bins,
+            origin='upper',
+            colors='k',
+            linestyles='solid',
+            levels=np.unique(mesh_bins),
+            extent=data_bounds
+        )
+
     def plotSourceSites(self):
         if not self.model.sourceSitesVisible or self.model.sourceSites is None:
             return
@@ -681,10 +705,7 @@ def add_outlines(self):
         # draw outlines as isocontours
         if cv.outlines:
             # set data extents for automatic reporting of pointer location
-            data_bounds = [cv.origin[self.main_window.xBasis] - cv.width/2.,
-                           cv.origin[self.main_window.xBasis] + cv.width/2.,
-                           cv.origin[self.main_window.yBasis] - cv.height/2.,
-                           cv.origin[self.main_window.yBasis] + cv.height/2.]
+            data_bounds = self.current_view_data_bounds()
             levels = np.unique(self.model.ids)
             self.contours = self.ax.contour(self.model.ids,
                                             origin='upper',

openmc_plotter/plotmodel.py

@@ -678,6 +678,30 @@ def _create_distribcell_image(self, tally, tally_value, scores, nuclides, cellin
 
         return image_data, None, data_min, data_max
 
+    def cpp_mesh_ids(self):
+        return list(openmc.lib.meshes.keys())
+
+    def mesh_plot_bins(self, mesh_id, view: PlotView = None, translation: tuple[float, float, float] = None):
+        mesh = openmc.lib.meshes[mesh_id]
+
+        if view is None:
+            view = self.currentView
+
+        if translation is None:
+            origin = view.origin
+        else:
+            origin = (view.origin[0] - translation[0],
+                      view.origin[1] - translation[1],
+                      view.origin[2] - translation[2])
+
+        mesh_bins = mesh.get_plot_bins(
+            origin=origin,
+            width=(view.width, view.height),
+            basis=view.basis,
+            pixels=(view.h_res, view.v_res),
+        )
+        return mesh_bins
+
     def _create_tally_mesh_image(
             self, tally: openmc.Tally, tally_value: TallyValueType,
             scores: Tuple[str], nuclides: Tuple[str], view: PlotView = None
@@ -737,21 +761,9 @@ def _do_op(array, tally_value, ax=0):
                 selected_scores.append(idx)
         data = _do_op(data[np.array(selected_scores)], tally_value)
 
-        # Account for mesh filter translation
-        if mesh_filter.translation is not None:
-            t = mesh_filter.translation
-            origin = (view.origin[0] - t[0], view.origin[1] - t[1], view.origin[2] - t[2])
-        else:
-            origin = view.origin
-
         # Get mesh bins from openmc.lib
         mesh_cpp = openmc.lib.meshes[mesh.id]
-        mesh_bins = mesh_cpp.get_plot_bins(
-            origin=origin,
-            width=(view.width, view.height),
-            basis=view.basis,
-            pixels=(view.h_res, view.v_res),
-        )
+        mesh_bins = self.mesh_plot_bins(mesh.id, view, mesh_filter.translation)
 
         # Apply volume normalization
         if view.tallyVolumeNorm:
@@ -1017,14 +1029,16 @@ class PlotView:
         Label of the currently selected tally
     """
 
-    attrs = ('view_ind', 'view_params', 'cells', 'materials', 'selectedTally')
+    attrs = ('view_ind', 'view_params', 'cells', 'materials', 'selectedTally', 'mesh_annotations')
     plotbase_attrs = ('level', 'origin', 'width', 'height',
                       'h_res', 'v_res', 'basis', 'llc', 'urc', 'color_overlaps')
 
     def __init__(self, origin=(0, 0, 0), width=10, height=10, restore_view=None,
                  restore_domains=False, default_res=None):
         """Initialize PlotView attributes"""
 
+        self.mesh_annotations = []
+
         if restore_view is not None:
             self.view_ind = copy.copy(restore_view.view_ind)
             self.view_params = copy.copy(restore_view.view_params)