Skip to content

Commit ec7ef1b

Browse files
pshriwisepaulromano
authored and
church89
committed
Apply memoization in get_all_universes (openmc-dev#2995)
Co-authored-by: Paul Romano <paul.k.romano@gmail.com>
1 parent 86d7fe6 commit ec7ef1b

5 files changed

Lines changed: 72 additions & 46 deletions

File tree

‎openmc/cell.py‎

Lines changed: 18 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -426,15 +426,13 @@ def get_all_cells(self, memo=None):
426426
instances
427427
428428
"""
429+
if memo is None:
430+
memo = set()
431+
elif self in memo:
432+
return {}
433+
memo.add(self)
429434

430435
cells = {}
431-
432-
if memo and self in memo:
433-
return cells
434-
435-
if memo is not None:
436-
memo.add(self)
437-
438436
if self.fill_type in ('universe', 'lattice'):
439437
cells.update(self.fill.get_all_cells(memo))
440438

@@ -465,7 +463,7 @@ def get_all_materials(self, memo=None):
465463

466464
return materials
467465

468-
def get_all_universes(self):
466+
def get_all_universes(self, memo=None):
469467
"""Return all universes that are contained within this one if any of
470468
its cells are filled with a universe or lattice.
471469
@@ -476,18 +474,22 @@ def get_all_universes(self):
476474
:class:`Universe` instances
477475
478476
"""
477+
if memo is None:
478+
memo = set()
479+
if self in memo:
480+
return {}
481+
memo.add(self)
479482

480483
universes = {}
481-
482484
if self.fill_type == 'universe':
483485
universes[self.fill.id] = self.fill
484-
universes.update(self.fill.get_all_universes())
486+
universes.update(self.fill.get_all_universes(memo))
485487
elif self.fill_type == 'lattice':
486-
universes.update(self.fill.get_all_universes())
488+
universes.update(self.fill.get_all_universes(memo))
487489

488490
return universes
489491

490-
def clone(self, clone_materials=True, clone_regions=True, memo=None):
492+
def clone(self, clone_materials=True, clone_regions=True, memo=None):
491493
"""Create a copy of this cell with a new unique ID, and clones
492494
the cell's region and fill.
493495
@@ -678,10 +680,11 @@ def create_xml_subelement(self, xml_element, memo=None):
678680
# thus far.
679681
def create_surface_elements(node, element, memo=None):
680682
if isinstance(node, Halfspace):
681-
if memo and node.surface in memo:
683+
if memo is None:
684+
memo = set()
685+
elif node.surface in memo:
682686
return
683-
if memo is not None:
684-
memo.add(node.surface)
687+
memo.add(node.surface)
685688
xml_element.append(node.surface.to_xml_element())
686689

687690
elif isinstance(node, Complement):

‎openmc/geometry.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -134,7 +134,7 @@ def to_xml_element(self, remove_surfs=False) -> ET.Element:
134134

135135
# Create XML representation
136136
element = ET.Element("geometry")
137-
self.root_universe.create_xml_subelement(element, memo=set())
137+
self.root_universe.create_xml_subelement(element)
138138

139139
# Sort the elements in the file
140140
element[:] = sorted(element, key=lambda x: (
@@ -373,7 +373,7 @@ def get_all_cells(self) -> typing.Dict[int, openmc.Cell]:
373373
374374
"""
375375
if self.root_universe is not None:
376-
return self.root_universe.get_all_cells(memo=set())
376+
return self.root_universe.get_all_cells()
377377
else:
378378
return {}
379379

@@ -417,7 +417,7 @@ def get_all_materials(self) -> typing.Dict[int, openmc.Material]:
417417
418418
"""
419419
if self.root_universe is not None:
420-
return self.root_universe.get_all_materials(memo=set())
420+
return self.root_universe.get_all_materials()
421421
else:
422422
return {}
423423

‎openmc/lattice.py‎

Lines changed: 23 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -170,11 +170,11 @@ def get_all_cells(self, memo=None):
170170
"""
171171
cells = {}
172172

173-
if memo and self in memo:
173+
if memo is None:
174+
memo = set()
175+
elif self in memo:
174176
return cells
175-
176-
if memo is not None:
177-
memo.add(self)
177+
memo.add(self)
178178

179179
unique_universes = self.get_unique_universes()
180180

@@ -194,6 +194,9 @@ def get_all_materials(self, memo=None):
194194
195195
"""
196196

197+
if memo is None:
198+
memo = set()
199+
197200
materials = {}
198201

199202
# Append all Cells in each Cell in the Universe to the dictionary
@@ -203,7 +206,7 @@ def get_all_materials(self, memo=None):
203206

204207
return materials
205208

206-
def get_all_universes(self):
209+
def get_all_universes(self, memo=None):
207210
"""Return all universes that are contained within the lattice
208211
209212
Returns
@@ -213,11 +216,16 @@ def get_all_universes(self):
213216
:class:`Universe` instances
214217
215218
"""
216-
217219
# Initialize a dictionary of all Universes contained by the Lattice
218220
# in each nested Universe level
219221
all_universes = {}
220222

223+
if memo is None:
224+
memo = set()
225+
elif self in memo:
226+
return all_universes
227+
memo.add(self)
228+
221229
# Get all unique Universes contained in each of the lattice cells
222230
unique_universes = self.get_unique_universes()
223231

@@ -226,7 +234,7 @@ def get_all_universes(self):
226234

227235
# Append all Universes containing each cell to the dictionary
228236
for universe in unique_universes.values():
229-
all_universes.update(universe.get_all_universes())
237+
all_universes.update(universe.get_all_universes(memo))
230238

231239
return all_universes
232240

@@ -846,10 +854,11 @@ def create_xml_subelement(self, xml_element, memo=None):
846854
847855
"""
848856
# If the element already contains the Lattice subelement, then return
849-
if memo and self in memo:
857+
if memo is None:
858+
memo = set()
859+
elif self in memo:
850860
return
851-
if memo is not None:
852-
memo.add(self)
861+
memo.add(self)
853862

854863
# Make sure universes have been assigned
855864
if self.universes is None:
@@ -1417,10 +1426,11 @@ def is_valid_index(self, idx):
14171426

14181427
def create_xml_subelement(self, xml_element, memo=None):
14191428
# If this subelement has already been written, return
1420-
if memo and self in memo:
1429+
if memo is None:
1430+
memo = set()
1431+
elif self in memo:
14211432
return
1422-
if memo is not None:
1423-
memo.add(self)
1433+
memo.add(self)
14241434

14251435
lattice_subelement = ET.Element("hex_lattice")
14261436
lattice_subelement.set("id", str(self._id))

‎openmc/tallies.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3209,6 +3209,7 @@ def _create_derivative_subelements(self, root_element):
32093209
def to_xml_element(self, memo=None):
32103210
"""Creates a 'tallies' element to be written to an XML file.
32113211
"""
3212+
memo = memo if memo is not None else set()
32123213
element = ET.Element("tallies")
32133214
self._create_mesh_subelements(element, memo)
32143215
self._create_filter_subelements(element)

‎openmc/universe.py‎

Lines changed: 27 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -90,7 +90,7 @@ def add_volume_information(self, volume_calc):
9090
else:
9191
raise ValueError('No volume information found for this universe.')
9292

93-
def get_all_universes(self):
93+
def get_all_universes(self, memo=None):
9494
"""Return all universes that are contained within this one.
9595
9696
Returns
@@ -100,10 +100,16 @@ def get_all_universes(self):
100100
:class:`Universe` instances
101101
102102
"""
103+
if memo is None:
104+
memo = set()
105+
elif self in memo:
106+
return {}
107+
memo.add(self)
108+
103109
# Append all Universes within each Cell to the dictionary
104110
universes = {}
105111
for cell in self.get_all_cells().values():
106-
universes.update(cell.get_all_universes())
112+
universes.update(cell.get_all_universes(memo))
107113

108114
return universes
109115

@@ -639,15 +645,14 @@ def get_all_cells(self, memo=None):
639645
640646
"""
641647

642-
cells = {}
643-
644-
if memo and self in memo:
645-
return cells
646-
647-
if memo is not None:
648-
memo.add(self)
648+
if memo is None:
649+
memo = set()
650+
elif self in memo:
651+
return {}
652+
memo.add(self)
649653

650654
# Add this Universe's cells to the dictionary
655+
cells = {}
651656
cells.update(self._cells)
652657

653658
# Append all Cells in each Cell in the Universe to the dictionary
@@ -667,6 +672,9 @@ def get_all_materials(self, memo=None):
667672
668673
"""
669674

675+
if memo is None:
676+
memo = set()
677+
670678
materials = {}
671679

672680
# Append all Cells in each Cell in the Universe to the dictionary
@@ -677,15 +685,17 @@ def get_all_materials(self, memo=None):
677685
return materials
678686

679687
def create_xml_subelement(self, xml_element, memo=None):
688+
if memo is None:
689+
memo = set()
690+
680691
# Iterate over all Cells
681692
for cell in self._cells.values():
682693

683694
# If the cell was already written, move on
684-
if memo and cell in memo:
695+
if cell in memo:
685696
continue
686697

687-
if memo is not None:
688-
memo.add(cell)
698+
memo.add(cell)
689699

690700
# Create XML subelement for this Cell
691701
cell_element = cell.create_xml_subelement(xml_element, memo)
@@ -928,11 +938,13 @@ def n_surfaces(self):
928938
return self._n_geom_elements('surface')
929939

930940
def create_xml_subelement(self, xml_element, memo=None):
931-
if memo and self in memo:
941+
if memo is None:
942+
memo = set()
943+
944+
if self in memo:
932945
return
933946

934-
if memo is not None:
935-
memo.add(self)
947+
memo.add(self)
936948

937949
# Set xml element values
938950
dagmc_element = ET.Element('dagmc_universe')

0 commit comments

Comments
 (0)