diff --git a/openmc/source.py b/openmc/source.py index 3863bb826a..b5b53167c5 100644 --- a/openmc/source.py +++ b/openmc/source.py @@ -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) diff --git a/tests/unit_tests/test_source_file.py b/tests/unit_tests/test_source_file.py index e9bcebc92e..ae6eff97f7 100644 --- a/tests/unit_tests/test_source_file.py +++ b/tests/unit_tests/test_source_file.py @@ -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))