added main method to voxel script

This commit is contained in:
Jonathan Shimwell 2022-08-24 10:50:40 +01:00
parent 2805060ddd
commit fb8b4e0690

View file

@ -1,54 +1,46 @@
#!/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()
def voxel_to_vtk(voxel_file: str, output: str):
# Read data from voxel file
fh = h5py.File(args.voxel_file, 'r')
fh = h5py.File(voxel_file, "r")
# check version
version = tuple(fh.attrs['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)
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']
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.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'][...]
h5data = fh["data"][...]
data = vtk.vtkIntArray()
data.SetName("id")
@ -62,11 +54,23 @@ def main():
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))
if not output.endswith(".vti"):
output += ".vti"
writer.SetFileName(output)
print("Writing VTK file {}...".format(output))
writer.Write()
if __name__ == '__main__':
main()
if __name__ == "__main__":
# Process command line arguments
parser = ArgumentParser('Converts a voxel HDF5 file to a VTK file')
parser.add_argument("voxel_file", help="Path to voxel h5 file")
parser.add_argument(
"-o",
"--output",
action="store",
default="plot",
help="Path to output VTK file.",
)
args = parser.parse_args()
voxel_to_vtk(args.voxel_file, args.output)