From 05afc55a88c11200d387519de6a71be9063dbbbc Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Thu, 15 Feb 2018 15:34:00 -0600 Subject: [PATCH] Give deplete.Chain some special methods to make it more Pythonic --- openmc/deplete/chain.py | 34 +++++++++----------------- openmc/deplete/openmc_wrapper.py | 9 +++---- tests/unit_tests/test_deplete_chain.py | 25 ++++++++++--------- 3 files changed, 29 insertions(+), 39 deletions(-) diff --git a/openmc/deplete/chain.py b/openmc/deplete/chain.py index 03decb72ab..846b3af564 100644 --- a/openmc/deplete/chain.py +++ b/openmc/deplete/chain.py @@ -113,8 +113,6 @@ class Chain(object): Attributes ---------- - n_nuclides : int - Number of nuclides in chain. nuclides : list of Nuclide List of nuclides in chain. nuclide_dict : OrderedDict of str to int @@ -132,8 +130,14 @@ class Chain(object): self.nuc_to_react_ind = OrderedDict() self.react_to_ind = OrderedDict() - @property - def n_nuclides(self): + def __contains__(self, nuclide): + return nuclide in self.nuclide_dict + + def __getitem__(self, name): + """Get a Nuclide by name.""" + return self.nuclides[self.nuclide_dict[name]] + + def __len__(self): """Number of nuclides in chain.""" return len(self.nuclides) @@ -381,14 +385,14 @@ class Chain(object): Parameters ---------- rates : numpy.ndarray - 2D array indexed by nuclide then by cell. + 2D array indexed by (nuclide, reaction) Returns ------- scipy.sparse.csr_matrix Sparse matrix representing depletion. - """ + """ matrix = defaultdict(float) reactions = set() @@ -450,21 +454,7 @@ class Chain(object): reactions.clear() # Use DOK matrix as intermediate representation, then convert to CSR and return - matrix_dok = sp.dok_matrix((self.n_nuclides, self.n_nuclides)) + n = len(self) + matrix_dok = sp.dok_matrix((n, n)) dict.update(matrix_dok, matrix) return matrix_dok.tocsr() - - def nuc_by_ind(self, ind): - """Extracts nuclides from the list by dictionary key. - - Parameters - ---------- - ind : str - Name of nuclide. - - Returns - ------- - Nuclide - Nuclide object that corresponds to ind. - """ - return self.nuclides[self.nuclide_dict[ind]] diff --git a/openmc/deplete/openmc_wrapper.py b/openmc/deplete/openmc_wrapper.py index 65ed4793dd..b8b9065276 100644 --- a/openmc/deplete/openmc_wrapper.py +++ b/openmc/deplete/openmc_wrapper.py @@ -328,7 +328,7 @@ class OpenMCOperator(Operator): i += 1 n_mat_burn = len(mat_burn) - n_nuc_burn = len(self.chain.nuclide_dict) + n_nuc_burn = len(self.chain) self.number = AtomNumber(mat_dict, nuc_dict, volume, n_mat_burn, n_nuc_burn) @@ -500,8 +500,7 @@ class OpenMCOperator(Operator): # Store list of tally nuclides on each process nuc_list = comm.bcast(nuc_list, root=0) - tally_nuclides = [nuc for nuc in nuc_list - if nuc in self.chain.nuclide_dict] + tally_nuclides = [nuc for nuc in nuc_list if nuc in self.chain] return tally_nuclides @@ -690,14 +689,14 @@ class OpenMCOperator(Operator): # and nuclides in depletion chain. if name not in self.participating_nuclides: self.participating_nuclides.add(name) - if name in self.chain.nuclide_dict: + if name in self.chain: self.burn_nuc_to_ind[name] = nuc_ind nuc_ind += 1 @property def n_nuc(self): """Number of nuclides considered in the decay chain.""" - return len(self.chain.nuclides) + return len(self.chain) def get_results_info(self): """Returns volume list, cell lists, and nuc lists. diff --git a/tests/unit_tests/test_deplete_chain.py b/tests/unit_tests/test_deplete_chain.py index 064b878afe..3a065fb46f 100644 --- a/tests/unit_tests/test_deplete_chain.py +++ b/tests/unit_tests/test_deplete_chain.py @@ -20,12 +20,12 @@ def test_init(): assert isinstance(dep.react_to_ind, Mapping) -def test_n_nuclides(): - """Test depletion chain n_nuclides parameter.""" +def test_len(): + """Test depletion chain length.""" dep = Chain() dep.nuclides = ["NucA", "NucB", "NucC"] - assert dep.n_nuclides == 3 + assert len(dep) == 3 def test_from_endf(): @@ -43,10 +43,10 @@ def test_from_xml(): dep = Chain.from_xml(_test_filename) # Basic checks - assert dep.n_nuclides == 3 + assert len(dep) == 3 # A tests - nuc = dep.nuclides[dep.nuclide_dict["A"]] + nuc = dep["A"] assert nuc.name == "A" assert nuc.half_life == 2.36520E+04 @@ -61,7 +61,7 @@ def test_from_xml(): assert [r.branching_ratio for r in nuc.reactions] == [1.0] # B tests - nuc = dep.nuclides[dep.nuclide_dict["B"]] + nuc = dep["B"] assert nuc.name == "B" assert nuc.half_life == 3.29040E+04 @@ -76,7 +76,7 @@ def test_from_xml(): assert [r.branching_ratio for r in nuc.reactions] == [1.0] # C tests - nuc = dep.nuclides[dep.nuclide_dict["C"]] + nuc = dep["C"] assert nuc.name == "C" assert nuc.n_decay_modes == 0 @@ -183,12 +183,13 @@ def test_form_matrix(): assert mat[2, 2] == mat22 -def test_nuc_by_ind(): +def test_getitem(): """ Test nuc_by_ind converter function. """ dep = Chain() dep.nuclides = ["NucA", "NucB", "NucC"] - dep.nuclide_dict = {"NucA" : 0, "NucB" : 1, "NucC" : 2} + dep.nuclide_dict = {nuc: dep.nuclides.index(nuc) + for nuc in dep.nuclides} - assert "NucA" == dep.nuc_by_ind("NucA") - assert "NucB" == dep.nuc_by_ind("NucB") - assert "NucC" == dep.nuc_by_ind("NucC") + assert "NucA" == dep["NucA"] + assert "NucB" == dep["NucB"] + assert "NucC" == dep["NucC"]