Parallelize sampling external sources and threadsafe rejection counters (#3830)

Co-authored-by: Paul Romano <paul.k.romano@gmail.com>

openmc/lib/core.py

@@ -33,7 +33,6 @@ class _SourceSite(Structure):
                 ('parent_id', c_int64),
                 ('progeny_id', c_int64)]
 
-
 # Define input type for numpy arrays that will be passed into C++ functions
 # Must be an int or double array, with single dimension that is contiguous
 _array_1d_int = np.ctypeslib.ndpointer(dtype=np.int32, ndim=1,
@@ -494,8 +493,9 @@ def run_random_ray(output=True):
 
 def sample_external_source(
         n_samples: int = 1000,
-        prn_seed: int | None = None
-) -> openmc.ParticleList:
+        prn_seed: int | None = None,
+        as_array: bool = False
+) -> openmc.ParticleList | np.ndarray:
     """Sample external source and return source particles.
 
     .. versionadded:: 0.13.1
@@ -507,11 +507,20 @@ def sample_external_source(
     prn_seed : int
         Pseudorandom number generator (PRNG) seed; if None, one will be
         generated randomly.
+    as_array : bool
+        If True, return a numpy structured array instead of a
+        :class:`~openmc.ParticleList`.  The array has fields ``'r'`` (float64,
+        shape 3), ``'u'`` (float64, shape 3), ``'E'`` (float64), ``'time'``
+        (float64), ``'wgt'`` (float64), ``'delayed_group'`` (int32),
+        ``'surf_id'`` (int32), and ``'particle'`` (int32).  This avoids the
+        overhead of constructing individual :class:`~openmc.SourceParticle`
+        objects and is substantially faster for large sample counts.
 
     Returns
     -------
-    openmc.ParticleList
-        List of sampled source particles
+    openmc.ParticleList or numpy.ndarray
+        List of sampled source particles, or a structured array when
+        *as_array* is True.
 
     """
     if n_samples <= 0:
@@ -519,18 +528,28 @@ def sample_external_source(
     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)
+    # Pre-allocate output array and sample all particles in a single C call
+    result = np.empty(n_samples, dtype=_SourceSite)
+    sites_array = (_SourceSite * n_samples).from_buffer(result)
+    _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.ParticleList([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)
+    if as_array:
+        return result
+
+    particles = [
+        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
-    ])
+    ]
+    return openmc.ParticleList(particles)
 
 
 def simulation_init():

openmc/model/model.py

@@ -1294,8 +1294,9 @@ class Model:
         self,
         n_samples: int = 1000,
         prn_seed: int | None = None,
+        as_array: bool = False,
         **init_kwargs
-    ) -> openmc.ParticleList:
+    ) -> openmc.ParticleList | np.ndarray:
         """Sample external source and return source particles.
 
         .. versionadded:: 0.15.1
@@ -1307,13 +1308,17 @@ class Model:
         prn_seed : int
             Pseudorandom number generator (PRNG) seed; if None, one will be
             generated randomly.
+        as_array : bool
+            If True, return a numpy structured array instead of a
+            :class:`~openmc.ParticleList`.
         **init_kwargs
             Keyword arguments passed to :func:`openmc.lib.init`
 
         Returns
         -------
-        openmc.ParticleList
-            List of samples source particles
+        openmc.ParticleList or numpy.ndarray
+            List of sampled source particles, or a structured array when
+            *as_array* is True.
         """
         import openmc.lib
 
@@ -1324,7 +1329,7 @@ class Model:
 
         with openmc.lib.TemporarySession(self, **init_kwargs):
             return openmc.lib.sample_external_source(
-                n_samples=n_samples, prn_seed=prn_seed
+                n_samples=n_samples, prn_seed=prn_seed, as_array=as_array
             )
 
     def apply_tally_results(self, statepoint: PathLike | openmc.StatePoint):
@@ -2588,7 +2593,7 @@ class Model:
                     # This mode doesn't require
                     # valid transport settings like particles/batches
                     original_run_mode = self.settings.run_mode
-                    self.settings.run_mode = 'volume' 
+                    self.settings.run_mode = 'volume'
                     self.init_lib(directory=tmpdir)
                     self.sync_dagmc_universes()
                     self.finalize_lib()