Move deplete function to dedicated pool module; removed timed_deplete

docs/source/pythonapi/deplete.rst

@@ -143,8 +143,7 @@ with :func:`cram.CRAM48` being the default.
 
    cram.CRAM16
    cram.CRAM48
-   cram.deplete
-   cram.timed_deplete
+   pool.deplete
 
 The following classes are used to help the :class:`openmc.deplete.Operator`
 compute quantities like effective fission yields, reaction rates, and

openmc/deplete/cram.py

@@ -4,9 +4,6 @@ Implements two different forms of CRAM for use in openmc.deplete.
 """
 
 import numbers
-from itertools import repeat
-from multiprocessing import Pool
-import time
 
 import numpy as np
 import scipy.sparse as sp
@@ -15,79 +12,7 @@ import scipy.sparse.linalg as sla
 from openmc.checkvalue import check_type, check_length
 from .abc import DepSystemSolver
 
-__all__ = [
-    "deplete", "timed_deplete", "CRAM16", "CRAM48",
-    "Cram16Solver", "Cram48Solver", "IPFCramSolver"]
-
-
-def deplete(func, chain, x, rates, dt, matrix_func=None):
-    """Deplete materials using given reaction rates for a specified time
-
-    Parameters
-    ----------
-    func : callable
-        Function to use to get new compositions. Expected to have the
-        signature ``func(A, n0, t) -> n1``
-    chain : openmc.deplete.Chain
-        Depletion chain
-    x : list of numpy.ndarray
-        Atom number vectors for each material
-    rates : openmc.deplete.ReactionRates
-        Reaction rates (from transport operator)
-    dt : float
-        Time in [s] to deplete for
-    maxtrix_func : Callable, optional
-        Function to form the depletion matrix after calling
-        ``matrix_func(chain, rates, fission_yields)``, where
-        ``fission_yields = {parent: {product: yield_frac}}``
-        Expected to return the depletion matrix required by
-        :func:`CRAM48`.
-
-    Returns
-    -------
-    x_result : list of numpy.ndarray
-        Updated atom number vectors for each material
-    """
-
-    fission_yields = chain.fission_yields
-    if len(fission_yields) == 1:
-        fission_yields = repeat(fission_yields[0])
-    elif len(fission_yields) != len(x):
-        raise ValueError(
-            "Number of material fission yield distributions {} is not equal "
-            "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:
-        inputs = zip(matrices, x, repeat(dt))
-        x_result = list(pool.starmap(func, inputs))
-
-    return x_result
-
-
-def timed_deplete(*args, **kwargs):
-    """Wrapper over :func:`deplete` that also returns process time
-
-    All arguments and keyword arguments are passed onto
-    :func:`deplete` directly.
-
-    Returns
-    -------
-    proc_time: float
-        Process time required to return from deplete
-    results: list of numpy arrays
-        Output from :func:`deplete` call
-    """
-
-    start = time.time()
-    results = deplete(*args, **kwargs)
-    return time.time() - start, results
+__all__ = ["CRAM16", "CRAM48", "Cram16Solver", "Cram48Solver", "IPFCramSolver"]
 
 
 class IPFCramSolver(DepSystemSolver):

openmc/deplete/pool.py

@@ -0,0 +1,57 @@
+"""Dedicated module containing depletion function
+
+Provided to avoid some circular imports
+"""
+from itertools import repeat
+from multiprocessing import Pool
+
+
+def deplete(func, chain, x, rates, dt, matrix_func=None):
+    """Deplete materials using given reaction rates for a specified time
+
+    Parameters
+    ----------
+    func : callable
+        Function to use to get new compositions. Expected to have the
+        signature ``func(A, n0, t) -> n1``
+    chain : openmc.deplete.Chain
+        Depletion chain
+    x : list of numpy.ndarray
+        Atom number vectors for each material
+    rates : openmc.deplete.ReactionRates
+        Reaction rates (from transport operator)
+    dt : float
+        Time in [s] to deplete for
+    maxtrix_func : Callable, optional
+        Function to form the depletion matrix after calling
+        ``matrix_func(chain, rates, fission_yields)``, where
+        ``fission_yields = {parent: {product: yield_frac}}``
+        Expected to return the depletion matrix required by
+        :func:`CRAM48`.
+
+    Returns
+    -------
+    x_result : list of numpy.ndarray
+        Updated atom number vectors for each material
+    """
+
+    fission_yields = chain.fission_yields
+    if len(fission_yields) == 1:
+        fission_yields = repeat(fission_yields[0])
+    elif len(fission_yields) != len(x):
+        raise ValueError(
+            "Number of material fission yield distributions {} is not equal "
+            "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:
+        inputs = zip(matrices, x, repeat(dt))
+        x_result = list(pool.starmap(func, inputs))
+
+    return x_result

tests/unit_tests/test_deplete_chain.py

@@ -7,7 +7,7 @@ import os
 from pathlib import Path
 
 import numpy as np
-from openmc.deplete import comm, Chain, reaction_rates, nuclide, cram
+from openmc.deplete import comm, Chain, reaction_rates, nuclide, cram, pool
 import pytest
 
 from tests import cdtemp
@@ -414,7 +414,7 @@ def test_fission_yield_attribute(simple_chain):
     dummy_conc = [[1, 2]] * (len(empty_chain.fission_yields) + 1)
     with pytest.raises(
             ValueError, match="fission yield.*not equal.*compositions"):
-        cram.deplete(cram.CRAM48, empty_chain, dummy_conc, None, 0.5)
+        pool.deplete(cram.CRAM48, empty_chain, dummy_conc, None, 0.5)
 
 
 def test_validate(simple_chain):