Give deplete.Chain some special methods to make it more Pythonic

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]]

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.

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"]