CelerisLab/src/CelerisLab/lbm/blocks.py

# CelerisLab/lbm/blocks.py
"""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). 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
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.
The one-cell same-level halo here is unrelated to the four-cell fine halo of
the fixed refinement patch (refinement.py), whose semantics are unchanged.

Static cut links (map_cut_links) use RECEIVER-LOCAL staging: each physical
link belongs once to the block whose core holds its fluid receiver. Order per
step: gather halos -> BFL writes the reflected source population into that
block's own core/halo copy of the solid cell -> pull/collide core cells, with
NO halo copy between BFL and pull. Reflected populations are temporary,
receiver-local boundary values: the solid's owner-core copy of that slot may
be stale, which is distinct from the unique core ownership of stored state.
No moving geometry is implied.
"""
from dataclasses import dataclass, replace
from numbers import Integral
from types import SimpleNamespace

import numpy as np

from .descriptors import D2Q9_EX, D2Q9_EY

NQ = 9
FACE, CORNER = 0, 1


def _readonly(values, dtype):
    array = np.ascontiguousarray(values, dtype=dtype)
    array.setflags(write=False)
    return array


def _index(value, upper, name, lower=0):
    """Integer index or index array in [lower, upper); rejects bools and floats."""
    array = np.asarray(value)
    if array.dtype.kind not in "iu" or (array.size and (array.min() < lower or array.max() >= upper)):
        raise ValueError(f"{name} must be integer index(es) in [{lower}, {upper})")
    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)


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).

    ``block`` owns the receiver core cell; ``receiver`` and ``donor`` are block
    cell ids in that block (donor = NO_DONOR unless q<.5 with a legal donor);
    ``write`` is the flat blocked slot of the reflected population.
    """
    block: np.ndarray
    receiver: np.ndarray
    donor: np.ndarray
    write: np.ndarray


NO_DONOR = np.iinfo(np.uint64).max


@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, 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.
    """
    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), 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:
            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 origins is not None:
            selected = _selected_blocks(origins, ny, nx, by, bx)
            origins = np.array(origins, dtype=np.int64).reshape(-1, 2)  # own copy, never caller storage
        # Capacity from arithmetic before any grid-sized allocation; hole =
        # nblocks * cells is below this bound too.
        if (nby * nbx if origins is None else len(origins)) * NQ * cells > np.iinfo(np.int64).max // 2:
            raise ValueError("blocked storage exceeds 64-bit indexing")
        if origins is None:
            selected = np.arange(nby * nbx)
            jy, jx = np.divmod(selected, 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]
        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]
        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)
            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])
        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 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

    @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))
        block_cell = _index(block_cell, self.nblocks * cells, "block_cell").astype(np.uint64)
        q = _index(q, NQ, "q").astype(np.uint64)
        b, c = np.divmod(block_cell, np.uint64(cells))
        return (b * np.uint64(NQ) + q) * 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;
        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
        ny, nx = self.shape
        block = int(_index(block, self.nblocks, "block"))
        y, x = int(_index(y, ny, "y")), int(_index(x, nx, "x"))
        y0, x0 = self.origins[block]
        ry, rx = y - int(y0) + self.halo, 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):
            raise ValueError("cell is not represented in this block's storage")
        return block * ly * lx + ry * lx + rx

    def _neighbor(self, y, x, dy, dx):
        ny, nx = self.shape
        y, x = y + dy, x + dx
        if self.periodic[1]:
            y %= ny
        if self.periodic[0]:
            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:
        """Map dense static cut links (fluid_idx, dir, q, fallback) to blocks.

        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). 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]
        cells = int(np.prod(self.local_shape))
        rows, seen, writers = [], set(), set()
        for link in links:
            r = int(_index(link.fluid_idx, ny * nx, "fluid_idx"))
            q = int(_index(link.dir, NQ, "dir", lower=1))
            if (r, q) in seen:
                raise ValueError("cut link is repeated")
            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.
            # Physical neighbors are formed only to reject non-periodic exteriors.
            step = D2Q9_EY[q] * lx + D2Q9_EX[q]
            self._neighbor(y, x, D2Q9_EY[q], D2Q9_EX[q])
            opp = q + 1 if q % 2 else q - 1
            write = int(self.slot(receiver + step, opp))
            if write in writers:
                raise ValueError("duplicate local boundary writer")
            writers.add(write)
            donor = NO_DONOR
            if float(link.q) < .5 and int(link.fallback) == 0:
                self._neighbor(y, x, -D2Q9_EY[q], -D2Q9_EX[q])
                donor = receiver - step
            rows.append((block, receiver, donor, write))
        table = np.array(rows, dtype=np.uint64).reshape(-1, 4)
        return BlockCutLinks(_readonly(table[:, 0], np.uint32),
                             *(_readonly(table[:, i], np.uint64) for i in (1, 2, 3)))

    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.

        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)
        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, holes are NaN.

        Full layouts accept any dtype; layouts with holes require floats.
        """
        blocked = self._check_blocked(blocked)
        flat = blocked.reshape(self.nblocks, NQ, -1)
        owned = self.owned.ravel()
        if not owned.all() and not np.issubdtype(blocked.dtype, np.floating):
            raise ValueError("a layout with holes unpacks floating arrays only (holes are NaN)")
        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():  # even an empty NaN assignment fails for integer dtypes
            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 \
                or not blocked.flags.c_contiguous:
            raise ValueError(f"expected C-contiguous blocked array {self.blocked_shape}")
        return blocked


def blocked_link_tables(layout, flags, fluid_idx, dir, q, fallback):
    """Storage adapter: dense static cut-link SoA + dense flags -> blocked tables.

    Legality (geometry, coverage, transfer isolation) stays with the caller's
    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. 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:
        raise ValueError("flags must be dense uint16 with one entry per grid cell")
    columns = [np.asarray(fluid_idx), np.asarray(dir), np.asarray(q), np.asarray(fallback)]
    if len({len(c) for c in columns}) != 1:
        raise ValueError("cut-link columns must have equal length")
    mapped = layout.map_cut_links(SimpleNamespace(fluid_idx=r, dir=d, q=v, fallback=f)
                                  for r, d, v, f in zip(*columns))
    cells = int(np.prod(layout.local_shape))
    blocked_flags = np.zeros(layout.nblocks * cells, np.uint16)
    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),
                  fallback=columns[3].astype(np.uint8))
    return tables, mapped


# One blocked level as FixedPatchStepper builds and launches it (any depth). Raw argument orders of the blocked.cu kernels live only here.

def level_config(cfg, layout, tau):
    """Kernel config of one blocked level: NX/NY = block window incl. halo.

    Physical wall/slot/inlet fields are neutralized; a blocked root gets them
    from its module globals (d_block_wall/d_block_channel) instead.
    """
    ly, lx = layout.local_shape
    local = replace(cfg, nx=lx, ny=ly, y_wall_bc="bounce_back", top_wall_bc="velocity",
                    wall_slot_width=0, wall_normal_velocity=0.0, velocity=0.0)
    local.validate()
    if not np.isclose(local.omega, 1 / tau, rtol=1e-14, atol=0):
        raise RuntimeError("block level omega/tau mismatch")
    return local


def launch_halo_copy(kernel, buffer, dst, src, count, *, block, stream):
    """BlockHaloCopy over ``count`` uploaded (dst, src) block-cell pairs."""
    kernel(buffer, dst, src, np.uint32(count),
           grid=((count + block[0] - 1) // block[0], 1), block=block, stream=stream)


def launch_curved_boundary(kernel, buffer, tables, count, *, block, stream,
                           velocity=None, rx=None, ry=None, raw=None):
    """BlockCurvedBoundary on uploaded link tables; None = no actuation/observation."""
    none = np.uintp(0)
    kernel(buffer, tables["block"], tables["cell"], tables["dir"], tables["q"], tables["fallback"],
           np.uint32(count), none if velocity is None else velocity,
           none if rx is None else rx, none if ry is None else ry, none if raw is None else raw,
           grid=((count + block[0] - 1) // block[0], 1), block=block, stream=stream)


def static_circle_tables(layout, objects):
    """Host tables for stationary circles; rejects before any CUDA work.

    Every periodic neighbor of every fluid core is inspected, including domain
    edges that CircleGeometry does not enumerate, so clipped or incomplete
    geometry cannot pass. Each physical link has one receiver-owning block.
    """
    from ..body.geometry.circle import CircleGeometry
    from .descriptors import FLUID, SOLID

    ny, nx = layout.shape
    flags = np.full(ny * nx, FLUID, np.uint16)
    for obj in objects:
        if (type(obj.geometry) is not CircleGeometry or obj.is_sensor or obj.is_force_region
                or obj.surface_actuation is not None
                or any(getattr(obj.state, name) != 0 for name in ("vx", "vy", "omega"))
                or any(vars(obj.control).values())):
            raise ValueError("Block bodies must be stationary passive circles")
        cx, cy = obj.geometry.center
        radius = obj.geometry._radius
        if not np.isfinite([cx, cy, radius]).all() or radius <= 0:
            raise ValueError("Block circle center/radius must be finite with radius > 0")
        mask = obj.geometry.build_flag_mask(nx, ny)
        solid = (mask & SOLID) != 0
        if not solid.any() or np.any(flags[solid] & SOLID):
            raise ValueError("Block circles must be resolved and disjoint")
        flags[solid] = mask[solid]
    solid = ((flags & SOLID) != 0).reshape(ny, nx)
    if solid[[0, -1]].any() or solid[:, [0, -1]].any():
        raise ValueError("Block circles must not touch the domain edge")
    links = [link for obj in objects
             for link in obj.geometry.build_cut_links(nx, ny, domain_flags=flags)]
    def source(r, q):
        y, x = divmod(r, nx)
        return (y + D2Q9_EY[q]) % ny, (x + D2Q9_EX[q]) % nx
    receivers = {(int(link.fluid_idx), int(link.dir)) for link in links}
    required = {(int(r), q) for r in np.flatnonzero(~solid.ravel())
                for q in range(1, NQ) if solid[source(r, q)]}
    if len(receivers) != len(links) or receivers != required:
        raise ValueError("Cut links must cover every fluid pull from a solid exactly once")
    tables, _ = blocked_link_tables(
        layout, flags, np.array([link.fluid_idx for link in links], np.uint32),
        np.array([link.dir for link in links], np.uint8),
        np.array([link.q for link in links], np.float32),
        np.array([link.fallback for link in links], np.uint8))
    return dict(solid=solid, **tables)