External sources alias sampler (#3201)

Co-authored-by: Paul Romano <[email protected]>
openmc-dev · Nov 23, 2024 · 2d988a6 · 2d988a6
1 parent dd01c40
commit 2d988a6
Show file tree

Hide file tree

Showing 6 changed files with 30 additions and 16 deletions.
diff --git a/include/openmc/source.h b/include/openmc/source.h
@@ -36,6 +36,9 @@ namespace model {
 
 extern vector<unique_ptr<Source>> external_sources;
 
+// Probability distribution for selecting external sources
+extern DiscreteIndex external_sources_probability;
+
 } // namespace model
 
 //==============================================================================

diff --git a/openmc/model/model.py b/openmc/model/model.py
@@ -885,7 +885,7 @@ def sample_external_source(
             n_samples: int = 1000,
             prn_seed: int | None = None,
             **init_kwargs
-    ) -> list[openmc.SourceParticle]:
+    ) -> openmc.ParticleList:
         """Sample external source and return source particles.
 
         .. versionadded:: 0.15.1
@@ -902,7 +902,7 @@ def sample_external_source(
 
         Returns
         -------
-        list of openmc.SourceParticle
+        openmc.ParticleList
             List of samples source particles
         """
         import openmc.lib

diff --git a/src/settings.cpp b/src/settings.cpp
@@ -563,6 +563,13 @@ void read_settings_xml(pugi::xml_node root)
     model::external_sources.push_back(make_unique<FileSource>(path));
   }
 
+  // Build probability mass function for sampling external sources
+  vector<double> source_strengths;
+  for (auto& s : model::external_sources) {
+    source_strengths.push_back(s->strength());
+  }
+  model::external_sources_probability.assign(source_strengths);
+
   // If no source specified, default to isotropic point source at origin with
   // Watt spectrum. No default source is needed in random ray mode.
   if (model::external_sources.empty() &&

diff --git a/src/source.cpp b/src/source.cpp
@@ -44,7 +44,10 @@ namespace openmc {
 namespace model {
 
 vector<unique_ptr<Source>> external_sources;
-}
+
+DiscreteIndex external_sources_probability;
+
+} // namespace model
 
 //==============================================================================
 // Source implementation
@@ -608,24 +611,13 @@ void initialize_source()
 
 SourceSite sample_external_source(uint64_t* seed)
 {
-  // Determine total source strength
-  double total_strength = 0.0;
-  for (auto& s : model::external_sources)
-    total_strength += s->strength();
-
   // Sample from among multiple source distributions
   int i = 0;
   if (model::external_sources.size() > 1) {
     if (settings::uniform_source_sampling) {
       i = prn(seed) * model::external_sources.size();
     } else {
-      double xi = prn(seed) * total_strength;
-      double c = 0.0;
-      for (; i < model::external_sources.size(); ++i) {
-        c += model::external_sources[i]->strength();
-        if (xi < c)
-          break;
-      }
+      i = model::external_sources_probability.sample(seed);
     }
   }
 

diff --git a/tests/regression_tests/source/results_true.dat b/tests/regression_tests/source/results_true.dat
@@ -1,2 +1,2 @@
 k-combined:
-3.054101E-01 1.167865E-03
+2.959436E-01 2.782384E-03
diff --git a/tests/unit_tests/test_uniform_source_sampling.py b/tests/unit_tests/test_uniform_source_sampling.py
@@ -61,3 +61,15 @@ def test_tally_mean(run_in_tmpdir, sphere_model):
 
     # Check that tally means match
     assert mean == pytest.approx(reference_mean)
+
+
+def test_multiple_sources(sphere_model):
+    low_strength_src = openmc.IndependentSource(
+        energy=openmc.stats.delta_function(1.0e6), strength=1e-7)
+    sphere_model.settings.source.append(low_strength_src)
+    sphere_model.settings.uniform_source_sampling = True
+
+    # Sample particles from source and make sure 1 MeV shows up despite
+    # negligible strength
+    particles = sphere_model.sample_external_source(100)
+    assert {p.E for p in particles} == {1.0e3, 1.0e6}