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
    for i in range(n):
        source.append(openmc.SourceParticle(
            r=(random(), i, 0),
            u=(0., 0., 1.),
            E=float(n - i),
        ))

    # Create source file
    openmc.write_source_file(source, 'test_source.h5')

    # 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))
    assert np.all(r['z'] == 0.0)
    u = arr['u']
    assert np.all(u['x'] == 0.0)
    assert np.all(u['y'] == 0.0)
    assert np.all(u['z'] == 1.0)
    assert np.all(arr['E'] == n - np.arange(n))
    assert np.all(arr['wgt'] == 1.0)
    assert np.all(arr['delayed_group'] == 0)
    assert np.all(arr['particle'] == 0)

@@ -18,16 +18,34 @@
 t random
+%0Aimport subprocess
 %0A%0Aimport
@@ -1270,16 +1270,1231 @@
 article'%5D == 0)%0A
+%0A%0Adef test_wrong_source_attributes(run_in_tmpdir):%0A    # Create a source file with animal attributes%0A    source_dtype = np.dtype(%5B%0A        ('platypus', '%3Cf8'),%0A        ('axolotl', '%3Cf8'),%0A        ('narwhal', '%3Ci4'),%0A    %5D)%0A    arr = np.array(%5B(1.0, 2.0, 3), (4.0, 5.0, 6), (7.0, 8.0, 9)%5D, dtype=source_dtype)%0A    with h5py.File('animal_source.h5', 'w') as fh:%0A        fh.attrs%5B'filetype'%5D = np.string_(%22source%22)%0A        fh.create_dataset('source_bank', data=arr)%0A%0A    # Create a simple model that uses this lovely animal source%0A    m = openmc.Material()%0A    m.add_nuclide('U235', 0.02)%0A    openmc.Materials(%5Bm%5D).export_to_xml()%0A    s = openmc.Sphere(r=10.0, boundary_type='vacuum')%0A    c = openmc.Cell(fill=m, region=-s)%0A    openmc.Geometry(%5Bc%5D).export_to_xml()%0A    settings =  openmc.Settings()%0A    settings.particles = 100%0A    settings.batches = 10%0A    settings.source = openmc.Source(filename='animal_source.h5')%0A    settings.export_to_xml()%0A%0A    # When we run the model, it should error out with a message that includes%0A    # the names of the wrong attributes%0A    with pytest.raises(subprocess.CalledProcessError) as excinfo:%0A        openmc.run()%0A    assert 'platypus, axolotl, narwhal' in excinfo.value.output%0A

Add test for wrong source attributes