Restrict depletion restart to schemes with common stages

Each scheme is capable of different amounts of concentration
and reaction data, depending on the number of intermediate
transport solutions of interest. The PredictorIntegrator
stores data from a single solution, while the CF4Integrator and
EPC_RK4_Integrator store four stages. This discrepancy can cause
issues when saving data in the depletion output file, as
either some data is calculated and thrown away [more stages going
to fewer stages] or excess space is not utilized [fewer stages
going in to a restart that previously used more stages].

Each subclass of Integrator is expected to declare a
_N_STAGES attribute that provides information on the number of
stages expected. Iterative schemes like SI_CELI_Integrator
only report two stages, for the initial concentration and
first prediction, rather than at each iteration point.

Logic that checks for consistent sizes is included in the
initialization of the Integrators, and has been removed
from Results.save

openmc/deplete/integrator/abc.py

@@ -34,6 +34,16 @@ class Integrator(ABC):
             initial heavy metal inventory to get total power if ``power``
             is not speficied.
         """
+        # Check number of stages previously used
+        if operator.prev_res is not None:
+            res = operator.prev_res[-1]
+            if res.data.shape[0] != self._N_STAGES:
+                raise ValueError(
+                    "{} incompatible with previous restart calculation. "
+                    "Previous scheme used {} intermediate solutions, while "
+                    "this uses {}".format(
+                        self.__class__.__name__, res.data.shape[0],
+                        self._N_STAGES))
         self.operator = operator
         self.chain = operator.chain
         if not isinstance(timesteps, Iterable):
@@ -84,6 +94,14 @@ class Integrator(ABC):
             simulations
         """
 
+    @property
+    @abstractmethod
+    def _N_STAGES(self):
+        """Number of intermediate transport solutions
+
+        Needed to ensure schemes are consistent with restarts
+        """
+
     def __iter__(self):
         return zip(self.timesteps, self.power)
 

openmc/deplete/integrator/cecm.py

@@ -22,6 +22,7 @@ class CECMIntegrator(Integrator):
         y_{n+1} &= \text{expm}(A_c h) y_n
         \end{aligned}
     """
+    _N_STAGES = 2
 
     def __call__(self, conc, rates, dt, power, _i=-1):
         """Integrate using CE/CM

openmc/deplete/integrator/celi.py

@@ -23,6 +23,7 @@ class CELIIntegrator(Integrator):
                    \text{expm}(\frac{5h}{12} A_0 + \frac{h}{12} A1) y_n
         \end{aligned}
     """
+    _N_STAGES = 2
 
     def __call__(self, bos_conc, rates, dt, power, _i=-1):
         """Perform the integration across one time step

openmc/deplete/integrator/cf4.py

@@ -28,6 +28,7 @@ class CF4Integrator(Integrator):
                \text{expm}(-1/12 F_1 + 1/6 F_2 + 1/6 F_3 + 1/4  F_4) y_0
         \end{aligned}
     """
+    _N_STAGES = 4
 
     def __call__(self, bos_conc, bos_rates, dt, power, i):
         """Perform the integration across one time step

openmc/deplete/integrator/epc_rk4.py

@@ -23,6 +23,7 @@ class EPC_RK4_Integrator(Integrator):
         y_4 &= \text{expm}(1/6 F_1 + 1/3 F_2 + 1/3 F_3 + 1/6 F_4) y_0
         \end{aligned}
     """
+    _N_STAGES = 4
 
     def __call__(self, conc, rates, dt, power, _i):
         """Perform the integration across one time step

openmc/deplete/integrator/leqi.py

@@ -37,6 +37,7 @@ class LEQIIntegrator(Integrator):
 
     It is initialized using the CE/LI algorithm.
     """
+    _N_STAGES = 2
 
     def __call__(self, bos_conc, bos_rates, dt, power, i):
         """Perform the integration across one time step

openmc/deplete/integrator/predictor.py

@@ -37,6 +37,7 @@ class PredictorIntegrator(Integrator):
         initial heavy metal inventory to get total power if ``power``
         is not speficied.
     """
+    _N_STAGES = 1
 
     def __call__(self, conc, rates, dt, power, _i=None):
         """Perform the integration across one time step

openmc/deplete/integrator/si_celi.py

@@ -18,6 +18,8 @@ class SI_CELI_Integrator(SI_Integrator):
     detailed algorithm can be found in section 3.2 in `colin josey's thesis
     <http://hdl.handle.net/1721.1/113721>`_.
     """
+    _N_STAGES = 2
+
     def __call__(self, bos_conc, bos_rates, dt, power, _i):
         """Perform the integration across one time step
 

openmc/deplete/integrator/si_leqi.py

@@ -21,6 +21,7 @@ class SI_LEQI_Integrator(SI_Integrator):
     Detailed algorithm can be found in Section 3.2 in `Colin Josey's thesis
     <http://hdl.handle.net/1721.1/113721>`_.
     """
+    _N_STAGES = 2
 
     def __call__(self, bos_conc, bos_rates, dt, power, i):
         """Perform the integration across one time step

openmc/deplete/results.py

@@ -425,16 +425,7 @@ class Results(object):
         # Get indexing terms
         vol_dict, nuc_list, burn_list, full_burn_list = op.get_results_info()
 
-        # For a restart calculation, limit number of stages saved to meet the
-        # format of the hdf5 file
         stages = len(x)
-        offset = 0
-        if op.prev_res is not None and op.prev_res[0].n_stages < stages:
-            offset = stages - op.prev_res[0].n_stages
-            stages = min(stages, op.prev_res[0].n_stages)
-            warn("Number of restart integrator stages saved limited by initial"
-                 " depletion integrator choice to {}"
-                 .format(op.prev_res[0].n_stages))
 
         # Create results
         results = Results()
@@ -444,7 +435,7 @@ class Results(object):
 
         for i in range(stages):
             for mat_i in range(n_mat):
-                results[i, mat_i, :] = x[offset + i][mat_i][:]
+                results[i, mat_i, :] = x[i][mat_i]
 
         results.k = [(r.k.nominal_value, r.k.std_dev) for r in op_results]
         results.rates = [r.rates for r in op_results]

tests/unit_tests/test_deplete_restart.py

@@ -4,7 +4,7 @@ These tests run in two steps, a first run then a restart run, a simple test
 problem described in dummy_geometry.py.
 """
 
-from pytest import approx
+from pytest import approx, raises
 import openmc.deplete
 from openmc.deplete import (
     CECMIntegrator, PredictorIntegrator, CELIIntegrator, LEQIIntegrator,
@@ -95,8 +95,7 @@ def test_restart_cecm(run_in_tmpdir):
 
 
 def test_restart_predictor_cecm(run_in_tmpdir):
-    """Integral regression test of integrator algorithm using predictor
-     for the first run then CE/CM for the restart run."""
+    """Test to ensure that schemes with different stages are not compatible"""
 
     op = dummy_operator.DummyOperator()
     output_dir = "test_restart_predictor_cecm"
@@ -114,25 +113,9 @@ def test_restart_predictor_cecm(run_in_tmpdir):
     op = dummy_operator.DummyOperator(prev_res)
     op.output_dir = output_dir
 
-    # Perform restarts simulation using the MCNPX/MCNP6 algorithm
-    cecm = CECMIntegrator(op, dt, power)
-    cecm.integrate()
-
-    # Load the files
-    res = openmc.deplete.ResultsList(op.output_dir / "depletion_results.h5")
-
-    _, y1 = res.get_atoms("1", "1")
-    _, y2 = res.get_atoms("1", "2")
-
-    # Test solution
-    s1 = [2.46847546272295, 0.986431226850467]
-    s2 = [3.09106948392, 0.607102912398]
-
-    assert y1[1] == approx(s1[0])
-    assert y2[1] == approx(s1[1])
-
-    assert y1[2] == approx(s2[0])
-    assert y2[2] == approx(s2[1])
+    # check ValueError is raised, indicating previous and current stages
+    with raises(ValueError, match="incompatible.* 1.*2"):
+        CECMIntegrator(op, dt, power)
 
 
 def test_restart_cecm_predictor(run_in_tmpdir):
@@ -156,25 +139,9 @@ def test_restart_cecm_predictor(run_in_tmpdir):
     op = dummy_operator.DummyOperator(prev_res)
     op.output_dir = output_dir
 
-    # Perform restarts simulation using the predictor algorithm
-    PredictorIntegrator(op, dt, power).integrate()
-
-    # Load the files
-    res = openmc.deplete.ResultsList(op.output_dir / "depletion_results.h5")
-
-    _, y1 = res.get_atoms("1", "1")
-    _, y2 = res.get_atoms("1", "2")
-
-    # Test solution
-    s1 = [1.86872629872102, 1.395525772416039]
-    s2 = [3.32776806576, 2.391425905]
-
-    assert y1[1] == approx(s1[0])
-    assert y2[1] == approx(s1[1])
-
-    assert y1[2] == approx(s2[0])
-    assert y2[2] == approx(s2[1])
-
+    # check ValueError is raised, indicating previous and current stages
+    with raises(ValueError, match="incompatible.* 2.*1"):
+        PredictorIntegrator(op, dt, power)
 
 def test_restart_cf4(run_in_tmpdir):
     """Integral regression test of integrator algorithm using CF4."""