Singleton approach to tracking openmc xml elements with sets.

openmc/cell.py

@@ -16,6 +16,7 @@ from openmc.region import Region, Intersection, Complement
 from openmc._xml import get_text
 from .mixin import IDManagerMixin
 
+from .geomtrack import ElementTracker
 
 class Cell(IDManagerMixin):
     r"""A region of space defined as the intersection of half-space created by
@@ -464,6 +465,7 @@ class Cell(IDManagerMixin):
         return memo[self]
 
     def create_xml_subelement(self, xml_element):
+        et = ElementTracker()
         element = ET.Element("cell")
         element.set("id", str(self.id))
 
@@ -499,10 +501,13 @@ class Cell(IDManagerMixin):
             # element for the corresponding surface if none has been created
             # thus far.
             def create_surface_elements(node, element):
+                et = ElementTracker()
                 if isinstance(node, Halfspace):
                     path = "./surface[@id='{}']".format(node.surface.id)
-                    if xml_element.find(path) is None:
+                    if node.surface.id not in et.surfaces:
+                        et.add_surface(node.surface.id)
                         xml_element.append(node.surface.to_xml_element())
+
                 elif isinstance(node, Complement):
                     create_surface_elements(node.node, element)
                 else:

openmc/geomtrack.py

@@ -0,0 +1,40 @@
+
+class ElementTracker:
+    class __ElementTracker:
+        def __init__(self):
+            self.cells = set()
+            self.surfaces = set()
+            self.lattices = set()
+            self.universes = set()
+
+        def update(self, element_type, element_id):
+            if element_type == "cell":
+                self.cells.add(element_id)
+            elif element_type == "surface":
+                self.surfaces.add(element_id)
+            elif element_type == "lattice":
+                self.lattices.add(element_id)
+            elif element_type == "universe":
+                self.universes.add(element_id)
+
+        def add_cell(self, cell_id):
+            self.update("cell", cell_id)
+
+        def add_surface(self, surface_id):
+            self.update("surface", surface_id)
+
+        def add_lattice(self, lattice_id):
+            self.update("lattice", lattice_id)
+
+        def add_universe(self, universe_id):
+            self.update("universe", universe_id)
+
+    instance = None
+
+    def __init__(self):
+        if not ElementTracker.instance:
+            ElementTracker.instance = ElementTracker.__ElementTracker()
+        else:
+            pass
+    def __getattr__(self, name):
+        return getattr(self.instance, name)

openmc/lattice.py

@@ -13,6 +13,7 @@ import openmc
 from openmc._xml import get_text
 from openmc.mixin import IDManagerMixin
 
+from .geomtrack import ElementTracker
 
 class Lattice(IDManagerMixin, metaclass=ABCMeta):
     """A repeating structure wherein each element is a universe.
@@ -754,14 +755,16 @@ class RectLattice(Lattice):
                     0 <= idx[2] < self.shape[2])
 
     def create_xml_subelement(self, xml_element):
-
+        et = ElementTracker()
         # Determine if XML element already contains subelement for this Lattice
         path = './lattice[@id=\'{0}\']'.format(self._id)
         test = xml_element.find(path)
 
         # If the element does contain the Lattice subelement, then return
-        if test is not None:
+        if self._id in et.lattices:
             return
+        else:
+            et.add_lattice(self._id)
 
         lattice_subelement = ET.Element("lattice")
         lattice_subelement.set("id", str(self._id))
@@ -1275,13 +1278,16 @@ class HexLattice(Lattice):
             return g < self.num_rings and 0 <= idx[2] < self.num_axial
 
     def create_xml_subelement(self, xml_element):
+        et = ElementTracker()
         # Determine if XML element already contains subelement for this Lattice
         path = './hex_lattice[@id=\'{0}\']'.format(self._id)
         test = xml_element.find(path)
 
         # If the element does contain the Lattice subelement, then return
-        if test is not None:
+        if self._id in et.lattices:
             return
+        else:
+            et.add_lattice(self._id)
 
         lattice_subelement = ET.Element("hex_lattice")
         lattice_subelement.set("id", str(self._id))

openmc/universe.py

@@ -12,6 +12,7 @@ import openmc.checkvalue as cv
 from openmc.plots import _SVG_COLORS
 from openmc.mixin import IDManagerMixin
 
+from .geomtrack import ElementTracker
 
 class Universe(IDManagerMixin):
     """A collection of cells that can be repeated.
@@ -513,12 +514,14 @@ class Universe(IDManagerMixin):
         return memo[self]
 
     def create_xml_subelement(self, xml_element):
+        et = ElementTracker()
         # Iterate over all Cells
         for cell_id, cell in self._cells.items():
             path = "./cell[@id='{}']".format(cell_id)
 
             # If the cell was not already written, write it
-            if xml_element.find(path) is None:
+            if cell_id not in et.cells:
+                et.add_cell(cell_id)
                 # Create XML subelement for this Cell
                 cell_element = cell.create_xml_subelement(xml_element)