Source code for openmc.deplete.pool

"""Dedicated module containing depletion function

Provided to avoid some circular imports
"""
from itertools import repeat, starmap
from multiprocessing import Pool


# Configurable switch that enables / disables the use of
# multiprocessing routines during depletion
USE_MULTIPROCESSING = True


[docs]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`` 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) inputs = zip(matrices, x, repeat(dt)) if USE_MULTIPROCESSING: with Pool() as pool: x_result = list(pool.starmap(func, inputs)) else: x_result = list(starmap(func, inputs)) return x_result