Set n_batches and n_max_batches through C API

docs/source/capi/index.rst

@@ -452,6 +452,16 @@ Functions
    :return: Return status (negative if an error occurs)
    :rtype: int
 
+.. c:function:: int openmc_set_n_batches(int32_t n_batches, int32_t n_max_batches, int32_t add_statepoint_batch)
+
+   Set number of batches and number of max batches
+
+   :param int32_t n_batches: Number of batches to simulate
+   :param int32_t n_batches: Maximum number of batches to simulate (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

@@ -89,6 +89,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, int32_t n_max_batches,
+                           bool add_statepoint_batch);
   int openmc_simulation_finalize();
   int openmc_simulation_init();
   int openmc_source_bank(void** ptr, int64_t* n);

openmc/lib/settings.py

@@ -2,6 +2,7 @@ from ctypes import c_int, c_int32, c_int64, c_double, c_char_p, c_bool
 
 from . import _dll
 from .core import _DLLGlobal
+from .error import _error_handler
 
 _RUN_MODES = {1: 'fixed source',
               2: 'eigenvalue',
@@ -11,11 +12,13 @@ _RUN_MODES = {1: 'fixed source',
 
 _dll.openmc_set_seed.argtypes = [c_int64]
 _dll.openmc_get_seed.restype = c_int64
+_dll.openmc_set_n_batches.argtypes = [c_int32, c_int32, 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 +62,23 @@ class _Settings:
     def seed(self, seed):
         _dll.openmc_set_seed(seed)
 
+    def set_n_batches(self, n_batches, n_max_batches=None, add_sp_batch=True):
+        """Set n_batches and n_max_batches
+
+        Parameters
+        ----------
+        n_batches : int
+            Number of batches to simulate
+        n_max_batches : int
+            Maximum number of batches. Only has an effect when triggers are used
+        add_sp_batch : bool
+            Whether to add `n_batches` as a statepoint batch
+
+        """
+        if not n_max_batches:
+            _dll.openmc_set_n_batches(n_batches, n_batches, add_sp_batch)
+        else:
+            _dll.openmc_set_n_batches(n_batches, n_max_batches, add_sp_batch)
+
 
 settings = _Settings()

src/settings.cpp

@@ -827,4 +827,48 @@ void free_memory_settings() {
   settings::res_scat_nuclides.clear();
 }
 
+//==============================================================================
+// C API functions
+//==============================================================================
+
+extern "C" int
+openmc_set_n_batches(int32_t n_batches, int32_t n_max_batches,
+                     bool add_statepoint_batch)
+{
+  if (settings::n_inactive >= n_batches ||
+      settings::n_inactive >= n_max_batches) {
+    set_errmsg("Number of active batches must be greater than zero.");
+    return OPENMC_E_INVALID_ARGUMENT;
+  }
+
+  if (simulation::current_batch >= n_batches ||
+      simulation::current_batch >= n_max_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 separately
+    settings::n_batches = n_batches;
+    settings::n_max_batches = n_max_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;
+}
+
 } // namespace openmc