Implement user-configurable random number stride (#3067)

Co-authored-by: Paul Romano <[email protected]>

Author: ahman24

docs/source/io_formats/settings.rst

@@ -538,6 +538,15 @@ pseudo-random number generator.
 
   *Default*: 1
 
+--------------------
+``<stride>`` Element
+--------------------
+
+The ``stride`` element is used to specify how many random numbers are allocated
+for each source particle history.
+
+  *Default*: 152,917
+
 .. _source_element:
 
 --------------------

docs/source/io_formats/statepoint.rst

@@ -23,6 +23,7 @@ The current version of the statepoint file format is 18.1.
                bank is present (1) or not (0).
 
 :Datasets: - **seed** (*int8_t*) -- Pseudo-random number generator seed.
+           - **stride** (*uint64_t*) -- Pseudo-random number generator stride.
            - **energy_mode** (*char[]*) -- Energy mode of the run, either
              'continuous-energy' or 'multi-group'.
            - **run_mode** (*char[]*) -- Run mode used, either 'eigenvalue' or

include/openmc/capi.h

@@ -71,6 +71,7 @@ int openmc_get_nuclide_index(const char name[], int* index);
 int openmc_add_unstructured_mesh(
   const char filename[], const char library[], int* id);
 int64_t openmc_get_seed();
+uint64_t openmc_get_stride();
 int openmc_get_tally_index(int32_t id, int32_t* index);
 void openmc_get_tally_next_id(int32_t* id);
 int openmc_global_tallies(double** ptr);
@@ -137,6 +138,7 @@ int openmc_reset_timers();
 int openmc_run();
 int openmc_sample_external_source(size_t n, uint64_t* seed, void* sites);
 void openmc_set_seed(int64_t new_seed);
+void openmc_set_stride(uint64_t new_stride);
 int openmc_set_n_batches(
   int32_t n_batches, bool set_max_batches, bool add_statepoint_batch);
 int openmc_simulation_finalize();

include/openmc/random_lcg.h

@@ -15,6 +15,7 @@ constexpr int STREAM_SOURCE {1};
 constexpr int STREAM_URR_PTABLE {2};
 constexpr int STREAM_VOLUME {3};
 constexpr int64_t DEFAULT_SEED {1};
+constexpr uint64_t DEFAULT_STRIDE {152917ULL};
 
 //==============================================================================
 //! Generate a pseudo-random number using a linear congruential generator.
@@ -98,5 +99,18 @@ extern "C" int64_t openmc_get_seed();
 
 extern "C" void openmc_set_seed(int64_t new_seed);
 
+//==============================================================================
+//! Get OpenMC's stride.
+//==============================================================================
+
+extern "C" uint64_t openmc_get_stride();
+
+//==============================================================================
+//! Set OpenMC's stride.
+//! @param new_stride Stride.
+//==============================================================================
+
+extern "C" void openmc_set_stride(uint64_t new_stride);
+
 } // namespace openmc
 #endif // OPENMC_RANDOM_LCG_H

openmc/lib/settings.py

@@ -12,6 +12,8 @@
 
 _dll.openmc_set_seed.argtypes = [c_int64]
 _dll.openmc_get_seed.restype = c_int64
+_dll.openmc_set_stride.argtypes = [c_int64]
+_dll.openmc_get_stride.restype = c_int64
 _dll.openmc_get_n_batches.argtypes = [POINTER(c_int), c_bool]
 _dll.openmc_get_n_batches.restype = c_int
 _dll.openmc_get_n_batches.errcheck = _error_handler
@@ -68,6 +70,14 @@ def seed(self):
     def seed(self, seed):
         _dll.openmc_set_seed(seed)
 
+    @property
+    def stride(self):
+        return _dll.openmc_get_stride()
+
+    @stride.setter
+    def stride(self, stride):
+        _dll.openmc_set_stride(stride)
+
     def set_batches(self, n_batches, set_max_batches=True, add_sp_batch=True):
         """Set number of batches or maximum number of batches
 

openmc/settings.py

@@ -193,6 +193,8 @@ class Settings:
         The type of calculation to perform (default is 'eigenvalue')
     seed : int
         Seed for the linear congruential pseudorandom number generator
+    stride : int
+        Number of random numbers allocated for each source particle history
     source : Iterable of openmc.SourceBase
         Distribution of source sites in space, angle, and energy
     sourcepoint : dict
@@ -338,6 +340,7 @@ def __init__(self, **kwargs):
         self._ptables = None
         self._uniform_source_sampling = None
         self._seed = None
+        self._stride = None
         self._survival_biasing = None
 
         # Shannon entropy mesh
@@ -614,6 +617,16 @@ def seed(self, seed: int):
         cv.check_greater_than('random number generator seed', seed, 0)
         self._seed = seed
 
+    @property
+    def stride(self) -> int:
+        return self._stride
+
+    @stride.setter
+    def stride(self, stride: int):
+        cv.check_type('random number generator stride', stride, Integral)
+        cv.check_greater_than('random number generator stride', stride, 0)
+        self._stride = stride
+
     @property
     def survival_biasing(self) -> bool:
         return self._survival_biasing
@@ -1336,6 +1349,11 @@ def _create_seed_subelement(self, root):
             element = ET.SubElement(root, "seed")
             element.text = str(self._seed)
 
+    def _create_stride_subelement(self, root):
+        if self._stride is not None:
+            element = ET.SubElement(root, "stride")
+            element.text = str(self._stride)
+
     def _create_survival_biasing_subelement(self, root):
         if self._survival_biasing is not None:
             element = ET.SubElement(root, "survival_biasing")
@@ -1763,6 +1781,11 @@ def _seed_from_xml_element(self, root):
         if text is not None:
             self.seed = int(text)
 
+    def _stride_from_xml_element(self, root):
+        text = get_text(root, 'stride')
+        if text is not None:
+            self.stride = int(text)
+
     def _survival_biasing_from_xml_element(self, root):
         text = get_text(root, 'survival_biasing')
         if text is not None:
@@ -2014,6 +2037,7 @@ def to_xml_element(self, mesh_memo=None):
         self._create_plot_seed_subelement(element)
         self._create_ptables_subelement(element)
         self._create_seed_subelement(element)
+        self._create_stride_subelement(element)
         self._create_survival_biasing_subelement(element)
         self._create_cutoff_subelement(element)
         self._create_entropy_mesh_subelement(element, mesh_memo)
@@ -2122,6 +2146,7 @@ def from_xml_element(cls, elem, meshes=None):
         settings._plot_seed_from_xml_element(elem)
         settings._ptables_from_xml_element(elem)
         settings._seed_from_xml_element(elem)
+        settings._stride_from_xml_element(elem)
         settings._survival_biasing_from_xml_element(elem)
         settings._cutoff_from_xml_element(elem)
         settings._entropy_mesh_from_xml_element(elem, meshes)

openmc/statepoint.py

@@ -99,6 +99,8 @@ class StatePoint:
         and whose values are time values in seconds.
     seed : int
         Pseudorandom number generator seed
+    stride : int
+        Number of random numbers allocated for each particle history
     source : numpy.ndarray of compound datatype
         Array of source sites. The compound datatype has fields 'r', 'u',
         'E', 'wgt', 'delayed_group', 'surf_id', and 'particle', corresponding to
@@ -356,6 +358,10 @@ def runtime(self):
     def seed(self):
         return self._f['seed'][()]
 
+    @property
+    def stride(self):
+        return self._f['stride'][()]
+
     @property
     def source(self):
         return self._f['source_bank'][()] if self.source_present else None

src/finalize.cpp

@@ -159,6 +159,7 @@ int openmc_finalize()
   model::root_universe = -1;
   model::plotter_seed = 1;
   openmc::openmc_set_seed(DEFAULT_SEED);
+  openmc::openmc_set_stride(DEFAULT_STRIDE);
 
   // Deallocate arrays
   free_memory();
@@ -221,5 +222,6 @@ int openmc_hard_reset()
 
   // Reset the random number generator state
   openmc::openmc_set_seed(DEFAULT_SEED);
+  openmc::openmc_set_stride(DEFAULT_STRIDE);
   return 0;
 }

src/initialize.cpp

@@ -99,9 +99,10 @@ int openmc_init(int argc, char* argv[], const void* intracomm)
   }
 #endif
 
-  // Initialize random number generator -- if the user specifies a seed, it
-  // will be re-initialized later
+  // Initialize random number generator -- if the user specifies a seed and/or
+  // stride, it will be re-initialized later
   openmc::openmc_set_seed(DEFAULT_SEED);
+  openmc::openmc_set_stride(DEFAULT_STRIDE);
 
   // Copy previous locale and set locale to C. This is a workaround for an issue
   // whereby when openmc_init is called from the plotter, the Qt application

src/random_lcg.cpp

@@ -10,7 +10,7 @@ int64_t master_seed {1};
 // LCG parameters
 constexpr uint64_t prn_mult {6364136223846793005ULL}; // multiplication
 constexpr uint64_t prn_add {1442695040888963407ULL};  // additive factor, c
-constexpr uint64_t prn_stride {152917LL}; // stride between particles
+uint64_t prn_stride {DEFAULT_STRIDE}; // stride between particles
 
 //==============================================================================
 // PRN
@@ -133,4 +133,14 @@ extern "C" void openmc_set_seed(int64_t new_seed)
   master_seed = new_seed;
 }
 
+extern "C" uint64_t openmc_get_stride()
+{
+  return prn_stride;
+}
+
+extern "C" void openmc_set_stride(uint64_t new_stride)
+{
+  prn_stride = new_stride;
+}
+
 } // namespace openmc

src/settings.cpp

@@ -514,6 +514,12 @@ void read_settings_xml(pugi::xml_node root)
     openmc_set_seed(seed);
   }
 
+  // Copy random number stride if specified
+  if (check_for_node(root, "stride")) {
+    auto stride = std::stoull(get_node_value(root, "stride"));
+    openmc_set_stride(stride);
+  }
+
   // Check for electron treatment
   if (check_for_node(root, "electron_treatment")) {
     auto temp_str = get_node_value(root, "electron_treatment", true, true);

src/state_point.cpp

@@ -89,6 +89,9 @@ extern "C" int openmc_statepoint_write(const char* filename, bool* write_source)
     // Write out random number seed
     write_dataset(file_id, "seed", openmc_get_seed());
 
+    // Write out random number stride
+    write_dataset(file_id, "stride", openmc_get_stride());
+
     // Write run information
     write_dataset(file_id, "energy_mode",
       settings::run_CE ? "continuous-energy" : "multi-group");
@@ -399,6 +402,11 @@ extern "C" int openmc_statepoint_load(const char* filename)
   read_dataset(file_id, "seed", seed);
   openmc_set_seed(seed);
 
+  // Read and overwrite random number stride
+  uint64_t stride;
+  read_dataset(file_id, "stride", stride);
+  openmc_set_stride(stride);
+
   // It is not impossible for a state point to be generated from a CE run but
   // to be loaded in to an MG run (or vice versa), check to prevent that.
   read_dataset(file_id, "energy_mode", word);

tests/regression_tests/stride/__init__.py

@@ -0,0 +1,25 @@
+<?xml version='1.0' encoding='utf-8'?>
+<model>
+  <materials>
+    <material depletable="true" id="1">
+      <density units="g/cm3" value="4.5"/>
+      <nuclide ao="1.0" name="U235"/>
+    </material>
+  </materials>
+  <geometry>
+    <cell id="1" material="1" region="-1" universe="1"/>
+    <surface boundary="vacuum" coeffs="0.0 0.0 0.0 10.0" id="1" type="sphere"/>
+  </geometry>
+  <settings>
+    <run_mode>eigenvalue</run_mode>
+    <particles>1000</particles>
+    <batches>10</batches>
+    <inactive>5</inactive>
+    <source particle="neutron" strength="1.0" type="independent">
+      <space type="box">
+        <parameters>-4 -4 -4 4 4 4</parameters>
+      </space>
+    </source>
+    <stride>1529170</stride>
+  </settings>
+  </model>

tests/regression_tests/stride/inputs_true.dat

@@ -0,0 +1,2 @@
+k-combined:
+2.978080E-01 6.106774E-03

tests/regression_tests/stride/results_true.dat

@@ -0,0 +1,29 @@
+import pytest
+import openmc
+
+from tests.testing_harness import PyAPITestHarness
+
+
+@pytest.fixture
+def model():
+    u = openmc.Material()
+    u.add_nuclide('U235', 1.0)
+    u.set_density('g/cm3', 4.5)
+    sph = openmc.Sphere(r=10.0, boundary_type='vacuum')
+    cell = openmc.Cell(fill=u, region=-sph)
+    model = openmc.Model()
+    model.geometry = openmc.Geometry([cell])
+
+    model.settings.batches = 10
+    model.settings.inactive = 5
+    model.settings.particles = 1000
+    model.settings.stride = 1_529_170
+    model.settings.source = openmc.IndependentSource(
+        space=openmc.stats.Box([-4, -4, -4], [4, 4, 4])
+    )
+    return model
+
+
+def test_seed(model):
+    harness = PyAPITestHarness('statepoint.10.h5', model)
+    harness.main()