Allow user to override num processes used for depletion

This commit is contained in:
Paul Romano 2022-03-22 06:53:47 -05:00
parent 73ba908f1f
commit dc18f59d82
2 changed files with 11 additions and 1 deletions

View file

@ -166,6 +166,12 @@ with :func:`cram.CRAM48` being the default.
:type: bool
.. data:: pool.NUM_PROCESSES
Number of worker processes used for depletion calculations, which rely on the
:class:`multiprocessing.pool.Pool` class. If set to ``None`` (default), the
number returned by :func:`os.cpu_count` is used.
The following classes are used to help the :class:`openmc.deplete.Operator`
compute quantities like effective fission yields, reaction rates, and
total system energy.

View file

@ -10,6 +10,10 @@ from multiprocessing import Pool
# multiprocessing routines during depletion
USE_MULTIPROCESSING = True
# Allow user to override the number of worker processes to use for depletion
# calculations
NUM_PROCESSES = None
def deplete(func, chain, x, rates, dt, matrix_func=None):
"""Deplete materials using given reaction rates for a specified time
@ -58,7 +62,7 @@ def deplete(func, chain, x, rates, dt, matrix_func=None):
inputs = zip(matrices, x, repeat(dt))
if USE_MULTIPROCESSING:
with Pool() as pool:
with Pool(NUM_PROCESSES) as pool:
x_result = list(pool.starmap(func, inputs))
else:
x_result = list(starmap(func, inputs))