From c0566e0a2d269d0749a98f13464ec665cedbfe90 Mon Sep 17 00:00:00 2001
From: Andrew Johnson <drewej@protonmail.com>
Date: Thu, 1 Aug 2019 09:24:52 -0500
Subject: [PATCH] Remove cram wrapper with maps of depletion matrices

openmc.deplete.cram.deplete now pre-generates a map object
that produces depletion matrices given the chain, rates,
and matrix_func. The internal cram wrapper function has
been removed, as the full inputs to CRAM48 are now known,
in zipped and mapped objects.

Some basic timing comparisons indicate that this approach
is at least as fast, if not marginally faster, for a
quarter PWR assembly and an SFR assembly.
---
 openmc/deplete/cram.py | 42 +++++++-----------------------------------
 1 file changed, 7 insertions(+), 35 deletions(-)

diff --git a/openmc/deplete/cram.py b/openmc/deplete/cram.py
index 184c411bf3..67a455cb35 100644
--- a/openmc/deplete/cram.py
+++ b/openmc/deplete/cram.py
@@ -51,11 +51,15 @@ def deplete(chain, x, rates, dt, matrix_func=None):
             "to the number of compositions {}".format(len(fission_yields),
                 len(x)))
 
+    if matrix_func is None:
+        matrices = map(chain.form_matrix, rates, fission_yields)
+    else:
+        matrices = map(matrix_func, repeat(chain), rates, fission_yields)
+
     # Use multiprocessing pool to distribute work
     with Pool() as pool:
-        iters = zip(repeat(chain), x, rates, repeat(dt),
-                    fission_yields, repeat(matrix_func))
-        x_result = list(pool.starmap(_cram_wrapper, iters))
+        inputs = zip(matrices, x, repeat(dt))
+        x_result = list(pool.starmap(CRAM48, inputs))
 
     return x_result
 
@@ -79,38 +83,6 @@ def timed_deplete(*args, **kwargs):
     return time.time() - start, results
 
 
-def _cram_wrapper(chain, n0, rates, dt, fission_yields, matrix_func=None):
-    """Wraps depletion matrix creation / CRAM solve for multiprocess execution
-
-    Parameters
-    ----------
-    chain : openmc.deplete.Chain
-        Depletion chain used to construct the burnup matrix
-    n0 : numpy.array
-        Vector to operate a matrix exponent on.
-    rates : numpy.ndarray
-        2D array indexed by nuclide then by cell.
-    dt : float
-        Time to integrate to.
-    maxtrix_func : function, optional
-        Function to form the depletion matrix
-    fission_yields : dict
-        Single-energy fission yields of the form
-        ``{parent: {product: fission_yield}}``
-
-    Returns
-    -------
-    numpy.array
-        Results of the matrix exponent.
-    """
-
-    if matrix_func is None:
-        A = chain.form_matrix(rates, fission_yields)
-    else:
-        A = matrix_func(chain, rates, fission_yields)
-    return CRAM48(A, n0, dt)
-
-
 def CRAM16(A, n0, dt):
     """Chebyshev Rational Approximation Method, order 16