diff --git a/include/openmc/source.h b/include/openmc/source.h index 2f32aa2a0..1ef7eba2b 100644 --- a/include/openmc/source.h +++ b/include/openmc/source.h @@ -4,6 +4,7 @@ #ifndef OPENMC_SOURCE_H #define OPENMC_SOURCE_H +#include #include #include @@ -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 source_n_accept; +extern std::atomic 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 diff --git a/openmc/lib/core.py b/openmc/lib/core.py index 2764ddff1..8e5cb56dc 100644 --- a/openmc/lib/core.py +++ b/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(): diff --git a/openmc/model/model.py b/openmc/model/model.py index 0920b79e4..67a8495ee 100644 --- a/openmc/model/model.py +++ b/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() diff --git a/src/simulation.cpp b/src/simulation.cpp index d0d6f037f..4fad196a6 100644 --- a/src/simulation.cpp +++ b/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 diff --git a/src/source.cpp b/src/source.cpp index 5300a685f..f0b8f48ed 100644 --- a/src/source.cpp +++ b/src/source.cpp @@ -37,6 +37,9 @@ namespace openmc { +std::atomic source_n_accept {0}; +std::atomic 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(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; } diff --git a/tests/unit_tests/test_lib.py b/tests/unit_tests/test_lib.py index 8ef8d5927..51e648dcf 100644 --- a/tests/unit_tests/test_lib.py +++ b/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