-
-
Save pshriwise/da30da3daf08594dddf1a58b9f10dcc8 to your computer and use it in GitHub Desktop.
#!/usr/env/python3 | |
""" | |
Script for converting OpenMC mesh in an OpenMC | |
statepoint file to other mesh formats for visualization | |
""" | |
import argparse | |
import sys | |
import math | |
import openmc | |
import numpy as np | |
def write_moab(xs, ys, zs, tally_label, tally_data, error_data, outfile): | |
# attempt to import pymoab | |
try: | |
from pymoab import core | |
from pymoab.hcoord import HomCoord | |
from pymoab.scd import ScdInterface | |
from pymoab import types | |
except (ImportError, ModuleNotFoundError): | |
msg = "Conversion to MOAB .h5m file requested," \ | |
"but PyMOAB is not installed" | |
raise ImportError(msg) | |
mb = core.Core() | |
scd = ScdInterface(mb) | |
coords = [] | |
for k in zs: | |
for j in ys: | |
for i in xs: | |
coords += [i,j,k] | |
low = HomCoord([0, 0, 0, 0]) | |
high = HomCoord([len(xs) - 1, len(ys) - 1, len(zs) - 1, 0]) | |
scdbox = scd.construct_box(low, high, coords) | |
hexes = mb.get_entities_by_type(0, types.MBHEX) | |
tally_tag = mb.tag_get_handle(tally_label,1,types.MB_TYPE_DOUBLE,types.MB_TAG_DENSE,True) | |
error_tag = mb.tag_get_handle("error_tag",1,types.MB_TYPE_DOUBLE,types.MB_TAG_DENSE,True) | |
mb.tag_set_data(tally_tag,hexes,tally_data) | |
mb.tag_set_data(error_tag,hexes,error_data) | |
print('Writing %s' % outfile) | |
mb.write_file(outfile) | |
def write_vtk(xs, ys, zs, tally_label, tally_data, error_data, outfile): | |
try: | |
import vtk | |
except (ImportError, ModuleNotFoundError) as e: | |
msg = "Conversion to VTK requested," \ | |
"but the Python VTK module is not installed." | |
raise ImportError(msg) | |
vtk_box = vtk.vtkRectilinearGrid() | |
vtk_box.SetDimensions(len(xs), len(ys), len(zs)) | |
vtk_x_array = vtk.vtkDoubleArray() | |
vtk_x_array.SetName('x-coords') | |
vtk_x_array.SetArray(xs, len(xs), True) | |
print(vtk_x_array) | |
vtk_box.SetXCoordinates(vtk_x_array) | |
vtk_y_array = vtk.vtkDoubleArray() | |
vtk_y_array.SetName('y-coords') | |
vtk_y_array.SetArray(ys, len(ys), True) | |
vtk_box.SetYCoordinates(vtk_y_array) | |
vtk_z_array = vtk.vtkDoubleArray() | |
vtk_z_array.SetName('z-coords') | |
vtk_z_array.SetArray(zs, len(zs), True) | |
vtk_box.SetZCoordinates(vtk_z_array) | |
tally = np.array(tally_data) | |
tally_data = vtk.vtkDoubleArray() | |
tally_data.SetName(tally_label) | |
tally_data.SetArray(tally, tally.size, True) | |
error = np.array(error_data) | |
error_data = vtk.vtkDoubleArray() | |
error_data.SetName("error_tag") | |
error_data.SetArray(error, error.size, True) | |
vtk_box.GetCellData().AddArray(tally_data) | |
vtk_box.GetCellData().AddArray(error_data) | |
writer = vtk.vtkRectilinearGridWriter() | |
writer.SetFileName(outfile) | |
writer.SetInputData(vtk_box) | |
print('Writing %s' % outfile) | |
writer.Write() | |
def main(): | |
ap = argparse.ArgumentParser(description=__doc__) | |
ap.add_argument('-i','--input', | |
required=True, | |
help='Path to statepoint h5 file') | |
ap.add_argument('-n','--tally-name', | |
help='Tally name to add to mesh') | |
ap.add_argument('-t','--tally-id', | |
type=int, | |
required=True, | |
help='Tally name to add to mesh') | |
ap.add_argument('-m','--mesh-id', | |
type=int, | |
required=True, | |
help='ID of the mesh to convert') | |
ap.add_argument('-o', '--output', | |
action='store', | |
default='meshtally.vtk', | |
help='Name of outputfile (.h5m for MOAB, .vtk for VTK)') | |
args = ap.parse_args() | |
print('Loading file %s' % args.input) | |
sp = openmc.StatePoint(args.input) | |
print('Loading mesh with ID of %s' % args.mesh_id) | |
mesh = sp.meshes[args.mesh_id] | |
xs = np.linspace(mesh.lower_left[0], | |
mesh.upper_right[0], | |
mesh.dimension[0] + 1) | |
ys = np.linspace(mesh.lower_left[1], | |
mesh.upper_right[1], | |
mesh.dimension[1] + 1) | |
zs = np.linspace(mesh.lower_left[2], | |
mesh.upper_right[2], | |
mesh.dimension[2] + 1) | |
tally_args = {'name' : args.tally_name, | |
'id' : args.tally_id} | |
msg = "Loading retrieving tally with \n" | |
if args.tally_name is not None: | |
msg += "Name: {}\n".format(args.tally_name) | |
if args.tally_id is not None: | |
msg += "ID: {}\n".format(args.tally_id) | |
try: | |
tally = sp.get_tally(**tally_args) | |
except LookupError as e: | |
raise e | |
data = tally.mean[:,0,0] | |
error = tally.std_dev[:,0,0] | |
data = data.tolist() | |
error = error.tolist() | |
for c, i in enumerate(data): | |
if math.isnan(i): | |
data[c] = 0. | |
for c, i in enumerate(error): | |
if math.isnan(i): | |
error[c] = 0. | |
if args.tally_name is None: | |
tally_label = "tally_{}".format(args.tally_id) | |
else: | |
tally_label = args.tally_name | |
if args.output.endswith(".vtk"): | |
write_vtk(xs, ys, zs, tally_label, data, error, args.output) | |
else: | |
write_moab(xs, ys, zs, tally_label, data, error, args.output) | |
if __name__ == "__main__": | |
main() |
pshriwise
commented
Sep 18, 2019
via email
•
Sometimes tally_label ends up being undefined as in the case where args.tally_name is != None then tally_label is undefined but called upon in lines 180 and 182
Indeed. I’ll patch that up. Thanks for catching it!
I have been stress testing this for large meshes, in the case where the mesh tally is large then making both the tally mesh and the error mesh pushes the computer over its RAM limits and I get a crash. I was wondering if we can have a flag that allows the script include either the tally or the tally and the error or just the error. This will allow us to make the largest VTK files within the RAM limits
@shimwell, were you seeing an issue with memory limits when writing a VTK or MOAB file? Maybe both?
I was using the VTK option, I removed the error tally and writing of the tally and this resulted in lower memory usage.
So I was thinking there could be some user option to allow more control of writing tally , error or both (default)
Another argument for the ap, perhaps something like this
ap.add_argument('-e', '--error',
choices=['error', 'tally'],
nargs='+',
default=['error', 'tally'],
help="Error and tally data included by default. Use this option to write just the error or just the tally data")
But this is not so important and can always be added later