Unit test for Mesh.from_rect_lattice()

tests/unit_tests/test_mesh_from_lattice.py

@@ -0,0 +1,116 @@
+import numpy as np
+import openmc
+import pytest
+
+
+@pytest.fixture(scope='module')
+def pincell1(uo2, water):
+    cyl = openmc.ZCylinder(R=0.35)
+    fuel = openmc.Cell(fill=uo2, region=-cyl)
+    moderator = openmc.Cell(fill=water, region=+cyl)
+
+    univ = openmc.Universe(cells=[fuel, moderator])
+    univ.fuel = fuel
+    univ.moderator = moderator
+    return univ
+
+
+@pytest.fixture(scope='module')
+def pincell2(uo2, water):
+    cyl = openmc.ZCylinder(R=0.4)
+    fuel = openmc.Cell(fill=uo2, region=-cyl)
+    moderator = openmc.Cell(fill=water, region=+cyl)
+
+    univ = openmc.Universe(cells=[fuel, moderator])
+    univ.fuel = fuel
+    univ.moderator = moderator
+    return univ
+
+
+@pytest.fixture(scope='module')
+def zr():
+    zr = openmc.Material()
+    zr.add_element('Zr', 1.0)
+    zr.set_density('g/cm3', 1.0)
+    return zr
+
+
+@pytest.fixture(scope='module')
+def rlat2(pincell1, pincell2, uo2, water, zr):
+    """2D Rectangular lattice for testing."""
+    all_zr = openmc.Cell(fill=zr)
+    pitch = 1.2
+    n = 3
+    u1, u2 = pincell1, pincell2
+    lattice = openmc.RectLattice()
+    lattice.lower_left = (-pitch*n/2, -pitch*n/2)
+    lattice.pitch = (pitch, pitch)
+    lattice.outer = openmc.Universe(cells=[all_zr])
+    lattice.universes = [
+        [u1, u2, u1],
+        [u2, u1, u2],
+        [u2, u1, u1]
+    ]
+
+    return lattice
+
+
+@pytest.fixture(scope='module')
+def rlat3(pincell1, pincell2, uo2, water, zr):
+    """3D Rectangular lattice for testing."""
+
+    # Create another universe for top layer
+    hydrogen = openmc.Material()
+    hydrogen.add_element('H', 1.0)
+    hydrogen.set_density('g/cm3', 0.09)
+    h_cell = openmc.Cell(fill=hydrogen)
+    u3 = openmc.Universe(cells=[h_cell])
+
+    all_zr = openmc.Cell(fill=zr)
+    pitch = 1.2
+    n = 3
+    u1, u2 = pincell1, pincell2
+    lattice = openmc.RectLattice()
+    lattice.lower_left = (-pitch*n/2, -pitch*n/2, -10.0)
+    lattice.pitch = (pitch, pitch, 10.0)
+    lattice.outer = openmc.Universe(cells=[all_zr])
+    lattice.universes = [
+        [[u1, u2, u1],
+         [u2, u1, u2],
+         [u2, u1, u1]],
+        [[u3, u1, u2],
+         [u1, u3, u2],
+         [u2, u1, u1]]
+    ]
+
+    return lattice
+
+
+def test_mesh2d(rlat2d):
+    shape = np.array(rlat2d.shape)
+    width = shape*rlat2d.pitch
+    
+    mesh1 = openmc.Mesh().from_rect_lattice(rlat2d)
+    assert np.array_equal(mesh1.dimension, (3, 3))
+    assert np.array_equal(mesh1.lower_left, rlat2d.lower_left)
+    assert np.array_equal(mesh1.upper_right, rlat2d.lower_left + width)
+    
+    mesh2 = openmc.Mesh().from_rect_lattice(rlat2d, division=3)
+    assert np.array_equal(mesh2.dimension, (9, 9))
+    assert np.array_equal(mesh2.lower_left, rlat2d.lower_left)
+    assert np.array_equal(mesh2.upper_right, rlat2d.lower_left + width)
+
+
+def test_mesh3d(rlat3d):
+    shape = np.array(rlat3d.shape)
+    width = shape*rlat3d.pitch
+
+    mesh1 = openmc.Mesh().from_rect_lattice(rlat3d)
+    assert np.array_equal(mesh1.dimension, (3, 3, 2))
+    assert np.array_equal(mesh1.lower_left, rlat3d.lower_left)
+    assert np.array_equal(mesh1.upper_right, rlat3d.lower_left + width)
+
+    mesh2 = openmc.Mesh().from_rect_lattice(rlat3d, division=3)
+    assert np.array_equal(mesh2.dimension, (9, 9, 6))
+    assert np.array_equal(mesh2.lower_left, rlat3d.lower_left)
+    assert np.array_equal(mesh2.upper_right, rlat3d.lower_left + width)