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

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

Author: eepeterson

include/openmc/source.h

@@ -4,6 +4,7 @@
 #ifndef OPENMC_SOURCE_H
 #define OPENMC_SOURCE_H
 
+#include <atomic>
 #include <limits>
 #include <unordered_set>
 
@@ -25,10 +26,18 @@ namespace openmc {
 // source_rejection_fraction
 constexpr int EXTSRC_REJECT_THRESHOLD {10000};
 
+// Maximum number of source rejections allowed while sampling a single site
+constexpr int64_t MAX_SOURCE_REJECTIONS_PER_SAMPLE {1'000'000};
+
 //==============================================================================
 // 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 +274,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,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,30 +507,49 @@ 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:
         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.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)
+    # 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,
+    )
+
+    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 @@ def sample_external_source(
         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 @@ 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`.
         **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 @@ 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, as_array=as_array
             )
 
     def apply_tally_results(self, statepoint: PathLike | openmc.StatePoint):
@@ -2588,7 +2593,7 @@ def convert_to_multigroup(
                     # 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()

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

src/source.cpp

@@ -37,6 +37,9 @@
 
 namespace openmc {
 
+std::atomic<int64_t> source_n_accept {0};
+std::atomic<int64_t> source_n_reject {0};
+
 namespace {
 
 void validate_particle_type(ParticleType type, const std::string& context)
@@ -191,9 +194,8 @@ void check_rejection_fraction(int64_t n_reject, int64_t n_accept)
 SourceSite Source::sample_with_constraints(uint64_t* seed) const
 {
   bool accepted = false;
-  static int64_t n_reject = 0;
-  static int64_t n_accept = 0;
-  SourceSite site;
+  int64_t n_local_reject = 0;
+  SourceSite site {};
 
   while (!accepted) {
     // Sample a source site without considering constraints yet
@@ -207,9 +209,13 @@ SourceSite Source::sample_with_constraints(uint64_t* seed) const
                  satisfies_energy_constraints(site.E) &&
                  satisfies_time_constraints(site.time);
       if (!accepted) {
-        // Increment number of rejections and check against minimum fraction
-        ++n_reject;
-        check_rejection_fraction(n_reject, n_accept);
+        ++n_local_reject;
+
+        // Check per-particle rejection limit
+        if (n_local_reject >= MAX_SOURCE_REJECTIONS_PER_SAMPLE) {
+          fatal_error("Exceeded maximum number of source rejections per "
+                      "sample. Please check your source definition.");
+        }
 
         // For the "kill" strategy, accept particle but set weight to 0 so that
         // it is terminated immediately
@@ -221,8 +227,13 @@ SourceSite Source::sample_with_constraints(uint64_t* seed) const
     }
   }
 
-  // Increment number of accepted samples
-  ++n_accept;
+  // Flush local rejection count, update accept counter, and check overall
+  // rejection fraction
+  if (n_local_reject > 0) {
+    source_n_reject += n_local_reject;
+  }
+  ++source_n_accept;
+  check_rejection_fraction(source_n_reject, source_n_accept);
 
   return site;
 }
@@ -361,15 +372,14 @@ IndependentSource::IndependentSource(pugi::xml_node node) : Source(node)
 
 SourceSite IndependentSource::sample(uint64_t* seed) const
 {
-  SourceSite site;
+  SourceSite site {};
   site.particle = particle_;
   double r_wgt = 1.0;
   double E_wgt = 1.0;
 
   // Repeat sampling source location until a good site has been accepted
   bool accepted = false;
-  static int64_t n_reject = 0;
-  static int64_t n_accept = 0;
+  int64_t n_local_reject = 0;
 
   while (!accepted) {
 
@@ -383,8 +393,11 @@ SourceSite IndependentSource::sample(uint64_t* seed) const
 
     // Check for rejection
     if (!accepted) {
-      ++n_reject;
-      check_rejection_fraction(n_reject, n_accept);
+      ++n_local_reject;
+      if (n_local_reject >= MAX_SOURCE_REJECTIONS_PER_SAMPLE) {
+        fatal_error("Exceeded maximum number of source rejections per "
+                    "sample. Please check your source definition.");
+      }
     }
   }
 
@@ -419,8 +432,11 @@ SourceSite IndependentSource::sample(uint64_t* seed) const
           (satisfies_energy_constraints(site.E)))
         break;
 
-      n_reject++;
-      check_rejection_fraction(n_reject, n_accept);
+      ++n_local_reject;
+      if (n_local_reject >= MAX_SOURCE_REJECTIONS_PER_SAMPLE) {
+        fatal_error("Exceeded maximum number of source rejections per "
+                    "sample. Please check your source definition.");
+      }
     }
 
     // Sample particle creation time
@@ -430,8 +446,10 @@ SourceSite IndependentSource::sample(uint64_t* seed) const
     site.wgt *= (E_wgt * time_wgt);
   }
 
-  // Increment number of accepted samples
-  ++n_accept;
+  // Flush local rejection count into global counter
+  if (n_local_reject > 0) {
+    source_n_reject += n_local_reject;
+  }
 
   return site;
 }
@@ -692,6 +710,13 @@ SourceSite sample_external_source(uint64_t* seed)
 void free_memory_source()
 {
   model::external_sources.clear();
+  reset_source_rejection_counters();
+}
+
+void reset_source_rejection_counters()
+{
+  source_n_accept = 0;
+  source_n_reject = 0;
 }
 
 //==============================================================================
@@ -712,8 +737,15 @@ extern "C" int openmc_sample_external_source(
   }
 
   auto sites_array = static_cast<SourceSite*>(sites);
+
+  // Derive independent per-particle seeds from the base seed so that
+  // each iteration has its own RNG state for thread-safe parallel sampling.
+  uint64_t base_seed = *seed;
+
+#pragma omp parallel for schedule(static)
   for (size_t i = 0; i < n; ++i) {
-    sites_array[i] = sample_external_source(seed);
+    uint64_t particle_seed = init_seed(base_seed + i, STREAM_SOURCE);
+    sites_array[i] = sample_external_source(&particle_seed);
   }
   return 0;
 }

tests/unit_tests/test_lib.py

@@ -1114,6 +1114,14 @@ def test_sample_external_source(run_in_tmpdir, mpi_intracomm):
         assert p1.time == p2.time
         assert p1.wgt == p2.wgt
 
+    # as_array should return a numpy structured array with matching values
+    arr = openmc.lib.sample_external_source(10, prn_seed=3, as_array=True)
+    assert isinstance(arr, np.ndarray)
+    assert len(arr) == 10
+    for p, row in zip(particles, arr):
+        assert p.r == pytest.approx(row['r'])
+        assert p.E == pytest.approx(row['E'])
+
     openmc.lib.finalize()
 
     # Make sure sampling works in volume calculation mode