Have Operator() return a namedtuple (simplifies integrators quite a bit)

This commit is contained in:
Paul Romano 2018-02-15 09:47:37 -06:00
parent 484a023888
commit 998a562a33
7 changed files with 55 additions and 92 deletions

View file

@ -4,6 +4,7 @@ This module contains the Operator class, which is then passed to an integrator
to run a full depletion simulation.
"""
from collections import namedtuple
import os
from pathlib import Path
@ -53,6 +54,9 @@ class Settings(object):
self._output_dir = Path(output_dir)
OperatorResult = namedtuple('OperatorResult', ['k', 'rates', 'seed'])
class Operator(metaclass=ABCMeta):
"""Abstract class defining all methods needed for the integrator.

View file

@ -12,7 +12,7 @@ from .save_results import save_results
def cecm(operator, print_out=True):
"""The CE/CM integrator.
r"""The CE/CM integrator.
Implements the second order CE/CM Predictor-Corrector algorithm [ref]_.
This algorithm is mathematically defined as:
@ -22,11 +22,11 @@ def cecm(operator, print_out=True):
A_p &= A(y_n, t_n)
y_m &= \\text{expm}(A_p h/2) y_n
y_m &= \text{expm}(A_p h/2) y_n
A_c &= A(y_m, t_n + h/2)
y_{n+1} &= \\text{expm}(A_c h) y_n
y_{n+1} &= \text{expm}(A_c h) y_n
.. [ref]
Isotalo, Aarno. "Comparison of Neutronics-Depletion Coupling Schemes
@ -35,88 +35,64 @@ def cecm(operator, print_out=True):
Parameters
----------
operator : Operator
operator : openmc.deplete.Operator
The operator object to simulate on.
print_out : bool, optional
Whether or not to print out time.
"""
"""
# Generate initial conditions
with operator as vec:
n_mats = len(vec)
t = 0.0
for i, dt in enumerate(operator.settings.dt_vec):
# Create vectors
# Get beginning-of-timestep reaction rates
x = [copy.deepcopy(vec)]
seeds = []
eigvls = []
rates_array = []
eigvl, rates, seed = operator(x[0])
eigvls.append(eigvl)
seeds.append(seed)
rates_array.append(rates)
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 = (rates_array[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)
# Get middle-of-timestep reaction rates
x.append(x_result)
results.append(operator(x_result))
eigvl, rates, seed = operator(x[1])
eigvls.append(eigvl)
seeds.append(seed)
rates_array.append(rates)
# 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 = (rates_array[1][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)
# Create results, write to disk
save_results(operator, x, rates_array, eigvls, seeds, [t, t + dt], i)
save_results(operator, x, results, [t, t + dt], i)
# Advance time, update vector
t += dt
vec = copy.deepcopy(x_result)
# Perform one last simulation
x = [copy.deepcopy(vec)]
seeds = []
eigvls = []
rates_array = []
eigvl, rates, seed = operator(x[0])
eigvls.append(eigvl)
seeds.append(seed)
rates_array.append(rates)
results = [operator(x[0])]
# Create results, write to disk
save_results(operator, x, rates_array, eigvls, seeds, [t, t],
len(operator.settings.dt_vec))
save_results(operator, x, results, [t, t], len(operator.settings.dt_vec))

View file

@ -12,7 +12,7 @@ from .save_results import save_results
def predictor(operator, print_out=True):
"""The basic predictor integrator.
r"""The basic predictor integrator.
Implements the first order predictor algorithm. This algorithm is
mathematically defined as:
@ -22,68 +22,50 @@ def predictor(operator, print_out=True):
A_p &= A(y_n, t_n)
y_{n+1} &= \\text{expm}(A_p h) y_n
y_{n+1} &= \text{expm}(A_p h) y_n
Parameters
----------
operator : Operator
operator : openmc.deplete.Operator
The operator object to simulate on.
print_out : bool, optional
Whether or not to print out time.
"""
"""
# Generate initial conditions
with operator as vec:
n_mats = len(vec)
t = 0.0
for i, dt in enumerate(operator.settings.dt_vec):
# Create vectors
# Get beginning-of-timestep reaction rates
x = [copy.deepcopy(vec)]
seeds = []
eigvls = []
rates_array = []
eigvl, rates, seed = operator(x[0])
eigvls.append(eigvl)
seeds.append(seed)
rates_array.append(rates)
results = [operator(x[0])]
# Create results, write to disk
save_results(operator, x, rates_array, eigvls, seeds, [t, t + dt], i)
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 = (rates_array[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)
# Advance time, update vector
t += dt
vec = copy.deepcopy(x_result)
# Perform one last simulation
x = [copy.deepcopy(vec)]
seeds = []
eigvls = []
rates_array = []
eigvl, rates, seed = operator(x[0])
eigvls.append(eigvl)
seeds.append(seed)
rates_array.append(rates)
results = [operator(x[0])]
# Create results, write to disk
save_results(operator, x, rates_array, eigvls, seeds, [t, t],
len(operator.settings.dt_vec))
save_results(operator, x, results, [t, t], len(operator.settings.dt_vec))

View file

@ -4,8 +4,8 @@
from ..results import Results, write_results
def save_results(op, x, rates, eigvls, seeds, t, step_ind):
""" Creates and writes results to disk
def save_results(op, x, op_results, t, step_ind):
"""Creates and writes depletion results to disk
Parameters
----------
@ -13,18 +13,14 @@ def save_results(op, x, rates, eigvls, seeds, t, step_ind):
The operator used to generate these results.
x : list of list of numpy.array
The prior x vectors. Indexed [i][cell] using the above equation.
rates : list of ReactionRates
The reaction rates for each substep.
eigvls : list of float
Eigenvalue for each substep
seeds : list of int
Seeds for each substep.
op_results : list of openmc.deplete.OperatorResult
Results of applying transport operator
t : list of float
Time indices.
step_ind : int
Step index.
"""
"""
# Get indexing terms
vol_list, nuc_list, burn_list, full_burn_list = op.get_results_info()
@ -39,9 +35,9 @@ def save_results(op, x, rates, eigvls, seeds, t, step_ind):
for mat_i in range(n_mat):
results[i, mat_i, :] = x[i][mat_i][:]
results.k = eigvls
results.seeds = seeds
results.k = [r.k for r in op_results]
results.seeds = [r.seed for r in op_results]
results.rates = [r.rates for r in op_results]
results.time = t
results.rates = rates
write_results(results, "depletion_results.h5", step_ind)

View file

@ -23,7 +23,7 @@ import openmc
import openmc.capi
from openmc.data import JOULE_PER_EV
from . import comm
from .abc import Settings, Operator
from .abc import Settings, Operator, OperatorResult
from .atom_number import AtomNumber
from .chain import Chain
from .reaction_rates import ReactionRates
@ -255,7 +255,7 @@ class OpenMCOperator(Operator):
print("Time to openmc: ", time_openmc - time_start)
print("Time to unpack: ", time_unpack - time_openmc)
return k, copy.deepcopy(self.reaction_rates), self.seed
return OperatorResult(k, copy.deepcopy(self.reaction_rates), self.seed)
def extract_mat_ids(self):
"""Extracts materials and assigns them to processes.

View file

@ -1,7 +1,7 @@
import numpy as np
import scipy.sparse as sp
from openmc.deplete.reaction_rates import ReactionRates
from openmc.deplete.abc import Operator
from openmc.deplete.abc import Operator, OperatorResult
class DummyGeometry(Operator):
@ -51,7 +51,7 @@ class DummyGeometry(Operator):
reaction_rates[0, 1, 0] = vec[0][1]
# Create a fake rates object
return 0.0, reaction_rates, 0
return OperatorResult(0.0, reaction_rates, 0)
@property
def chain(self):

View file

@ -11,7 +11,8 @@ import os
from unittest.mock import MagicMock
import numpy as np
from openmc.deplete import integrator, ReactionRates, results, comm
from openmc.deplete import (integrator, ReactionRates, results, comm,
OperatorResult)
def test_save_results(run_in_tmpdir):
@ -49,8 +50,8 @@ def test_save_results(run_in_tmpdir):
x2.append([np.random.rand(2), np.random.rand(2)])
# Construct r
cell_dict = {s:i for i, s in enumerate(burn_list)}
r1 = ReactionRates(cell_dict, {"na":0, "nb":1}, {"ra":0, "rb":1})
cell_dict = {s: i for i, s in enumerate(burn_list)}
r1 = ReactionRates(cell_dict, {"na": 0, "nb": 1}, {"ra": 0, "rb": 1})
r1.rates = np.random.rand(2, 2, 2)
rate1 = []
@ -76,8 +77,12 @@ def test_save_results(run_in_tmpdir):
t1 = [0.0, 1.0]
t2 = [1.0, 2.0]
integrator.save_results(op, x1, rate1, eigvl1, seed1, t1, 0)
integrator.save_results(op, x2, rate2, eigvl2, seed2, t2, 1)
op_result1 = [OperatorResult(k, rates, seed)
for k, rates, seed in zip(eigvl1, rate1, seed1)]
op_result2 = [OperatorResult(k, rates, seed)
for k, rates, seed in zip(eigvl2, rate2, seed2)]
integrator.save_results(op, x1, op_result1, t1, 0)
integrator.save_results(op, x2, op_result2, t2, 1)
# Load the files
res = results.read_results("depletion_results.h5")