Merge pull request #1578 from shikhar413/lib_batches_setter

Python bindings for setting n_batches through C API

docs/source/capi/index.rst

@@ -266,6 +266,15 @@ Functions
    :return: Return status (negative if an error occurs)
    :rtype: int
 
+.. c:function:: int openmc_get_n_batches(int* n_batches, bool get_max_batches)
+
+   Get number of batches to simulate
+
+   :param int* n_batches: Number of batches to simulate
+   :param bool get_max_batches: Whether to return `n_batches` or `n_max_batches` (only relevant when triggers are used)
+   :return: Return status (negative if an error occurred)
+   :rtype: int
+
 .. c:function:: int openmc_get_nuclide_index(const char name[], int* index)
 
    Get the index in the nuclides array for a nuclide with a given name
@@ -452,6 +461,16 @@ Functions
    :return: Return status (negative if an error occurs)
    :rtype: int
 
+.. c:function:: int openmc_set_n_batches(int32_t n_batches, bool set_max_batches, bool add_statepoint_batch)
+
+   Set number of batches and number of max batches
+
+   :param int32_t n_batches: Number of batches to simulate
+   :param bool set_max_batches: Whether to set `settings::n_max_batches` or `settings::n_batches` (only relevant when triggers are used)
+   :param bool add_statepoint_batch: Whether to add `n_batches` to `settings::statepoint_batch`
+   :return: Return status (negative if an error occurred)
+   :rtype: int
+
 .. c:function:: int openmc_simulation_finalize()
 
    Finalize a simulation.

include/openmc/capi.h

@@ -44,6 +44,7 @@ extern "C" {
   int openmc_get_keff(double k_combined[]);
   int openmc_get_material_index(int32_t id, int32_t* index);
   int openmc_get_mesh_index(int32_t id, int32_t* index);
+  int openmc_get_n_batches(int* n_batches, bool get_max_batches);
   int openmc_get_nuclide_index(const char name[], int* index);
   int64_t openmc_get_seed();
   int openmc_get_tally_index(int32_t id, int32_t* index);
@@ -89,6 +90,8 @@ extern "C" {
   int openmc_reset_timers();
   int openmc_run();
   void openmc_set_seed(int64_t new_seed);
+  int openmc_set_n_batches(int32_t n_batches, bool set_max_batches,
+                           bool add_statepoint_batch);
   int openmc_simulation_finalize();
   int openmc_simulation_init();
   int openmc_source_bank(void** ptr, int64_t* n);

include/openmc/settings.h

@@ -64,7 +64,6 @@ extern std::string path_source_library;   //!< path to the source shared object
 extern std::string path_sourcepoint;      //!< path to a source file
 extern "C" std::string path_statepoint;   //!< path to a statepoint file
 
-extern "C" int32_t n_batches;                //!< number of (inactive+active) batches
 extern "C" int32_t n_inactive;               //!< number of inactive batches
 extern "C" int32_t max_lost_particles;     //!< maximum number of lost particles
 extern double rel_max_lost_particles;   //!< maximum number of lost particles, relative to the total number of particles
@@ -79,6 +78,7 @@ extern std::array<double, 4> energy_cutoff;  //!< Energy cutoff in [eV] for each
 extern int legendre_to_tabular_points; //!< number of points to convert Legendres
 extern int max_order;                //!< Maximum Legendre order for multigroup data
 extern int n_log_bins;               //!< number of bins for logarithmic energy grid
+extern int n_batches;                //!< number of (inactive+active) batches
 extern int n_max_batches;            //!< Maximum number of batches
 extern ResScatMethod res_scat_method; //!< resonance upscattering method
 extern double res_scat_energy_min;   //!< Min energy in [eV] for res. upscattering

openmc/cmfd.py

@@ -830,7 +830,7 @@ class CMFDRun:
 
         """
         if filename is None:
-            batch_str_len = len(str(openmc.lib.settings.batches))
+            batch_str_len = len(str(openmc.lib.settings.get_batches()))
             batch_str = str(openmc.lib.current_batch()).zfill(batch_str_len)
             filename = 'statepoint.{}.h5'.format(batch_str)
 
@@ -1164,8 +1164,8 @@ class CMFDRun:
                 self._k_cmfd.append(self._keff)
 
                 # Check to perform adjoint on last batch
-                if (openmc.lib.current_batch() == openmc.lib.settings.batches
-                        and self._run_adjoint):
+                batches = openmc.lib.settings.get_batches()
+                if openmc.lib.current_batch() == batches and self._run_adjoint:
                     self._cmfd_solver_execute(adjoint=True)
 
                 # Calculate fission source

openmc/lib/settings.py

@@ -1,7 +1,8 @@
-from ctypes import c_int, c_int32, c_int64, c_double, c_char_p, c_bool
+from ctypes import c_int, c_int32, c_int64, c_double, c_char_p, c_bool, POINTER
 
 from . import _dll
 from .core import _DLLGlobal
+from .error import _error_handler
 
 _RUN_MODES = {1: 'fixed source',
               2: 'eigenvalue',
@@ -11,11 +12,16 @@ _RUN_MODES = {1: 'fixed source',
 
 _dll.openmc_set_seed.argtypes = [c_int64]
 _dll.openmc_get_seed.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
+_dll.openmc_set_n_batches.argtypes = [c_int32, c_bool, c_bool]
+_dll.openmc_set_n_batches.restype = c_int
+_dll.openmc_set_n_batches.errcheck = _error_handler
 
 
 class _Settings:
     # Attributes that are accessed through a descriptor
-    batches = _DLLGlobal(c_int32, 'n_batches')
     cmfd_run = _DLLGlobal(c_bool, 'cmfd_run')
     entropy_on = _DLLGlobal(c_bool, 'entropy_on')
     generations_per_batch = _DLLGlobal(c_int32, 'gen_per_batch')
@@ -59,5 +65,42 @@ class _Settings:
     def seed(self, seed):
         _dll.openmc_set_seed(seed)
 
+    def set_batches(self, n_batches, set_max_batches=True, add_sp_batch=True):
+        """Set number of batches or maximum number of batches
+
+        Parameters
+        ----------
+        n_batches : int
+            Number of batches to simulate
+        set_max_batches : bool
+            Whether to set maximum number of batches. If True, the value of
+            `n_max_batches` is overridden, otherwise the value of `n_batches`
+            is overridden. Only has an effect when triggers are used
+        add_sp_batch : bool
+            Whether to add `n_batches` as a statepoint batch
+
+        """
+        _dll.openmc_set_n_batches(n_batches, set_max_batches, add_sp_batch)
+
+    def get_batches(self, get_max_batches=True):
+        """Get number of batches or maximum number of batches
+
+        Parameters
+        ----------
+        get_max_batches : bool
+            Return `n_max_batches` if true, else return `n_batches`. Difference
+            arises only if triggers are used.
+
+        Returns
+        -------
+        int
+            Number of batches to simulate
+
+        """
+        n_batches = c_int()
+        _dll.openmc_get_n_batches(n_batches, get_max_batches)
+
+        return n_batches.value
+
 
 settings = _Settings()

src/finalize.cpp

@@ -162,6 +162,8 @@ int openmc_reset()
   simulation::k_abs_tra = 0.0;
   simulation::k_sum = {0.0, 0.0};
 
+  settings::cmfd_run = false;
+
   return 0;
 }
 

src/settings.cpp

@@ -78,7 +78,6 @@ std::string path_source_library;
 std::string path_sourcepoint;
 std::string path_statepoint;
 
-int32_t n_batches;
 int32_t n_inactive {0};
 int32_t max_lost_particles {10};
 double rel_max_lost_particles {1.0e-6};
@@ -92,6 +91,7 @@ std::array<double, 4> energy_cutoff {0.0, 1000.0, 0.0, 0.0};
 int legendre_to_tabular_points {C_NONE};
 int max_order {0};
 int n_log_bins {8000};
+int n_batches;
 int n_max_batches;
 ResScatMethod res_scat_method {ResScatMethod::rvs};
 double res_scat_energy_min {0.01};
@@ -827,4 +827,56 @@ void free_memory_settings() {
   settings::res_scat_nuclides.clear();
 }
 
+//==============================================================================
+// C API functions
+//==============================================================================
+
+extern "C" int
+openmc_set_n_batches(int32_t n_batches, bool set_max_batches,
+                     bool add_statepoint_batch)
+{
+  if (settings::n_inactive >= n_batches) {
+    set_errmsg("Number of active batches must be greater than zero.");
+    return OPENMC_E_INVALID_ARGUMENT;
+  }
+
+  if (simulation::current_batch >= n_batches) {
+    set_errmsg("Number of batches must be greater than current batch.");
+    return OPENMC_E_INVALID_ARGUMENT;
+  }
+
+  if (!settings::trigger_on) {
+    // Set n_batches and n_max_batches to same value
+    settings::n_batches = n_batches;
+    settings::n_max_batches = n_batches;
+  } else {
+    // Set n_batches and n_max_batches based on value of set_max_batches
+    if (set_max_batches) {
+      settings::n_max_batches = n_batches;
+    } else {
+      settings::n_batches = n_batches;
+    }
+  }
+
+  // Update size of k_generation and entropy
+  int m = settings::n_max_batches * settings::gen_per_batch;
+  simulation::k_generation.reserve(m);
+  simulation::entropy.reserve(m);
+
+  // Add value of n_batches to statepoint_batch
+  if (add_statepoint_batch &&
+      !(contains(settings::statepoint_batch, n_batches)))
+    settings::statepoint_batch.insert(n_batches);
+
+  return 0;
+}
+
+extern "C" int
+openmc_get_n_batches(int* n_batches, bool get_max_batches)
+{
+  *n_batches = get_max_batches ? settings::n_max_batches : settings::n_batches;
+
+  return 0;
+}
+
 } // namespace openmc

tests/unit_tests/test_lib.py

@@ -49,6 +49,37 @@ def pincell_model():
         yield
 
 
+@pytest.fixture(scope='module')
+def uo2_trigger_model():
+    """Set up a simple UO2 model with k-eff trigger"""
+    model = openmc.model.Model()
+    m = openmc.Material(name='UO2')
+    m.add_nuclide('U235', 1.0)
+    m.add_nuclide('O16', 2.0)
+    m.set_density('g/cm3', 10.0)
+    model.materials.append(m)
+
+    cyl = openmc.ZCylinder(r=1.0, boundary_type='vacuum')
+    c = openmc.Cell(fill=m, region=-cyl)
+    model.geometry.root_universe = openmc.Universe(cells=[c])
+
+    model.settings.batches = 10
+    model.settings.inactive = 5
+    model.settings.particles = 100
+    model.settings.source = openmc.Source(space=openmc.stats.Box(
+        [-0.5, -0.5, -1], [0.5, 0.5, 1], only_fissionable=True))
+    model.settings.verbosity = 1
+    model.settings.keff_trigger = {'type': 'std_dev', 'threshold': 0.001}
+    model.settings.trigger_active = True
+    model.settings.trigger_max_batches = 10
+    model.settings.trigger_batch_interval = 1
+
+    # Write XML files in tmpdir
+    with cdtemp():
+        model.export_to_xml()
+        yield
+
+
 @pytest.fixture(scope='module')
 def lib_init(pincell_model, mpi_intracomm):
     openmc.lib.init(intracomm=mpi_intracomm)
@@ -162,8 +193,6 @@ def test_nuclide_mapping(lib_init):
 
 def test_settings(lib_init):
     settings = openmc.lib.settings
-    assert settings.batches == 10
-    settings.batches = 10
     assert settings.inactive == 5
     assert settings.generations_per_batch == 1
     assert settings.particles == 100
@@ -318,6 +347,41 @@ def test_by_batch(lib_run):
         openmc.lib.simulation_finalize()
 
 
+def test_set_n_batches(lib_run):
+    # Run simulation_init so that current_batch reset to 0
+    openmc.lib.hard_reset()
+    openmc.lib.simulation_init()
+
+    settings = openmc.lib.settings
+    assert settings.get_batches() == 10
+
+    # Setting n_batches less than n_inactive should raise error
+    with pytest.raises(exc.InvalidArgumentError):
+        settings.set_batches(3)
+    # n_batches should stay the same
+    assert settings.get_batches() == 10
+
+    for i in range(7):
+        openmc.lib.next_batch()
+    # Setting n_batches less than current_batch should raise error
+    with pytest.raises(exc.InvalidArgumentError):
+        settings.set_batches(6)
+    # n_batches should stay the same
+    assert settings.get_batches() == 10
+
+    # Change n_batches from 10 to 20
+    settings.set_batches(20)
+    for _ in openmc.lib.iter_batches():
+        pass
+    openmc.lib.simulation_finalize()
+
+    # n_active should have been overwritten from 5 to 15
+    assert openmc.lib.num_realizations() == 15
+
+    # Ensure statepoint created at new value of n_batches
+    assert os.path.exists('statepoint.20.h5')
+
+
 def test_reset(lib_run):
     # Init and run 10 batches.
     openmc.lib.hard_reset()
@@ -515,3 +579,28 @@ def test_global_bounding_box(lib_init):
 
     assert tuple(llc) == expected_llc
     assert tuple(urc) == expected_urc
+
+
+def test_trigger_set_n_batches(uo2_trigger_model, mpi_intracomm):
+    openmc.lib.finalize()
+    openmc.lib.init(intracomm=mpi_intracomm)
+    openmc.lib.simulation_init()
+
+    settings = openmc.lib.settings
+    # Change n_batches to 12 and n_max_batches to 20
+    settings.set_batches(12, set_max_batches=False, add_sp_batch=False)
+    settings.set_batches(20, set_max_batches=True, add_sp_batch=True)
+
+    assert settings.get_batches(get_max_batches=False) == 12
+    assert settings.get_batches(get_max_batches=True) == 20
+
+    for _ in openmc.lib.iter_batches():
+        pass
+    openmc.lib.simulation_finalize()
+
+    # n_active should have been overwritten from 5 to 15
+    assert openmc.lib.num_realizations() == 15
+
+    # Ensure statepoint was created only at batch 20 when calling set_batches
+    assert not os.path.exists('statepoint.12.h5')
+    assert os.path.exists('statepoint.20.h5')