Adding nascent source site plot capability. (#131)

* Adding nascent source site plot capability.

* Polishing source site plotting capability

* Plot sites before setting axes extents

* Updating README

* Clean up interface for source site sampling

* Update openmc_plotter/plotgui.py

---------

Co-authored-by: Paul Romano <[email protected]>

Author: pshriwise

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
 from .overlays import ShortcutsOverlay
-from .tools import ExportDataDialog
+from .tools import ExportDataDialog, SourceSitesDialog
 
 
 def _openmcReload(threads=None, model_path='.'):
@@ -97,6 +97,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)
@@ -206,12 +207,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()
 
@@ -525,13 +532,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)
@@ -616,6 +624,11 @@ def updateDataMenu(self):
         elif hasattr(self, "closeStatePointAction"):
             self.dataMenu.removeAction(self.closeStatePointAction)
 
+    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()
 
         # always make sure the data bounds are set correctly
         self.ax.set_xbound(data_bounds[0], data_bounds[1])
@@ -651,6 +652,30 @@ def updatePixmap(self):
         self.draw()
         return "Done"
 
+    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))
@@ -786,6 +810,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
@@ -806,6 +832,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':

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