Frank14f/CelerisLab
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

BlockLayout2D: explicit selected block origins with hole owners

Browse Source 
build(..., origins=None|[(y0, x0), ...]) stores only grid-aligned selected
blocks (unique, strictly row-major, no silent reordering). Cells of absent
blocks own layout.hole (= nblocks*L, one past the last block cell) so an
unguarded index fails rather than wrapping like -1; layout.owned is the
dense mask. Halo pairs never read holes, to_blocked ignores hole values,
to_dense returns NaN there, map_cut_links rejects receivers, solid
neighbors or q<.5 donors in absent blocks, blocked_link_tables stores
owned flags only. origins=None or every origin reproduces the regular
layout exactly. Storage only: no CUDA or stepper changes.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Claude
2026-10-16 03:09:15 +00:00
co-authored by Claude Opus 5.5
parent 99ddb5fc27
commit c397e08565
2 changed files with 226 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/CelerisLab/lbm/blocks.py
+90 -20
View File
@@ -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),
tests/unit/test_block_layout.py
+136 -1
View File
@@ -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()