Make sure write_source_file creates filetype attribute

This commit is contained in:
Paul Romano 2020-10-20 14:07:45 -05:00
parent 2d8c9aed81
commit 490eadd864
2 changed files with 9 additions and 0 deletions

View file

@ -308,9 +308,11 @@ def write_source_file(source_particles, filename, **kwargs):
])
# Create array of source particles
cv.check_iterable_type("source particles", source_particles, SourceParticle)
arr = np.array([s.to_tuple() for s in source_particles], dtype=source_dtype)
# Write array to file
kwargs.setdefault('mode', 'w')
with h5py.File(filename, **kwargs) as fh:
fh.attrs['filetype'] = np.string_("source")
fh.create_dataset('source_bank', data=arr, dtype=source_dtype)

View file

@ -3,9 +3,14 @@ from random import random
import h5py
import numpy as np
import openmc
import pytest
def test_source_file(run_in_tmpdir):
# write_source_file shouldn't accept non-SourceParticle items
with pytest.raises(TypeError):
openmc.write_source_file([1, 2, 3], 'test_source.h5')
# Create source particles
source = []
n = 1000
@ -21,9 +26,11 @@ def test_source_file(run_in_tmpdir):
# Get array of source particles from file
with h5py.File('test_source.h5', 'r') as fh:
filetype = fh.attrs['filetype']
arr = fh['source_bank'][...]
# Ensure data is consistent
assert filetype == b'source'
r = arr['r']
assert np.all((r['x'] > 0.0) & (r['x'] < 1.0))
assert np.all(r['y'] == np.arange(1000))