Statepoint file loading refactor and CAPI function (#2886)

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

docs/source/pythonapi/capi.rst

@@ -43,6 +43,7 @@ Functions
    simulation_finalize
    simulation_init
    source_bank
+   statepoint_load
    statepoint_write
 
 Classes

docs/source/usersguide/settings.rst

@@ -628,3 +628,37 @@ instance, whereas the :meth:`openmc.Track.filter` method returns a new
           .. code-block:: sh
 
             openmc-track-combine tracks_p*.h5 --out tracks.h5
+
+-----------------------
+Restarting a Simulation
+-----------------------
+
+OpenMC can be run in a mode where it reads in a statepoint file and continues a
+simulation from the ending point of the statepoint file. A restart simulation
+can be performed by passing the path to the statepoint file to the OpenMC
+executable:
+
+.. code-block:: sh
+
+    openmc -r statepoint.100.h5
+
+From the Python API, the `restart_file` argument provides the same behavior:
+
+.. code-block:: python
+
+    openmc.run(restart_file='statepoint.100.h5')
+
+or if using the :class:`~openmc.Model` class:
+
+.. code-block:: python
+
+    model.run(restart_file='statepoint.100.h5')
+
+The restart simulation will execute until the number of batches specified in the
+:class:`~openmc.Settings` object on a model (or in the :ref:`settings XML file
+<io_settings>`) is satisfied. Note that if the number of batches in the
+statepoint file is the same as that specified in the settings object (i.e., if
+the inputs were not modified before the restart run), no particles will be
+transported and OpenMC will exit immediately.
+
+.. note:: A statepoint file must match the input model to be successfully used in a restart simulation.

include/openmc/capi.h

@@ -150,6 +150,7 @@ int openmc_sphharm_filter_get_cosine(int32_t index, char cosine[]);
 int openmc_sphharm_filter_set_order(int32_t index, int order);
 int openmc_sphharm_filter_set_cosine(int32_t index, const char cosine[]);
 int openmc_statepoint_write(const char* filename, bool* write_source);
+int openmc_statepoint_load(const char* filename);
 int openmc_tally_allocate(int32_t index, const char* type);
 int openmc_tally_get_active(int32_t index, bool* active);
 int openmc_tally_get_estimator(int32_t index, int* estimator);

openmc/lib/core.py

@@ -95,6 +95,11 @@ _dll.openmc_simulation_finalize.errcheck = _error_handler
 _dll.openmc_statepoint_write.argtypes = [c_char_p, POINTER(c_bool)]
 _dll.openmc_statepoint_write.restype = c_int
 _dll.openmc_statepoint_write.errcheck = _error_handler
+_dll.openmc_statepoint_load.argtypes = [c_char_p]
+_dll.openmc_statepoint_load.restype = c_int
+_dll.openmc_statepoint_load.errcheck = _error_handler
+_dll.openmc_statepoint_write.restype = c_int
+_dll.openmc_statepoint_write.errcheck = _error_handler
 _dll.openmc_global_bounding_box.argtypes = [POINTER(c_double),
                                             POINTER(c_double)]
 _dll.openmc_global_bounding_box.restype = c_int
@@ -568,6 +573,19 @@ def statepoint_write(filename=None, write_source=True):
     _dll.openmc_statepoint_write(filename, c_bool(write_source))
 
 
+def statepoint_load(filename: PathLike):
+    """Load a statepoint file.
+
+    Parameters
+    ----------
+    filename : path-like
+        Path to the statepoint to load.
+
+    """
+    filename = c_char_p(str(filename).encode())
+    _dll.openmc_statepoint_load(filename)
+
+
 @contextmanager
 def run_in_memory(**kwargs):
     """Provides context manager for calling OpenMC shared library functions.

src/simulation.cpp

@@ -54,8 +54,14 @@ int openmc_run()
   openmc::simulation::time_total.start();
   openmc_simulation_init();
 
-  int err = 0;
+  // Ensure that a batch isn't executed in the case that the maximum number of
+  // batches has already been run in a restart statepoint file
   int status = 0;
+  if (openmc::simulation::current_batch >= openmc::settings::n_max_batches) {
+    status = openmc::STATUS_EXIT_MAX_BATCH;
+  }
+
+  int err = 0;
   while (status == 0 && err == 0) {
     err = openmc_next_batch(&status);
   }

src/state_point.cpp

@@ -364,11 +364,27 @@ void restart_set_keff()
 
 void load_state_point()
 {
-  // Write message
-  write_message("Loading state point " + settings::path_statepoint + "...", 5);
+  write_message(
+    fmt::format("Loading state point {}...", settings::path_statepoint_c), 5);
+  openmc_statepoint_load(settings::path_statepoint.c_str());
+}
 
+void statepoint_version_check(hid_t file_id)
+{
+  // Read revision number for state point file and make sure it matches with
+  // current version
+  array<int, 2> version_array;
+  read_attribute(file_id, "version", version_array);
+  if (version_array != VERSION_STATEPOINT) {
+    fatal_error(
+      "State point version does not match current version in OpenMC.");
+  }
+}
+
+extern "C" int openmc_statepoint_load(const char* filename)
+{
   // Open file for reading
-  hid_t file_id = file_open(settings::path_statepoint.c_str(), 'r', true);
+  hid_t file_id = file_open(filename, 'r', true);
 
   // Read filetype
   std::string word;
@@ -377,14 +393,7 @@ void load_state_point()
     fatal_error("OpenMC tried to restart from a non-statepoint file.");
   }
 
-  // Read revision number for state point file and make sure it matches with
-  // current version
-  array<int, 2> array;
-  read_attribute(file_id, "version", array);
-  if (array != VERSION_STATEPOINT) {
-    fatal_error(
-      "State point version does not match current version in OpenMC.");
-  }
+  statepoint_version_check(file_id);
 
   // Read and overwrite random number seed
   int64_t seed;
@@ -421,9 +430,10 @@ void load_state_point()
   read_dataset(file_id, "current_batch", simulation::restart_batch);
 
   if (simulation::restart_batch >= settings::n_max_batches) {
-    fatal_error(fmt::format(
-      "The number of batches specified for simulation ({}) is smaller"
-      " than the number of batches in the restart statepoint file ({})",
+    warning(fmt::format(
+      "The number of batches specified for simulation ({}) is smaller "
+      "than or equal to the number of batches in the restart statepoint file "
+      "({})",
       settings::n_max_batches, simulation::restart_batch));
   }
 
@@ -489,7 +499,6 @@ void load_state_point()
         if (internal) {
           tally->writable_ = false;
         } else {
-
           auto& results = tally->results_;
           read_tally_results(tally_group, results.shape()[0],
             results.shape()[1], results.data());
@@ -497,7 +506,6 @@ void load_state_point()
           close_group(tally_group);
         }
       }
-
       close_group(tallies_group);
     }
   }
@@ -525,6 +533,8 @@ void load_state_point()
 
   // Close file
   file_close(file_id);
+
+  return 0;
 }
 
 hid_t h5banktype()

tests/regression_tests/statepoint_restart/test.py

@@ -1,7 +1,6 @@
 from pathlib import Path
 
 import openmc
-import pytest
 
 from tests.testing_harness import TestHarness
 from tests.regression_tests import config
@@ -61,24 +60,35 @@ def test_statepoint_restart():
     harness.main()
 
 
-def test_batch_check(request):
+def test_batch_check(request, capsys):
     xmls = list(request.path.parent.glob('*.xml'))
 
     with cdtemp(xmls):
         model = openmc.Model.from_xml()
         model.settings.particles = 100
+
         # run the model
-        sp_file = model.run()
+        sp_file = model.run(export_model_xml=False)
+        assert sp_file is not None
 
         # run a restart with the resulting statepoint
         # and the settings unchanged
-        with pytest.raises(RuntimeError, match='is smaller than the number of batches'):
-            model.run(restart_file=sp_file)
+        model.settings.batches = 6
+        # ensure we capture output only from the next run
+        capsys.readouterr()
+        sp_file = model.run(export_model_xml=False, restart_file=sp_file)
+        # indicates that a new statepoint file was not created
+        assert sp_file is None
 
-        # update the number of batches and run again
+        output = capsys.readouterr().out
+        assert "WARNING" in output
+        assert "The number of batches specified for simulation" in output
+
+        # update the number of batches and run again,
+        # this restart run should be successful
         model.settings.batches = 15
         model.settings.statepoint = {}
-        sp_file = model.run(restart_file=sp_file)
+        sp_file = model.run(export_model_xml=False, restart_file=sp_file)
 
         sp = openmc.StatePoint(sp_file)
         assert sp.n_batches == 15