From c9abcfc0a1890da0263642e22b3cd9675f28242c Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Sat, 17 Feb 2018 10:28:24 -0600 Subject: [PATCH] Make ReactionRates a subclass of ndarray, simplifying mapping dictionaries --- openmc/deplete/chain.py | 41 +++++++-------- openmc/deplete/openmc_wrapper.py | 20 ++++---- openmc/deplete/reaction_rates.py | 86 +++++++++++--------------------- openmc/deplete/results.py | 12 ++--- 4 files changed, 63 insertions(+), 96 deletions(-) diff --git a/openmc/deplete/chain.py b/openmc/deplete/chain.py index 846b3af564..2187b04860 100644 --- a/openmc/deplete/chain.py +++ b/openmc/deplete/chain.py @@ -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 diff --git a/openmc/deplete/openmc_wrapper.py b/openmc/deplete/openmc_wrapper.py index b8b9065276..46d07a23f7 100644 --- a/openmc/deplete/openmc_wrapper.py +++ b/openmc/deplete/openmc_wrapper.py @@ -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) diff --git a/openmc/deplete/reaction_rates.py b/openmc/deplete/reaction_rates.py index e8bab101bf..b437773826 100644 --- a/openmc/deplete/reaction_rates.py +++ b/openmc/deplete/reaction_rates.py @@ -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) diff --git a/openmc/deplete/results.py b/openmc/deplete/results.py index 37c4b6e921..4a5ca24352 100644 --- a/openmc/deplete/results.py +++ b/openmc/deplete/results.py @@ -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