Merge pull request #2009 from paulromano/depletion-num-processes

Allow user to override number of processes used for depletion
This commit is contained in:
Patrick Shriwise 2022-03-25 13:06:28 -05:00 committed by GitHub
commit 4df659d0c2
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
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))