Add simulation_init/finalize, iter_batches, source_bank in openmc.capi

openmc/capi/core.py

@@ -1,11 +1,23 @@
 from contextlib import contextmanager
-from ctypes import CDLL, c_int, c_int32, c_double, POINTER
+from ctypes import CDLL, c_int, c_int32, c_int64, c_double, POINTER, Structure
 from warnings import warn
 
+import numpy as np
+from numpy.ctypeslib import as_array
+
 from . import _dll
 from .error import _error_handler
 
 
+class _Bank(Structure):
+    _fields_ = [('wgt', c_double),
+                ('xyz', c_double*3),
+                ('uvw', c_double*3),
+                ('E', c_double),
+                ('delayed_group', c_int)]
+
+
+
 _dll.openmc_calculate_volumes.restype = None
 _dll.openmc_finalize.restype = None
 _dll.openmc_find.argtypes = [POINTER(c_double*3), c_int, POINTER(c_int32),
@@ -18,9 +30,16 @@ _dll.openmc_init.restype = None
 _dll.openmc_get_keff.argtypes = [POINTER(c_double*2)]
 _dll.openmc_get_keff.restype = c_int
 _dll.openmc_get_keff.errcheck = _error_handler
+_dll.openmc_next_batch.restype = c_int
 _dll.openmc_plot_geometry.restype = None
 _dll.openmc_run.restype = None
 _dll.openmc_reset.restype = None
+_dll.openmc_source_bank.argtypes = [POINTER(POINTER(_Bank)), POINTER(c_int64)]
+_dll.openmc_source_bank.restype = c_int
+_dll.openmc_source_bank.errcheck = _error_handler
+_dll.openmc_simulation_init.restype = None
+_dll.openmc_simulation_finalize.restype = None
+_dll.openmc_statepoint_write.restype = None
 
 
 def calculate_volumes():
@@ -102,6 +121,20 @@ def init(intracomm=None):
         _dll.openmc_init(None)
 
 
+def iter_batches():
+    """Iterator over batches."""
+    while True:
+        # Run next batch
+        retval = _dll.openmc_next_batch()
+
+        # Provide opportunity for user to perform action between batches
+        yield
+
+        # End the iteration
+        if retval < 0:
+            break
+
+
 def keff():
     """Return the calculated k-eigenvalue and its standard deviation.
 
@@ -115,6 +148,10 @@ def keff():
     _dll.openmc_get_keff(k)
     return tuple(k)
 
+def next_batch():
+    """Run next batch."""
+    return _dll.openmc_next_batch()
+
 
 def plot_geometry():
     """Plot geometry"""
@@ -131,6 +168,40 @@ def run():
     _dll.openmc_run()
 
 
+def simulation_init():
+    """Initialize simulation"""
+    _dll.openmc_simulation_init()
+
+
+def simulation_finalize():
+    """Finalize simulation"""
+    _dll.openmc_simulation_finalize()
+
+
+def source_bank():
+    """Return source bank as NumPy array
+
+    Returns
+    -------
+    numpy.ndarray
+        Source sites
+
+    """
+    # Get pointer to source bank
+    ptr = POINTER(_Bank)()
+    n = c_int64()
+    _dll.openmc_source_bank(ptr, n)
+
+    # Convert to numpy array with appropriate datatype
+    bank_dtype = np.dtype(_Bank)
+    return as_array(ptr, (n.value,)).view(bank_dtype)
+
+
+def statepoint_write():
+    """Write a statepoint."""
+    _dll.openmc_statepoint_write()
+
+
 @contextmanager
 def run_in_memory(intracomm=None):
     """Provides context manager for calling OpenMC shared library functions.

src/bank_header.F90

@@ -2,6 +2,8 @@ module bank_header
 
   use, intrinsic :: ISO_C_BINDING
 
+  use error, only: E_ALLOCATE, set_errmsg
+
   implicit none
 
 !===============================================================================
@@ -48,4 +50,24 @@ contains
 
   end subroutine free_memory_bank
 
+!===============================================================================
+!                               C API FUNCTIONS
+!===============================================================================
+
+  function openmc_source_bank(ptr, n) result(err) bind(C)
+    ! Return a pointer to the source bank
+    type(C_PTR), intent(out) :: ptr
+    integer(C_INT64_T), intent(out) :: n
+    integer(C_INT) :: err
+
+    if (.not. allocated(source_bank)) then
+      err = E_ALLOCATE
+      call set_errmsg("Source bank has not been allocated.")
+    else
+      err = 0
+      ptr = C_LOC(source_bank)
+      n = size(source_bank)
+    end if
+  end function openmc_source_bank
+
 end module bank_header

src/simulation.F90

@@ -29,7 +29,7 @@ module simulation
   use settings
   use simulation_header
   use source,          only: initialize_source, sample_external_source
-  use state_point,     only: write_state_point, write_source_point, load_state_point
+  use state_point,     only: openmc_statepoint_write, write_source_point, load_state_point
   use string,          only: to_str
   use tally,           only: accumulate_tallies, setup_active_tallies, &
                              init_tally_routines
@@ -349,7 +349,7 @@ contains
 
     ! Write out state point if it's been specified for this batch
     if (statepoint_batch % contains(current_batch)) then
-      call write_state_point()
+      call openmc_statepoint_write()
     end if
 
     ! Write out source point if it's been specified for this batch

src/state_point.F90

@@ -40,10 +40,10 @@ module state_point
 contains
 
 !===============================================================================
-! WRITE_STATE_POINT
+! OPENMC_STATEPOINT_WRITE writes an HDF5 statepoint file to disk
 !===============================================================================
 
-  subroutine write_state_point()
+  subroutine openmc_statepoint_write() bind(C)
 
     integer :: i, j, k
     integer :: i_xs
@@ -433,7 +433,7 @@ contains
 
       call file_close(file_id)
     end if
-  end subroutine write_state_point
+  end subroutine openmc_statepoint_write
 
 !===============================================================================
 ! WRITE_SOURCE_POINT

src/tallies/tally_header.F90

@@ -22,6 +22,7 @@ module tally_header
   public :: free_memory_tally
   public :: openmc_extend_tallies
   public :: openmc_get_tally_index
+  public :: openmc_global_tallies
   public :: openmc_tally_get_id
   public :: openmc_tally_get_filters
   public :: openmc_tally_get_nuclides
@@ -476,6 +477,7 @@ contains
 
     if (.not. allocated(global_tallies)) then
       err = E_ALLOCATE
+      call set_errmsg("Global tallies have not been allocated yet.")
     else
       err = 0
       ptr = C_LOC(global_tallies)