diff --git a/src/CelerisLab/lbm/blocks.py b/src/CelerisLab/lbm/blocks.py index 57d8269..1bc665a 100644 --- a/src/CelerisLab/lbm/blocks.py +++ b/src/CelerisLab/lbm/blocks.py @@ -3,8 +3,14 @@ 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. +(nblocks, q, by+2*halo, bx+2*halo). A full layout stores every grid block, +block b = jy*nbx + jx with core origin (jy*by, jx*bx), so every dense cell, +fluid or solid, has exactly one core owner. A sparse layout stores only the +selected grid-aligned blocks (row-major order, b = position in that list); +cells of absent blocks are HOLES: owner = layout.hole (one past the last +block cell, so any unguarded use fails instead of wrapping like -1), no halo +pair reads them, packing ignores them and unpacking returns NaN there. This +is physical-block storage, not arbitrary cell refinement. 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 @@ -57,6 +63,24 @@ def _positive_int(value, name): return int(value) +def _selected_blocks(origins, ny, nx, by, bx): + """Row-major grid block indices of explicit (y0, x0) origins, validated.""" + array = np.asarray(origins) + if array.dtype.kind not in "iu" or array.ndim != 2 or array.shape[1] != 2 or not len(array): + raise ValueError("origins must be a non-empty integer (n, 2) array of (y0, x0)") + y0, x0 = array[:, 0].astype(np.int64), array[:, 1].astype(np.int64) + if (y0 % by).any() or (x0 % bx).any(): + raise ValueError("origins must be multiples of the block size") + if (y0 < 0).any() or (y0 >= ny).any() or (x0 < 0).any() or (x0 >= nx).any(): + raise ValueError("origins must lie inside the grid") + index = (y0 // by) * (nx // bx) + x0 // bx + if (np.diff(index) == 0).any(): + raise ValueError("origins must be unique") + if (np.diff(index) < 0).any(): + raise ValueError("origins must be in row-major order") + return index + + @dataclass(frozen=True, eq=False) class BlockCutLinks: """One row per input link, in input order (so force is counted once). @@ -79,8 +103,9 @@ 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; + to its owning block cell id, or ``hole`` if its block is absent. + ``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. """ @@ -95,8 +120,14 @@ class BlockLayout2D: 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.""" + def build(cls, ny, nx, by, bx, *, halo=1, periodic=(True, True), origins=None) -> "BlockLayout2D": + """``periodic`` is (periodic_x, periodic_y); block sizes must divide the grid. + + ``origins`` (optional) selects stored blocks by (y0, x0) core origin: + integer multiples of (by, bx) inside the grid, at least one, unique + and strictly increasing in row-major order (no silent reordering). + None stores all blocks; passing every origin gives the same layout. + """ 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: @@ -109,14 +140,23 @@ class BlockLayout2D: 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: + if origins is None: + jy, jx = np.divmod(np.arange(nby * nbx), nbx) + origins = np.stack([jy * by, jx * bx], axis=1) + selected = _selected_blocks(origins, ny, nx, by, bx) + origins = np.asarray(origins, dtype=np.int64).reshape(-1, 2) + if (len(origins) + 1) * 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) - + # Grid block -> stored block, or len(origins) for an absent block, so + # absent cells land exactly on hole = nblocks * cells. + stored = np.full(nby * nbx, len(origins), dtype=np.int64) + stored[selected] = np.arange(len(origins)) gy, gx = np.mgrid[0:ny, 0:nx] - owner = ((gy // by) * nbx + gx // bx) * cells + (gy % by + halo) * lx + gx % bx + halo + gb = stored[(gy // by) * nbx + gx // bx] + hole = len(origins) * cells + owner = np.where(gb < len(origins), + gb * cells + (gy % by + halo) * lx + gx % bx + halo, hole) # Halo cells of one block in local row-major order, then all blocks. hy, hx = np.mgrid[0:ly, 0:lx] @@ -132,6 +172,7 @@ class BlockLayout2D: if px: sx %= nx keep = (sy >= 0) & (sy < ny) & (sx >= 0) & (sx < nx) + keep[keep] = owner[sy[keep], sx[keep]] != hole dst.append(b * cells + hy[keep] * lx + hx[keep]) src.append(owner[sy[keep], sx[keep]]) kinds.append(kind[keep]) @@ -145,6 +186,16 @@ class BlockLayout2D: def nblocks(self) -> int: return len(self.origins) + @property + def hole(self) -> int: + """Owner value of cells in absent blocks: one past the last block cell.""" + return self.nblocks * int(np.prod(self.local_shape)) + + @property + def owned(self) -> np.ndarray: + """Dense (ny, nx) mask of cells stored in some block core.""" + return self.owner != self.hole + @property def local_shape(self) -> tuple[int, int]: return self.block[0] + 2 * self.halo, self.block[1] + 2 * self.halo @@ -168,6 +219,8 @@ class BlockLayout2D: if two images fit, the core-side one is returned (use local offsets, not this lookup, when a stencil must follow a specific image). Raises if the cell is not represented there or lies across an exterior. + In a sparse layout a represented halo cell whose source is a hole has + no pair and stays invalid (NaN after packing). """ by, bx = self.block ly, lx = self.local_shape @@ -193,6 +246,8 @@ class BlockLayout2D: x %= nx if not (0 <= y < ny and 0 <= x < nx): raise ValueError("cut-link stencil crosses a non-periodic exterior") + if self.owner[y, x] == self.hole: + raise ValueError("cut-link stencil touches an absent block") return y, x def map_cut_links(self, links) -> BlockCutLinks: @@ -200,7 +255,9 @@ class BlockLayout2D: Rejects repeated physical links and repeated local writers. Addresses are geometry-free: q, fallback and BFL math stay as built on the dense - grid. Halo must be at least one cell (always true here). + grid. Halo must be at least one cell (always true here). In a sparse + layout the receiver, its solid neighbor and any q<.5 donor must all lie + in stored blocks. """ ny, nx = self.shape lx = self.local_shape[1] @@ -214,6 +271,8 @@ class BlockLayout2D: seen.add((r, q)) y, x = divmod(r, nx) receiver = int(self.owner[y, x]) + if receiver == self.hole: + raise ValueError("cut-link receiver lies in an absent block") block = receiver // cells # Local stencil offsets from the core receiver: the window image the # next local pull reads, even where a periodic seam wraps physically. @@ -244,23 +303,32 @@ class BlockLayout2D: flat[db, :, dc] = flat[sb, :, sc] def to_blocked(self, dense): - """Pack (NQ, ny, nx) floats; halo cells without a pair stay NaN.""" + """Pack (NQ, ny, nx) floats; halo cells without a pair stay NaN. + + Values at holes (absent blocks) are ignored. + """ 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 + owned = self.owned.ravel() + ob, oc = np.divmod(self.owner.ravel()[owned].astype(np.intp), flat.shape[2]) + flat[ob, :, oc] = dense.reshape(NQ, -1)[:, owned].T self.exchange(blocked) return blocked def to_dense(self, blocked): - """Unpack owner cores only; halo contents are ignored.""" + """Unpack owner cores only; halo contents are ignored, holes are NaN.""" 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) + owned = self.owned.ravel() + ob, oc = np.divmod(self.owner.ravel()[owned].astype(np.intp), flat.shape[2]) + dense = np.empty((NQ, owned.size), dtype=blocked.dtype) + if not owned.all(): # full layouts keep any dtype; holes need a float NaN + dense[:, ~owned] = np.nan + dense[:, owned] = flat[ob, :, oc].T + return dense.reshape(NQ, *self.shape) def _check_blocked(self, blocked): if not isinstance(blocked, np.ndarray) or blocked.shape != self.blocked_shape \ @@ -276,7 +344,8 @@ def blocked_link_tables(layout, flags, fluid_idx, dir, q, fallback): own builder; this only maps storage via map_cut_links (unique receiver owner and local writer, input row order kept). Returns (tables, mapped): tables holds device-ready arrays, with flags sized nblocks*L (unused halo - entries 0) and cell = the receiver's local core cell. + entries 0) and cell = the receiver's local core cell. Flags at holes + (absent blocks) are not stored. """ flags = np.asarray(flags) if flags.shape != (int(np.prod(layout.shape)),) or flags.dtype != np.uint16: @@ -288,7 +357,8 @@ def blocked_link_tables(layout, flags, fluid_idx, dir, q, fallback): for r, d, v, f in zip(*columns)) cells = int(np.prod(layout.local_shape)) blocked_flags = np.zeros(layout.nblocks * cells, np.uint16) - blocked_flags[layout.owner.ravel().astype(np.intp)] = flags + owned = layout.owned.ravel() + blocked_flags[layout.owner.ravel()[owned].astype(np.intp)] = flags[owned] tables = dict(flags=blocked_flags, block=mapped.block, cell=(mapped.receiver % cells).astype(np.uint32), dir=columns[1].astype(np.uint8), q=columns[2].astype(np.float32), diff --git a/tests/unit/test_block_layout.py b/tests/unit/test_block_layout.py index 1c4c45c..3d7f482 100644 --- a/tests/unit/test_block_layout.py +++ b/tests/unit/test_block_layout.py @@ -1,12 +1,13 @@ """CPU block layout contracts: ownership, halo pairs, packing and tagged pull.""" import sys import unittest +from types import SimpleNamespace 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.blocks import CORNER, FACE, NO_DONOR, NQ, BlockLayout2D, blocked_link_tables from CelerisLab.lbm.descriptors import D2Q9_EX, D2Q9_EY CASES = ( # (ny, nx, by, bx, halo) @@ -285,5 +286,139 @@ class BlockLayoutTests(unittest.TestCase): self.assertFalse(array.flags.writeable) +def all_origins(ny, nx, by, bx): + return [(y, x) for y in range(0, ny, by) for x in range(0, nx, bx)] + + +def link(y, x, d, q, nx=12, fallback=0): + return SimpleNamespace(fluid_idx=y * nx + x, dir=d, q=q, fallback=fallback) + + +# 6x12 grid of 3x4 blocks (2x3 block grid); stored L: (0,0), (0,4), (3,4). +L_ORIGINS = [(0, 0), (0, 4), (3, 4)] +SPARSE = ( # (ny, nx, by, bx, halo, origins) + (6, 12, 3, 4, 1, L_ORIGINS), + (6, 12, 3, 4, 1, [(3, 8)]), # N=1 corner block + (8, 8, 4, 4, 2, [(0, 4), (4, 0)]), # diagonal pair, wider halo + (6, 8, 3, 4, 1, [(0, 0), (0, 4), (3, 0)]), # missing one of 2x2 +) + + +class SparseBlockLayoutTests(unittest.TestCase): + """Explicit selected block origins: absent blocks are holes, not wrapped -1.""" + + def test_all_origins_reproduce_full_layout(self): + for ny, nx, by, bx, h in CASES + ((4, 8, 4, 8, 1),): # last: N=1 full grid + for periodic in PERIODIC: + with self.subTest(case=(ny, nx, by, bx, h), periodic=periodic): + full = BlockLayout2D.build(ny, nx, by, bx, halo=h, periodic=periodic) + explicit = BlockLayout2D.build(ny, nx, by, bx, halo=h, periodic=periodic, + origins=all_origins(ny, nx, by, bx)) + for name in ("origins", "owner", "halo_dst", "halo_src", "halo_kind"): + np.testing.assert_array_equal(getattr(explicit, name), getattr(full, name)) + self.assertEqual(getattr(explicit, name).dtype, getattr(full, name).dtype) + self.assertTrue(full.owned.all()) + self.assertEqual(explicit.hole, full.hole) + + def test_holes_and_ownership(self): + layout = BlockLayout2D.build(6, 12, 3, 4, periodic=(False, False), origins=L_ORIGINS) + np.testing.assert_array_equal(layout.origins, L_ORIGINS) + self.assertEqual(layout.nblocks, 3) + cells = int(np.prod(layout.local_shape)) + self.assertEqual(layout.hole, 3 * cells) + expected = np.zeros((6, 12), dtype=bool) + expected[0:3, 0:8] = expected[3:6, 4:8] = True + np.testing.assert_array_equal(layout.owned, expected) + self.assertTrue((layout.owner[~expected] == layout.hole).all()) + for b, (y0, x0) in enumerate(L_ORIGINS): + for y in range(y0, y0 + 3): + for x in range(x0, x0 + 4): + self.assertEqual(layout.block_cell(b, y, x), layout.owner[y, x]) + owned = layout.owner[expected].tolist() + self.assertEqual(len(owned), len(set(owned))) + with self.assertRaises(ValueError): # the hole value is never a valid slot + layout.slot(layout.hole, 0) + + def test_halo_pairs_skip_absent_blocks(self): + layout = BlockLayout2D.build(6, 12, 3, 4, periodic=(False, False), origins=L_ORIGINS) + # Block 0: 3 east faces + SE corner; block 1: 3 west + 4 south faces; + # block 2: 4 north faces + NW corner. Nothing reads (3..5, 0..3) or x>=8. + self.assertEqual(len(layout.halo_dst), 16) + np.testing.assert_array_equal(np.bincount(layout.halo_kind, minlength=2), [14, 2]) + for ny, nx, by, bx, h, origins in SPARSE: + for periodic in PERIODIC: + with self.subTest(origins=origins, periodic=periodic): + layout = BlockLayout2D.build(ny, nx, by, bx, halo=h, periodic=periodic, + origins=origins) + core = set(layout.owner[layout.owned].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.assertLess(max(dst + [-1]), layout.hole) + + def test_pack_unpack_and_pull_with_holes(self): + for ny, nx, by, bx, h, origins in SPARSE: + for periodic in PERIODIC: + with self.subTest(origins=origins, periodic=periodic): + layout = BlockLayout2D.build(ny, nx, by, bx, halo=h, periodic=periodic, + origins=origins) + owned = layout.owned + a = tags(ny, nx) + holed = a.copy() + holed[:, ~owned] = np.nan + blocked = layout.to_blocked(a) + np.testing.assert_array_equal(layout.to_dense(blocked), holed) + poisoned = a.copy() + poisoned[:, ~owned] = np.inf # hole values are ignored by packing + np.testing.assert_array_equal(layout.to_blocked(poisoned), blocked) + self.assertEqual(np.isfinite(blocked[:, 0]).sum(), + owned.sum() + len(layout.halo_dst)) + expected = dense_pull(holed, periodic) + expected[:, ~owned] = np.nan + got = layout.to_dense(blocked_pull(layout, blocked)) + np.testing.assert_array_equal(got, expected) + + def test_cut_links_in_or_next_to_absent_blocks_are_rejected(self): + layout = BlockLayout2D.build(6, 12, 3, 4, periodic=(False, False), origins=L_ORIGINS) + ok = [link(1, 1, 3, .6), link(1, 1, 3, .3), link(1, 7, 2, .6), link(1, 4, 2, .3)] + mapped = layout.map_cut_links(ok) # fully owned stencils, incl. across a seam + cells = int(np.prod(layout.local_shape)) + np.testing.assert_array_equal(mapped.block, [0, 0, 1, 1]) + np.testing.assert_array_equal(mapped.receiver, [layout.owner[1, 1]] * 2 + + [layout.owner[1, 7], layout.owner[1, 4]]) + self.assertEqual(int(mapped.donor[1]), int(layout.owner[0, 1])) + self.assertTrue((mapped.write < np.uint64(layout.nblocks * NQ * cells)).all()) + for bad in (link(4, 1, 1, .6), # receiver in an absent block + link(2, 1, 3, .6), # solid neighbor (3,1) absent + link(1, 7, 2, .3), # q<.5 donor (1,8) absent + link(2, 7, 1, .6)): # solid neighbor (2,8) absent + with self.subTest(fluid=bad.fluid_idx, dir=bad.dir, q=bad.q): + with self.assertRaises(ValueError): + layout.map_cut_links([bad]) + + def test_link_tables_store_owned_flags_only(self): + layout = BlockLayout2D.build(6, 12, 3, 4, periodic=(False, False), origins=L_ORIGINS) + flags = np.arange(1, 6 * 12 + 1, dtype=np.uint16) + tables, _ = blocked_link_tables(layout, flags, [], [], [], []) + owned = layout.owned.ravel() + stored = tables["flags"] + self.assertEqual(stored.shape, (layout.hole,)) + np.testing.assert_array_equal(stored[layout.owner.ravel()[owned].astype(np.intp)], flags[owned]) + self.assertEqual(np.count_nonzero(stored), owned.sum()) + + def test_invalid_origins_rejected(self): + for origins in ([(0, 4), (0, 0)], # not row-major + [(3, 0), (0, 4)], # not row-major across rows + [(0, 0), (0, 0)], # repeated + [(1, 0)], [(0, 2)], # misaligned + [(6, 0)], [(0, 12)], [(-3, 0)], [(0, -4)], # outside + np.zeros((0, 2), dtype=np.int64), # empty + [(0.0, 0.0)], [(True, False)], [0, 0], [(0, 0, 0)]): + with self.subTest(origins=origins), self.assertRaises(ValueError): + BlockLayout2D.build(6, 12, 3, 4, origins=origins) + + if __name__ == "__main__": unittest.main()