Simplify integrator implementations by separating out function for depletion

openmc/deplete/integrator/cecm.py

@@ -1,13 +1,8 @@
-""" The CE/CM integrator."""
+"""The CE/CM integrator."""
 
 import copy
-from itertools import repeat
-import os
-from multiprocessing import Pool
-import time
 
-from .. import comm
-from .cram import CRAM48, cram_wrapper
+from .cram import deplete
 from .save_results import save_results
 
 
@@ -43,8 +38,7 @@ def cecm(operator, print_out=True):
     """
     # Generate initial conditions
     with operator as vec:
-        n_mats = len(vec)
-
+        chain = operator.chain
         t = 0.0
         for i, dt in enumerate(operator.settings.dt_vec):
             # Get beginning-of-timestep reaction rates
@@ -52,43 +46,21 @@ def cecm(operator, print_out=True):
             results = [operator(x[0])]
 
             # Deplete for first half of timestep
-            t_start = time.time()
-            chains = repeat(operator.chain, n_mats)
-            vecs = (x[0][i] for i in range(n_mats))
-            rates = (results[0].rates[i, :, :] for i in range(n_mats))
-            dts = repeat(dt/2, n_mats)
-            with Pool() as pool:
-                iters = zip(chains, vecs, rates, dts)
-                x_result = list(pool.starmap(cram_wrapper, iters))
-            t_end = time.time()
-            if comm.rank == 0:
-                if print_out:
-                    print("Time to matexp: ", t_end - t_start)
+            x_middle = deplete(chain, x[0], results[0], dt/2, print_out)
 
             # Get middle-of-timestep reaction rates
-            x.append(x_result)
-            results.append(operator(x_result))
+            x.append(x_middle)
+            results.append(operator(x_middle))
 
-            # Deplete for second half of timestep
-            t_start = time.time()
-            chains = repeat(operator.chain, n_mats)
-            vecs = (x[0][i] for i in range(n_mats))
-            rates = (results[1].rates[i, :, :] for i in range(n_mats))
-            dts = repeat(dt, n_mats)
-            with Pool() as pool:
-                iters = zip(chains, vecs, rates, dts)
-                x_result = list(pool.starmap(cram_wrapper, iters))
-            t_end = time.time()
-            if comm.rank == 0:
-                if print_out:
-                    print("Time to matexp: ", t_end - t_start)
+            # Deplete for full timestep using beginning-of-step materials
+            x_end = deplete(chain, x[0], results[1], dt, print_out)
 
             # Create results, write to disk
             save_results(operator, x, results, [t, t + dt], i)
 
             # Advance time, update vector
             t += dt
-            vec = copy.deepcopy(x_result)
+            vec = copy.deepcopy(x_end)
 
         # Perform one last simulation
         x = [copy.deepcopy(vec)]

openmc/deplete/integrator/cram.py

@@ -3,10 +3,60 @@
 Implements two different forms of CRAM for use in openmc.deplete.
 """
 
+from itertools import repeat
+from multiprocessing import Pool
+import time
+
 import numpy as np
 import scipy.sparse as sp
 import scipy.sparse.linalg as sla
 
+from .. import comm
+
+
+def deplete(chain, x, op_result, dt, print_out):
+    """Deplete materials using given reaction rates for a specified time
+
+    Parameters
+    ----------
+    chain : openmc.deplete.Chain
+        Depletion chain
+    x : list of numpy.ndarray
+        Atom number vectors for each material
+    op_result : openmc.deplete.OperatorResult
+        Result of applying transport operator (contains reaction rates)
+    dt : float
+        Time in [s] to deplete for
+    print_out : bool
+        Whether to show elapsed time
+
+    Returns
+    -------
+    x_result : list of numpy.ndarray
+        Updated atom number vectors for each material
+
+    """
+    t_start = time.time()
+
+    # Set up iterators
+    n_mats = len(x)
+    chains = repeat(chain, n_mats)
+    vecs = (x[i] for i in range(n_mats))
+    rates = (op_result.rates[i, :, :] for i in range(n_mats))
+    dts = repeat(dt, n_mats)
+
+    # Use multiprocessing pool to distribute work
+    with Pool() as pool:
+        iters = zip(chains, vecs, rates, dts)
+        x_result = list(pool.starmap(cram_wrapper, iters))
+
+    t_end = time.time()
+    if comm.rank == 0:
+        if print_out:
+            print("Time to matexp: ", t_end - t_start)
+
+    return x_result
+
 
 def cram_wrapper(chain, n0, rates, dt):
     """Wraps depletion matrix creation / CRAM solve for multiprocess execution

openmc/deplete/integrator/predictor.py

@@ -1,20 +1,15 @@
-""" The Predictor algorithm."""
+"""The Predictor algorithm."""
 
 import copy
-from itertools import repeat
-import os
-from multiprocessing import Pool
-import time
 
-from .. import comm
-from .cram import CRAM48, cram_wrapper
+from .cram import deplete
 from .save_results import save_results
 
 
 def predictor(operator, print_out=True):
     r"""The basic predictor integrator.
 
-    Implements the first order predictor algorithm. This algorithm is
+    Implements the first-order predictor algorithm. This algorithm is
     mathematically defined as:
 
     .. math::
@@ -34,8 +29,7 @@ def predictor(operator, print_out=True):
     """
     # Generate initial conditions
     with operator as vec:
-        n_mats = len(vec)
-
+        chain = operator.chain
         t = 0.0
         for i, dt in enumerate(operator.settings.dt_vec):
             # Get beginning-of-timestep reaction rates
@@ -46,22 +40,11 @@ def predictor(operator, print_out=True):
             save_results(operator, x, results, [t, t + dt], i)
 
             # Deplete for full timestep
-            t_start = time.time()
-            chains = repeat(operator.chain, n_mats)
-            vecs = (x[0][i] for i in range(n_mats))
-            rates = (results[0].rates[i, :, :] for i in range(n_mats))
-            dts = repeat(dt, n_mats)
-            with Pool() as pool:
-                iters = zip(chains, vecs, rates, dts)
-                x_result = list(pool.starmap(cram_wrapper, iters))
-            t_end = time.time()
-            if comm.rank == 0:
-                if print_out:
-                    print("Time to matexp: ", t_end - t_start)
+            x_end = deplete(chain, x[0], results[0], dt, print_out)
 
             # Advance time, update vector
             t += dt
-            vec = copy.deepcopy(x_result)
+            vec = copy.deepcopy(x_end)
 
         # Perform one last simulation
         x = [copy.deepcopy(vec)]