Converted plot_mesh_tally.py to use Tkinter rather than PyQt.

This commit is contained in:
Paul Romano 2014-01-20 19:00:12 -05:00
parent b38af09b82
commit 55d257351c

View file

@ -1,269 +1,241 @@
#!/usr/bin/env python2
#!/usr/bin/env python
'''Python script to plot tally data generated by OpenMC.'''
"""Python script to plot tally data generated by OpenMC."""
import os
import sys
from statepoint import *
# Color intensity dependent on individual score?
from PyQt4.QtCore import *
from PyQt4.QtGui import *
import matplotlib.pyplot as plt
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg as FigureCanvas
from matplotlib.backends.backend_tkagg import NavigationToolbar2TkAgg as NavigationToolbar
from matplotlib.figure import Figure
from matplotlib.backends.backend_qt4agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qt4agg import NavigationToolbar2QTAgg as NavigationToolbar
import matplotlib.pyplot as plt
import numpy as np
class AppForm(QMainWindow):
def __init__(self, argv, parent=None):
QMainWindow.__init__(self, parent)
from statepoint import *
if sys.version_info[0] < 3:
import Tkinter as tk
else:
import tkinter as tk
import tkFileDialog
import tkFont
import tkMessageBox
import ttk
class MeshPlotter(tk.Frame):
def __init__(self, parent, filename):
tk.Frame.__init__(self, parent)
self.labels = {'cell': 'Cell:', 'cellborn': 'Cell born:',
'surface': 'Surface:', 'material': 'Material:',
'universe': 'Universe:', 'energyin': 'Energy in:',
'energyout': 'Energy out:'}
self.filterBoxes = {}
# Read data from source or leakage fraction file
self.all_good = False
while not self.all_good:
if len(argv) > 1:
cl_file = str(argv[1])
else:
cl_file = None
self.get_file_data(cl_file)
# Check that there are any mesh tallies at all
if len(self.tally_ids) != 0:
self.all_good = True
else:
# if there are not, the user will be given the choice to choose
# another file (but only if using interactive chooser)
if cl_file is None:
choice = QMessageBox.critical(None, "Invalid StatePoint File",
"File Does Not Contain Mesh " +
"Tallies!" +
"\nSelect Another File Or Quit",
QMessageBox.Retry,
QMessageBox.Abort)
if choice == QMessageBox.Abort:
self.all_good = False
break
else:
print("Invalid StatePoint File; File Does Not Contain " +
"Mesh Tallies!")
self.all_good = False
break
self.get_file_data(filename)
# Set up top-level window
top = self.winfo_toplevel()
top.title('Mesh Tally Plotter: ' + filename)
top.rowconfigure(0, weight=1)
top.columnconfigure(0, weight=1)
self.grid(sticky=tk.W+tk.N)
if self.all_good:
# Set maximum colorbar value by maximum tally data value
self.maxvalue = self.datafile.tallies[0].results.max()
# Create widgets and draw to screen
self.create_widgets()
self.update()
self.main_frame = QWidget()
self.setCentralWidget(self.main_frame)
def create_widgets(self):
figureFrame = tk.Frame(self)
figureFrame.grid(row=0, column=0)
# Create the Figure, Canvas, and Axes
self.dpi = 100
self.fig = Figure((5.0, 15.0), dpi=self.dpi)
self.canvas = FigureCanvas(self.fig)
self.canvas.setParent(self.main_frame)
self.axes = self.fig.add_subplot(111)
# Create the Figure and Canvas
self.dpi = 100
self.fig = Figure((5.0, 5.0), dpi=self.dpi)
self.canvas = FigureCanvas(self.fig, master=figureFrame)
self.canvas.get_tk_widget().pack(side=tk.TOP, fill=tk.BOTH, expand=1)
# Create the navigation toolbar, tied to the canvas
self.mpl_toolbar = NavigationToolbar(self.canvas, self.main_frame)
# Create the navigation toolbar, tied to the canvas
self.mpl_toolbar = NavigationToolbar(self.canvas, figureFrame)
self.mpl_toolbar.update()
self.canvas._tkcanvas.pack(side=tk.TOP, fill=tk.BOTH, expand=1)
# Grid layout at bottom
self.grid = QGridLayout()
# Create frame for comboboxes
self.selectFrame = tk.Frame(self)
self.selectFrame.grid(row=1, column=0, sticky=tk.W+tk.E)
# Overall layout
self.vbox = QVBoxLayout()
self.vbox.addWidget(self.canvas)
self.vbox.addWidget(self.mpl_toolbar)
self.vbox.addLayout(self.grid)
self.main_frame.setLayout(self.vbox)
# Tally selection
labelTally = tk.Label(self.selectFrame, text='Tally:')
labelTally.grid(row=0, column=0, sticky=tk.W)
self.tallyBox = ttk.Combobox(self.selectFrame, state='readonly')
self.tallyBox['values'] = [self.datafile.tallies[i].id
for i in self.meshTallies]
self.tallyBox.current(0)
self.tallyBox.grid(row=0, column=1, sticky=tk.W+tk.E)
self.tallyBox.bind('<<ComboboxSelected>>', self.update)
# Tally selections
label_tally = QLabel("Tally:")
self.tally = QComboBox()
# Only show options for the tallies with meshes
self.tally.addItems([str(i + 1) for i in self.tally_ids])
self.connect(self.tally, SIGNAL('activated(int)'),
self._update)
self.connect(self.tally, SIGNAL('activated(int)'),
self.populate_boxes)
self.connect(self.tally, SIGNAL('activated(int)'),
self.on_draw)
# Planar basis selection
labelBasis = tk.Label(self.selectFrame, text='Basis:')
labelBasis.grid(row=1, column=0, sticky=tk.W)
self.basisBox = ttk.Combobox(self.selectFrame, state='readonly')
self.basisBox['values'] = ('xy', 'yz', 'xz')
self.basisBox.current(0)
self.basisBox.grid(row=1, column=1, sticky=tk.W+tk.E)
self.basisBox.bind('<<ComboboxSelected>>', self.update)
# Planar basis
label_basis = QLabel("Basis:")
self.basis = QComboBox()
self.basis.addItems(['xy', 'yz', 'xz'])
# Axial level
labelAxial = tk.Label(self.selectFrame, text='Axial level:')
labelAxial.grid(row=2, column=0, sticky=tk.W)
self.axialBox = ttk.Combobox(self.selectFrame, state='readonly')
self.axialBox.grid(row=2, column=1, sticky=tk.W+tk.E)
self.axialBox.bind('<<ComboboxSelected>>', self.redraw)
# Update window when 'Basis' selection is changed
self.connect(self.basis, SIGNAL('activated(int)'),
self._update)
self.connect(self.basis, SIGNAL('activated(int)'),
self.populate_boxes)
self.connect(self.basis, SIGNAL('activated(int)'),
self.on_draw)
# Option for mean/uncertainty
labelMean = tk.Label(self.selectFrame, text='Mean/Uncertainty:')
labelMean.grid(row=3, column=0, sticky=tk.W)
self.meanBox = ttk.Combobox(self.selectFrame, state='readonly')
self.meanBox['values'] = ('Mean', 'Absolute uncertainty',
'Relative uncertainty')
self.meanBox.current(0)
self.meanBox.grid(row=3, column=1, sticky=tk.W+tk.E)
self.meanBox.bind('<<ComboboxSelected>>', self.update)
# Axial level within selected basis
label_axial_level = QLabel("Axial Level:")
self.axial_level = QComboBox()
self.connect(self.axial_level, SIGNAL('activated(int)'),
self.on_draw)
# Scores
labelScore = tk.Label(self.selectFrame, text='Score:')
labelScore.grid(row=4, column=0, sticky=tk.W)
self.scoreBox = ttk.Combobox(self.selectFrame, state='readonly')
self.scoreBox.grid(row=4, column=1, sticky=tk.W+tk.E)
self.scoreBox.bind('<<ComboboxSelected>>', self.redraw)
# Add Option to plot mean or uncertainty
label_mean = QLabel("Mean or Uncertainty:")
self.mean = QComboBox()
self.mean.addItems(['Mean','Absolute Uncertainty',
'Relative Uncertainty'])
# Filter label
font = tkFont.Font(weight='bold')
labelFilters = tk.Label(self.selectFrame, text='Filters:', font=font)
labelFilters.grid(row=5, column=0, sticky=tk.W)
# Update window when mean selection is changed
self.connect(self.mean, SIGNAL('activated(int)'),
self.on_draw)
self.label_filters = QLabel("Filter options:")
# Labels for all possible filters
self.labels = {'cell': 'Cell: ', 'cellborn': 'Cell born: ',
'surface': 'Surface: ', 'material': 'Material',
'universe': 'Universe: ', 'energyin': 'Energy in: ',
'energyout': 'Energy out: '}
# Empty reusable labels
self.qlabels = {}
for j in range(8):
self.nextLabel = QLabel
self.qlabels[j] = self.nextLabel
# Reusable comboboxes labelled with filter names
self.boxes = {}
for key in self.labels.keys():
self.nextBox = QComboBox()
self.connect(self.nextBox, SIGNAL('activated(int)'),
self.on_draw)
self.boxes[key] = self.nextBox
# Combobox to select among scores
self.score_label = QLabel("Score:")
self.scoreBox = QComboBox()
for item in self.tally_scores[0]:
self.scoreBox.addItems(str(item))
self.connect(self.scoreBox, SIGNAL('activated(int)'),
self.on_draw)
# Fill layout
self.grid.addWidget(label_tally, 0, 0)
self.grid.addWidget(self.tally, 0, 1)
self.grid.addWidget(label_basis, 1, 0)
self.grid.addWidget(self.basis, 1, 1)
self.grid.addWidget(label_axial_level, 2, 0)
self.grid.addWidget(self.axial_level, 2, 1)
self.grid.addWidget(label_mean, 3, 0)
self.grid.addWidget(self.mean, 3, 1)
self.grid.addWidget(self.label_filters, 4, 0)
self._update()
self.populate_boxes()
self.on_draw()
def get_file_data(self, cl_file=None):
# Get data file name from "open file" browser
if cl_file is None:
filename = QFileDialog.getOpenFileName(self,
'Select statepoint file', '.')
def update(self, event=None):
if not event:
widget = None
else:
filename = cl_file
widget = event.widget
# Create StatePoint object and read in data
self.datafile = StatePoint(str(filename))
self.datafile.read_results()
self.datafile.generate_stdev()
tally_id = self.meshTallies[self.tallyBox.current()]
selectedTally = self.datafile.tallies[tally_id]
self.setWindowTitle('Core Map Tool : ' + str(self.datafile.path))
# Get mesh for selected tally
self.mesh = self.datafile.meshes[selectedTally.filters['mesh'].bins[0] - 1]
self.labelList = []
# Get mesh dimensions
self.nx, self.ny, self.nz = self.mesh.dimension
# Read mesh dimensions
# for mesh in self.datafile.meshes:
# self.nx, self.ny, self.nz = mesh.dimension
# Repopulate comboboxes baesd on current basis selection
text = self.basisBox['values'][self.basisBox.current()]
if text == 'xy':
self.axialBox['values'] = [str(i+1) for i in range(self.nz)]
elif text == 'yz':
self.axialBox['values'] = [str(i+1) for i in range(self.nx)]
else:
self.axialBox['values'] = [str(i+1) for i in range(self.ny)]
self.axialBox.current(0)
# Find which tallies have meshes so the rest can be ignored,
# and for these tallies read the filter and score types
self.tally_ids = []
self.n_tallies = len(self.datafile.tallies)
self.tally_list = []
self.tally_scores = []
for itally, tally in enumerate(self.datafile.tallies):
if 'mesh' in tally.filters:
# Then we have a good tally, store the ID, filters and
# scores
self.tally_ids.append(itally)
self.filter_types = []
for f in tally.filters:
self.filter_types.append(f)
self.tally_list.append(self.filter_types)
self.score_types = []
for s in tally.scores:
self.score_types.append(s)
self.tally_scores.append(self.score_types)
# If update() was called by a change in the basis combobox, we don't
# need to repopulate the filters
if widget == self.basisBox:
self.redraw()
return
def on_draw(self):
""" Redraws the figure
"""
# Update scores
self.scoreBox['values'] = selectedTally.scores # self.tally_scores[self.tallyBox.current()]
self.scoreBox.current(0)
# print 'Calling on_draw...'
# Get selected basis, axial_level and stage
basis = self.basis.currentIndex() + 1
axial_level = self.axial_level.currentIndex() + 1
is_mean = self.mean.currentIndex()
# Remove any filter labels/comboboxes that exist
for row in range(6, self.selectFrame.grid_size()[1]):
for w in self.selectFrame.grid_slaves(row=row):
w.grid_forget()
w.destroy()
# get current tally index
tally_id = self.tally_ids[self.tally.currentIndex()]
# create a label/combobox for each filter in selected tally
count = 0
for filterType in selectedTally.filters:
if filterType == 'mesh':
continue
count += 1
# Create label and combobox for this filter
label = tk.Label(self.selectFrame, text=self.labels[filterType])
label.grid(row=count+6, column=0, sticky=tk.W)
combobox = ttk.Combobox(self.selectFrame, state='readonly')
self.filterBoxes[filterType] = combobox
# Set combobox items
f = selectedTally.filters[filterType]
if filterType in ['energyin', 'energyout']:
combobox['values'] = ['{0} to {1}'.format(*f.bins[i:i+2])
for i in range(f.length)]
else:
combobox['values'] = [str(i) for i in f.bins]
combobox.current(0)
combobox.grid(row=count+6, column=1, sticky=tk.W+tk.E)
combobox.bind('<<ComboboxSelected>>', self.redraw)
self.redraw()
def redraw(self, event=None):
basis = self.basisBox.current() + 1
axial_level = self.axialBox.current() + 1
is_mean = self.meanBox.current()
# Get selected tally
tally_id = self.meshTallies[self.tallyBox.current()]
selectedTally = self.datafile.tallies[tally_id]
# Create spec_list
spec_list = []
for tally in self.datafile.tallies[tally_id].filters.values():
if tally.type == 'mesh':
for f in selectedTally.filters.values():
if f.type == 'mesh':
continue
index = self.boxes[tally.type].currentIndex()
spec_list.append((tally.type, index))
index = self.filterBoxes[f.type].current()
spec_list.append((f.type, index))
# Take is_mean and convert it to an index of the score
score_loc = is_mean
if score_loc > 1:
score_loc = 1
if self.basis.currentText() == 'xy':
text = self.basisBox['values'][self.basisBox.current()]
if text == 'xy':
matrix = np.zeros((self.nx, self.ny))
for i in range(self.nx):
for j in range(self.ny):
matrix[i,j] = self.datafile.get_value(tally_id,
spec_list + [('mesh', (i + 1, j + 1, axial_level))],
self.scoreBox.currentIndex())[score_loc]
# Calculate relative uncertainty from absolute, if
# requested
self.scoreBox.current())[score_loc]
# Calculate relative uncertainty from absolute, if requested
if is_mean == 2:
# Take care to handle zero means when normalizing
mean_val = self.datafile.get_value(tally_id,
spec_list + [('mesh', (i + 1, j + 1, axial_level))],
self.scoreBox.currentIndex())[0]
self.scoreBox.current())[0]
if mean_val > 0.0:
matrix[i,j] = matrix[i,j] / mean_val
else:
matrix[i,j] = 0.0
elif self.basis.currentText() == 'yz':
elif text == 'yz':
matrix = np.zeros((self.ny, self.nz))
for i in range(self.ny):
for j in range(self.nz):
matrix[i,j] = self.datafile.get_value(tally_id,
spec_list + [('mesh', (axial_level, i + 1, j + 1))],
self.scoreBox.currentIndex())[score_loc]
# Calculate relative uncertainty from absolute, if
# requested
self.scoreBox.current())[score_loc]
# Calculate relative uncertainty from absolute, if requested
if is_mean == 2:
# Take care to handle zero means when normalizing
mean_val = self.datafile.get_value(tally_id,
spec_list + [('mesh', (axial_level, i + 1, j + 1))],
self.scoreBox.currentIndex())[0]
self.scoreBox.current())[0]
if mean_val > 0.0:
matrix[i,j] = matrix[i,j] / mean_val
else:
@ -275,20 +247,18 @@ class AppForm(QMainWindow):
for j in range(self.nz):
matrix[i,j] = self.datafile.get_value(tally_id,
spec_list + [('mesh', (i + 1, axial_level, j + 1))],
self.scoreBox.currentIndex())[score_loc]
# Calculate relative uncertainty from absolute, if
# requested
self.scoreBox.current())[score_loc]
# Calculate relative uncertainty from absolute, if requested
if is_mean == 2:
# Take care to handle zero means when normalizing
mean_val = self.datafile.get_value(tally_id,
spec_list + [('mesh', (i + 1, axial_level, j + 1))],
self.scoreBox.currentIndex())[0]
self.scoreBox.current())[0]
if mean_val > 0.0:
matrix[i,j] = matrix[i,j] / mean_val
else:
matrix[i,j] = 0.0
# print spec_list
# Clear the figure
self.fig.clear()
@ -296,7 +266,7 @@ class AppForm(QMainWindow):
# Make figure, set up color bar
self.axes = self.fig.add_subplot(111)
cax = self.axes.imshow(matrix.transpose(), vmin=0.0, vmax=matrix.max(),
interpolation="nearest", origin='lower')
interpolation='none', origin='lower')
self.fig.colorbar(cax)
self.axes.set_xticks([])
@ -305,100 +275,44 @@ class AppForm(QMainWindow):
# Draw canvas
self.canvas.draw()
def get_file_data(self, filename):
# Create StatePoint object and read in data
self.datafile = StatePoint(filename)
self.datafile.read_results()
self.datafile.generate_stdev()
def _update(self):
'''Updates widget to display new relevant comboboxes and figure data
'''
# print 'Calling _update...'
# Find which tallies are mesh tallies
self.meshTallies = []
for itally, tally in enumerate(self.datafile.tallies):
if 'mesh' in tally.filters:
self.meshTallies.append(itally)
# get current tally index
tally_id = self.tally_ids[self.tally.currentIndex()]
if not self.meshTallies:
tkMessageBox.showerror("Invalid StatePoint File",
"File does not contain mesh tallies!")
sys.exit(1)
self.mesh = self.datafile.meshes[
self.datafile.tallies[tally_id].filters['mesh'].bins[0] - 1]
if __name__ == '__main__':
# Hide root window
root = tk.Tk()
root.withdraw()
self.nx, self.ny, self.nz = self.mesh.dimension
# If no filename given as command-line argument, open file dialog
if len(sys.argv) < 2:
filename = tkFileDialog.askopenfilename(title='Select statepoint file',
initialdir='.')
else:
filename = sys.argv[1]
# Clear axial level combobox
self.axial_level.clear()
if filename:
# Check to make sure file exists
if not os.path.isfile(filename):
tkMessageBox.showerror("File not found",
"Could not find regular file: " + filename)
sys.exit(1)
# Repopulate axial level combobox based on current basis selection
if (self.basis.currentText() == 'xy'):
self.axial_level.addItems([str(i+1) for i in range(self.nz)])
elif (self.basis.currentText() == 'yz'):
self.axial_level.addItems([str(i+1) for i in range(self.nx)])
else:
self.axial_level.addItems([str(i+1) for i in range(self.ny)])
# Determine maximum value from current tally data set
self.maxvalue = self.datafile.tallies[tally_id].results.max()
# print self.maxvalue
# Clear and hide old filter labels
for item in self.labelList:
item.clear()
# Clear and hide old filter boxes
for j in self.labels:
self.boxes[j].clear()
self.boxes[j].setParent(None)
self.update()
def populate_boxes(self):
# print 'Calling populate_boxes...'
# get current tally index
tally_id = self.tally_ids[self.tally.currentIndex()]
n = 5
labels = {'cell': 'Cell : ',
'cellborn': 'Cell born: ',
'surface': 'Surface: ',
'material': 'Material: ',
'universe': 'Universe: '}
# For each filter in newly-selected tally, name a label and fill the
# relevant combobox with options
for element in self.tally_list[self.tally.currentIndex()]:
nextFilter = self.datafile.tallies[tally_id].filters[element]
if element == 'mesh':
continue
label = QLabel(self.labels[element])
self.labelList.append(label)
combobox = self.boxes[element]
self.grid.addWidget(label, n, 0)
self.grid.addWidget(combobox, n, 1)
n += 1
# print element
if element in ['cell', 'cellborn', 'surface', 'material', 'universe']:
combobox.addItems([str(i) for i in nextFilter.bins])
# for i in nextFilter.bins:
# print i
elif element == 'energyin' or element == 'energyout':
for i in range(nextFilter.length):
text = (str(nextFilter.bins[i]) + ' to ' +
str(nextFilter.bins[i+1]))
combobox.addItem(text)
self.scoreBox.clear()
for item in self.tally_scores[self.tally.currentIndex()]:
self.scoreBox.addItem(str(item))
self.grid.addWidget(self.score_label, n, 0)
self.grid.addWidget(self.scoreBox, n, 1)
def main():
app = QApplication(sys.argv)
form = AppForm(app.arguments())
if form.all_good:
form.show()
app.exec_()
if __name__ == "__main__":
main()
app = MeshPlotter(root, filename)
root.deiconify()
root.mainloop()