diff --git a/openmc/deplete/atom_number.py b/openmc/deplete/atom_number.py index 5179224c2..8f6a41911 100644 --- a/openmc/deplete/atom_number.py +++ b/openmc/deplete/atom_number.py @@ -2,6 +2,7 @@ An ndarray to store atom densities with string, integer, or slice indexing. """ +from collections import OrderedDict import numpy as np @@ -44,8 +45,8 @@ class AtomNumber(object): """ def __init__(self, local_mats, nuclides, volume, n_nuc_burn): - self.index_mat = {mat: i for i, mat in enumerate(local_mats)} - self.index_nuc = {nuc: i for i, nuc in enumerate(nuclides)} + self.index_mat = OrderedDict((mat, i) for i, mat in enumerate(local_mats)) + self.index_nuc = OrderedDict((nuc, i) for i, nuc in enumerate(nuclides)) self.volume = np.ones(len(local_mats)) for mat, val in volume.items(): diff --git a/openmc/deplete/integrator/save_results.py b/openmc/deplete/integrator/save_results.py index 40d2f70b6..98245fdea 100644 --- a/openmc/deplete/integrator/save_results.py +++ b/openmc/deplete/integrator/save_results.py @@ -22,12 +22,12 @@ def save_results(op, x, op_results, t, step_ind): """ # Get indexing terms - vol_list, nuc_list, burn_list, full_burn_list = op.get_results_info() + vol_dict, nuc_list, burn_list, full_burn_list = op.get_results_info() # Create results stages = len(x) results = Results() - results.allocate(vol_list, nuc_list, burn_list, full_burn_list, stages) + results.allocate(vol_dict, nuc_list, burn_list, full_burn_list, stages) n_mat = len(burn_list) diff --git a/openmc/deplete/operator.py b/openmc/deplete/operator.py index af08727ea..d0a64a06d 100644 --- a/openmc/deplete/operator.py +++ b/openmc/deplete/operator.py @@ -111,9 +111,8 @@ class Operator(TransportOperator): self._extract_number(self.local_mats, volume, nuclides) # Create reaction rates array - index_rx = {rx: i for i, rx in enumerate(self.chain.reactions)} self.reaction_rates = ReactionRates( - self.local_mats, self._burnable_nucs, index_rx) + self.local_mats, self._burnable_nucs, self.chain.reactions) def __call__(self, vec, power, print_out=True): """Runs a simulation. diff --git a/openmc/deplete/reaction_rates.py b/openmc/deplete/reaction_rates.py index c66003530..482c357a2 100644 --- a/openmc/deplete/reaction_rates.py +++ b/openmc/deplete/reaction_rates.py @@ -2,6 +2,7 @@ An ndarray to store reaction rates with string, integer, or slice indexing. """ +from collections import OrderedDict import numpy as np @@ -19,8 +20,8 @@ class ReactionRates(np.ndarray): Material IDs nuclides : list of str Depletable nuclides - index_rx : OrderedDict of str to int - A dictionary mapping reaction name as string to index. + reactions : list of str + Transmutation reactions being tracked Attributes ---------- @@ -46,16 +47,16 @@ class ReactionRates(np.ndarray): # the __array_finalize__ method (discussed here: # https://docs.scipy.org/doc/numpy/user/basics.subclassing.html) - def __new__(cls, local_mats, nuclides, index_rx): + def __new__(cls, local_mats, nuclides, reactions): # Create appropriately-sized zeroed-out ndarray - shape = (len(local_mats), len(nuclides), len(index_rx)) + shape = (len(local_mats), len(nuclides), len(reactions)) obj = super().__new__(cls, shape) obj[:] = 0.0 # Add mapping attributes obj.index_mat = {mat: i for i, mat in enumerate(local_mats)} obj.index_nuc = {nuc: i for i, nuc in enumerate(nuclides)} - obj.index_rx = index_rx + obj.index_rx = {rx: i for i, rx in enumerate(reactions)} return obj diff --git a/openmc/deplete/results.py b/openmc/deplete/results.py index 177158dbe..307d18a40 100644 --- a/openmc/deplete/results.py +++ b/openmc/deplete/results.py @@ -76,16 +76,10 @@ class Results(object): """ self.volume = copy.deepcopy(volume) - self.nuc_to_ind = OrderedDict() - self.mat_to_ind = OrderedDict() + self.nuc_to_ind = {nuc: i for i, nuc in enumerate(nuc_list)} + self.mat_to_ind = {mat: i for i, mat in enumerate(burn_list)} self.mat_to_hdf5_ind = {mat: i for i, mat in enumerate(full_burn_list)} - for i, mat in enumerate(burn_list): - self.mat_to_ind[mat] = i - - for i, nuc in enumerate(nuc_list): - self.nuc_to_ind[nuc] = i - # Create storage array self.data = np.zeros((stages, self.n_mat, self.n_nuc)) @@ -123,8 +117,8 @@ class Results(object): ------- float The atoms for stage, mat, nuc - """ + """ stage, mat, nuc = pos if isinstance(mat, str): mat = self.mat_to_ind[mat] @@ -145,8 +139,8 @@ class Results(object): val : float The value to set data to. - """ + """ stage, mat, nuc = pos if isinstance(mat, str): mat = self.mat_to_ind[mat] @@ -162,8 +156,8 @@ class Results(object): ---------- handle : h5py.File or h5py.Group An hdf5 file or group type to store this in. - """ + """ # Create and save the 5 dictionaries: # quantities # self.mat_to_ind -> self.volume (TODO: support for changing volumes) @@ -234,8 +228,8 @@ class Results(object): An hdf5 file or group type to store this in. index : int What step is this? - """ + """ if "/number" not in handle: comm.barrier() self.create_hdf5(handle) @@ -299,6 +293,7 @@ class Results(object): An hdf5 file or group type to load from. index : int What step is this? + """ results = cls() @@ -357,8 +352,8 @@ def write_results(result, filename, index): Target filename. index : int What step is this? - """ + """ if have_mpi and h5py.get_config().mpi: kwargs = {'driver': 'mpio', 'comm': comm} else: diff --git a/tests/unit_tests/test_deplete_chain.py b/tests/unit_tests/test_deplete_chain.py index ae900e62d..7ac3e2f8d 100644 --- a/tests/unit_tests/test_deplete_chain.py +++ b/tests/unit_tests/test_deplete_chain.py @@ -136,11 +136,10 @@ def test_form_matrix(): chain = Chain.from_xml(_test_filename) - mat_ind = {"10000": 0, "10001": 1} - nuc_ind = {"A": 0, "B": 1, "C": 2} - react_ind = {rx: i for i, rx in enumerate(chain.reactions)} + mats = ["10000", "10001"] + nuclides = ["A", "B", "C"] - react = reaction_rates.ReactionRates(mat_ind, nuc_ind, react_ind) + react = reaction_rates.ReactionRates(mats, nuclides, chain.reactions) react.set("10000", "C", "fission", 1.0) react.set("10000", "A", "(n,gamma)", 2.0) diff --git a/tests/unit_tests/test_deplete_integrator.py b/tests/unit_tests/test_deplete_integrator.py index 78dd6bcef..b20990088 100644 --- a/tests/unit_tests/test_deplete_integrator.py +++ b/tests/unit_tests/test_deplete_integrator.py @@ -26,20 +26,18 @@ def test_save_results(run_in_tmpdir): op = MagicMock() vol_dict = {} - full_burn_dict = {} + full_burn_list = [] - j = 0 for i in range(comm.size): vol_dict[str(2*i)] = 1.2 vol_dict[str(2*i + 1)] = 1.2 - full_burn_dict[str(2*i)] = j - full_burn_dict[str(2*i + 1)] = j + 1 - j += 2 + full_burn_list.append(str(2*i)) + full_burn_list.append(str(2*i + 1)) - burn_list = [str(i) for i in range(2*comm.rank, 2*comm.rank + 2)] + burn_list = full_burn_list[2*comm.rank : 2*comm.rank + 2] nuc_list = ["na", "nb"] - op.get_results_info.return_value = vol_dict, nuc_list, burn_list, full_burn_dict + op.get_results_info.return_value = vol_dict, nuc_list, burn_list, full_burn_list # Construct x x1 = [] @@ -50,18 +48,17 @@ 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}) - r1.rates = np.random.rand(2, 2, 2) + r1 = ReactionRates(burn_list, ["na", "nb"], ["ra", "rb"]) + r1[:] = np.random.rand(2, 2, 2) rate1 = [] rate2 = [] for i in range(stages): rate1.append(copy.deepcopy(r1)) - r1.rates = np.random.rand(2, 2, 2) + r1[:] = np.random.rand(2, 2, 2) rate2.append(copy.deepcopy(r1)) - r1.rates = np.random.rand(2, 2, 2) + r1[:] = np.random.rand(2, 2, 2) # Create global terms eigvl1 = np.random.rand(stages) diff --git a/tests/unit_tests/test_deplete_reaction.py b/tests/unit_tests/test_deplete_reaction.py index e46a7b13d..18639a27a 100644 --- a/tests/unit_tests/test_deplete_reaction.py +++ b/tests/unit_tests/test_deplete_reaction.py @@ -9,9 +9,9 @@ def test_get_set(): local_mats = ["10000", "10001"] nuclides = ["U238", "U235"] - react_to_ind = {"fission": 0, "(n,gamma)": 1} + reactions = ["fission", "(n,gamma)"] - rates = ReactionRates(local_mats, nuclides, react_to_ind) + rates = ReactionRates(local_mats, nuclides, reactions) assert rates.shape == (2, 2, 2) assert np.all(rates == 0.0) @@ -54,9 +54,9 @@ def test_properties(): """Test number of materials property.""" local_mats = ["10000", "10001"] nuclides = ["U238", "U235", "Gd157"] - react_to_ind = {"fission": 0, "(n,gamma)": 1, "(n,2n)": 2, "(n,3n)": 3} + reactions = ["fission", "(n,gamma)", "(n,2n)", "(n,3n)"] - rates = ReactionRates(local_mats, nuclides, react_to_ind) + rates = ReactionRates(local_mats, nuclides, reactions) assert rates.n_mat == 2 assert rates.n_nuc == 3