OpenMC/scripts/openmc-voxel-to-vtk
2019-10-28 11:55:45 -05:00

72 lines
2.1 KiB
Python
Executable file

#!/usr/bin/env python3
import struct
import sys
from argparse import ArgumentParser
import numpy as np
import h5py
import vtk
_min_version = (2,0)
def main():
# Process command line arguments
parser = ArgumentParser()
parser.add_argument('voxel_file', help='Path to voxel file')
parser.add_argument('-o', '--output', action='store',
default='plot', help='Path to output VTK file.')
args = parser.parse_args()
# Read data from voxel file
fh = h5py.File(args.voxel_file, 'r')
# check version
version = tuple(fh.attrs['version'])
if version < _min_version:
old_version = ".".join(map(str,version))
min_version = ".".join(map(str,_min_version))
err_msg = "This voxel file's version is {}. This script " \
"only supports voxel files with version {} or " \
"higher. Please generate a new voxel file using " \
"a newer version of OpenMC.".format(old_version, min_version)
raise ValueError(err_msg)
dimension = fh.attrs['num_voxels']
width = fh.attrs['voxel_width']
lower_left = fh.attrs['lower_left']
nx, ny, nz = dimension
upper_right = lower_left + width*dimension
grid = vtk.vtkImageData()
grid.SetDimensions(nx+1, ny+1, nz+1)
grid.SetOrigin(*lower_left)
grid.SetSpacing(*width)
# transpose data from OpenMC ordering (zyx) to VTK ordering (xyz)
# and flatten to 1-D array
print("Reading and translating data...")
h5data = fh['data'][...]
data = vtk.vtkIntArray()
data.SetName("id")
# set the array using the h5data array
data.SetArray(h5data, h5data.size, True)
# add data to image grid
grid.GetCellData().AddArray(data)
writer = vtk.vtkXMLImageDataWriter()
if vtk.vtkVersion.GetVTKMajorVersion() > 5:
writer.SetInputData(grid)
else:
writer.SetInput(grid)
if not args.output.endswith(".vti"):
args.output += ".vti"
writer.SetFileName(args.output)
print("Writing VTK file {}...".format(args.output))
writer.Write()
if __name__ == '__main__':
main()