Fix weight modification for uniform source sampling (#3395)

Co-authored-by: Patrick Shriwise <pshriwise@gmail.com>

Author: paulromano

openmc/deplete/d1s.py

@@ -141,7 +141,9 @@ def apply_time_correction(
     time_correction_factors : dict
         Time correction factors as returned by :func:`time_correction_factors`
     index : int, optional
-        Index to use for the correction factors
+        Index of the time of interest. If N timesteps are provided in
+        :func:`time_correction_factors`, there are N + 1 times to select from.
+        The default is -1 which corresponds to the final time.
     sum_nuclides : bool
         Whether to sum over the parent nuclides
 

openmc/source.py

@@ -110,7 +110,7 @@ def constraints(self, constraints: dict[str, Any] | None):
                 cv.check_value('rejection strategy', value, ('resample', 'kill'))
                 self._constraints['rejection_strategy'] = value
             else:
-                raise ValueError('Unknown key in constraints dictionary: {key}')
+                raise ValueError(f'Unknown key in constraints dictionary: {key}')
 
     @abstractmethod
     def populate_xml_element(self, element):

src/settings.cpp

@@ -620,13 +620,6 @@ 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() &&
@@ -639,6 +632,13 @@ void read_settings_xml(pugi::xml_node root)
       UPtrDist {new Discrete(T, p, 1)}));
   }
 
+  // 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);
+
   // Check if we want to write out source
   if (check_for_node(root, "write_initial_source")) {
     write_initial_source = get_node_value_bool(root, "write_initial_source");

src/source.cpp

@@ -613,9 +613,10 @@ SourceSite sample_external_source(uint64_t* seed)
 {
   // Sample from among multiple source distributions
   int i = 0;
-  if (model::external_sources.size() > 1) {
+  int n_sources = model::external_sources.size();
+  if (n_sources > 1) {
     if (settings::uniform_source_sampling) {
-      i = prn(seed) * model::external_sources.size();
+      i = prn(seed) * n_sources;
     } else {
       i = model::external_sources_probability.sample(seed);
     }
@@ -624,9 +625,13 @@ SourceSite sample_external_source(uint64_t* seed)
   // Sample source site from i-th source distribution
   SourceSite site {model::external_sources[i]->sample_with_constraints(seed)};
 
-  // Set particle creation weight
-  if (settings::uniform_source_sampling) {
-    site.wgt *= model::external_sources[i]->strength();
+  // For uniform source sampling, multiply the weight by the ratio of the actual
+  // probability of sampling source i to the biased probability of sampling
+  // source i, which is (strength_i / total_strength) / (1 / n)
+  if (n_sources > 1 && settings::uniform_source_sampling) {
+    double total_strength = model::external_sources_probability.integral();
+    site.wgt *=
+      model::external_sources[i]->strength() * n_sources / total_strength;
   }
 
   // If running in MG, convert site.E to group

src/tallies/tally.cpp

@@ -815,11 +815,8 @@ void Tally::accumulate()
   if (mpi::master || !settings::reduce_tallies) {
     // Calculate total source strength for normalization
     double total_source = 0.0;
-    if (settings::run_mode == RunMode::FIXED_SOURCE &&
-        !settings::uniform_source_sampling) {
-      for (const auto& s : model::external_sources) {
-        total_source += s->strength();
-      }
+    if (settings::run_mode == RunMode::FIXED_SOURCE) {
+      total_source = model::external_sources_probability.integral();
     } else {
       total_source = 1.0;
     }

tests/unit_tests/test_uniform_source_sampling.py

@@ -14,10 +14,15 @@ def sphere_model():
 
     model.settings.particles = 100
     model.settings.batches = 1
-    model.settings.source = openmc.IndependentSource(
+    src1 = openmc.IndependentSource(
         energy=openmc.stats.delta_function(1.0e3),
-        strength=100.0
+        strength=75.0
     )
+    src2 = openmc.IndependentSource(
+        energy=openmc.stats.delta_function(1.0e3),
+        strength=25.0
+    )
+    model.settings.source = [src1, src2]
     model.settings.run_mode = "fixed source"
     model.settings.surf_source_write = {
         "max_particles": 100,
@@ -42,11 +47,13 @@ def test_source_weight(run_in_tmpdir, sphere_model):
     sphere_model.settings.uniform_source_sampling = True
     sphere_model.run()
     particles = openmc.ParticleList.from_hdf5('surface_source.h5')
-    strength = sphere_model.settings.source[0].strength
-    assert set(p.wgt for p in particles) == {strength}
+    assert set(p.wgt for p in particles) == {0.5, 1.5}
 
 
 def test_tally_mean(run_in_tmpdir, sphere_model):
+    # Use only one source
+    sphere_model.settings.source.pop()
+
     # Run without uniform source sampling
     sphere_model.settings.uniform_source_sampling = False
     sp_file = sphere_model.run()