diff --git a/openmc/trackfile.py b/openmc/trackfile.py index 8c2d02cd1b..64ed5b7076 100644 --- a/openmc/trackfile.py +++ b/openmc/trackfile.py @@ -130,3 +130,27 @@ class TrackFile(list): for track in self: track.plot(ax) plt.show() + + @staticmethod + def combine(track_files, path='tracks.h5'): + """Combine multiple track files into a single track file + + Parameters + ---------- + track_files : list of path-like + Paths to track files to combine + path : path-like + Path of combined track file to create + + """ + with h5py.File(path, 'w') as h5_out: + for i, fname in enumerate(track_files): + with h5py.File(fname, 'r') as h5_in: + # Copy file attributes for first file + if i == 0: + h5_out.attrs['filetype'] = h5_in.attrs['filetype'] + h5_out.attrs['version'] = h5_in.attrs['version'] + + # Copy each 'track_*' dataset from input file + for dset in h5_in: + h5_in.copy(dset, h5_out) diff --git a/scripts/openmc-track-combine b/scripts/openmc-track-combine index 4bd2fb0820..155fe9b541 100755 --- a/scripts/openmc-track-combine +++ b/scripts/openmc-track-combine @@ -1,12 +1,10 @@ #!/usr/bin/env python3 -"""Combine multiple HDF5 particle track files. - -""" +"""Combine multiple HDF5 particle track files.""" import argparse -import h5py +import openmc def main(): @@ -19,18 +17,7 @@ def main(): help='Output HDF5 particle track file.') args = parser.parse_args() - - with h5py.File(args.out, 'w') as h5_out: - for i, fname in enumerate(args.input): - with h5py.File(fname, 'r') as h5_in: - # Copy file attributes for first file - if i == 0: - h5_out.attrs['filetype'] = h5_in.attrs['filetype'] - h5_out.attrs['version'] = h5_in.attrs['version'] - - # Copy each 'track_*' dataset from input file - for dset in h5_in: - h5_in.copy(dset, h5_out) + openmc.TrackFile.combine(args.input, args.out) if __name__ == '__main__':