OpenMC/scripts/openmc-plot-mesh-tally

344 lines
12 KiB
Python
Executable file

#!/usr/bin/env python3
"""Python script to plot tally data generated by OpenMC."""
import os
import sys
import argparse
import tkinter as tk
import tkinter.filedialog as filedialog
import tkinter.font as font
import tkinter.messagebox as messagebox
import tkinter.ttk as ttk
import matplotlib
matplotlib.use("TkAgg")
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
from matplotlib.backends.backend_tkagg import NavigationToolbar2Tk
from matplotlib.figure import Figure
import matplotlib.pyplot as plt
import numpy as np
from openmc import StatePoint, MeshFilter, UnstructuredMesh
_COMBOBOX_SELECTED = '<<ComboboxSelected>>'
def mesh_filter_check(filter):
"""
Check that the filter is a usable mesh filter
"""
return isinstance(filter, MeshFilter) and not isinstance(filter.mesh, UnstructuredMesh)
class MeshPlotter(tk.Frame):
def __init__(self, parent, filename):
super().__init__(parent)
self.labels = {
'Cell': 'Cell:',
'Cellborn': 'Cell born:',
'Surface': 'Surface:',
'Material': 'Material:',
'Universe': 'Universe:',
'Energy': 'Energy in:',
'Energyout': 'Energy out:'
}
self.filterBoxes = {}
# Read data from source or leakage fraction file
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)
# Create widgets and draw to screen
self.create_widgets()
self.update()
def create_widgets(self):
figureFrame = tk.Frame(self)
figureFrame.grid(row=0, column=0)
# Create the Figure and Canvas
self.dpi = 100
self.fig = Figure((5.0, 5.0), dpi=self.dpi)
self.canvas = FigureCanvasTkAgg(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 = NavigationToolbar2Tk(self.canvas, figureFrame)
self.mpl_toolbar.update()
self.canvas._tkcanvas.pack(side=tk.TOP, fill=tk.BOTH, expand=1)
# Create frame for comboboxes
self.selectFrame = tk.Frame(self)
self.selectFrame.grid(row=1, column=0, sticky=tk.W+tk.E)
# 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(_COMBOBOX_SELECTED, self.update)
# 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(_COMBOBOX_SELECTED, self.update)
# 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(_COMBOBOX_SELECTED, self.redraw)
# 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(_COMBOBOX_SELECTED, self.update)
# 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(_COMBOBOX_SELECTED, self.redraw)
# Filter label
boldfont = font.Font(weight='bold')
labelFilters = tk.Label(self.selectFrame, text='Filters:',
font=boldfont)
labelFilters.grid(row=5, column=0, sticky=tk.W)
def update(self, event=None):
widget = event.widget if event else None
tally_id = self.meshTallies[self.tallyBox.current()]
selectedTally = self.datafile.tallies[tally_id]
# Get mesh for selected tally
self.mesh = selectedTally.find_filter(MeshFilter).mesh
# Get mesh dimensions
if len(self.mesh.dimension) == 2:
self.nx, self.ny = self.mesh.dimension
self.nz = 1
else:
self.nx, self.ny, self.nz = self.mesh.dimension
# Repopulate comboboxes baesd on current basis selection
text = self.basisBox.get()
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)
# 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
# Update scores
self.scoreBox['values'] = selectedTally.scores
self.scoreBox.current(0)
# 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()
# create a label/combobox for each filter in selected tally
count = 0
for f in selectedTally.filters:
filterType = f.short_name
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
if filterType in ['Energy', 'Energyout']:
combobox['values'] = ['{} to {}'.format(*ebin)
for ebin in f.bins]
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(_COMBOBOX_SELECTED, self.redraw)
# If There are no filters, leave a 'None available' message
if count == 0:
count += 1
label = tk.Label(self.selectFrame, text="None Available")
label.grid(row=count+6, column=0, sticky=tk.W)
self.redraw()
def redraw(self, event=None):
basis = self.basisBox.current() + 1
axial_level = self.axialBox.current() + 1
mbvalue = self.meanBox.get()
# Get selected tally
tally_id = self.meshTallies[self.tallyBox.current()]
selectedTally = self.datafile.tallies[tally_id]
# Create spec_list
spec_list = []
for f in selectedTally.filters:
if f.short_name == 'Mesh':
mesh_filter = f
continue
elif f.short_name in ['Energy', 'Energyout']:
index = self.filterBoxes[f.short_name].current()
ebin = f.bins[index]
spec_list.append((type(f), (ebin,)))
else:
index = self.filterBoxes[f.short_name].current()
spec_list.append((type(f), (index,)))
dims = (self.nx, self.ny, self.nz)
text = self.basisBox.get()
if text == 'xy':
h_ind = 0
v_ind = 1
elif text == 'yz':
h_ind = 1
v_ind = 2
else:
h_ind = 0
v_ind = 2
axial_ind = 3 - (h_ind + v_ind)
dims = (dims[h_ind], dims[v_ind])
mesh_dim = len(self.mesh.dimension)
if mesh_dim == 3:
mesh_indices = [0,0,0]
else:
mesh_indices = [0,0]
matrix = np.zeros(dims)
for i in range(dims[0]):
for j in range(dims[1]):
if mesh_dim == 3:
mesh_indices[h_ind] = i + 1
mesh_indices[v_ind] = j + 1
mesh_indices[axial_ind] = axial_level
else:
mesh_indices[0] = i + 1
mesh_indices[1] = j + 1
filters, filter_bins = zip(*spec_list + [
(type(mesh_filter), (tuple(mesh_indices),))])
mean = selectedTally.get_values(
[self.scoreBox.get()], filters, filter_bins)
stdev = selectedTally.get_values(
[self.scoreBox.get()], filters, filter_bins,
value='std_dev')
if mbvalue == 'Mean':
matrix[i, j] = mean
elif mbvalue == 'Absolute uncertainty':
matrix[i, j] = stdev
else:
if mean > 0.:
matrix[i, j] = stdev/mean
else:
matrix[i, j] = 0.
# Clear the figure
self.fig.clear()
# 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='none', origin='lower')
self.fig.colorbar(cax)
self.axes.set_xticks([])
self.axes.set_yticks([])
self.axes.set_aspect('equal')
# Draw canvas
self.canvas.draw()
def get_file_data(self, filename):
# Create StatePoint object and read in data
self.datafile = StatePoint(filename)
meshes = self.datafile.meshes
if any(isinstance(m, UnstructuredMesh) for m in meshes.values()):
warn_msg = "Unstructured meshes are present in the" \
" statepoint file but are not currently" \
" supported by this script"
messagebox.showwarning("Unstructured Meshes", message=warn_msg)
# Find which tallies are mesh tallies
self.meshTallies = []
for itally, tally in self.datafile.tallies.items():
if any([mesh_filter_check(f) for f in tally.filters]):
self.meshTallies.append(itally)
if not self.meshTallies:
messagebox.showerror("Invalid StatePoint File",
"File does not contain mesh tallies!")
sys.exit(1)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('statepoint', nargs='?', help='Statepoint file')
args = parser.parse_args()
# Hide root window
root = tk.Tk()
root.withdraw()
# If no filename given as command-line argument, open file dialog
if args.statepoint is None:
filename = filedialog.askopenfilename(title='Select statepoint file',
initialdir='.')
else:
filename = args.statepoint
if filename:
# Check to make sure file exists
if not os.path.isfile(filename):
messagebox.showerror("File not found",
"Could not find regular file: " + filename)
sys.exit(1)
app = MeshPlotter(root, filename)
root.deiconify()
root.mainloop()