openmc-track-to-vtk

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