Adding nascent source site plot capability.

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='.'):
@@ -95,6 +95,8 @@ 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)
+
         # Restore Window Settings
         self.restoreWindowSettings()
 
@@ -203,12 +205,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()
 
@@ -522,13 +530,15 @@ def loadViewFile(self, filename):
         except Exception:
             message = 'Error loading plot settings'
             saved = {'version': None,
-                        'current': None}
+                     'current': None}
+
+        print(saved['current'])
         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)
@@ -613,6 +623,12 @@ def updateDataMenu(self):
         elif hasattr(self, "closeStatePointAction"):
             self.dataMenu.removeAction(self.closeStatePointAction)
 
+    def plotSourceSites(self):
+        self.sourceSitesDialog.show()
+        self.sourceSitesDialog.raise_()
+        self.sourceSitesDialog.activateWindow()
+        self.applyChanges()
+
     def applyChanges(self):
         if self.model.activeView != self.model.currentView:
             self.statusBar().showMessage('Generating Plot...')

openmc_plotter/plotgui.py

@@ -644,9 +644,30 @@ def updatePixmap(self):
         self.ax.dataLim.y0 = data_bounds[2]
         self.ax.dataLim.y1 = data_bounds[3]
 
+        self.plotSites()
+
         self.draw()
         return "Done"
 
+    def plotSites(self):
+        if self.model.source_sites is None:
+            return
+
+        cv = self.model.currentView
+        basis = cv.view_params.basis
+
+        print(self.model.source_sites[:]['xyz'.index(basis[0])],
+              self.model.source_sites[:]['xyz'.index(basis[1])])
+
+        h_idx = 'xyz'.index(basis[0])
+        v_idx = 'xyz'.index(basis[1])
+
+        self.ax.scatter([s[h_idx] for s in self.model.source_sites],
+                        [s[v_idx] for s in self.model.source_sites],
+                        marker='o',
+
+                        color='blue')
+
     def add_outlines(self):
         cv = self.model.currentView
         # draw outlines as isocontours

openmc_plotter/plotmodel.py

@@ -123,6 +123,8 @@ class PlotModel:
     subsequentViews : list of PlotView instances
         List of undone plot view settings used to redo changes made
         in plot explorer
+    sourceSites :  Source sites to plot
+        Set of source locations to plot
     defaultView : PlotView
         Default settings for given geometry
     currentView : PlotView
@@ -164,6 +166,8 @@ def __init__(self, use_settings_pkl, model_path):
 
         self.defaultView = self.getDefaultView()
 
+        self.source_sites = None
+
         if model_path.is_file():
             settings_pkl = model_path.with_name('plot_settings.pkl')
         else:
@@ -378,6 +382,12 @@ def redo(self):
             self.activeView = self.subsequentViews.pop()
             self.generatePlot()
 
+    def getExternalSourceSites(self, n=100):
+        """Plot source sites from a source file
+        """
+        sites = openmc.lib.sample_external_source(n)
+        self.source_sites = 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))

openmc_plotter/tools.py

@@ -8,6 +8,40 @@
 from .scientific_spin_box import ScientificDoubleSpinBox
 
 
+class SourceSitesDialog(QtWidgets.QDialog):
+    def __init__(self, model, font_metric, parent=None):
+        super().__init__(parent)
+
+        self.model = model
+        self.font_metric = font_metric
+        self.parent = parent
+
+        self.layout = QtWidgets.QGridLayout()
+        self.setLayout(self.layout)
+
+        # disable interaction with main window while this is open
+        self.setModal(True)
+
+    def show(self):
+        self.populate()
+        super().show()
+
+    def populate(self):
+        self.layout.addWidget(QtWidgets.QLabel("# Source Sites:"), 0, 0)
+        self.nSitesBox = QtWidgets.QSpinBox()
+        self.nSitesBox.setMaximum(10000)
+        self.nSitesBox.setMinimum(0)
+        self.layout.addWidget(self.nSitesBox, 0, 1)
+
+        self.sampleButton = QtWidgets.QPushButton("Sample Sites")
+        self.sampleButton.clicked.connect(self._sample_sites)
+
+        self.layout.addWidget(self.sampleButton, 1, 0, 1, 2)
+
+    def _sample_sites(self):
+        self.model.getExternalSourceSites(self.nSitesBox.value())
+        self.close()
+
 class ExportDataDialog(QtWidgets.QDialog):
     """
     A dialog to facilitate generation of VTK files for