Merge pull request #2119 from paulromano/sample-source-capi

Add C API function for sampling external source and Python binding

docs/source/pythonapi/capi.rst

@@ -29,6 +29,7 @@ Functions
    reset
    run
    run_in_memory
+   sample_external_source
    simulation_init
    simulation_finalize
    source_bank

include/openmc/capi.h

@@ -121,6 +121,7 @@ int openmc_regular_mesh_set_params(int32_t index, int n, const double* ll,
 int openmc_reset();
 int openmc_reset_timers();
 int openmc_run();
+int openmc_sample_external_source(size_t n, uint64_t* seed, void* sites);
 void openmc_set_seed(int64_t new_seed);
 int openmc_set_n_batches(
   int32_t n_batches, bool set_max_batches, bool add_statepoint_batch);

openmc/lib/core.py

@@ -1,8 +1,10 @@
 from contextlib import contextmanager
 from ctypes import (c_bool, c_int, c_int32, c_int64, c_double, c_char_p,
-                    c_char, POINTER, Structure, c_void_p, create_string_buffer)
+                    c_char, POINTER, Structure, c_void_p, create_string_buffer,
+                    c_uint64, c_size_t)
 import sys
 import os
+from random import getrandbits
 
 import numpy as np
 from numpy.ctypeslib import as_array
@@ -10,6 +12,7 @@ from numpy.ctypeslib import as_array
 from . import _dll
 from .error import _error_handler
 import openmc.lib
+import openmc
 
 
 class _SourceSite(Structure):
@@ -95,6 +98,9 @@ _dll.openmc_global_bounding_box.argtypes = [POINTER(c_double),
                                             POINTER(c_double)]
 _dll.openmc_global_bounding_box.restype = c_int
 _dll.openmc_global_bounding_box.errcheck = _error_handler
+_dll.openmc_sample_external_source.argtypes = [c_size_t, POINTER(c_uint64), POINTER(_SourceSite)]
+_dll.openmc_sample_external_source.restype = c_int
+_dll.openmc_sample_external_source.errcheck = _error_handler
 
 
 def global_bounding_box():
@@ -415,6 +421,45 @@ def run(output=True):
         _dll.openmc_run()
 
 
+def sample_external_source(n_samples=1, prn_seed=None):
+    """Sample external source
+
+    .. versionadded:: 0.13.1
+
+    Parameters
+    ----------
+    n_samples : int
+        Number of samples
+    prn_seed : int
+        Pseudorandom number generator (PRNG) seed; if None, one will be
+        generated randomly.
+
+    Returns
+    -------
+    list of openmc.SourceParticle
+        List of samples source particles
+
+    """
+    if n_samples <= 0:
+        raise ValueError("Number of samples must be positive")
+    if prn_seed is None:
+        prn_seed = getrandbits(63)
+
+    # Call into C API to sample source
+    sites_array = (_SourceSite * n_samples)()
+    _dll.openmc_sample_external_source(c_size_t(n_samples), c_uint64(prn_seed), sites_array)
+
+    # Convert to list of SourceParticle and return
+    return [
+        openmc.SourceParticle(
+            r=site.r, u=site.u, E=site.E, time=site.time, wgt=site.wgt,
+            delayed_group=site.delayed_group, surf_id=site.surf_id,
+            particle=openmc.ParticleType(site.particle)
+        )
+        for site in sites_array
+    ]
+
+
 def simulation_init():
     """Initialize simulation"""
     _dll.openmc_simulation_init()

openmc/lib/math.py

@@ -1,12 +1,11 @@
 from ctypes import c_int, c_double, POINTER, c_uint64
+from random import getrandbits
 
 import numpy as np
 from numpy.ctypeslib import ndpointer
 
 from . import _dll
 
-from random import getrandbits
-
 
 _dll.t_percentile.restype = c_double
 _dll.t_percentile.argtypes = [c_double, c_int]
@@ -200,7 +199,8 @@ def rotate_angle(uvw0, mu, phi, prn_seed=None):
     phi : float
         Azimuthal angle; if None, one will be sampled uniformly
     prn_seed : int
-        PRNG seed; if None, one will be generated randomly
+        Pseudorandom number generator (PRNG) seed; if None, one will be
+        generated randomly.
 
     Returns
     -------
@@ -232,7 +232,8 @@ def maxwell_spectrum(T, prn_seed=None):
     T : float
         Spectrum parameter
     prn_seed : int
-        PRNG seed; if None, one will be generated randomly
+        Pseudorandom number generator (PRNG) seed; if None, one will be
+        generated randomly.
 
     Returns
     -------
@@ -240,10 +241,10 @@ def maxwell_spectrum(T, prn_seed=None):
         Sampled outgoing energy
 
     """
-    
+
     if prn_seed is None:
         prn_seed = getrandbits(63)
-	
+
     return _dll.maxwell_spectrum(T, c_uint64(prn_seed))
 
 
@@ -257,7 +258,8 @@ def watt_spectrum(a, b, prn_seed=None):
     b : float
         Spectrum parameter b
     prn_seed : int
-        PRNG seed; if None, one will be generated randomly
+        Pseudorandom number generator (PRNG) seed; if None, one will be
+        generated randomly.
 
     Returns
     -------
@@ -265,7 +267,7 @@ def watt_spectrum(a, b, prn_seed=None):
         Sampled outgoing energy
 
     """
-    
+
     if prn_seed is None:
         prn_seed = getrandbits(63)
 
@@ -282,7 +284,8 @@ def normal_variate(mean_value, std_dev, prn_seed=None):
     std_dev : float
         Standard deviation of the normal distribution
     prn_seed : int
-        PRNG seed; if None, one will be generated randomly
+        Pseudorandom number generator (PRNG) seed; if None, one will be
+        generated randomly.
 
     Returns
     -------
@@ -290,7 +293,7 @@ def normal_variate(mean_value, std_dev, prn_seed=None):
         Sampled outgoing normally distributed value
 
     """
-    
+
     if prn_seed is None:
         prn_seed = getrandbits(63)
 

src/source.cpp

@@ -399,4 +399,23 @@ void free_memory_source()
   model::external_sources.clear();
 }
 
+//==============================================================================
+// C API
+//==============================================================================
+
+extern "C" int openmc_sample_external_source(
+  size_t n, uint64_t* seed, void* sites)
+{
+  if (!sites || !seed) {
+    set_errmsg("Received null pointer.");
+    return OPENMC_E_INVALID_ARGUMENT;
+  }
+
+  auto sites_array = static_cast<SourceSite*>(sites);
+  for (size_t i = 0; i < n; ++i) {
+    sites_array[i] = sample_external_source(seed);
+  }
+  return 0;
+}
+
 } // namespace openmc

tests/unit_tests/test_lib.py

@@ -807,3 +807,46 @@ def test_cell_rotation(pincell_model_w_univ, mpi_intracomm):
     cell.rotation = (180., 0., 0.)
     assert cell.rotation == pytest.approx([180., 0., 0.])
     openmc.lib.finalize()
+
+
+def test_sample_external_source(run_in_tmpdir, mpi_intracomm):
+    # Define a simple model and export
+    mat = openmc.Material()
+    mat.add_nuclide('U235', 1.0e-2)
+    sph = openmc.Sphere(r=100.0, boundary_type='vacuum')
+    cell = openmc.Cell(fill=mat, region=-sph)
+    model = openmc.Model()
+    model.geometry = openmc.Geometry([cell])
+    model.settings.source = openmc.Source(
+        space=openmc.stats.Box([-5., -5., -5.], [5., 5., 5.]),
+        angle=openmc.stats.Monodirectional((0., 0., 1.)),
+        energy=openmc.stats.Discrete([1.0e5], [1.0])
+    )
+    model.settings.particles = 1000
+    model.settings.batches = 10
+    model.export_to_xml()
+
+    # Sample some particles and make sure they match specified source
+    openmc.lib.init()
+    particles = openmc.lib.sample_external_source(10, prn_seed=3)
+    assert len(particles) == 10
+    for p in particles:
+        assert -5. < p.r[0] < 5.
+        assert -5. < p.r[1] < 5.
+        assert -5. < p.r[2] < 5.
+        assert p.u[0] == 0.0
+        assert p.u[1] == 0.0
+        assert p.u[2] == 1.0
+        assert p.E == 1.0e5
+
+    # Using the same seed should produce the same particles
+    other_particles = openmc.lib.sample_external_source(10, prn_seed=3)
+    assert len(other_particles) == 10
+    for p1, p2 in zip(particles, other_particles):
+        assert p1.r == p2.r
+        assert p1.u == p2.u
+        assert p1.E == p2.E
+        assert p1.time == p2.time
+        assert p1.wgt == p2.wgt
+
+    openmc.lib.finalize()