Parallelize sampling external sources and threadsafe rejection counters

docs/source/io_formats/settings.rst

@@ -1062,6 +1062,16 @@ variable and whose sub-elements/attributes are as follows:
 
   *Default*: None
 
+------------------------------------------------
+``<max_source_rejections_per_sample>`` Element
+------------------------------------------------
+
+The ``<max_source_rejections_per_sample>`` element specifies the maximum number
+of rejection attempts allowed when sampling a single source site with
+constraints. If this limit is exceeded, the simulation will abort with an error.
+
+   *Default*: 1000000
+
 ---------------------------------------
 ``<source_rejection_fraction>`` Element
 ---------------------------------------

include/openmc/capi.h

@@ -184,7 +184,8 @@ int openmc_reset();
 int openmc_reset_timers();
 int openmc_run();
 void openmc_run_random_ray();
-int openmc_sample_external_source(size_t n, uint64_t* seed, void* sites);
+int openmc_sample_external_source(
+  size_t n, uint64_t* seed, void* sites, int threads);
 void openmc_set_seed(int64_t new_seed);
 void openmc_set_stride(uint64_t new_stride);
 int openmc_set_n_batches(

include/openmc/settings.h

@@ -167,8 +167,10 @@ extern std::unordered_set<int>
 extern CollisionTrackConfig collision_track_config;
 extern double source_rejection_fraction; //!< Minimum fraction of source sites
                                          //!< that must be accepted
-extern double free_gas_threshold;        //!< Threshold multiplier for free gas
-                                         //!< scattering treatment
+extern int64_t max_source_rejections_per_sample; //!< Maximum number of
+                                                 //!< rejections per sample
+extern double free_gas_threshold; //!< Threshold multiplier for free gas
+                                  //!< scattering treatment
 
 extern int
   max_history_splits; //!< maximum number of particle splits for weight windows

include/openmc/source.h

@@ -4,6 +4,7 @@
 #ifndef OPENMC_SOURCE_H
 #define OPENMC_SOURCE_H
 
+#include <atomic>
 #include <limits>
 #include <unordered_set>
 
@@ -29,6 +30,11 @@ constexpr int EXTSRC_REJECT_THRESHOLD {10000};
 // Global variables
 //==============================================================================
 
+// Cumulative counters for source rejection diagnostics. These are atomic to
+// allow thread-safe concurrent sampling of external sources.
+extern std::atomic<int64_t> source_n_accept;
+extern std::atomic<int64_t> source_n_reject;
+
 class Source;
 
 namespace model {
@@ -265,6 +271,9 @@ SourceSite sample_external_source(uint64_t* seed);
 
 void free_memory_source();
 
+//! Reset cumulative source rejection counters
+void reset_source_rejection_counters();
+
 } // namespace openmc
 
 #endif // OPENMC_SOURCE_H

openmc/lib/core.py

@@ -33,6 +33,26 @@ class _SourceSite(Structure):
                 ('parent_id', c_int64),
                 ('progeny_id', c_int64)]
 
+# Numpy dtype matching _SourceSite memory layout (104 bytes) for zero-copy
+# interpretation of the ctypes buffer as a numpy structured array.
+_source_site_dtype = np.dtype([
+    ('r', '<f8', (3,)),
+    ('u', '<f8', (3,)),
+    ('E', '<f8'),
+    ('time', '<f8'),
+    ('wgt', '<f8'),
+    ('delayed_group', '<i4'),
+    ('surf_id', '<i4'),
+    ('particle', '<i4'),
+    ('parent_nuclide', '<i4'),
+    ('parent_id', '<i8'),
+    ('progeny_id', '<i8'),
+])
+
+# Maximum number of source sites to sample per C API call.  Requests for
+# more sites are automatically split into batches of this size so that the
+# intermediate ctypes buffer stays bounded (~104 MB at the default value).
+_SOURCE_SAMPLE_BATCH_SIZE: int = 1_000_000
 
 # 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
@@ -110,7 +130,7 @@ class _SourceSite(Structure):
                                             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.argtypes = [c_size_t, POINTER(c_uint64), POINTER(_SourceSite), c_int]
 _dll.openmc_sample_external_source.restype = c_int
 _dll.openmc_sample_external_source.errcheck = _error_handler
 
@@ -494,8 +514,10 @@ 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,
+        n_threads: int = 1,
+        as_array: bool = False
+) -> openmc.ParticleList | np.ndarray:
     """Sample external source and return source particles.
 
     .. versionadded:: 0.13.1
@@ -507,30 +529,80 @@ def sample_external_source(
     prn_seed : int
         Pseudorandom number generator (PRNG) seed; if None, one will be
         generated randomly.
+    n_threads : int
+        Number of OpenMP threads to use for parallel sampling. Defaults to 1
+        (serial). Each sample gets an independent RNG stream derived from
+        the base seed, so results are deterministic regardless of thread count.
+    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:
         raise ValueError("Number of samples must be positive")
+    if n_threads < 1:
+        raise ValueError("Number of threads must be at least 1")
     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)
+    batch_size = min(n_samples, _SOURCE_SAMPLE_BATCH_SIZE)
 
-    # 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)
+    # Allocate the ctypes buffer once; it is reused for every batch.
+    sites_array = (_SourceSite * batch_size)()
+
+    # Pre-allocate the output container.  For ``as_array`` mode we create
+    # a single numpy array and fill it in slices; for the ParticleList
+    # path we accumulate SourceParticle objects across batches.
+    if as_array:
+        result = np.empty(n_samples, dtype=_source_site_dtype)
+    else:
+        particles = []
+
+    for offset in range(0, n_samples, batch_size):
+        n_batch = min(batch_size, n_samples - offset)
+
+        # Each batch's base seed is shifted by ``offset`` so that
+        # particle *i* within a batch gets
+        #   init_seed(prn_seed + offset + i, STREAM_SOURCE)
+        # which is identical to the seed it would receive in an
+        # unbatched call.  Results are therefore deterministic
+        # regardless of batch size.
+        _dll.openmc_sample_external_source(
+            c_size_t(n_batch),
+            c_uint64(prn_seed + offset),
+            sites_array,
+            c_int(n_threads),
         )
-        for site in sites_array
-    ])
+
+        if as_array:
+            result[offset:offset + n_batch] = np.frombuffer(
+                sites_array, dtype=_source_site_dtype, count=n_batch
+            )
+        else:
+            particles.extend(
+                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[:n_batch]
+            )
+
+    if as_array:
+        return result
+    return openmc.ParticleList(particles)
 
 
 def simulation_init():

openmc/model/model.py

@@ -1294,8 +1294,10 @@ def sample_external_source(
         self,
         n_samples: int = 1000,
         prn_seed: int | None = None,
+        n_threads: int = 1,
+        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 +1309,19 @@ def sample_external_source(
         prn_seed : int
             Pseudorandom number generator (PRNG) seed; if None, one will be
             generated randomly.
+        n_threads : int
+            Number of OpenMP threads to use for parallel sampling.
+        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 +1332,8 @@ def sample_external_source(
 
         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,
+                n_threads=n_threads, as_array=as_array
             )
 
     def apply_tally_results(self, statepoint: PathLike | openmc.StatePoint):

openmc/settings.py

@@ -254,6 +254,10 @@ class Settings:
         Minimum fraction of source sites that must be accepted when applying
         rejection sampling based on constraints. If not specified, the default
         value is 0.05.
+    max_source_rejections_per_sample : int
+        Maximum number of rejections allowed when sampling a single source
+        particle before raising an error. If not specified, the default
+        value is 1,000,000.
     sourcepoint : dict
         Options for writing source points. Acceptable keys are:
 
@@ -390,6 +394,7 @@ def __init__(self, **kwargs):
         # Source subelement
         self._source = cv.CheckedList(SourceBase, 'source distributions')
         self._source_rejection_fraction = None
+        self._max_source_rejections_per_sample = None
 
         self._confidence_intervals = None
         self._electron_treatment = None
@@ -1389,6 +1394,18 @@ def source_rejection_fraction(self, source_rejection_fraction: float):
                            source_rejection_fraction, 1)
         self._source_rejection_fraction = source_rejection_fraction
 
+    @property
+    def max_source_rejections_per_sample(self) -> int | None:
+        return self._max_source_rejections_per_sample
+
+    @max_source_rejections_per_sample.setter
+    def max_source_rejections_per_sample(self, max_source_rejections_per_sample: int):
+        cv.check_type('max_source_rejections_per_sample',
+                      max_source_rejections_per_sample, Integral)
+        cv.check_greater_than('max_source_rejections_per_sample',
+                              max_source_rejections_per_sample, 0)
+        self._max_source_rejections_per_sample = max_source_rejections_per_sample
+
     @property
     def free_gas_threshold(self) -> float | None:
         return self._free_gas_threshold
@@ -1927,6 +1944,11 @@ def _create_source_rejection_fraction_subelement(self, root):
             element = ET.SubElement(root, "source_rejection_fraction")
             element.text = str(self._source_rejection_fraction)
 
+    def _create_max_source_rejections_per_sample_subelement(self, root):
+        if self._max_source_rejections_per_sample is not None:
+            element = ET.SubElement(root, "max_source_rejections_per_sample")
+            element.text = str(self._max_source_rejections_per_sample)
+
     def _create_free_gas_threshold_subelement(self, root):
         if self._free_gas_threshold is not None:
             element = ET.SubElement(root, "free_gas_threshold")
@@ -2393,6 +2415,11 @@ def _source_rejection_fraction_from_xml_element(self, root):
         if text is not None:
             self.source_rejection_fraction = float(text)
 
+    def _max_source_rejections_per_sample_from_xml_element(self, root):
+        text = get_text(root, 'max_source_rejections_per_sample')
+        if text is not None:
+            self.max_source_rejections_per_sample = int(text)
+
     def _free_gas_threshold_from_xml_element(self, root):
         text = get_text(root, 'free_gas_threshold')
         if text is not None:
@@ -2469,6 +2496,7 @@ def to_xml_element(self, mesh_memo=None):
         self._create_random_ray_subelement(element, mesh_memo)
         self._create_use_decay_photons_subelement(element)
         self._create_source_rejection_fraction_subelement(element)
+        self._create_max_source_rejections_per_sample_subelement(element)
         self._create_free_gas_threshold_subelement(element)
 
         # Clean the indentation in the file to be user-readable
@@ -2582,6 +2610,7 @@ def from_xml_element(cls, elem, meshes=None):
         settings._random_ray_from_xml_element(elem, meshes)
         settings._use_decay_photons_from_xml_element(elem)
         settings._source_rejection_fraction_from_xml_element(elem)
+        settings._max_source_rejections_per_sample_from_xml_element(elem)
         settings._free_gas_threshold_from_xml_element(elem)
 
         return settings

src/settings.cpp

@@ -129,6 +129,7 @@ SolverType solver_type {SolverType::MONTE_CARLO};
 std::unordered_set<int> sourcepoint_batch;
 std::unordered_set<int> statepoint_batch;
 double source_rejection_fraction {0.05};
+int64_t max_source_rejections_per_sample {1'000'000};
 double free_gas_threshold {400.0};
 std::unordered_set<int> source_write_surf_id;
 CollisionTrackConfig collision_track_config {};
@@ -674,6 +675,11 @@ void read_settings_xml(pugi::xml_node root)
       std::stod(get_node_value(root, "source_rejection_fraction"));
   }
 
+  if (check_for_node(root, "max_source_rejections_per_sample")) {
+    max_source_rejections_per_sample =
+      std::stoll(get_node_value(root, "max_source_rejections_per_sample"));
+  }
+
   if (check_for_node(root, "free_gas_threshold")) {
     free_gas_threshold = std::stod(get_node_value(root, "free_gas_threshold"));
   }

src/simulation.cpp

@@ -122,6 +122,7 @@ int openmc_simulation_init()
   simulation::ssw_current_file = 1;
   simulation::k_generation.clear();
   simulation::entropy.clear();
+  reset_source_rejection_counters();
   openmc_reset();
 
   // If this is a restart run, load the state point data and binary source