allowing to pass additiinal args to pool.deplete when using user specified maxtrix_func

This commit is contained in:
Baptiste Mouginot 2023-01-18 14:04:41 +01:00
parent 303812202d
commit bbacb383cd

View file

@ -15,7 +15,7 @@ USE_MULTIPROCESSING = True
NUM_PROCESSES = None
def deplete(func, chain, x, rates, dt, matrix_func=None):
def deplete(func, chain, x, rates, dt, matrix_func=None, **func_args):
"""Deplete materials using given reaction rates for a specified time
Parameters
@ -37,6 +37,8 @@ def deplete(func, chain, x, rates, dt, matrix_func=None):
``fission_yields = {parent: {product: yield_frac}}``
Expected to return the depletion matrix required by
``func``
func_args : dict
Remaining keyword arguments passed to the matrix_func
Returns
-------
@ -57,7 +59,8 @@ def deplete(func, chain, x, rates, dt, matrix_func=None):
if matrix_func is None:
matrices = map(chain.form_matrix, rates, fission_yields)
else:
matrices = map(matrix_func, repeat(chain), rates, fission_yields)
matrices = map(matrix_func, repeat(chain), rates, fission_yields,
**func_args)
inputs = zip(matrices, x, repeat(dt))