ClaudeandClaude Opus 5.5
a64af5bece
Fixed-size 2D blocks, block-contiguous Q-major storage, explicit per-axis periodicity; non-periodic exterior halo cells stay unpaired (NaN). Provides owner/halo tables for a later periodic GPU consumer. Cut-link mapping deliberately omitted until the BFL write policy is settled. Tests: ownership, pair invariants, pack round trip, tagged pull vs dense. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
140 lines
6.3 KiB
Python
140 lines
6.3 KiB
Python
"""CPU block layout contracts: ownership, halo pairs, packing and tagged pull."""
|
|
import sys
|
|
import unittest
|
|
|
|
import numpy as np
|
|
|
|
from CelerisLab.lbm.blocks import CORNER, FACE, 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)
|
|
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
|
|
|
|
|
|
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_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()
|