Convert OpenMC Meshtally to MOAB/VTK mesh

openmc-statepoint-3d.py

#!/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_name, 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)

    flux_tag = mb.tag_get_handle(tally_name,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(flux_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_name, 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)

    flux = np.array(list(tally_data))
    flux_data = vtk.vtkDoubleArray()
    flux_data.SetName(tally_name)
    flux_data.SetArray(flux, flux.size, True)

    error = np.array(list(error_data))
    error_data = vtk.vtkDoubleArray()
    error_data.SetName("error_tag")
    error_data.SetArray(error, error.size, True)

    vtk_box.GetCellData().AddArray(flux_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('-t','--tally_name', 
                    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)

    print('Loading tally with name of %s' % args.tally_name)
    tally = sp.get_tally(name=args.tally_name)

    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.output.endswith(".vtk"):
        write_vtk(xs, ys, zs, args.tally_name, data, error, args.output)
    else:
        write_moab(xs, ys, zs, args.tally_name, data, error, args.output)

if __name__ == "__main__":
    main()

Currently getting errors with mesh tallies that don't score in the mesh element

Loading file statepoint.10.h5
Loading mesh with ID of 1
Loading tally with score of flux
/home/shimwell/openmc/openmc/tallies.py:289: RuntimeWarning: invalid value encountered in greater
  nonzero = np.abs(self.mean) > 0
/home/shimwell/openmc/openmc/tallies.py:292: RuntimeWarning: invalid value encountered in sqrt
  self.mean[nonzero]**2)/(n - 1))
Writing meshtally.vtk