Statistical weights in IndependentSource (#3195)

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

docs/source/usersguide/settings.rst

@@ -183,6 +183,7 @@ source distributions and has four main attributes that one can set:
 :attr:`IndependentSource.energy`, which defines the energy distribution, and
 :attr:`IndependentSource.time`, which defines the time distribution.
 
+
 The spatial distribution can be set equal to a sub-class of
 :class:`openmc.stats.Spatial`; common choices are :class:`openmc.stats.Point` or
 :class:`openmc.stats.Box`. To independently specify distributions in the
@@ -225,6 +226,7 @@ distribution. This could be a probability mass function
 (:class:`openmc.stats.Tabular`). By default, if no time distribution is
 specified, particles are started at :math:`t=0`.
 
+
 As an example, to create an isotropic, 10 MeV monoenergetic source uniformly
 distributed over a cube centered at the origin with an edge length of 10 cm, and
 emitting a pulse of particles from 0 to 10 µs, one
@@ -252,6 +254,24 @@ sampled 70% of the time and another that should be sampled 30% of the time::
 
   settings.source = [src1, src2]
 
+When the relative strengths are several orders of magnitude different, it may
+happen that not enough statistics are obtained from the lower strength source.
+This can be improved by sampling among the sources with equal probability,
+applying the source strength as a weight on the sampled source particles. The
+:attr:`Settings.uniform_source_sampling` attribute can be used to enable this
+option::
+
+  src1 = openmc.IndependentSource()
+  src1.strength = 100.0
+  ...
+
+  src2 = openmc.IndependentSource()
+  src2.strength = 1.0
+  ...
+
+  settings.source = [src1, src2]
+  settings.uniform_source_sampling = True
+
 Finally, the :attr:`IndependentSource.particle` attribute can be used to
 indicate the source should be composed of particles other than neutrons. For
 example, the following would generate a photon source::

include/openmc/settings.h

@@ -44,29 +44,30 @@ extern "C" bool entropy_on; //!< calculate Shannon entropy?
 extern "C" bool
   event_based; //!< use event-based mode (instead of history-based)
 extern bool legendre_to_tabular; //!< convert Legendre distributions to tabular?
-extern bool material_cell_offsets; //!< create material cells offsets?
-extern "C" bool output_summary;    //!< write summary.h5?
-extern bool output_tallies;        //!< write tallies.out?
-extern bool particle_restart_run;  //!< particle restart run?
-extern "C" bool photon_transport;  //!< photon transport turned on?
-extern "C" bool reduce_tallies;    //!< reduce tallies at end of batch?
-extern bool res_scat_on;           //!< use resonance upscattering method?
-extern "C" bool restart_run;       //!< restart run?
-extern "C" bool run_CE;            //!< run with continuous-energy data?
-extern bool source_latest;         //!< write latest source at each batch?
-extern bool source_separate;       //!< write source to separate file?
-extern bool source_write;          //!< write source in HDF5 files?
-extern bool source_mcpl_write;     //!< write source in mcpl files?
-extern bool surf_source_write;     //!< write surface source file?
-extern bool surf_mcpl_write;       //!< write surface mcpl file?
-extern bool surf_source_read;      //!< read surface source file?
-extern bool survival_biasing;      //!< use survival biasing?
-extern bool temperature_multipole; //!< use multipole data?
-extern "C" bool trigger_on;        //!< tally triggers enabled?
-extern bool trigger_predict;       //!< predict batches for triggers?
-extern bool ufs_on;                //!< uniform fission site method on?
-extern bool urr_ptables_on;        //!< use unresolved resonance prob. tables?
-extern "C" bool weight_windows_on; //!< are weight windows are enabled?
+extern bool material_cell_offsets;   //!< create material cells offsets?
+extern "C" bool output_summary;      //!< write summary.h5?
+extern bool output_tallies;          //!< write tallies.out?
+extern bool particle_restart_run;    //!< particle restart run?
+extern "C" bool photon_transport;    //!< photon transport turned on?
+extern "C" bool reduce_tallies;      //!< reduce tallies at end of batch?
+extern bool res_scat_on;             //!< use resonance upscattering method?
+extern "C" bool restart_run;         //!< restart run?
+extern "C" bool run_CE;              //!< run with continuous-energy data?
+extern bool source_latest;           //!< write latest source at each batch?
+extern bool source_separate;         //!< write source to separate file?
+extern bool source_write;            //!< write source in HDF5 files?
+extern bool source_mcpl_write;       //!< write source in mcpl files?
+extern bool surf_source_write;       //!< write surface source file?
+extern bool surf_mcpl_write;         //!< write surface mcpl file?
+extern bool surf_source_read;        //!< read surface source file?
+extern bool survival_biasing;        //!< use survival biasing?
+extern bool temperature_multipole;   //!< use multipole data?
+extern "C" bool trigger_on;          //!< tally triggers enabled?
+extern bool trigger_predict;         //!< predict batches for triggers?
+extern bool uniform_source_sampling; //!< sample sources uniformly?
+extern bool ufs_on;                  //!< uniform fission site method on?
+extern bool urr_ptables_on;          //!< use unresolved resonance prob. tables?
+extern "C" bool weight_windows_on;   //!< are weight windows are enabled?
 extern bool weight_window_checkpoint_surface;   //!< enable weight window check
                                                 //!< upon surface crossing?
 extern bool weight_window_checkpoint_collision; //!< enable weight window check

openmc/settings.py

@@ -266,6 +266,9 @@ class Settings:
         Maximum number of batches simulated. If this is set, the number of
         batches specified via ``batches`` is interpreted as the minimum number
         of batches
+    uniform_source_sampling : bool
+        Whether to sampling among multiple sources uniformly, applying their
+        strengths as weights to sampled particles.
     ufs_mesh : openmc.RegularMesh
         Mesh to be used for redistributing source sites via the uniform fission
         site (UFS) method.
@@ -328,6 +331,7 @@ class Settings:
         self._photon_transport = None
         self._plot_seed = None
         self._ptables = None
+        self._uniform_source_sampling = None
         self._seed = None
         self._survival_biasing = None
 
@@ -575,6 +579,15 @@ class Settings:
         cv.check_type('photon transport', photon_transport, bool)
         self._photon_transport = photon_transport
 
+    @property
+    def uniform_source_sampling(self) -> bool:
+        return self._uniform_source_sampling
+
+    @uniform_source_sampling.setter
+    def uniform_source_sampling(self, uniform_source_sampling: bool):
+        cv.check_type('strength as weights', uniform_source_sampling, bool)
+        self._uniform_source_sampling = uniform_source_sampling
+
     @property
     def plot_seed(self):
         return self._plot_seed
@@ -1221,6 +1234,11 @@ class Settings:
                 subelement.text = ' '.join(
                     str(x) for x in self._statepoint['batches'])
 
+    def _create_uniform_source_sampling_subelement(self, root):
+        if self._uniform_source_sampling is not None:
+            element = ET.SubElement(root, "uniform_source_sampling")
+            element.text = str(self._uniform_source_sampling).lower()
+
     def _create_sourcepoint_subelement(self, root):
         if self._sourcepoint:
             element = ET.SubElement(root, "source_point")
@@ -1702,6 +1720,11 @@ class Settings:
         if text is not None:
             self.photon_transport = text in ('true', '1')
 
+    def _uniform_source_sampling_from_xml_element(self, root):
+        text = get_text(root, 'uniform_source_sampling')
+        if text is not None:
+            self.uniform_source_sampling = text in ('true', '1')
+
     def _plot_seed_from_xml_element(self, root):
         text = get_text(root, 'plot_seed')
         if text is not None:
@@ -1957,6 +1980,7 @@ class Settings:
         self._create_energy_mode_subelement(element)
         self._create_max_order_subelement(element)
         self._create_photon_transport_subelement(element)
+        self._create_uniform_source_sampling_subelement(element)
         self._create_plot_seed_subelement(element)
         self._create_ptables_subelement(element)
         self._create_seed_subelement(element)
@@ -2063,6 +2087,7 @@ class Settings:
         settings._energy_mode_from_xml_element(elem)
         settings._max_order_from_xml_element(elem)
         settings._photon_transport_from_xml_element(elem)
+        settings._uniform_source_sampling_from_xml_element(elem)
         settings._plot_seed_from_xml_element(elem)
         settings._ptables_from_xml_element(elem)
         settings._seed_from_xml_element(elem)

src/finalize.cpp

@@ -133,6 +133,7 @@ int openmc_finalize()
   settings::trigger_on = false;
   settings::trigger_predict = false;
   settings::trigger_batch_interval = 1;
+  settings::uniform_source_sampling = false;
   settings::ufs_on = false;
   settings::urr_ptables_on = true;
   settings::verbosity = 7;

src/settings.cpp

@@ -72,6 +72,7 @@ bool survival_biasing {false};
 bool temperature_multipole {false};
 bool trigger_on {false};
 bool trigger_predict {false};
+bool uniform_source_sampling {false};
 bool ufs_on {false};
 bool urr_ptables_on {true};
 bool weight_windows_on {false};
@@ -786,6 +787,12 @@ void read_settings_xml(pugi::xml_node root)
     sourcepoint_batch = statepoint_batch;
   }
 
+  // Check is the user specified to convert strength to statistical weight
+  if (check_for_node(root, "uniform_source_sampling")) {
+    uniform_source_sampling =
+      get_node_value_bool(root, "uniform_source_sampling");
+  }
+
   // Check if the user has specified to write surface source
   if (check_for_node(root, "surf_source_write")) {
     surf_source_write = true;

src/source.cpp

@@ -616,18 +616,27 @@ SourceSite sample_external_source(uint64_t* seed)
   // Sample from among multiple source distributions
   int i = 0;
   if (model::external_sources.size() > 1) {
-    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;
+    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;
+      }
     }
   }
 
   // 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();
+  }
+
   // If running in MG, convert site.E to group
   if (!settings::run_CE) {
     site.E = lower_bound_index(data::mg.rev_energy_bins_.begin(),

src/tallies/tally.cpp

@@ -751,7 +751,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) {
+    if (settings::run_mode == RunMode::FIXED_SOURCE &&
+        !settings::uniform_source_sampling) {
       for (const auto& s : model::external_sources) {
         total_source += s->strength();
       }

tests/unit_tests/test_uniform_source_sampling.py

@@ -0,0 +1,63 @@
+import openmc
+import pytest
+
+
+@pytest.fixture
+def sphere_model():
+    mat = openmc.Material()
+    mat.add_nuclide('Li6', 1.0)
+    mat.set_density('g/cm3', 1.0)
+    sphere = openmc.Sphere(r=1.0, boundary_type='vacuum')
+    cell = openmc.Cell(region=-sphere, fill=mat)
+    model = openmc.Model()
+    model.geometry = openmc.Geometry([cell])
+
+    model.settings.particles = 100
+    model.settings.batches = 1
+    model.settings.source = openmc.IndependentSource(
+        energy=openmc.stats.delta_function(1.0e3),
+        strength=100.0
+    )
+    model.settings.run_mode = "fixed source"
+    model.settings.surf_source_write = {
+        "max_particles": 100,
+    }
+
+    tally = openmc.Tally()
+    tally.scores = ['flux']
+    model.tallies = [tally]
+    return model
+
+
+def test_source_weight(run_in_tmpdir, sphere_model):
+    # Run OpenMC without uniform source sampling and check that banked particles
+    # have weight 1
+    sphere_model.settings.uniform_source_sampling = False
+    sphere_model.run()
+    particles = openmc.ParticleList.from_hdf5('surface_source.h5')
+    assert set(p.wgt for p in particles) == {1.0}
+
+    # Run with uniform source sampling and check that banked particles have
+    # weight == strength
+    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}
+
+
+def test_tally_mean(run_in_tmpdir, sphere_model):
+    # Run without uniform source sampling
+    sphere_model.settings.uniform_source_sampling = False
+    sp_file = sphere_model.run()
+    with openmc.StatePoint(sp_file) as sp:
+        reference_mean = sp.tallies[sphere_model.tallies[0].id].mean
+
+    # Run with uniform source sampling
+    sphere_model.settings.uniform_source_sampling = True
+    sp_file = sphere_model.run()
+    with openmc.StatePoint(sp_file) as sp:
+        mean = sp.tallies[sphere_model.tallies[0].id].mean
+
+    # Check that tally means match
+    assert mean == pytest.approx(reference_mean)