Hexagonal Lattice Roundtrip (openmc-dev#3003)

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

openmc/lattice.py

@@ -1574,6 +1574,7 @@ def from_xml_element(cls, elem, get_universe):
                             alpha -= 1
                             if not lat.is_valid_index((x, alpha, z)):
                                 # Reached the bottom
+                                j += 1
                                 break
                     j += 1
             else:
@@ -1593,6 +1594,7 @@ def from_xml_element(cls, elem, get_universe):
 
                         # Check if we've reached the bottom
                         if y == -n_rings:
+                            j += 1
                             break
 
                         while not lat.is_valid_index((alpha, y, z)):

tests/unit_tests/test_lattice.py

@@ -377,3 +377,46 @@ def test_unset_universes():
     hex_lattice.pitch = (1.,)
     with pytest.raises(ValueError):
         hex_lattice.create_xml_subelement(elem)
+
+
+@pytest.mark.parametrize("orientation", ['x', 'y'])
+def test_hex_lattice_roundtrip(orientation):
+    openmc.reset_auto_ids()
+
+    # ensure that the lattice universes are the same on all axial levels
+    def check_lattice_universes(og_lattice, xml_lattice):
+        for axial_og, axial_rt in zip(og_lattice.universes, xml_lattice.universes):
+            for ring_og, ring_rt in zip(axial_og, axial_rt):
+                assert [u.id for u in ring_og] == [u.id for u in ring_rt]
+
+    latt = openmc.HexLattice()
+    latt.pitch = (1.0, 1.0)
+    latt.center = (0.0, 0.0, 0.0)
+    latt.orientation = orientation
+
+    # fill the lattice with universes in increasing order and repeat for
+    # the second actial level
+    lvl_one_univs = [openmc.Universe(cells=[openmc.Cell()]) for _ in range(19)]
+    lvl_one_univs = [lvl_one_univs[-12:], lvl_one_univs[1:7], lvl_one_univs[:1]]
+    latt.universes = [lvl_one_univs, lvl_one_univs]
+
+    geom = openmc.Geometry([openmc.Cell(fill=latt)])
+    geom.export_to_xml()
+
+    xml_geom = openmc.Geometry.from_xml(materials=openmc.Materials())
+
+    xml_latt = xml_geom.get_all_lattices()[latt.id]
+
+    check_lattice_universes(latt, xml_latt)
+
+    # same test but with unique universes for each axial level
+    lvl_two_univs = [openmc.Universe(cells=[openmc.Cell()]) for _ in range(19)]
+    lvl_two_univs = [lvl_two_univs[-12:], lvl_two_univs[1:7], lvl_two_univs[:1]]
+    latt.universes = [lvl_one_univs, lvl_two_univs]
+
+    geom.export_to_xml()
+
+    xml_geom = openmc.Geometry.from_xml(materials=openmc.Materials())
+    xml_latt = xml_geom.get_all_lattices()[latt.id]
+
+    check_lattice_universes(latt, xml_latt)