Make ReactionRates a subclass of ndarray, simplifying mapping dictionaries

This commit is contained in:
Paul Romano 2018-02-17 10:28:24 -06:00
parent a6c095c4e9
commit c9abcfc0a1
4 changed files with 63 additions and 96 deletions

View file

@ -117,9 +117,7 @@ class Chain(object):
List of nuclides in chain.
nuclide_dict : OrderedDict of str to int
Maps a nuclide name to an index in nuclides.
nuc_to_react_ind : OrderedDict of str to int
Dictionary mapping a nuclide name to an index in ReactionRates.
react_to_ind : OrderedDict of str to int
index_reaction : OrderedDict of str to int
Dictionary mapping a reaction name to an index in ReactionRates.
"""
@ -127,8 +125,7 @@ class Chain(object):
def __init__(self):
self.nuclides = []
self.nuclide_dict = OrderedDict()
self.nuc_to_react_ind = OrderedDict()
self.react_to_ind = OrderedDict()
self.index_reaction = OrderedDict()
def __contains__(self, nuclide):
return nuclide in self.nuclide_dict
@ -155,7 +152,7 @@ class Chain(object):
List of ENDF neutron reaction sub-library files
"""
depl_chain = cls()
chain = cls()
# Create dictionary mapping target to filename
reactions = {}
@ -200,8 +197,8 @@ class Chain(object):
nuclide = Nuclide()
nuclide.name = parent
depl_chain.nuclides.append(nuclide)
depl_chain.nuclide_dict[parent] = idx
chain.nuclides.append(nuclide)
chain.nuclide_dict[parent] = idx
if not data.nuclide['stable'] and data.half_life.nominal_value != 0.0:
nuclide.half_life = data.half_life.nominal_value
@ -235,8 +232,8 @@ class Chain(object):
Z = data.nuclide['atomic_number'] + delta_Z
daughter = '{}{}'.format(openmc.data.ATOMIC_SYMBOL[Z], A)
if name not in depl_chain.react_to_ind:
depl_chain.react_to_ind[name] = reaction_index
if name not in chain.index_reaction:
chain.index_reaction[name] = reaction_index
reaction_index += 1
if daughter not in decay_data:
@ -259,8 +256,8 @@ class Chain(object):
nuclide.reactions.append(
ReactionTuple('fission', 0, q_value, 1.0))
if 'fission' not in depl_chain.react_to_ind:
depl_chain.react_to_ind['fission'] = reaction_index
if 'fission' not in chain.index_reaction:
chain.index_reaction['fission'] = reaction_index
reaction_index += 1
else:
missing_fpy.append(parent)
@ -316,7 +313,7 @@ class Chain(object):
for vals in missing_fp:
print(' {}, E={} eV (total yield={})'.format(*vals))
return depl_chain
return chain
@classmethod
def from_xml(cls, filename):
@ -331,7 +328,7 @@ class Chain(object):
----
Allow for branching on capture, etc.
"""
depl_chain = cls()
chain = cls()
# Load XML tree
try:
@ -346,17 +343,17 @@ class Chain(object):
reaction_index = 0
for i, nuclide_elem in enumerate(root.findall('nuclide_table')):
nuc = Nuclide.from_xml(nuclide_elem)
depl_chain.nuclide_dict[nuc.name] = i
chain.nuclide_dict[nuc.name] = i
# Check for reaction paths
for rx in nuc.reactions:
if rx.type not in depl_chain.react_to_ind:
depl_chain.react_to_ind[rx.type] = reaction_index
if rx.type not in chain.index_reaction:
chain.index_reaction[rx.type] = reaction_index
reaction_index += 1
depl_chain.nuclides.append(nuc)
chain.nuclides.append(nuc)
return depl_chain
return chain
def export_to_xml(self, filename):
"""Writes a depletion chain XML file.
@ -416,14 +413,14 @@ class Chain(object):
k = self.nuclide_dict[target]
matrix[k, i] += branch_val
if nuc.name in self.nuc_to_react_ind:
if nuc.name in rates.index_nuc:
# Extract all reactions for this nuclide in this cell
nuc_ind = self.nuc_to_react_ind[nuc.name]
nuc_ind = rates.index_nuc[nuc.name]
nuc_rates = rates[nuc_ind, :]
for r_type, target, _, br in nuc.reactions:
# Extract reaction index, and then final reaction rate
r_id = self.react_to_ind[r_type]
r_id = rates.index_rx[r_type]
path_rate = nuc_rates[r_id]
# Loss term -- make sure we only count loss once for

View file

@ -369,9 +369,7 @@ class OpenMCOperator(Operator):
self.reaction_rates = ReactionRates(
self.burn_mat_to_ind,
self.burn_nuc_to_ind,
self.chain.react_to_ind)
self.chain.nuc_to_react_ind = self.burn_nuc_to_ind
self.chain.index_reaction)
def form_matrix(self, y, mat):
"""Forms the depletion matrix.
@ -522,7 +520,7 @@ class OpenMCOperator(Operator):
# transmutation. The nuclides for the tally are set later when eval() is
# called.
tally_dep = openmc.capi.Tally(1)
tally_dep.scores = self.chain.react_to_ind.keys()
tally_dep.scores = self.chain.index_reaction.keys()
tally_dep.filters = [mat_filter]
def total_density_list(self):
@ -579,11 +577,11 @@ class OpenMCOperator(Operator):
# Extract tally bins
materials = list(self.mat_tally_ind.keys())
nuclides = openmc.capi.tallies[1].nuclides
reactions = list(self.chain.react_to_ind.keys())
reactions = list(self.chain.index_reaction.keys())
# Form fast map
nuc_ind = [rates.nuc_to_ind[nuc] for nuc in nuclides]
react_ind = [rates.react_to_ind[react] for react in reactions]
nuc_ind = [rates.index_nuc[nuc] for nuc in nuclides]
react_ind = [rates.index_rx[react] for react in reactions]
# Compute fission power
# TODO : improve this calculation
@ -598,13 +596,13 @@ class OpenMCOperator(Operator):
rates_expanded = np.zeros((rates.n_nuc, rates.n_react))
number = np.zeros(rates.n_nuc)
fission_ind = rates.react_to_ind["fission"]
fission_ind = rates.index_rx["fission"]
for nuclide in self.chain.nuclides:
if nuclide.name in rates.nuc_to_ind:
if nuclide.name in rates.index_nuc:
for rx in nuclide.reactions:
if rx.type == 'fission':
ind = rates.nuc_to_ind[nuclide.name]
ind = rates.index_nuc[nuclide.name]
fission_Q[ind] = rx.Q
break
@ -637,7 +635,7 @@ class OpenMCOperator(Operator):
for react in react_ind:
rates_expanded[i_nuc_results, react] /= number[i_nuc_results]
rates.rates[i, :, :] = rates_expanded
rates[i, :, :] = rates_expanded
# Reduce energy produced from all processes
energy = comm.allreduce(energy)

View file

@ -6,7 +6,7 @@ An ndarray to store reaction rates with string, integer, or slice indexing.
import numpy as np
class ReactionRates(object):
class ReactionRates(np.ndarray):
"""ReactionRates class.
An ndarray to store reaction rates with string, integer, or slice indexing.
@ -38,76 +38,48 @@ class ReactionRates(object):
Array storing rates indexed by the above dictionaries.
"""
def __init__(self, mat_to_ind, nuc_to_ind, react_to_ind):
def __new__(cls, index_mat, index_nuc, index_rx):
# Create appropriately-sized zeroed-out ndarray
shape = (len(index_mat), len(index_nuc), len(index_rx))
obj = super().__new__(cls, shape)
obj[:] = 0.0
self.mat_to_ind = mat_to_ind
self.nuc_to_ind = nuc_to_ind
self.react_to_ind = react_to_ind
# Add mapping attributes
obj.index_mat = index_mat
obj.index_nuc = index_nuc
obj.index_rx = index_rx
self.rates = np.zeros((self.n_mat, self.n_nuc, self.n_react))
return obj
def __getitem__(self, pos):
"""Retrieves an item from reaction_rates.
def __array_finalize__(self, obj):
if obj is None:
return
self.index_mat = getattr(obj, 'index_mat', None)
self.index_nuc = getattr(obj, 'index_nuc', None)
self.index_rx = getattr(obj, 'index_rx', None)
Parameters
----------
pos : tuple
A three-length tuple containing a material index, a nuc index, and a
reaction index. These indexes can be strings (which get converted
to integers via the dictionaries), integers used directly, or
slices.
def __reduce__(self):
state = super().__reduce__()
new_state = state[2] + (self.index_mat, self.index_nuc, self.index_rx)
return (state[0], state[1], new_state)
Returns
-------
numpy.array
The value indexed from self.rates.
"""
mat, nuc, react = pos
if isinstance(mat, str):
mat = self.mat_to_ind[mat]
if isinstance(nuc, str):
nuc = self.nuc_to_ind[nuc]
if isinstance(react, str):
react = self.react_to_ind[react]
return self.rates[mat, nuc, react]
def __setitem__(self, pos, val):
"""Sets an item from reaction_rates.
Parameters
----------
pos : tuple
A three-length tuple containing a material index, a nuc index, and a
reaction index. These indexes can be strings (which get converted
to integers via the dictionaries), integers used directly, or
slices.
val : float
The value to set the array to.
"""
mat, nuc, react = pos
if isinstance(mat, str):
mat = self.mat_to_ind[mat]
if isinstance(nuc, str):
nuc = self.nuc_to_ind[nuc]
if isinstance(react, str):
react = self.react_to_ind[react]
self.rates[mat, nuc, react] = val
def __setstate__(self, state):
self.index_mat = state[-3]
self.index_nuc = state[-2]
self.index_rx = state[-1]
super().__setstate__(state[0:-3])
@property
def n_mat(self):
"""Number of cells."""
return len(self.mat_to_ind)
return len(self.index_mat)
@property
def n_nuc(self):
"""Number of nucs."""
return len(self.nuc_to_ind)
return len(self.index_nuc)
@property
def n_react(self):
"""Number of reactions."""
return len(self.react_to_ind)
return len(self.index_rx)

View file

@ -180,11 +180,11 @@ class Results(object):
mat_int = sorted([int(mat) for mat in self.mat_to_hdf5_ind])
mat_list = [str(mat) for mat in mat_int]
nuc_list = sorted(self.nuc_to_ind.keys())
rxn_list = sorted(self.rates[0].react_to_ind.keys())
rxn_list = sorted(self.rates[0].index_rx.keys())
n_mats = self.n_hdf5_mats
n_nuc_number = len(nuc_list)
n_nuc_rxn = len(self.rates[0].nuc_to_ind)
n_nuc_rxn = len(self.rates[0].index_nuc)
n_rxn = len(rxn_list)
n_stages = self.n_stages
@ -200,14 +200,14 @@ class Results(object):
for nuc in nuc_list:
nuc_single_group = nuc_group.create_group(nuc)
nuc_single_group.attrs["atom number index"] = self.nuc_to_ind[nuc]
if nuc in self.rates[0].nuc_to_ind:
nuc_single_group.attrs["reaction rate index"] = self.rates[0].nuc_to_ind[nuc]
if nuc in self.rates[0].index_nuc:
nuc_single_group.attrs["reaction rate index"] = self.rates[0].index_nuc[nuc]
rxn_group = handle.create_group("reactions")
for rxn in rxn_list:
rxn_single_group = rxn_group.create_group(rxn)
rxn_single_group.attrs["index"] = self.rates[0].react_to_ind[rxn]
rxn_single_group.attrs["index"] = self.rates[0].index_rx[rxn]
# Construct array storage
@ -341,7 +341,7 @@ class Results(object):
for i in range(results.n_stages):
rate = ReactionRates(results.mat_to_ind, rxn_nuc_to_ind, rxn_to_ind)
rate.rates = handle["/reaction rates"][index, i, :, :, :]
rate[:] = handle["/reaction rates"][index, i, :, :, :]
results.rates.append(rate)
return results