"""CPU block layout contracts: ownership, halo pairs, packing and tagged pull.""" import sys import unittest import numpy as np from CelerisLab.body.geometry.base import CutLink from CelerisLab.body.geometry.circle import CircleGeometry from CelerisLab.lbm.blocks import CORNER, FACE, NO_DONOR, NQ, BlockLayout2D from CelerisLab.lbm.descriptors import D2Q9_EX, D2Q9_EY CASES = ( # (ny, nx, by, bx, halo) (6, 8, 3, 4, 1), # 2x2 split (6, 12, 3, 4, 1), # 2x3 split (8, 8, 4, 4, 2), # wider halo (4, 8, 4, 4, 1), # one block along y ) PERIODIC = ((True, True), (False, False), (True, False), (False, True)) def tags(ny, nx): """Unique, exactly representable value per dense population slot.""" return np.arange(NQ * ny * nx, dtype=np.float64).reshape(NQ, ny, nx) def dense_pull(a, periodic): """f[q](x) = a[q](x - c_q); NaN where a non-periodic exterior is crossed.""" px, py = periodic _, ny, nx = a.shape gy, gx = np.mgrid[0:ny, 0:nx] out = np.empty_like(a) for q in range(NQ): sy, sx = gy - D2Q9_EY[q], gx - D2Q9_EX[q] valid = ((py | ((sy >= 0) & (sy < ny))) & (px | ((sx >= 0) & (sx < nx)))) out[q] = np.where(valid, a[q, sy % ny, sx % nx], np.nan) return out def blocked_pull(layout, blocked): """Exchange, then pull every core cell from its own block storage only.""" layout.exchange(blocked) return blocked_pull_core(layout, blocked) def blocked_pull_core(layout, blocked): """Pull core cells with no halo copy (a copy would undo BFL halo writes).""" out = np.full_like(blocked, np.nan) (by, bx), h = layout.block, layout.halo for q in range(NQ): ey, ex = D2Q9_EY[q], D2Q9_EX[q] out[:, q, h:h + by, h:h + bx] = blocked[:, q, h - ey:h - ey + by, h - ex:h - ex + bx] return out def opp(q): return q + 1 if q % 2 else q - 1 def reflect(q, toward, other): """Static linear Bouzidi stand-in; equivalence holds for any read->write map.""" if q < .5: return 2 * q * toward + (1 - 2 * q) * other return toward / (2 * q) + (2 * q - 1) / (2 * q) * other def dense_bfl(a, links): out = a.copy() _, ny, nx = a.shape for link in links: y, x = divmod(link.fluid_idx, nx) q, d = link.q, link.dir if q < .5 and link.fallback == 0: other = a[d, (y - D2Q9_EY[d]) % ny, (x - D2Q9_EX[d]) % nx] elif q < .5: other = 0.0 else: other = a[opp(d), y, x] out[opp(d), (y + D2Q9_EY[d]) % ny, (x + D2Q9_EX[d]) % nx] = reflect(q, a[d, y, x], other) return out def blocked_bfl(layout, blocked, links, mapped): """Reads and writes only the receiver block's storage; no exchange inside.""" flat = blocked.ravel() source = flat.copy() for link, receiver, donor, write in zip(links, mapped.receiver, mapped.donor, mapped.write): q, d = link.q, link.dir toward = source[int(layout.slot(receiver, d))] if donor != NO_DONOR: other = source[int(layout.slot(donor, d))] elif q < .5: other = 0.0 else: other = source[int(layout.slot(receiver, opp(d)))] flat[int(write)] = reflect(q, toward, other) def seam_circle(n=16): """Circle on the shared corner of a 2x2 split (integer-node frame).""" geometry = CircleGeometry(n / 2 - .5, n / 2 - .5, 3.3) links = geometry.build_cut_links(n, n) gy, gx = np.mgrid[0:n, 0:n] solid = (gx - (n / 2 - .5)) ** 2 + (gy - (n / 2 - .5)) ** 2 < 3.3 ** 2 return links, solid class BlockLayoutTests(unittest.TestCase): def test_cpu_import_boundary(self): self.assertFalse(any(name == "pycuda" or name.startswith("pycuda.") for name in sys.modules)) def test_origins_and_unique_core_ownership(self): layout = BlockLayout2D.build(6, 8, 3, 4) np.testing.assert_array_equal(layout.origins, [[0, 0], [0, 4], [3, 0], [3, 4]]) self.assertEqual(layout.blocked_shape, (4, NQ, 5, 6)) for ny, nx, by, bx, h in CASES: layout = BlockLayout2D.build(ny, nx, by, bx, halo=h) ly, lx = layout.local_shape core = [b * ly * lx + (y + h) * lx + x + h for b in range(layout.nblocks) for y in range(by) for x in range(bx)] self.assertEqual(sorted(layout.owner.ravel().tolist()), sorted(core)) for y in range(ny): for x in range(nx): b = (y // by) * (nx // bx) + x // bx self.assertEqual(layout.block_cell(b, y, x), layout.owner[y, x]) def test_halo_pairs_are_owner_core_to_unique_halo(self): for ny, nx, by, bx, h in CASES: for periodic in PERIODIC: layout = BlockLayout2D.build(ny, nx, by, bx, halo=h, periodic=periodic) core = set(layout.owner.ravel().tolist()) dst = layout.halo_dst.tolist() self.assertEqual(dst, sorted(dst)) self.assertEqual(len(dst), len(set(dst))) self.assertTrue(core.isdisjoint(dst)) self.assertTrue(set(layout.halo_src.tolist()) <= core) self.assertTrue(set(layout.halo_kind.tolist()) <= {FACE, CORNER}) full = BlockLayout2D.build(6, 8, 3, 4) per_block = np.bincount(full.halo_kind, minlength=2) // full.nblocks np.testing.assert_array_equal(per_block, [2 * (3 + 4), 4]) closed = BlockLayout2D.build(6, 8, 3, 4, periodic=(False, False)) # Each 2x2 block keeps one inner face column, one face row and one corner. self.assertEqual(len(closed.halo_dst), 4 * (3 + 4 + 1)) def test_pack_round_trip_and_exterior_stays_invalid(self): for ny, nx, by, bx, h in CASES: for periodic in PERIODIC: layout = BlockLayout2D.build(ny, nx, by, bx, halo=h, periodic=periodic) a = tags(ny, nx) blocked = layout.to_blocked(a) np.testing.assert_array_equal(layout.to_dense(blocked), a) filled = np.isfinite(blocked[:, 0]).sum() self.assertEqual(filled, ny * nx + len(layout.halo_dst)) flat = blocked.ravel() for q in (0, 5): np.testing.assert_array_equal(flat[layout.slot(layout.halo_dst, q).astype(np.intp)], flat[layout.slot(layout.halo_src, q).astype(np.intp)]) def test_tagged_pull_matches_dense_pull(self): for ny, nx, by, bx, h in CASES: for periodic in PERIODIC: with self.subTest(case=(ny, nx, by, bx, h), periodic=periodic): layout = BlockLayout2D.build(ny, nx, by, bx, halo=h, periodic=periodic) a = tags(ny, nx) got = layout.to_dense(blocked_pull(layout, layout.to_blocked(a))) np.testing.assert_array_equal(got, dense_pull(a, periodic)) def test_block_cell_resolves_periodic_halo_images(self): layout = BlockLayout2D.build(6, 8, 3, 4) ly, lx = layout.local_shape # Block 0 sees x=7 (wrapped) in its west halo and y=5 in its south halo. self.assertEqual(layout.block_cell(0, 0, 7), 1 * lx + 0) self.assertEqual(layout.block_cell(0, 5, 0), 0 * lx + 1) closed = BlockLayout2D.build(6, 8, 3, 4, periodic=(False, False)) with self.assertRaises(ValueError): closed.block_cell(0, 0, 7) with self.assertRaises(ValueError): layout.block_cell(0, 4, 0) # two rows past the one-cell halo def test_invalid_inputs_rejected(self): for args, kw in (((6, 8, 4, 4), {}), ((6, 8, 3, 4), {"halo": 0}), ((6, 8, 3, 4), {"halo": 4}), ((6, 8, 3, 4), {"periodic": (1, 1)}), ((6, 8, 3, True), {}), ((6, 8.0, 3, 4), {})): with self.assertRaises(ValueError): BlockLayout2D.build(*args, **kw) layout = BlockLayout2D.build(6, 8, 3, 4) with self.assertRaises(ValueError): layout.to_blocked(np.zeros((NQ, 6, 8), dtype=np.int64)) with self.assertRaises(ValueError): layout.to_dense(np.zeros((4, NQ, 5, 5))) def test_cut_links_are_receiver_local_and_unique(self): links, _ = seam_circle() for periodic in ((True, True), (False, False)): layout = BlockLayout2D.build(16, 16, 8, 8, periodic=periodic) mapped = layout.map_cut_links(links) self.assertEqual(len(mapped.write), len(links)) # one row, one force per link self.assertEqual(len(set(mapped.write.tolist())), len(links)) self.assertEqual(len(set(mapped.block.tolist())), 4) # seams are crossed cells = int(np.prod(layout.local_shape)) core = set(layout.owner.ravel().tolist()) a = tags(16, 16) flat = layout.to_blocked(a).ravel() halo_reads = 0 for link, b, r, d, w in zip(links, mapped.block, mapped.receiver, mapped.donor, mapped.write): y, x = divmod(link.fluid_idx, 16) self.assertEqual(r, layout.owner[y, x]) self.assertEqual(int(r) // cells, b) self.assertEqual(int(w) // (NQ * cells), b) sy, sx = y + D2Q9_EY[link.dir], x + D2Q9_EX[link.dir] self.assertEqual(flat[int(w)], a[opp(link.dir), sy, sx]) if d != NO_DONOR: self.assertEqual(int(d) // cells, b) self.assertEqual(flat[int(layout.slot(d, link.dir))], a[link.dir, y - D2Q9_EY[link.dir], x - D2Q9_EX[link.dir]]) halo_reads += int(d) not in core self.assertGreater(halo_reads, 0) # some q<.5 donors come from a neighbor's core def test_static_bfl_pull_matches_dense_on_fluid_cells(self): links, solid = seam_circle() a = tags(16, 16) / 7.0 expected = dense_pull(dense_bfl(a, links), (True, True)) layout = BlockLayout2D.build(16, 16, 8, 8) blocked = layout.to_blocked(a) # gather halos blocked_bfl(layout, blocked, links, layout.map_cut_links(links)) got = layout.to_dense(blocked_pull_core(layout, blocked)) np.testing.assert_array_equal(got[:, ~solid], expected[:, ~solid]) def test_periodic_seam_links_follow_the_pulled_halo_image(self): # Receivers on x=0 / x=nx-1 / a corner whose sources wrap physically. links = [CutLink(fluid_idx=1 * 6 + 0, dir=2, q=.6, body_id=0), CutLink(fluid_idx=2 * 6 + 0, dir=2, q=.3, body_id=0), CutLink(fluid_idx=3 * 6 + 5, dir=1, q=.3, body_id=0), CutLink(fluid_idx=0 * 6 + 0, dir=6, q=.4, body_id=0)] sources = np.zeros((8, 6), dtype=bool) for y, x in ((1, 5), (2, 5), (3, 0), (7, 5)): sources[y, x] = True a = tags(8, 6) / 7.0 expected = dense_pull(dense_bfl(a, links), (True, True)) # One block along x, along y, along both (periodic self-images). for by, bx in ((4, 6), (8, 3), (8, 6)): with self.subTest(block=(by, bx)): layout = BlockLayout2D.build(8, 6, by, bx) mapped = layout.map_cut_links(links) blocked = layout.to_blocked(a) blocked_bfl(layout, blocked, links, mapped) got = layout.to_dense(blocked_pull_core(layout, blocked)) np.testing.assert_array_equal(got[:, ~sources], expected[:, ~sources]) def test_cut_link_rejections(self): layout = BlockLayout2D.build(16, 16, 8, 8, periodic=(False, False)) link = CutLink(fluid_idx=5 * 16 + 5, dir=1, q=.6, body_id=0) with self.assertRaises(ValueError): layout.map_cut_links([link, link]) with self.assertRaises(ValueError): layout.map_cut_links([CutLink(fluid_idx=15, dir=1, q=.6, body_id=0)]) with self.assertRaises(ValueError): layout.map_cut_links([CutLink(fluid_idx=5 * 16 + 5, dir=0, q=.6, body_id=0)]) def test_tables_are_readonly(self): layout = BlockLayout2D.build(6, 8, 3, 4) for array in (layout.origins, layout.owner, layout.halo_dst, layout.halo_src, layout.halo_kind): self.assertFalse(array.flags.writeable) if __name__ == "__main__": unittest.main()