Merge pull request #1039 from smharper/capi_sourcepoint

Write statepoints with source banks from openmc.capi

openmc/capi/core.py

@@ -1,6 +1,6 @@
 from contextlib import contextmanager
-from ctypes import (CDLL, c_int, c_int32, c_int64, c_double, c_char_p, c_char,
-                    POINTER, Structure, c_void_p, create_string_buffer)
+from ctypes import (CDLL, c_bool, c_int, c_int32, c_int64, c_double, c_char_p,
+                    c_char, POINTER, Structure, c_void_p, create_string_buffer)
 from warnings import warn
 
 import numpy as np
@@ -52,7 +52,7 @@ _dll.openmc_simulation_init.restype = c_int
 _dll.openmc_simulation_init.errcheck = _error_handler
 _dll.openmc_simulation_finalize.restype = c_int
 _dll.openmc_simulation_finalize.errcheck = _error_handler
-_dll.openmc_statepoint_write.argtypes = [POINTER(c_char_p)]
+_dll.openmc_statepoint_write.argtypes = [POINTER(c_char_p), POINTER(c_bool)]
 _dll.openmc_statepoint_write.restype = c_int
 _dll.openmc_statepoint_write.errcheck = _error_handler
 
@@ -269,7 +269,7 @@ def source_bank():
     return as_array(ptr, (n.value,)).view(bank_dtype)
 
 
-def statepoint_write(filename=None):
+def statepoint_write(filename=None, write_source=True):
     """Write a statepoint file.
 
     Parameters
@@ -277,11 +277,13 @@ def statepoint_write(filename=None):
     filename : str or None
         Path to the statepoint to write. If None is passed, a default name that
         contains the current batch will be written.
+    write_source : bool
+        Whether or not to include the source bank in the statepoint.
 
     """
     if filename is not None:
         filename = c_char_p(filename.encode())
-    _dll.openmc_statepoint_write(filename)
+    _dll.openmc_statepoint_write(filename, c_bool(write_source))
 
 
 @contextmanager

src/simulation.F90

@@ -317,6 +317,7 @@ contains
 #ifdef OPENMC_MPI
     integer :: mpi_err ! MPI error code
 #endif
+    character(MAX_FILE_LEN) :: filename
 
     ! Reduce tallies onto master process and accumulate
     call time_tallies % start()
@@ -349,13 +350,22 @@ contains
 
     ! Write out state point if it's been specified for this batch
     if (statepoint_batch % contains(current_batch)) then
-      err = openmc_statepoint_write()
+      if (sourcepoint_batch % contains(current_batch) .and. source_write &
+           .and. .not. source_separate) then
+        err = openmc_statepoint_write(write_source=.true._C_BOOL)
+      else
+        err = openmc_statepoint_write(write_source=.false._C_BOOL)
+      end if
     end if
 
-    ! Write out source point if it's been specified for this batch
-    if ((sourcepoint_batch % contains(current_batch) .or. source_latest) .and. &
-         source_write) then
-      call write_source_point()
+    ! Write out a separate source point if it's been specified for this batch
+    if (sourcepoint_batch % contains(current_batch) .and. source_write &
+         .and. source_separate) call write_source_point()
+
+    ! Write a continously-overwritten source point if requested.
+    if (source_latest) then
+      filename = trim(path_output) // 'source' // '.h5'
+      call write_source_point(filename)
     end if
 
   end subroutine finalize_batch

src/state_point.F90

@@ -58,10 +58,12 @@ contains
 ! OPENMC_STATEPOINT_WRITE writes an HDF5 statepoint file to disk
 !===============================================================================
 
-  function openmc_statepoint_write(filename) result(err) bind(C)
-    type(C_PTR), intent(in), optional :: filename
+  function openmc_statepoint_write(filename, write_source) result(err) bind(C)
+    type(C_PTR),     intent(in), optional :: filename
+    logical(C_BOOL), intent(in), optional :: write_source
     integer(C_INT) :: err
 
+    logical :: write_source_
     integer :: i, j, k
     integer :: i_xs
     integer, allocatable :: id_array(:)
@@ -69,13 +71,17 @@ contains
     integer(HID_T) :: cmfd_group, tallies_group, tally_group, meshes_group, &
                       filters_group, filter_group, derivs_group, &
                       deriv_group, runtime_group
+    integer(C_INT) :: ignored_err
     real(C_DOUBLE) :: k_combined(2)
     character(MAX_WORD_LEN), allocatable :: str_array(:)
     character(C_CHAR), pointer :: string(:)
     character(len=:, kind=C_CHAR), allocatable :: filename_
     character(MAX_WORD_LEN, kind=C_CHAR) :: temp_name
+    logical :: parallel
 
     err = 0
+
+    ! Set the filename
     if (present(filename)) then
       call c_f_pointer(filename, string, [MAX_FILE_LEN])
       filename_ = to_f_string(string)
@@ -86,6 +92,13 @@ contains
       filename_ = trim(filename_) // '.h5'
     end if
 
+    ! Determine whether or not to write the source bank
+    if (present(write_source)) then
+      write_source_ = write_source
+    else
+      write_source_ = .true.
+    end if
+
     ! Write message
     call write_message("Creating state point " // trim(filename_) // "...", 5)
 
@@ -138,10 +151,10 @@ contains
       call write_dataset(file_id, "current_batch", current_batch)
 
       ! Indicate whether source bank is stored in statepoint
-      if (source_separate) then
-        call write_attribute(file_id, "source_present", 0)
-      else
+      if (write_source_) then
         call write_attribute(file_id, "source_present", 1)
+      else
+        call write_attribute(file_id, "source_present", 0)
       end if
 
       ! Write out information for eigenvalue run
@@ -156,7 +169,7 @@ contains
         call write_dataset(file_id, "k_col_abs", k_col_abs)
         call write_dataset(file_id, "k_col_tra", k_col_tra)
         call write_dataset(file_id, "k_abs_tra", k_abs_tra)
-        err = openmc_get_keff(k_combined)
+        ignored_err = openmc_get_keff(k_combined)
         call write_dataset(file_id, "k_combined", k_combined)
 
         ! Write out CMFD info
@@ -432,17 +445,33 @@ contains
 
       call file_close(file_id)
     end if
+
+#ifdef PHDF5
+    parallel = .true.
+#else
+    parallel = .false.
+#endif
+
+    ! Write the source bank if desired
+    if (write_source_) then
+      if (master .or. parallel) then
+        file_id = file_open(filename_, 'a', parallel=.true.)
+      end if
+      call write_source_bank(file_id, work_index, source_bank)
+      if (master .or. parallel) call file_close(file_id)
+    end if
   end function openmc_statepoint_write
 
 !===============================================================================
 ! WRITE_SOURCE_POINT
 !===============================================================================
 
-  subroutine write_source_point()
+  subroutine write_source_point(filename)
+    character(MAX_FILE_LEN), intent(in), optional :: filename
 
     logical :: parallel
     integer(HID_T) :: file_id
-    character(MAX_FILE_LEN) :: filename
+    character(MAX_FILE_LEN) :: filename_
 
     ! When using parallel HDF5, the file is written to collectively by all
     ! processes. With MPI-only, the file is opened and written by the master
@@ -454,47 +483,20 @@ contains
     parallel = .false.
 #endif
 
-    ! Check to write out source for a specified batch
-    if (sourcepoint_batch%contains(current_batch)) then
-      if (source_separate) then
-        filename = trim(path_output) // 'source.' // &
-             & zero_padded(current_batch, count_digits(n_max_batches))
-        filename = trim(filename) // '.h5'
-        call write_message("Creating source file " // trim(filename) &
-             // "...", 5)
-
-        ! Create separate source file
-        if (master .or. parallel) then
-          file_id = file_open(filename, 'w', parallel=.true.)
-          call write_attribute(file_id, "filetype", 'source')
-        end if
-      else
-        filename = trim(path_output) // 'statepoint.' // &
-             zero_padded(current_batch, count_digits(n_max_batches))
-        filename = trim(filename) // '.h5'
-
-        if (master .or. parallel) then
-          file_id = file_open(filename, 'a', parallel=.true.)
-        end if
-      end if
-
-      call write_source_bank(file_id, work_index, source_bank)
-      if (master .or. parallel) call file_close(file_id)
+    if (present(filename)) then
+      filename_ = filename
+    else
+      filename_ = trim(path_output) // 'source.' // &
+           & zero_padded(current_batch, count_digits(n_max_batches))
+      filename_ = trim(filename_) // '.h5'
     end if
 
-    ! Also check to write source separately in overwritten file
-    if (source_latest) then
-      filename = trim(path_output) // 'source' // '.h5'
-      call write_message("Creating source file " // trim(filename) // "...", 5)
-      if (master .or. parallel) then
-        file_id = file_open(filename, 'w', parallel=.true.)
-        call write_attribute(file_id, "filetype", 'source')
-      end if
-
-      call write_source_bank(file_id, work_index, source_bank)
-
-      if (master .or. parallel) call file_close(file_id)
+    if (master .or. parallel) then
+      file_id = file_open(filename_, 'w', parallel=.true.)
+      call write_attribute(file_id, "filetype", 'source')
     end if
+    call write_source_bank(file_id, work_index, source_bank)
+    if (master .or. parallel) call file_close(file_id)
 
   end subroutine write_source_point
 

tests/unit_tests/test_capi.py

@@ -311,3 +311,35 @@ def test_mesh(capi_init):
 
     msf = openmc.capi.MeshSurfaceFilter(mesh)
     assert msf.mesh == mesh
+
+
+def test_restart(capi_init):
+    # Finalize and re-init to make internal state consistent with XML.
+    openmc.capi.hard_reset()
+    openmc.capi.finalize()
+    openmc.capi.init()
+    openmc.capi.simulation_init()
+
+    # Run for 7 batches then write a statepoint.
+    for i in range(7):
+        openmc.capi.next_batch()
+    openmc.capi.statepoint_write('restart_test.h5', True)
+
+    # Run 3 more batches and copy the keff.
+    for i in range(3):
+        openmc.capi.next_batch()
+    keff0 = openmc.capi.keff()
+
+    # Restart the simulation from the statepoint and the 5 active batches.
+    openmc.capi.simulation_finalize()
+    openmc.capi.hard_reset()
+    openmc.capi.finalize()
+    openmc.capi.init(args=('-r', 'restart_test.h5'))
+    openmc.capi.simulation_init()
+    for i in range(5):
+        openmc.capi.next_batch()
+    keff1 = openmc.capi.keff()
+    openmc.capi.simulation_finalize()
+
+    # Compare the keff values.
+    assert keff0 == pytest.approx(keff1)