diff --git a/docs/source/pythonapi/base.rst b/docs/source/pythonapi/base.rst index d621637bb..544e7009c 100644 --- a/docs/source/pythonapi/base.rst +++ b/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 ---------------------- diff --git a/openmc/source.py b/openmc/source.py index 22952b4a5..3863bb826 100644 --- a/openmc/source.py +++ b/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,87 @@ 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', '