Adding memo for get_all_cells/get_all_materials.

openmc/cell.py

@@ -350,7 +350,7 @@ class Cell(IDManagerMixin):
 
         return nuclides
 
-    def get_all_cells(self):
+    def get_all_cells(self, memo=None):
         """Return all cells that are contained within this one if it is filled with a
         universe or lattice
 
@@ -364,12 +364,18 @@ class Cell(IDManagerMixin):
 
         cells = OrderedDict()
 
+        if memo and id(self) in memo:
+            return cells
+
+        if memo is not None:
+            memo.add(id(self))
+
         if self.fill_type in ('universe', 'lattice'):
-            cells.update(self.fill.get_all_cells())
+            cells.update(self.fill.get_all_cells(memo))
 
         return cells
 
-    def get_all_materials(self):
+    def get_all_materials(self, memo=None):
         """Return all materials that are contained within the cell
 
         Returns
@@ -388,9 +394,9 @@ class Cell(IDManagerMixin):
                     materials[m.id] = m
         else:
             # Append all Cells in each Cell in the Universe to the dictionary
-            cells = self.get_all_cells()
+            cells = self.get_all_cells(memo)
             for cell in cells.values():
-                materials.update(cell.get_all_materials())
+                materials.update(cell.get_all_materials(memo))
 
         return materials
 

openmc/geometry.py

@@ -273,8 +273,11 @@ class Geometry(object):
             Dictionary mapping cell IDs to :class:`openmc.Cell` instances
 
         """
+
+        memo = set()
+
         if self.root_universe is not None:
-            return self.root_universe.get_all_cells()
+            return self.root_universe.get_all_cells(memo)
         else:
             return []
 
@@ -303,7 +306,9 @@ class Geometry(object):
             instances
 
         """
-        return self.root_universe.get_all_materials()
+        memo = set()
+
+        return self.root_universe.get_all_materials(memo)
 
     def get_all_material_cells(self):
         """Return all cells filled by a material

openmc/lattice.py

@@ -332,7 +332,7 @@ class Lattice(IDManagerMixin, metaclass=ABCMeta):
 
         return nuclides
 
-    def get_all_cells(self):
+    def get_all_cells(self, memo=None):
         """Return all cells that are contained within the lattice
 
         Returns
@@ -344,14 +344,21 @@ class Lattice(IDManagerMixin, metaclass=ABCMeta):
         """
 
         cells = OrderedDict()
+
+        if memo and id(self) in memo:
+            return cells
+
+        if memo is not None:
+            memo.add(id(self))
+
         unique_universes = self.get_unique_universes()
 
         for universe_id, universe in unique_universes.items():
-            cells.update(universe.get_all_cells())
+            cells.update(universe.get_all_cells(memo))
 
         return cells
 
-    def get_all_materials(self):
+    def get_all_materials(self, memo=None):
         """Return all materials that are contained within the lattice
 
         Returns
@@ -365,9 +372,9 @@ class Lattice(IDManagerMixin, metaclass=ABCMeta):
         materials = OrderedDict()
 
         # Append all Cells in each Cell in the Universe to the dictionary
-        cells = self.get_all_cells()
+        cells = self.get_all_cells(memo)
         for cell_id, cell in cells.items():
-            materials.update(cell.get_all_materials())
+            materials.update(cell.get_all_materials(memo))
 
         return materials
 

openmc/universe.py

@@ -418,7 +418,7 @@ class Universe(IDManagerMixin):
 
         return nuclides
 
-    def get_all_cells(self):
+    def get_all_cells(self, memo=None):
         """Return all cells that are contained within the universe
 
         Returns
@@ -431,16 +431,22 @@ class Universe(IDManagerMixin):
 
         cells = OrderedDict()
 
+        if memo and id(self) in memo:
+            return cells
+
+        if memo is not None:
+            memo.add(id(self))
+
         # Add this Universe's cells to the dictionary
         cells.update(self._cells)
 
         # Append all Cells in each Cell in the Universe to the dictionary
         for cell in self._cells.values():
-            cells.update(cell.get_all_cells())
+            cells.update(cell.get_all_cells(memo))
 
         return cells
 
-    def get_all_materials(self):
+    def get_all_materials(self, memo=None):
         """Return all materials that are contained within the universe
 
         Returns
@@ -454,9 +460,9 @@ class Universe(IDManagerMixin):
         materials = OrderedDict()
 
         # Append all Cells in each Cell in the Universe to the dictionary
-        cells = self.get_all_cells()
+        cells = self.get_all_cells(memo)
         for cell in cells.values():
-            materials.update(cell.get_all_materials())
+            materials.update(cell.get_all_materials(memo))
 
         return materials