From 08bf32e1ae3ec0b8384dc809a77250bf4d9097b2 Mon Sep 17 00:00:00 2001
From: Paul Romano <paul.k.romano@gmail.com>
Date: Fri, 15 Dec 2017 11:38:28 +0700
Subject: [PATCH] Throw exception if next_batch() is called before simulation
 is initialized

---
 openmc/capi/core.py           | 10 +++++++---
 src/simulation.F90            | 18 ++++++++++++++++++
 src/simulation_header.F90     |  6 +++++-
 tests/unit_tests/test_capi.py | 17 ++++++++++-------
 4 files changed, 40 insertions(+), 11 deletions(-)

diff --git a/openmc/capi/core.py b/openmc/capi/core.py
index ced939043..415fa2e93 100644
--- a/openmc/capi/core.py
+++ b/openmc/capi/core.py
@@ -7,7 +7,7 @@ import numpy as np
 from numpy.ctypeslib import as_array
 
 from . import _dll
-from .error import _error_handler
+from .error import _error_handler, AllocationError
 import openmc.capi
 
 
@@ -147,7 +147,7 @@ def iter_batches():
     """
     while True:
         # Run next batch
-        retval = _dll.openmc_next_batch()
+        retval = next_batch()
 
         # Provide opportunity for user to perform action between batches
         yield
@@ -181,7 +181,11 @@ def keff():
 
 def next_batch():
     """Run next batch."""
-    return _dll.openmc_next_batch()
+    retval = _dll.openmc_next_batch()
+    if retval == -3:
+        raise AllocationError('Simulation has not been initialized. You must call '
+                              'openmc.capi.simulation_init() first.')
+    return retval
 
 
 def plot_geometry():
diff --git a/src/simulation.F90 b/src/simulation.F90
index e3eba0d51..5eacbd40b 100644
--- a/src/simulation.F90
+++ b/src/simulation.F90
@@ -73,6 +73,12 @@ contains
     type(Particle) :: p
     integer(8)     :: i_work
 
+    ! Make sure simulation has been initialized
+    if (.not. simulation_initialized) then
+      retval = -3
+      return
+    end if
+
     call initialize_batch()
 
     ! Handle restart runs
@@ -396,6 +402,9 @@ contains
 
   subroutine openmc_simulation_init() bind(C)
 
+    ! Skip if simulation has already been initialized
+    if (simulation_initialized) return
+
     ! Set up tally procedure pointers
     call init_tally_routines()
 
@@ -438,6 +447,9 @@ contains
     ! Reset current batch
     current_batch = 0
 
+    ! Set flag indicating initialization is done
+    simulation_initialized = .true.
+
   end subroutine openmc_simulation_init
 
 !===============================================================================
@@ -455,6 +467,9 @@ contains
     real(8)    :: tempr(3) ! temporary array for communication
 #endif
 
+    ! Skip if simulation was never run
+    if (.not. simulation_initialized) return
+
     ! Stop active batch timer
     call time_active % stop()
 
@@ -511,6 +526,9 @@ contains
       if (check_overlaps) call print_overlap_check()
     end if
 
+    ! Reset initialization flag
+    simulation_initialized = .false.
+
   end subroutine openmc_simulation_finalize
 
 !===============================================================================
diff --git a/src/simulation_header.F90 b/src/simulation_header.F90
index 62c66e5f9..5041f0793 100644
--- a/src/simulation_header.F90
+++ b/src/simulation_header.F90
@@ -18,11 +18,12 @@ module simulation_header
   real(8) :: log_spacing ! spacing on logarithmic grid
 
   ! ============================================================================
-  ! EIGENVALUE SIMULATION VARIABLES
+  ! SIMULATION VARIABLES
 
   integer    :: current_batch     ! current batch
   integer    :: current_gen       ! current generation within a batch
   integer    :: total_gen     = 0 ! total number of generations simulated
+  logical(C_BOOL), bind(C) :: simulation_initialized = .false.
 
   ! ============================================================================
   ! TALLY PRECISION TRIGGER VARIABLES
@@ -33,6 +34,9 @@ module simulation_header
   integer(8), allocatable :: work_index(:) ! starting index in source bank for each process
   integer(8) :: current_work ! index in source bank of current history simulated
 
+  ! ============================================================================
+  ! K-EIGENVALUE SIMULATION VARIABLES
+
   ! Temporary k-effective values
   type(VectorReal) :: k_generation ! single-generation estimates of k
   real(C_DOUBLE), bind(C) :: keff = ONE  ! average k over active batches
diff --git a/tests/unit_tests/test_capi.py b/tests/unit_tests/test_capi.py
index cde3d8724..ba9ac5c9e 100644
--- a/tests/unit_tests/test_capi.py
+++ b/tests/unit_tests/test_capi.py
@@ -206,17 +206,20 @@ def test_source_bank(capi_run):
     assert np.all(source['wgt'] == 1.0)
 
 
-def test_keff(capi_run):
-    mean, std_dev = openmc.capi.keff()
-    assert 0.0 < mean < 2.5
-    assert std_dev > 0.0
-
-
 def test_by_batch(capi_run):
     openmc.capi.hard_reset()
+
+    # Running next batch before simulation is initialized should raise an
+    # exception
+    with pytest.raises(openmc.capi.AllocationError):
+        openmc.capi.next_batch()
+
     openmc.capi.simulation_init()
     for _ in openmc.capi.iter_batches():
-        pass
+        # Make sure we can get k-effective during inactive/active batches
+        mean, std_dev = openmc.capi.keff()
+        assert 0.0 < mean < 2.5
+        assert std_dev > 0.0
     assert openmc.capi.num_realizations() == 5
 
     for i in range(3):