Merge pull request #1593 from drewejohnson/deplete-with-itertools

Allow depletion without multiprocessing

docs/source/pythonapi/deplete.rst

@@ -145,6 +145,17 @@ with :func:`cram.CRAM48` being the default.
    cram.CRAM48
    pool.deplete
 
+.. data:: pool.USE_MULTIPROCESSING
+
+   Boolean switch to enable or disable the use of :mod:`multiprocessing`
+   when solving the Bateman equations. The default is to use
+   :mod:`multiprocessing`, but can cause the simulation to hang in
+   some computing environments, namely due to MPI and networking
+   restrictions. Disabling this option will result in only a single
+   CPU core being used for depletion.
+
+   :type: bool
+
 The following classes are used to help the :class:`openmc.deplete.Operator`
 compute quantities like effective fission yields, reaction rates, and
 total system energy.

openmc/deplete/pool.py

@@ -2,10 +2,15 @@
 
 Provided to avoid some circular imports
 """
-from itertools import repeat
+from itertools import repeat, starmap
 from multiprocessing import Pool
 
 
+# Configurable switch that enables / disables the use of
+# multiprocessing routines during depletion
+USE_MULTIPROCESSING = True
+
+
 def deplete(func, chain, x, rates, dt, matrix_func=None):
     """Deplete materials using given reaction rates for a specified time
 
@@ -41,8 +46,8 @@ def deplete(func, chain, x, rates, dt, matrix_func=None):
         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(
+            "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:
@@ -50,9 +55,12 @@ def deplete(func, chain, x, rates, dt, matrix_func=None):
     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))
+    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

tests/regression_tests/deplete/test.py

@@ -5,6 +5,7 @@ import shutil
 from pathlib import Path
 
 import numpy as np
+import pytest
 import openmc
 from openmc.data import JOULE_PER_EV
 import openmc.deplete
@@ -13,7 +14,17 @@ from tests.regression_tests import config
 from example_geometry import generate_problem
 
 
-def test_full(run_in_tmpdir):
+@pytest.fixture(scope="module")
+def problem():
+    n_rings = 2
+    n_wedges = 4
+
+    # Load geometry from example
+    return generate_problem(n_rings, n_wedges)
+
+
+@pytest.mark.parametrize("multiproc", [True, False])
+def test_full(run_in_tmpdir, problem, multiproc):
     """Full system test suite.
 
     Runs an entire OpenMC simulation with depletion coupling and verifies
@@ -25,11 +36,7 @@ def test_full(run_in_tmpdir):
 
     """
 
-    n_rings = 2
-    n_wedges = 4
-
-    # Load geometry from example
-    geometry, lower_left, upper_right = generate_problem(n_rings, n_wedges)
+    geometry, lower_left, upper_right = problem
 
     # OpenMC-specific settings
     settings = openmc.Settings()
@@ -54,6 +61,7 @@ def test_full(run_in_tmpdir):
     power = 2.337e15*4*JOULE_PER_EV*1e6  # MeV/second cm from CASMO
 
     # Perform simulation using the predictor algorithm
+    openmc.deplete.pool.USE_MULTIPROCESSING = multiproc
     openmc.deplete.PredictorIntegrator(op, dt, power).integrate()
 
     # Get path to test and reference results