From 46e98e4a01d71ebf67cb65e4630df540b657e197 Mon Sep 17 00:00:00 2001 From: Andrew Johnson Date: Sat, 9 May 2020 14:38:46 -0400 Subject: [PATCH] Move deplete function to dedicated pool module; removed timed_deplete --- docs/source/pythonapi/deplete.rst | 3 +- openmc/deplete/cram.py | 77 +------------------------- openmc/deplete/pool.py | 57 +++++++++++++++++++ tests/unit_tests/test_deplete_chain.py | 4 +- 4 files changed, 61 insertions(+), 80 deletions(-) create mode 100644 openmc/deplete/pool.py diff --git a/docs/source/pythonapi/deplete.rst b/docs/source/pythonapi/deplete.rst index b9970ffb1..b01de477c 100644 --- a/docs/source/pythonapi/deplete.rst +++ b/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 diff --git a/openmc/deplete/cram.py b/openmc/deplete/cram.py index 939ff5432..6b6a909ec 100644 --- a/openmc/deplete/cram.py +++ b/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): diff --git a/openmc/deplete/pool.py b/openmc/deplete/pool.py new file mode 100644 index 000000000..ef92569b7 --- /dev/null +++ b/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 diff --git a/tests/unit_tests/test_deplete_chain.py b/tests/unit_tests/test_deplete_chain.py index 9506ad317..7162be6de 100644 --- a/tests/unit_tests/test_deplete_chain.py +++ b/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):