Put track file combining logic in TrackFile.combine staticmethod

This commit is contained in:
Paul Romano 2022-05-19 17:10:27 -05:00
parent 8ecfdf0796
commit 583d864f08
2 changed files with 27 additions and 16 deletions

View file

@ -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)

View file

@ -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__':