Merge pull request #1691 from paulromano/sourcesite-python

Generation of source files from Python

docs/source/io_formats/source.rst

@@ -15,7 +15,7 @@ is that documented here.
 :Datasets:
 
            - **source_bank** (Compound type) -- Source bank information for each
-             particle. The compound type has fields ``wgt``, ``xyz``, ``uvw``,
-             ``E``, ``delayed_group``, and ``particle``, which represent the
-             weight, position, direction, energy, energy group, delayed group,
-             and type of the source particle, respectively.
+             particle. The compound type has fields ``r``, ``u``, ``E``,
+             ``wgt``, ``delayed_group``, and ``particle``, which represent the
+             position, direction, energy, weight, delayed group, and particle
+             type (0=neutron, 1=photon, 2=electron, 3=positron), respectively.

docs/source/io_formats/statepoint.rst

@@ -51,11 +51,11 @@ The current version of the statepoint file format is 17.0.
            - **global_tallies** (*double[][2]*) -- Accumulated sum and
              sum-of-squares for each global tally.
            - **source_bank** (Compound type) -- Source bank information for each
-             particle. The compound type has fields ``wgt``, ``xyz``, ``uvw``,
-             ``E``, ``g``, and ``delayed_group``, which represent the weight,
-             position, direction, energy, energy group, and delayed_group of the
-             source particle, respectively. Only present when `run_mode` is
-             'eigenvalue'.
+             particle. The compound type has fields ``r``, ``u``, ``E``,
+             ``wgt``, ``delayed_group``, and ``particle``, which represent the
+             position, direction, energy, weight, delayed group, and particle
+             type (0=neutron, 1=photon, 2=electron, 3=positron), respectively.
+             Only present when `run_mode` is 'eigenvalue'.
 
 **/tallies/**
 

docs/source/pythonapi/base.rst

@@ -22,9 +22,19 @@ Simulation Settings
    :template: myclass.rst
 
    openmc.Source
+   openmc.SourceParticle
    openmc.VolumeCalculation
    openmc.Settings
 
+The following function can be used for generating a source file:
+
+.. autosummary::
+   :toctree: generated
+   :nosignatures:
+   :template: myfunction.rst
+
+   openmc.write_source_file
+
 Material Specification
 ----------------------
 

openmc/source.py

@@ -1,6 +1,10 @@
+from enum import Enum
 from numbers import Real
 from xml.etree import ElementTree as ET
 
+import numpy as np
+import h5py
+
 import openmc.checkvalue as cv
 from openmc.stats.multivariate import UnitSphere, Spatial
 from openmc.stats.univariate import Univariate
@@ -226,3 +230,89 @@ class Source:
             source.energy = Univariate.from_xml_element(energy)
 
         return source
+
+
+class ParticleType(Enum):
+    NEUTRON = 0
+    PHOTON = 1
+    ELECTRON = 2
+    POSITRON = 3
+
+
+class SourceParticle:
+    """Source particle
+
+    This class can be used to create source particles that can be written to a
+    file and used by OpenMC
+
+    Parameters
+    ----------
+    r : iterable of float
+        Position of particle in Cartesian coordinates
+    u : iterable of float
+        Directional cosines
+    E : float
+        Energy of particle in [eV]
+    wgt : float
+        Weight of the particle
+    delayed_group : int
+        Delayed group particle was created in (neutrons only)
+    particle : ParticleType
+        Type of the particle
+
+    """
+    def __init__(self, r=(0., 0., 0.), u=(0., 0., 1.), E=1.0e6, wgt=1.0,
+                 delayed_group=0, particle=ParticleType.NEUTRON):
+        self.r = tuple(r)
+        self.u = tuple(u)
+        self.E = float(E)
+        self.wgt = float(wgt)
+        self.delayed_group = delayed_group
+        self.particle = particle
+
+    def to_tuple(self):
+        """Return source particle attributes as a tuple
+
+        Returns
+        -------
+        tuple
+            Source particle attributes
+
+        """
+        return (self.r, self.u, self.E, self.wgt,
+                self.delayed_group, self.particle.value)
+
+
+def write_source_file(source_particles, filename, **kwargs):
+    """Write a source file using a collection of source particles
+
+    Parameters
+    ----------
+    source_particles : iterable of SourceParticle
+        Source particles to write to file
+    filename : str or path-like
+        Path to source file to write
+    **kwargs
+        Keyword arguments to pass to :class:`h5py.File`
+
+    """
+    # Create compound datatype for source particles
+    pos_dtype = np.dtype([('x', '<f8'), ('y', '<f8'), ('z', '<f8')])
+    source_dtype = np.dtype([
+        ('r', pos_dtype),
+        ('u', pos_dtype),
+        ('E', '<f8'),
+        ('wgt', '<f8'),
+        ('delayed_group', '<i4'),
+        ('particle', '<i4'),
+    ])
+
+    # 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)

openmc/statepoint.py

@@ -93,9 +93,10 @@ class StatePoint:
     seed : int
         Pseudorandom number generator seed
     source : numpy.ndarray of compound datatype
-        Array of source sites. The compound datatype has fields 'wgt', 'xyz',
-        'uvw', and 'E' corresponding to the weight, position, direction, and
-        energy of the source site.
+        Array of source sites. The compound datatype has fields 'r', 'u',
+        'E', 'wgt', 'delayed_group', and 'particle', corresponding to the
+        position, direction, energy, weight, delayed group, and particle type
+        of the source site, respectively.
     source_present : bool
         Indicate whether source sites are present
     sparse : bool

src/state_point.cpp

@@ -506,6 +506,12 @@ hid_t h5banktype() {
   H5Tinsert(postype, "z", HOFFSET(Position, z), H5T_NATIVE_DOUBLE);
 
   // Create bank datatype
+  //
+  // If you make changes to the compound datatype here, make sure you update:
+  // - openmc/source.py
+  // - openmc/statepoint.py
+  // - docs/source/io_formats/statepoint.rst
+  // - docs/source/io_formats/source.rst
   hid_t banktype = H5Tcreate(H5T_COMPOUND, sizeof(struct Particle::Bank));
   H5Tinsert(banktype, "r", HOFFSET(Particle::Bank, r), postype);
   H5Tinsert(banktype, "u", HOFFSET(Particle::Bank, u), postype);

tests/unit_tests/test_source_file.py

@@ -0,0 +1,45 @@
+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)