Apply Weight Window Upon Surface Crossing (#2670)

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

docs/source/io_formats/settings.rst

@@ -1211,6 +1211,25 @@ mesh-based weight windows.
 
         *Default*: 5.0
 
+---------------------------------------
+``<weight_window_checkpoints>`` Element
+---------------------------------------
+
+The ``<weight_window_checkpoints>`` element indicates the checkpoints for weight
+window split/roulette (surface, collision or both). This element has the
+following sub-elements/attributes:
+
+  :surface:
+    If set to "true", weight window checks will be performed at surface
+    crossings.
+
+    *Default*: False
+
+  :collision:
+    If set to "true", weight window checks will be performed at collisions.
+
+    *Default*: True
+
 --------------------------------------
 ``<weight_windows_file>`` Element
 --------------------------------------

include/openmc/settings.h

@@ -59,8 +59,12 @@ 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 write_all_tracks;      //!< write track files for every particle?
-extern bool write_initial_source;  //!< write out initial source file?
+extern bool weight_window_checkpoint_surface;   //!< enable weight window check
+                                                //!< upon surface crossing?
+extern bool weight_window_checkpoint_collision; //!< enable weight window check
+                                                //!< upon collision?
+extern bool write_all_tracks;     //!< write track files for every particle?
+extern bool write_initial_source; //!< write out initial source file?
 
 // Paths to various files
 extern std::string path_cross_sections; //!< path to cross_sections.xml

openmc/settings.py

@@ -236,6 +236,11 @@ class Settings:
         Weight windows to use for variance reduction
 
         .. versionadded:: 0.13
+    weight_window_checkpoints : dict
+        Indicates the checkpoints for weight window split/roulettes. Valid keys
+        include "collision" and "surface". Values must be of type bool.
+
+        .. versionadded:: 0.13.4
     weight_window_generators : WeightWindowGenerator or iterable of WeightWindowGenerator
         Weight windows generation parameters to apply during simulation
 
@@ -335,6 +340,7 @@ class Settings:
         self._weight_window_generators = cv.CheckedList(WeightWindowGenerator, 'weight window generators')
         self._weight_windows_on = None
         self._weight_windows_file = None
+        self._weight_window_checkpoints = {}
         self._max_splits = None
         self._max_tracks = None
 
@@ -961,6 +967,16 @@ class Settings:
         cv.check_type('weight windows on', value, bool)
         self._weight_windows_on = value
 
+    @property
+    def weight_window_checkpoints(self) -> dict:
+        return self._weight_window_checkpoints
+
+    @weight_window_checkpoints.setter
+    def weight_window_checkpoints(self, weight_window_checkpoints: dict):
+        for key in weight_window_checkpoints.keys():
+            cv.check_value('weight_window_checkpoints', key, ('collision', 'surface'))
+        self._weight_window_checkpoints = weight_window_checkpoints
+
     @property
     def max_splits(self) -> int:
         return self._max_splits
@@ -1377,6 +1393,19 @@ class Settings:
             element.text = self.weight_windows_file
             root.append(element)
 
+    def _create_weight_window_checkpoints_subelement(self, root):
+        if not self._weight_window_checkpoints:
+            return
+        element = ET.SubElement(root, "weight_window_checkpoints")
+
+        if 'collision' in self._weight_window_checkpoints:
+            subelement = ET.SubElement(element, "collision")
+            subelement.text = str(self._weight_window_checkpoints['collision']).lower()
+
+        if 'surface' in self._weight_window_checkpoints:
+            subelement = ET.SubElement(element, "surface")
+            subelement.text = str(self._weight_window_checkpoints['surface']).lower()
+
     def _create_max_splits_subelement(self, root):
         if self._max_splits is not None:
             elem = ET.SubElement(root, "max_splits")
@@ -1718,6 +1747,16 @@ class Settings:
         if meshes is not None and self.weight_windows:
             meshes.update({ww.mesh.id: ww.mesh for ww in self.weight_windows})
 
+    def _weight_window_checkpoints_from_xml_element(self, root):
+        elem = root.find('weight_window_checkpoints')
+        if elem is None:
+            return
+        for key in ('collision', 'surface'):
+            value = get_text(elem, key)
+            if value is not None:
+                value = value in ('true', '1')
+                self.weight_window_checkpoints[key] = value
+
     def _max_splits_from_xml_element(self, root):
         text = get_text(root, 'max_splits')
         if text is not None:
@@ -1786,6 +1825,7 @@ class Settings:
         self._create_weight_windows_subelement(element, mesh_memo)
         self._create_weight_window_generators_subelement(element, mesh_memo)
         self._create_weight_windows_file_element(element)
+        self._create_weight_window_checkpoints_subelement(element)
         self._create_max_splits_subelement(element)
         self._create_max_tracks_subelement(element)
 
@@ -1882,6 +1922,7 @@ class Settings:
         settings._write_initial_source_from_xml_element(elem)
         settings._weight_windows_from_xml_element(elem, meshes)
         settings._weight_window_generators_from_xml_element(elem, meshes)
+        settings._weight_window_checkpoints_from_xml_element(elem)
         settings._max_splits_from_xml_element(elem)
         settings._max_tracks_from_xml_element(elem)
 

src/particle.cpp

@@ -30,6 +30,7 @@
 #include "openmc/tallies/tally.h"
 #include "openmc/tallies/tally_scoring.h"
 #include "openmc/track_output.h"
+#include "openmc/weight_windows.h"
 
 #ifdef DAGMC
 #include "DagMC.hpp"
@@ -272,6 +273,9 @@ void Particle::event_cross_surface()
   } else {
     // Particle crosses surface
     cross_surface();
+    if (settings::weight_window_checkpoint_surface) {
+      apply_weight_windows(*this);
+    }
     event() = TallyEvent::SURFACE;
   }
   // Score cell to cell partial currents

src/physics.cpp

@@ -47,15 +47,9 @@ void collision(Particle& p)
   switch (p.type()) {
   case ParticleType::neutron:
     sample_neutron_reaction(p);
-    if (settings::weight_windows_on) {
-      apply_weight_windows(p);
-    }
     break;
   case ParticleType::photon:
     sample_photon_reaction(p);
-    if (settings::weight_windows_on) {
-      apply_weight_windows(p);
-    }
     break;
   case ParticleType::electron:
     sample_electron_reaction(p);
@@ -65,6 +59,9 @@ void collision(Particle& p)
     break;
   }
 
+  if (settings::weight_window_checkpoint_collision)
+    apply_weight_windows(p);
+
   // Kill particle if energy falls below cutoff
   int type = static_cast<int>(p.type());
   if (p.E() < settings::energy_cutoff[type]) {

src/settings.cpp

@@ -74,6 +74,8 @@ bool trigger_predict {false};
 bool ufs_on {false};
 bool urr_ptables_on {true};
 bool weight_windows_on {false};
+bool weight_window_checkpoint_surface {false};
+bool weight_window_checkpoint_collision {true};
 bool write_all_tracks {false};
 bool write_initial_source {false};
 
@@ -974,6 +976,19 @@ void read_settings_xml(pugi::xml_node root)
       }
     }
   }
+
+  // Set up weight window checkpoints
+  if (check_for_node(root, "weight_window_checkpoints")) {
+    xml_node ww_checkpoints = root.child("weight_window_checkpoints");
+    if (check_for_node(ww_checkpoints, "collision")) {
+      weight_window_checkpoint_collision =
+        get_node_value_bool(ww_checkpoints, "collision");
+    }
+    if (check_for_node(ww_checkpoints, "surface")) {
+      weight_window_checkpoint_surface =
+        get_node_value_bool(ww_checkpoints, "surface");
+    }
+  }
 }
 
 void free_memory_settings()

src/weight_windows.cpp

@@ -53,6 +53,13 @@ openmc::vector<unique_ptr<WeightWindowsGenerator>> weight_windows_generators;
 
 void apply_weight_windows(Particle& p)
 {
+  if (!settings::weight_windows_on)
+    return;
+
+  // WW on photon and neutron only
+  if (p.type() != ParticleType::neutron && p.type() != ParticleType::photon)
+    return;
+
   // skip dead or no energy
   if (p.E() <= 0 || !p.alive())
     return;

tests/unit_tests/test_settings.py

@@ -57,6 +57,7 @@ def test_export_to_xml(run_in_tmpdir):
     s.photon_transport = False
     s.electron_treatment = 'led'
     s.write_initial_source = True
+    s.weight_window_checkpoints = {'surface': True, 'collision': False}
 
     # Make sure exporting XML works
     s.export_to_xml()
@@ -126,3 +127,4 @@ def test_export_to_xml(run_in_tmpdir):
     assert vol.samples == 1000
     assert vol.lower_left == (-10., -10., -10.)
     assert vol.upper_right == (10., 10., 10.)
+    assert s.weight_window_checkpoints == {'surface': True, 'collision': False}