Add CPU BlockLayout2D: core ownership, halo pairs, blocked packing

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>

src/CelerisLab/lbm/blocks.py

@@ -0,0 +1,173 @@
+"""CPU layout for fixed-size 2D blocks: core ownership, halo pairs and packing.
+
+Dense arrays are Q-major (q, y, x) with cell y*nx+x, as in core/layout.cuh.
+Blocked storage is block-contiguous and Q-major inside each block:
+(nblocks, q, by+2*halo, bx+2*halo). Block b = jy*nbx + jx has core origin
+(jy*by, jx*bx). Every dense cell, fluid or solid, has exactly one core owner.
+
+A "block cell" id is b*L + c with L=(by+2h)*(bx+2h) and c the local row-major
+cell; its population q lives at flat slot b*NQ*L + q*L + c. Halo pairs copy
+all NQ populations from an owner core cell into another block's halo cell.
+Periodicity is explicit per axis. Across a non-periodic exterior, halo cells
+have no pair and stay invalid (NaN after packing): no wrap, no implied BC.
+
+Preparation only: no CUDA, no boundary conditions, no refinement levels.
+"""
+from dataclasses import dataclass
+from numbers import Integral
+
+import numpy as np
+
+NQ = 9
+FACE, CORNER = 0, 1
+
+
+def _readonly(values, dtype):
+    array = np.ascontiguousarray(values, dtype=dtype)
+    array.setflags(write=False)
+    return array
+
+
+def _positive_int(value, name):
+    if isinstance(value, (bool, np.bool_)) or not isinstance(value, Integral) or value <= 0:
+        raise ValueError(f"{name} must be a positive integer")
+    return int(value)
+
+
+@dataclass(frozen=True, eq=False)
+class BlockLayout2D:
+    """Fixed-size block partition of an (ny, nx) grid; construct with build().
+
+    ``origins`` rows are (y0, x0) core origins. ``owner`` maps each dense cell
+    to its owning block cell id. ``halo_dst``/``halo_src`` are block cell ids
+    (destination halo cell, source owner core cell) sorted by destination;
+    ``halo_kind`` marks FACE or CORNER pairs. Sources are never halo cells, so
+    one gather in any order fills every valid halo cell.
+    """
+    shape: tuple[int, int]
+    block: tuple[int, int]
+    halo: int
+    periodic: tuple[bool, bool]
+    origins: np.ndarray
+    owner: np.ndarray
+    halo_dst: np.ndarray
+    halo_src: np.ndarray
+    halo_kind: np.ndarray
+
+    @classmethod
+    def build(cls, ny, nx, by, bx, *, halo=1, periodic=(True, True)) -> "BlockLayout2D":
+        """``periodic`` is (periodic_x, periodic_y); block sizes must divide the grid."""
+        ny, nx, by, bx, halo = (_positive_int(v, n) for v, n in
+                                ((ny, "ny"), (nx, "nx"), (by, "by"), (bx, "bx"), (halo, "halo")))
+        if ny % by or nx % bx:
+            raise ValueError("block sizes must divide the grid")
+        if halo > min(by, bx):
+            raise ValueError("halo may not exceed the block size")
+        if len(periodic) != 2 or not all(isinstance(p, (bool, np.bool_)) for p in periodic):
+            raise ValueError("periodic must be two booleans (x, y)")
+        px, py = map(bool, periodic)
+        nby, nbx = ny // by, nx // bx
+        ly, lx = by + 2 * halo, bx + 2 * halo
+        cells = ly * lx
+        if nby * nbx * NQ * cells > np.iinfo(np.int64).max // 2:
+            raise ValueError("blocked storage exceeds 64-bit indexing")
+
+        jy, jx = np.divmod(np.arange(nby * nbx), nbx)
+        origins = np.stack([jy * by, jx * bx], axis=1)
+
+        gy, gx = np.mgrid[0:ny, 0:nx]
+        owner = ((gy // by) * nbx + gx // bx) * cells + (gy % by + halo) * lx + gx % bx + halo
+
+        # Halo cells of one block in local row-major order, then all blocks.
+        hy, hx = np.mgrid[0:ly, 0:lx]
+        out_y = (hy < halo) | (hy >= by + halo)
+        out_x = (hx < halo) | (hx >= bx + halo)
+        ring = out_y | out_x
+        hy, hx, kind = hy[ring], hx[ring], np.where(out_y[ring] & out_x[ring], CORNER, FACE)
+        dst, src, kinds = [], [], []
+        for b, (y0, x0) in enumerate(origins):
+            sy, sx = y0 + hy - halo, x0 + hx - halo
+            if py:
+                sy %= ny
+            if px:
+                sx %= nx
+            keep = (sy >= 0) & (sy < ny) & (sx >= 0) & (sx < nx)
+            dst.append(b * cells + hy[keep] * lx + hx[keep])
+            src.append(owner[sy[keep], sx[keep]])
+            kinds.append(kind[keep])
+        return cls((ny, nx), (by, bx), halo, (px, py),
+                   _readonly(origins, np.int64), _readonly(owner, np.uint64),
+                   _readonly(np.concatenate(dst), np.uint64),
+                   _readonly(np.concatenate(src), np.uint64),
+                   _readonly(np.concatenate(kinds), np.uint8))
+
+    @property
+    def nblocks(self) -> int:
+        return len(self.origins)
+
+    @property
+    def local_shape(self) -> tuple[int, int]:
+        return self.block[0] + 2 * self.halo, self.block[1] + 2 * self.halo
+
+    @property
+    def blocked_shape(self) -> tuple[int, int, int, int]:
+        return (self.nblocks, NQ, *self.local_shape)
+
+    def slot(self, block_cell, q):
+        """Flat blocked population slot of block cell id(s) for direction q."""
+        cells = int(np.prod(self.local_shape))
+        b, c = np.divmod(np.asarray(block_cell, dtype=np.uint64), np.uint64(cells))
+        return (b * np.uint64(NQ) + np.asarray(q, dtype=np.uint64)) * np.uint64(cells) + c
+
+    def block_cell(self, block, y, x):
+        """Address of in-grid cell (y, x) in ``block``'s own storage, core or halo.
+
+        Periodic axes resolve to the image inside the block's halo window.
+        Raises if the cell is not represented there or lies across an exterior.
+        """
+        by, bx = self.block
+        ly, lx = self.local_shape
+        ny, nx = self.shape
+        y0, x0 = self.origins[block]
+        ry, rx = int(y) - int(y0) + self.halo, int(x) - int(x0) + self.halo
+        if self.periodic[1] and not 0 <= ry < ly:
+            ry = next((r for r in (ry - ny, ry + ny) if 0 <= r < ly), ry)
+        if self.periodic[0] and not 0 <= rx < lx:
+            rx = next((r for r in (rx - nx, rx + nx) if 0 <= r < lx), rx)
+        if not (0 <= ry < ly and 0 <= rx < lx and 0 <= y < ny and 0 <= x < nx):
+            raise ValueError("cell is not represented in this block's storage")
+        return int(block) * ly * lx + ry * lx + rx
+
+    def exchange(self, blocked) -> None:
+        """Fill every valid halo cell in place from its owner's core."""
+        blocked = self._check_blocked(blocked)
+        flat = blocked.reshape(self.nblocks, NQ, -1)
+        cells = flat.shape[2]
+        db, dc = np.divmod(self.halo_dst.astype(np.intp), cells)
+        sb, sc = np.divmod(self.halo_src.astype(np.intp), cells)
+        flat[db, :, dc] = flat[sb, :, sc]
+
+    def to_blocked(self, dense):
+        """Pack (NQ, ny, nx) floats; halo cells without a pair stay NaN."""
+        dense = np.asarray(dense)
+        if dense.shape != (NQ, *self.shape) or not np.issubdtype(dense.dtype, np.floating):
+            raise ValueError(f"expected floating (NQ, ny, nx) = {(NQ, *self.shape)}")
+        blocked = np.full(self.blocked_shape, np.nan, dtype=dense.dtype)
+        flat = blocked.reshape(self.nblocks, NQ, -1)
+        ob, oc = np.divmod(self.owner.ravel().astype(np.intp), flat.shape[2])
+        flat[ob, :, oc] = dense.reshape(NQ, -1).T
+        self.exchange(blocked)
+        return blocked
+
+    def to_dense(self, blocked):
+        """Unpack owner cores only; halo contents are ignored."""
+        blocked = self._check_blocked(blocked)
+        flat = blocked.reshape(self.nblocks, NQ, -1)
+        ob, oc = np.divmod(self.owner.ravel().astype(np.intp), flat.shape[2])
+        return np.ascontiguousarray(flat[ob, :, oc].T).reshape(NQ, *self.shape)
+
+    def _check_blocked(self, blocked):
+        if not isinstance(blocked, np.ndarray) or blocked.shape != self.blocked_shape \
+                or not blocked.flags.c_contiguous:
+            raise ValueError(f"expected C-contiguous blocked array {self.blocked_shape}")
+        return blocked

tests/unit/test_block_layout.py

@@ -0,0 +1,139 @@
+"""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()