Merge branch 'develop' into mesh_annotation

README.md

@@ -98,6 +98,15 @@ user's Matplotlib installation.
 
   - Tally and geometry data (material/cell IDs) can be exported to a VTK file under "File->Export"
 
+### Source Site Plotting
+
+Source locations from an externally defined source can be visualized in the plotter to verify
+source definitions. These source sites are gathered as generated by the transport model. A tolerance
+can be provided to filter out source sites that are too far from the slice plane, otherwise source
+locations are projected onto the slice plane.
+
+![Source plotting](./screenshots/source-sites.png)
+
 # Options/Functionality
 
 ## Menu Bar:

openmc_plotter/main_window.py

@@ -24,7 +24,7 @@
 from .plotgui import PlotImage, ColorDialog
 from .docks import DomainDock, TallyDock, MeshAnnotationDock
 from .overlays import ShortcutsOverlay
-from .tools import ExportDataDialog
+from .tools import ExportDataDialog, SourceSitesDialog
 
 
 def _openmcReload(threads=None, model_path='.'):
@@ -103,6 +103,7 @@ def loadGui(self, use_settings_pkl=True):
 
         # Tools
         self.exportDataDialog = ExportDataDialog(self.model, self.font_metric, self)
+        self.sourceSitesDialog = SourceSitesDialog(self.model, self.font_metric, self)
 
         # Keyboard overlay
         self.shortcutOverlay = ShortcutsOverlay(self)
@@ -212,12 +213,18 @@ def createMenuBar(self):
         self.openStatePointAction.setToolTip('Open statepoint file')
         self.openStatePointAction.triggered.connect(self.openStatePoint)
 
+        self.sourceSitesAction = QAction('&Sample source sites...', self)
+        self.sourceSitesAction.setToolTip('Add source sites to plot')
+        self.setStatusTip('Sample and add source sites to the plot')
+        self.sourceSitesAction.triggered.connect(self.plotSourceSites)
+
         self.importPropertiesAction = QAction("&Import properties...", self)
         self.importPropertiesAction.setToolTip("Import properties")
         self.importPropertiesAction.triggered.connect(self.importProperties)
 
         self.dataMenu = self.mainMenu.addMenu('D&ata')
         self.dataMenu.addAction(self.openStatePointAction)
+        self.dataMenu.addAction(self.sourceSitesAction)
         self.dataMenu.addAction(self.importPropertiesAction)
         self.updateDataMenu()
 
@@ -538,13 +545,14 @@ def loadViewFile(self, filename):
         except Exception:
             message = 'Error loading plot settings'
             saved = {'version': None,
-                        'current': None}
+                     'current': None}
+
         if saved['version'] == self.model.version:
             self.model.activeView = saved['current']
             self.dock.updateDock()
             self.colorDialog.updateDialogValues()
             self.applyChanges()
-            message = '{} settings loaded'.format(filename)
+            message = '{} loaded'.format(filename)
         else:
             message = 'Error loading plot settings. Incompatible model.'
         self.statusBar().showMessage(message, 5000)
@@ -629,9 +637,15 @@ 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_()
+        self.sourceSitesDialog.activateWindow()
+
     def applyChanges(self):
         if self.model.activeView != self.model.currentView:
             self.statusBar().showMessage('Generating Plot...')

openmc_plotter/plotgui.py

@@ -639,6 +639,7 @@ def updatePixmap(self):
 
         # annotate outlines
         self.add_outlines()
+        self.plotSourceSites()
 
         # annotate mesh boundaries
         for mid in cv.mesh_annotations:
@@ -674,6 +675,30 @@ def annotate_mesh(self, mesh_id):
                                             levels=np.unique(mesh_bins),
                                             extent=data_bounds)
 
+    def plotSourceSites(self):
+        if not self.model.sourceSitesVisible or self.model.sourceSites is None:
+            return
+
+        cv = self.model.currentView
+        basis = cv.view_params.basis
+
+        h_idx = 'xyz'.index(basis[0])
+        v_idx = 'xyz'.index(basis[1])
+
+        sites = self.model.sourceSites
+
+        slice_ax = cv.view_params.slice_axis
+
+        if self.model.sourceSitesApplyTolerance:
+            sites_to_plot = sites[np.abs(sites[:, slice_ax] - cv.origin[slice_ax]) <= self.model.sourceSitesTolerance]
+        else:
+            sites_to_plot = sites
+
+        self.ax.scatter([s[h_idx] for s in sites_to_plot],
+                        [s[v_idx] for s in sites_to_plot],
+                        marker='o',
+                        color=rgb_normalize(self.model.sourceSitesColor))
+
     def add_outlines(self):
         cv = self.model.currentView
         # draw outlines as isocontours

openmc_plotter/plotmodel.py

@@ -140,6 +140,14 @@ class PlotModel:
     subsequentViews : list of PlotView instances
         List of undone plot view settings used to redo changes made
         in plot explorer
+    sourceSitesTolerance : float
+        Tolerance for source site plotting (default 0.1 cm)
+    sourceSitesColor : tuple of 3 int
+        RGB color for source site plotting (default blue)
+    sourceSitesVisible : bool
+        Whether to plot source sites (default True)
+    sourceSites :  Source sites to plot
+        Set of source locations to plot
     defaultView : PlotView
         Default settings for given geometry
     currentView : PlotView
@@ -178,6 +186,13 @@ def __init__(self, use_settings_pkl, model_path, default_res):
 
         self.defaultView = self.getDefaultView(default_res)
 
+        # Source site defaults
+        self.sourceSitesApplyTolerance = False
+        self.sourceSitesTolerance = 0.1 # cm
+        self.sourceSitesColor = (0, 0, 255)
+        self.sourceSitesVisible = True
+        self.sourceSites = None
+
         if model_path.is_file():
             settings_pkl = model_path.with_name('plot_settings.pkl')
         else:
@@ -399,6 +414,15 @@ def redo(self):
             self.activeView = self.subsequentViews.pop()
             self.generatePlot()
 
+    def getExternalSourceSites(self, n=100):
+        """Plot source sites from a source file
+        """
+        if n == 0:
+            self.source_sites = None
+            return
+        sites = openmc.lib.sample_external_source(n)
+        self.sourceSites = np.array([s.r for s in sites[:n]], dtype=float)
+
     def storeCurrent(self):
         """ Add current view to previousViews list """
         self.previousViews.append(copy.deepcopy(self.currentView))
@@ -798,6 +822,8 @@ class ViewParam(openmc.lib.plot._PlotBase):
         Vertical resolution of plot image
     basis : {'xy', 'xz', 'yz'}
         The basis directions for the plot
+    slice_axis : int
+        The axis along which the plot is sliced
     color_overlaps : bool
         Indicator of whether or not overlaps will be shown
     level : int
@@ -818,6 +844,15 @@ def __init__(self, origin=(0, 0, 0), width=10, height=10, default_res=1000):
         self.basis = 'xy'
         self.color_overlaps = False
 
+    @property
+    def slice_axis(self):
+        if self.basis == 'xy':
+            return 2
+        elif self.basis == 'yz':
+            return 0
+        else:
+            return 1
+
     @property
     def llc(self):
         if self.basis == 'xy':
@@ -1155,6 +1190,12 @@ def __deepcopy__(self, memo):
         obj.defaults = self.defaults
         return obj
 
+    def get_defaults(self, key: int) -> DomainView:
+        return self.defaults[key]
+
+    def get_default_color(self, key: int):
+        return self.get_defaults(key).color
+
     def set_name(self, key: int, name: Optional[str]):
         domain = self[key]
         self[key] = DomainView(domain.id, name, domain.color, domain.masked, domain.highlight)
@@ -1244,7 +1285,8 @@ def data(self, index, role=Qt.DisplayRole):
             elif column == COLOR:
                 return '' if domain.color is not None else '+'
             elif column == COLORLABEL:
-                return str(tuple(domain.color)) if domain.color is not None else '--'
+                return (str(tuple(int(x) for x in domain.color))
+                        if domain.color is not None else '--')
             elif column == MASK:
                 return None
             elif column == HIGHLIGHT:
@@ -1322,10 +1364,11 @@ def setData(self, index, value, role=Qt.EditRole):
 
         if column == NAME:
             self.domains.set_name(key, value if value else None)
-        elif column == COLOR:
-            self.domains.set_color(key, value)
-        elif column == COLORLABEL:
-            self.domains.set_color(key, value)
+        elif column == COLOR or column == COLORLABEL:
+                # reset the color to the default value if the coloar value is None
+                if value is None:
+                    value = self.domains.get_default_color(key)
+                self.domains.set_color(key, value)
         elif column == MASK:
             if role == Qt.CheckStateRole:
                 self.domains.set_masked(key, Qt.CheckState(value) == Qt.Checked)
@@ -1381,7 +1424,7 @@ def setEditorData(self, editor, index):
     def editorEvent(self, event, model, option, index):
 
         if index.column() in (COLOR, COLORLABEL):
-            if not int(index.flags() & Qt.ItemIsEditable) > 0:
+            if (index.flags() & Qt.ItemFlag.ItemIsEditable) == Qt.ItemFlag.NoItemFlags:
                 return False
             if event.type() == QEvent.MouseButtonRelease \
                and event.button() == Qt.RightButton:

openmc_plotter/tools.py

@@ -8,6 +8,93 @@
 from .scientific_spin_box import ScientificDoubleSpinBox
 
 
+class SourceSitesDialog(QtWidgets.QDialog):
+    def __init__(self, model, font_metric, parent=None):
+        super().__init__(parent)
+
+        self.setWindowTitle('Sample Source Sites')
+        self.model = model
+        self.font_metric = font_metric
+        self.parent = parent
+
+        self.layout = QtWidgets.QFormLayout()
+        self.setLayout(self.layout)
+
+        self.populate()
+
+    def populate(self):
+        self.nSitesBox = QtWidgets.QSpinBox(self)
+        self.nSitesBox.setMaximum(1_000_000)
+        self.nSitesBox.setMinimum(0)
+        self.nSitesBox.setValue(1000)
+        self.nSitesBox.setToolTip('Number of source sites to sample from the OpenMC source')
+
+        self.sites_visible = QtWidgets.QCheckBox(self)
+        self.sites_visible.setChecked(self.model.sourceSitesVisible)
+        self.sites_visible.setToolTip('Toggle visibility of source sites on the slice plane')
+        self.sites_visible.stateChanged.connect(self._toggle_source_sites)
+
+        self.colorButton = QtWidgets.QPushButton(self)
+        self.colorButton.setToolTip('Select color for displaying source sites on the slice plane')
+        self.colorButton.setCursor(QtCore.Qt.PointingHandCursor)
+        self.colorButton.setFixedHeight(self.font_metric.height() * 1.5)
+        self.colorButton.clicked.connect(self._select_source_site_color)
+        rgb = self.model.sourceSitesColor
+        self.colorButton.setStyleSheet(
+            f"border-radius: 8px; background-color: rgb{rgb}")
+
+        self.toleranceBox = ScientificDoubleSpinBox()
+        self.toleranceBox.setToolTip('Slice axis tolerance for displaying source sites on the slice plane')
+        self.toleranceBox.setValue(self.model.sourceSitesTolerance)
+        self.toleranceBox.valueChanged.connect(self._set_source_site_tolerance)
+        self.toleranceBox.setEnabled(self.model.sourceSitesApplyTolerance)
+
+        self.toleranceToggle = QtWidgets.QCheckBox(self)
+        self.toleranceToggle.setChecked(self.model.sourceSitesApplyTolerance)
+        self.toleranceToggle.stateChanged.connect(self._toggle_tolerance)
+
+        self.sampleButton = QtWidgets.QPushButton("Sample New Sites")
+        self.sampleButton.setToolTip('Sample new source sites from the OpenMC source')
+        self.sampleButton.clicked.connect(self._sample_sites)
+
+        self.closeButton = QtWidgets.QPushButton("Close")
+        self.closeButton.clicked.connect(self.close)
+
+        self.layout.addRow("Source Sites:", self.nSitesBox)
+        self.layout.addRow("Visible:", self.sites_visible)
+        self.layout.addRow("Color:", self.colorButton)
+        self.layout.addRow('Tolerance:', self.toleranceBox)
+        self.layout.addRow('Apply tolerance:', self.toleranceToggle)
+        self.layout.addRow(HorizontalLine())
+        self.layout.addRow(self.sampleButton)
+        self.layout.addRow(self.closeButton)
+
+    def _sample_sites(self):
+        self.model.getExternalSourceSites(self.nSitesBox.value())
+        self.parent.applyChanges()
+
+    def _toggle_source_sites(self):
+        self.model.sourceSitesVisible = self.sites_visible.isChecked()
+        self.parent.applyChanges()
+
+    def _select_source_site_color(self):
+        color = QtWidgets.QColorDialog.getColor()
+        if color.isValid():
+            rgb = self.model.sourceSitesColor = color.getRgb()[:3]
+            self.colorButton.setStyleSheet(
+                f"border-radius: 8px; background-color: rgb{rgb}")
+            self.parent.applyChanges()
+
+    def _toggle_tolerance(self):
+        self.model.sourceSitesApplyTolerance = self.toleranceToggle.isChecked()
+        self.toleranceBox.setEnabled(self.toleranceToggle.isChecked())
+        self.parent.applyChanges()
+
+    def _set_source_site_tolerance(self):
+        self.model.sourceSitesTolerance = self.toleranceBox.value()
+        self.parent.applyChanges()
+
+
 class ExportDataDialog(QtWidgets.QDialog):
     """
     A dialog to facilitate generation of VTK files for