Apply memoization in get_all_universes (#2995)

Co-authored-by: Paul Romano <[email protected]>

openmc/cell.py

@@ -426,15 +426,13 @@ def get_all_cells(self, memo=None):
             instances
 
         """
+        if memo is None:
+            memo = set()
+        elif self in memo:
+            return {}
+        memo.add(self)
 
         cells = {}
-
-        if memo and self in memo:
-            return cells
-
-        if memo is not None:
-            memo.add(self)
-
         if self.fill_type in ('universe', 'lattice'):
             cells.update(self.fill.get_all_cells(memo))
 
@@ -465,7 +463,7 @@ def get_all_materials(self, memo=None):
 
         return materials
 
-    def get_all_universes(self):
+    def get_all_universes(self, memo=None):
         """Return all universes that are contained within this one if any of
         its cells are filled with a universe or lattice.
 
@@ -476,18 +474,22 @@ def get_all_universes(self):
             :class:`Universe` instances
 
         """
+        if memo is None:
+            memo = set()
+        if self in memo:
+            return {}
+        memo.add(self)
 
         universes = {}
-
         if self.fill_type == 'universe':
             universes[self.fill.id] = self.fill
-            universes.update(self.fill.get_all_universes())
+            universes.update(self.fill.get_all_universes(memo))
         elif self.fill_type == 'lattice':
-            universes.update(self.fill.get_all_universes())
+            universes.update(self.fill.get_all_universes(memo))
 
         return universes
 
-    def clone(self, clone_materials=True, clone_regions=True,  memo=None):
+    def clone(self, clone_materials=True, clone_regions=True, memo=None):
         """Create a copy of this cell with a new unique ID, and clones
         the cell's region and fill.
 
@@ -678,10 +680,11 @@ def create_xml_subelement(self, xml_element, memo=None):
             # thus far.
             def create_surface_elements(node, element, memo=None):
                 if isinstance(node, Halfspace):
-                    if memo and node.surface in memo:
+                    if memo is None:
+                        memo = set()
+                    elif node.surface in memo:
                         return
-                    if memo is not None:
-                        memo.add(node.surface)
+                    memo.add(node.surface)
                     xml_element.append(node.surface.to_xml_element())
 
                 elif isinstance(node, Complement):

openmc/geometry.py

@@ -134,7 +134,7 @@ def to_xml_element(self, remove_surfs=False) -> ET.Element:
 
         # Create XML representation
         element = ET.Element("geometry")
-        self.root_universe.create_xml_subelement(element, memo=set())
+        self.root_universe.create_xml_subelement(element)
 
         # Sort the elements in the file
         element[:] = sorted(element, key=lambda x: (
@@ -373,7 +373,7 @@ def get_all_cells(self) -> typing.Dict[int, openmc.Cell]:
 
         """
         if self.root_universe is not None:
-            return self.root_universe.get_all_cells(memo=set())
+            return self.root_universe.get_all_cells()
         else:
             return {}
 
@@ -417,7 +417,7 @@ def get_all_materials(self) -> typing.Dict[int, openmc.Material]:
 
         """
         if self.root_universe is not None:
-            return self.root_universe.get_all_materials(memo=set())
+            return self.root_universe.get_all_materials()
         else:
             return {}
 

openmc/lattice.py

@@ -170,11 +170,11 @@ def get_all_cells(self, memo=None):
         """
         cells = {}
 
-        if memo and self in memo:
+        if memo is None:
+            memo = set()
+        elif self in memo:
             return cells
-
-        if memo is not None:
-            memo.add(self)
+        memo.add(self)
 
         unique_universes = self.get_unique_universes()
 
@@ -194,6 +194,9 @@ def get_all_materials(self, memo=None):
 
         """
 
+        if memo is None:
+            memo = set()
+
         materials = {}
 
         # Append all Cells in each Cell in the Universe to the dictionary
@@ -203,7 +206,7 @@ def get_all_materials(self, memo=None):
 
         return materials
 
-    def get_all_universes(self):
+    def get_all_universes(self, memo=None):
         """Return all universes that are contained within the lattice
 
         Returns
@@ -213,11 +216,16 @@ def get_all_universes(self):
             :class:`Universe` instances
 
         """
-
         # Initialize a dictionary of all Universes contained by the Lattice
         # in each nested Universe level
         all_universes = {}
 
+        if memo is None:
+            memo = set()
+        elif self in memo:
+            return all_universes
+        memo.add(self)
+
         # Get all unique Universes contained in each of the lattice cells
         unique_universes = self.get_unique_universes()
 
@@ -226,7 +234,7 @@ def get_all_universes(self):
 
         # Append all Universes containing each cell to the dictionary
         for universe in unique_universes.values():
-            all_universes.update(universe.get_all_universes())
+            all_universes.update(universe.get_all_universes(memo))
 
         return all_universes
 
@@ -846,10 +854,11 @@ def create_xml_subelement(self, xml_element, memo=None):
 
         """
         # If the element already contains the Lattice subelement, then return
-        if memo and self in memo:
+        if memo is None:
+            memo = set()
+        elif self in memo:
             return
-        if memo is not None:
-            memo.add(self)
+        memo.add(self)
 
         # Make sure universes have been assigned
         if self.universes is None:
@@ -1417,10 +1426,11 @@ def is_valid_index(self, idx):
 
     def create_xml_subelement(self, xml_element, memo=None):
         # If this subelement has already been written, return
-        if memo and self in memo:
+        if memo is None:
+            memo = set()
+        elif self in memo:
             return
-        if memo is not None:
-            memo.add(self)
+        memo.add(self)
 
         lattice_subelement = ET.Element("hex_lattice")
         lattice_subelement.set("id", str(self._id))

openmc/tallies.py

@@ -3209,6 +3209,7 @@ def _create_derivative_subelements(self, root_element):
     def to_xml_element(self, memo=None):
         """Creates a 'tallies' element to be written to an XML file.
         """
+        memo = memo if memo is not None else set()
         element = ET.Element("tallies")
         self._create_mesh_subelements(element, memo)
         self._create_filter_subelements(element)

openmc/universe.py

@@ -90,7 +90,7 @@ def add_volume_information(self, volume_calc):
         else:
             raise ValueError('No volume information found for this universe.')
 
-    def get_all_universes(self):
+    def get_all_universes(self, memo=None):
         """Return all universes that are contained within this one.
 
         Returns
@@ -100,10 +100,16 @@ def get_all_universes(self):
             :class:`Universe` instances
 
         """
+        if memo is None:
+            memo = set()
+        elif self in memo:
+            return {}
+        memo.add(self)
+
         # Append all Universes within each Cell to the dictionary
         universes = {}
         for cell in self.get_all_cells().values():
-            universes.update(cell.get_all_universes())
+            universes.update(cell.get_all_universes(memo))
 
         return universes
 
@@ -639,15 +645,14 @@ def get_all_cells(self, memo=None):
 
         """
 
-        cells = {}
-
-        if memo and self in memo:
-            return cells
-
-        if memo is not None:
-            memo.add(self)
+        if memo is None:
+            memo = set()
+        elif self in memo:
+            return {}
+        memo.add(self)
 
         # Add this Universe's cells to the dictionary
+        cells = {}
         cells.update(self._cells)
 
         # Append all Cells in each Cell in the Universe to the dictionary
@@ -667,6 +672,9 @@ def get_all_materials(self, memo=None):
 
         """
 
+        if memo is None:
+            memo = set()
+
         materials = {}
 
         # Append all Cells in each Cell in the Universe to the dictionary
@@ -677,15 +685,17 @@ def get_all_materials(self, memo=None):
         return materials
 
     def create_xml_subelement(self, xml_element, memo=None):
+        if memo is None:
+            memo = set()
+
         # Iterate over all Cells
         for cell in self._cells.values():
 
             # If the cell was already written, move on
-            if memo and cell in memo:
+            if cell in memo:
                 continue
 
-            if memo is not None:
-                memo.add(cell)
+            memo.add(cell)
 
             # Create XML subelement for this Cell
             cell_element = cell.create_xml_subelement(xml_element, memo)
@@ -928,11 +938,13 @@ def n_surfaces(self):
         return self._n_geom_elements('surface')
 
     def create_xml_subelement(self, xml_element, memo=None):
-        if memo and self in memo:
+        if memo is None:
+            memo = set()
+
+        if self in memo:
             return
 
-        if memo is not None:
-            memo.add(self)
+        memo.add(self)
 
         # Set xml element values
         dagmc_element = ET.Element('dagmc_universe')