mirror of
https://github.com/openmc-dev/openmc.git
synced 2026-07-24 03:55:38 -04:00
78 lines
2.2 KiB
Python
Executable file
78 lines
2.2 KiB
Python
Executable file
#!/usr/bin/env python3
|
|
|
|
"""Convert HDF5 particle track to VTK poly data.
|
|
|
|
"""
|
|
|
|
import os
|
|
import argparse
|
|
import struct
|
|
|
|
import h5py
|
|
import vtk
|
|
|
|
|
|
def _parse_args():
|
|
# Create argument parser.
|
|
parser = argparse.ArgumentParser(
|
|
description='Convert particle track file to a .pvtp file.')
|
|
parser.add_argument('input', metavar='IN', type=str, nargs='+',
|
|
help='Input particle track data filename(s).')
|
|
parser.add_argument('-o', '--out', metavar='OUT', type=str, dest='out',
|
|
help='Output VTK poly data filename.')
|
|
|
|
# Parse and return commandline arguments.
|
|
return parser.parse_args()
|
|
|
|
|
|
def main():
|
|
# Parse commandline arguments.
|
|
args = _parse_args()
|
|
|
|
# Make sure that the output filename ends with '.pvtp'.
|
|
if not args.out:
|
|
args.out = 'tracks.pvtp'
|
|
elif not args.out.endswith('.pvtp'):
|
|
args.out += '.pvtp'
|
|
|
|
# Initialize data arrays and offset.
|
|
points = vtk.vtkPoints()
|
|
cells = vtk.vtkCellArray()
|
|
point_offset = 0
|
|
for fname in args.input:
|
|
# Write coordinate values to points array.
|
|
track = h5py.File(fname)
|
|
n_particles = track.attrs['n_particles']
|
|
n_coords = track.attrs['n_coords']
|
|
coords = []
|
|
for i in range(n_particles):
|
|
coords.append(track['coordinates_' + str(i + 1)].value)
|
|
for j in range(n_coords[i]):
|
|
points.InsertNextPoint(coords[i][j,:])
|
|
|
|
for i in range(n_particles):
|
|
# Create VTK line and assign points to line.
|
|
line = vtk.vtkPolyLine()
|
|
line.GetPointIds().SetNumberOfIds(n_coords[i])
|
|
for j in range(n_coords[i]):
|
|
line.GetPointIds().SetId(j, point_offset + j)
|
|
|
|
# Add line to cell array
|
|
cells.InsertNextCell(line)
|
|
point_offset += n_coords[i]
|
|
|
|
data = vtk.vtkPolyData()
|
|
data.SetPoints(points)
|
|
data.SetLines(cells)
|
|
|
|
writer = vtk.vtkXMLPPolyDataWriter()
|
|
if vtk.vtkVersion.GetVTKMajorVersion() > 5:
|
|
writer.SetInputData(data)
|
|
else:
|
|
writer.SetInput(data)
|
|
writer.SetFileName(args.out)
|
|
writer.Write()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|