diff --git a/openmc_plotter/main_window.py b/openmc_plotter/main_window.py index f7aa756..ae2b182 100755 --- a/openmc_plotter/main_window.py +++ b/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...') diff --git a/openmc_plotter/plotgui.py b/openmc_plotter/plotgui.py index 0c32173..4f27c67 100644 --- a/openmc_plotter/plotgui.py +++ b/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 diff --git a/openmc_plotter/plotmodel.py b/openmc_plotter/plotmodel.py index 9b822ae..48d6423 100644 --- a/openmc_plotter/plotmodel.py +++ b/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)) diff --git a/openmc_plotter/tools.py b/openmc_plotter/tools.py index 16194ab..7d27555 100644 --- a/openmc_plotter/tools.py +++ b/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