Adding mesh annotation dock

paulromano · Nov 23, 2024 · Jul 22, 2024 · Jul 22, 2024 · Jul 24, 2024 · Jul 26, 2024
commit 397b9204f2202361f378ae915ccbc5147bb9918c
diff --git a/openmc_plotter/__init__.py b/openmc_plotter/__init__.py
@@ -1 +1 @@
-__version__ = '0.4.0'
+__version__ = '0.4.1'
diff --git a/openmc_plotter/docks.py b/openmc_plotter/docks.py
@@ -32,6 +32,55 @@ 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.Checked)
+            self.mesh_items.append((mesh_id, mesh_item))
+            self.meshTree.addTopLevelItem(mesh_item)
+
+        # 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):
     """

diff --git a/openmc_plotter/main_window.py b/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
 
@@ -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()
@@ -379,6 +385,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+T")
+        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')
@@ -388,6 +400,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)
@@ -616,12 +629,16 @@ 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 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.tallyDock.updateModel()
+            self.updateMeshAnnotations()
             self.model.storeCurrent()
             self.model.subsequentViews = []
             self.plotIm.generatePixmap()
@@ -798,6 +815,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)

diff --git a/openmc_plotter/plotgui.py b/openmc_plotter/plotgui.py
@@ -641,7 +641,8 @@ def updatePixmap(self):
         self.add_outlines()
 
         # annotate mesh boundaries
-        self.annotate_mesh(3)
+        for mid in cv.mesh_annotations:
+            self.annotate_mesh(mid)
 
         # always make sure the data bounds are set correctly
         self.ax.set_xbound(data_bounds[0], data_bounds[1])

diff --git a/openmc_plotter/plotmodel.py b/openmc_plotter/plotmodel.py
@@ -654,6 +654,9 @@ 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):
         mesh = openmc.lib.meshes[mesh_id]
 
@@ -991,14 +994,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)