From b62e25bcf5f84e5178990519a5e6c2ec433cb7bd Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Thu, 15 Feb 2018 10:39:16 -0600 Subject: [PATCH] Simplify integrator implementations by separating out function for depletion --- openmc/deplete/integrator/cecm.py | 46 +++++------------------- openmc/deplete/integrator/cram.py | 50 ++++++++++++++++++++++++++ openmc/deplete/integrator/predictor.py | 29 ++++----------- 3 files changed, 65 insertions(+), 60 deletions(-) diff --git a/openmc/deplete/integrator/cecm.py b/openmc/deplete/integrator/cecm.py index 16fa299fd..760e3c89d 100644 --- a/openmc/deplete/integrator/cecm.py +++ b/openmc/deplete/integrator/cecm.py @@ -1,13 +1,8 @@ -""" The CE/CM integrator.""" +"""The CE/CM integrator.""" import copy -from itertools import repeat -import os -from multiprocessing import Pool -import time -from .. import comm -from .cram import CRAM48, cram_wrapper +from .cram import deplete from .save_results import save_results @@ -43,8 +38,7 @@ def cecm(operator, print_out=True): """ # Generate initial conditions with operator as vec: - n_mats = len(vec) - + chain = operator.chain t = 0.0 for i, dt in enumerate(operator.settings.dt_vec): # Get beginning-of-timestep reaction rates @@ -52,43 +46,21 @@ def cecm(operator, print_out=True): results = [operator(x[0])] # Deplete for first half of timestep - t_start = time.time() - chains = repeat(operator.chain, n_mats) - vecs = (x[0][i] for i in range(n_mats)) - rates = (results[0].rates[i, :, :] for i in range(n_mats)) - dts = repeat(dt/2, n_mats) - with Pool() as pool: - iters = zip(chains, vecs, rates, dts) - x_result = list(pool.starmap(cram_wrapper, iters)) - t_end = time.time() - if comm.rank == 0: - if print_out: - print("Time to matexp: ", t_end - t_start) + x_middle = deplete(chain, x[0], results[0], dt/2, print_out) # Get middle-of-timestep reaction rates - x.append(x_result) - results.append(operator(x_result)) + x.append(x_middle) + results.append(operator(x_middle)) - # Deplete for second half of timestep - t_start = time.time() - chains = repeat(operator.chain, n_mats) - vecs = (x[0][i] for i in range(n_mats)) - rates = (results[1].rates[i, :, :] for i in range(n_mats)) - dts = repeat(dt, n_mats) - with Pool() as pool: - iters = zip(chains, vecs, rates, dts) - x_result = list(pool.starmap(cram_wrapper, iters)) - t_end = time.time() - if comm.rank == 0: - if print_out: - print("Time to matexp: ", t_end - t_start) + # Deplete for full timestep using beginning-of-step materials + x_end = deplete(chain, x[0], results[1], dt, print_out) # Create results, write to disk save_results(operator, x, results, [t, t + dt], i) # Advance time, update vector t += dt - vec = copy.deepcopy(x_result) + vec = copy.deepcopy(x_end) # Perform one last simulation x = [copy.deepcopy(vec)] diff --git a/openmc/deplete/integrator/cram.py b/openmc/deplete/integrator/cram.py index 56476384c..09207fbc6 100644 --- a/openmc/deplete/integrator/cram.py +++ b/openmc/deplete/integrator/cram.py @@ -3,10 +3,60 @@ Implements two different forms of CRAM for use in openmc.deplete. """ +from itertools import repeat +from multiprocessing import Pool +import time + import numpy as np import scipy.sparse as sp import scipy.sparse.linalg as sla +from .. import comm + + +def deplete(chain, x, op_result, dt, print_out): + """Deplete materials using given reaction rates for a specified time + + Parameters + ---------- + chain : openmc.deplete.Chain + Depletion chain + x : list of numpy.ndarray + Atom number vectors for each material + op_result : openmc.deplete.OperatorResult + Result of applying transport operator (contains reaction rates) + dt : float + Time in [s] to deplete for + print_out : bool + Whether to show elapsed time + + Returns + ------- + x_result : list of numpy.ndarray + Updated atom number vectors for each material + + """ + t_start = time.time() + + # Set up iterators + n_mats = len(x) + chains = repeat(chain, n_mats) + vecs = (x[i] for i in range(n_mats)) + rates = (op_result.rates[i, :, :] for i in range(n_mats)) + dts = repeat(dt, n_mats) + + # Use multiprocessing pool to distribute work + with Pool() as pool: + iters = zip(chains, vecs, rates, dts) + x_result = list(pool.starmap(cram_wrapper, iters)) + + t_end = time.time() + if comm.rank == 0: + if print_out: + print("Time to matexp: ", t_end - t_start) + + return x_result + def cram_wrapper(chain, n0, rates, dt): """Wraps depletion matrix creation / CRAM solve for multiprocess execution diff --git a/openmc/deplete/integrator/predictor.py b/openmc/deplete/integrator/predictor.py index 872b5ebb1..064078331 100644 --- a/openmc/deplete/integrator/predictor.py +++ b/openmc/deplete/integrator/predictor.py @@ -1,20 +1,15 @@ -""" The Predictor algorithm.""" +"""The Predictor algorithm.""" import copy -from itertools import repeat -import os -from multiprocessing import Pool -import time -from .. import comm -from .cram import CRAM48, cram_wrapper +from .cram import deplete from .save_results import save_results def predictor(operator, print_out=True): r"""The basic predictor integrator. - Implements the first order predictor algorithm. This algorithm is + Implements the first-order predictor algorithm. This algorithm is mathematically defined as: .. math:: @@ -34,8 +29,7 @@ def predictor(operator, print_out=True): """ # Generate initial conditions with operator as vec: - n_mats = len(vec) - + chain = operator.chain t = 0.0 for i, dt in enumerate(operator.settings.dt_vec): # Get beginning-of-timestep reaction rates @@ -46,22 +40,11 @@ def predictor(operator, print_out=True): save_results(operator, x, results, [t, t + dt], i) # Deplete for full timestep - t_start = time.time() - chains = repeat(operator.chain, n_mats) - vecs = (x[0][i] for i in range(n_mats)) - rates = (results[0].rates[i, :, :] for i in range(n_mats)) - dts = repeat(dt, n_mats) - with Pool() as pool: - iters = zip(chains, vecs, rates, dts) - x_result = list(pool.starmap(cram_wrapper, iters)) - t_end = time.time() - if comm.rank == 0: - if print_out: - print("Time to matexp: ", t_end - t_start) + x_end = deplete(chain, x[0], results[0], dt, print_out) # Advance time, update vector t += dt - vec = copy.deepcopy(x_result) + vec = copy.deepcopy(x_end) # Perform one last simulation x = [copy.deepcopy(vec)]