CelerisLab/src/CelerisLab/lbm/refinement_tables.py

# CelerisLab/lbm/refinement_tables.py
"""Host-only (numpy) mapping of fixed-patch transfer and boundary tables to block storage.

No CUDA import: CPU tests import these directly. refinement_stepper re-exports the names.
"""
from dataclasses import dataclass, replace

import numpy as np

from .blocks import blocked_link_tables
from ..body.jets import _surface_links
from .descriptors import D2Q9_EX, D2Q9_EY
from .refinement import ACTIVE, OUTER_GHOST


def blocked_fine_tables(patch, layout, order):
    """Remap fixed-patch fine tables to owner-core blocked storage (host only).

    Returns (explosion target block cells, coalescence child slots in the same
    rows/child order as ``patch.coalesce_fine[order]``, per-block-cell kind with
    unused halo entries OUTER). Coarse donors/targets are unchanged.
    """
    owner = layout.owner.ravel().astype(np.intp)
    cells = patch.fine_shape[0] * patch.fine_shape[1]
    q, cell = np.divmod(patch.coalesce_fine[order], np.uint64(cells))
    kind = blocked_kind(layout, patch.fine_kind)
    return (owner[patch.ghost_fine.astype(np.intp)].astype(np.uint64),
            layout.slot(owner[cell.astype(np.intp)], q), kind)


def blocked_kind(layout, dense_kind):
    """Per-block-cell kind; unused halo entries and unstored holes stay OUTER."""
    owner = layout.owner.ravel().astype(np.intp)
    owned = layout.owned.ravel()  # sparse multipatch levels leave holes unstored
    kind = np.full(layout.nblocks * int(np.prod(layout.local_shape)), OUTER_GHOST, np.uint8)
    kind[owner[owned]] = np.asarray(dense_kind).ravel()[owned]
    return kind

def blocked_coarse_tables(patch, layout, order, wall_aware=False, *, composed_kind=None):
    """Parent donor cells and receiver-local arrival slots, in transfer row order.

    Stage arrivals only after child reads and parent halo gather, immediately
    before parent pull. Source *owner* slots are wrong at block seams.
    Nested pairs must pass composed_kind to retain inherited INNER/OUTER and
    sparse holes; donors and receivers must be owned, uncovered ACTIVE cells.
    """
    owner = layout.owner.ravel()
    q, receiver = np.divmod(patch.coalesce_coarse[order], np.uint64(owner.size))
    if composed_kind is None:
        dense_kind = np.where(patch.covered, OUTER_GHOST, ACTIVE).astype(np.uint8)
        if wall_aware:
            dense_kind[[0, -1]] = OUTER_GHOST
    else:
        dense_kind = np.asarray(composed_kind)
        if dense_kind.shape != patch.coarse_shape:
            raise ValueError("composed parent kind has the wrong shape")
    donors = patch.ghost_coarse.astype(np.intp)
    receivers = receiver.astype(np.intp)
    legal = layout.owned.ravel() & (dense_kind.ravel() == ACTIVE) & ~patch.covered.ravel()
    if not legal[donors].all() or not legal[receivers].all():
        raise ValueError("child donors/receivers need owned uncovered parent ACTIVE cells")
    offset = (np.asarray(D2Q9_EY)[q.astype(np.intp)] * layout.local_shape[1]
              + np.asarray(D2Q9_EX)[q.astype(np.intp)])
    source = owner[receivers].astype(np.int64) - offset
    return (owner[donors].astype(np.uint64), layout.slot(source, q),
            blocked_kind(layout, dense_kind))


def projected_explosion_tables(pair, parent_layout, composed_parent_kind):
    """Experimental transverse linear explosion, in the pair's existing row order.

    Schornbaum & Rüde (2016), arXiv:1508.07982, Algorithm 2, following
    Chen et al., doi:10.1016/j.physa.2005.09.036. Independently implemented
    from the mathematics: centered gradients, perpendicular projection and
    child offsets +/-1/4. The kernel applies the projection per population.

    Only original explosion donors that are owned, uncovered ACTIVE cells
    may supply gradients. Both axis neighbors must be legal; otherwise repeat
    the center twice for zero slope. Owner-core addresses avoid stale halos.
    The same rule applies to scalar windows, unions and every adjacent pair;
    nested layouts retain their own per-axis (usually nonperiodic) topology.
    """
    ny, nx = pair.coarse_shape
    kind = np.asarray(composed_parent_kind)
    if parent_layout.shape != pair.coarse_shape or kind.shape != pair.coarse_shape:
        raise ValueError("projected explosion parent layout/kind has the wrong shape")
    donor = pair.ghost_coarse.astype(np.intp)
    legal = np.zeros(ny * nx, bool)
    legal[donor] = True
    legal &= (parent_layout.owned.ravel() & ~pair.covered.ravel()
              & (kind.ravel() == ACTIVE))
    if not legal[donor].all():
        raise ValueError("projected explosion donors need owned uncovered parent ACTIVE cells")
    y, x = np.divmod(donor, nx)
    neighbors = np.repeat(donor[:, None], 4, axis=1)
    for axis, (dx, dy) in enumerate(((1, 0), (0, 1))):
        minus = ((y - dy) % ny) * nx + (x - dx) % nx
        plus = ((y + dy) % ny) * nx + (x + dx) % nx
        valid = legal[minus] & legal[plus]
        if not parent_layout.periodic[axis]:
            coordinate, extent = (x, nx) if axis == 0 else (y, ny)
            valid &= (coordinate > 0) & (coordinate < extent - 1)
        neighbors[valid, 2 * axis] = minus[valid]
        neighbors[valid, 2 * axis + 1] = plus[valid]
    fy, fx = np.divmod(pair.ghost_fine, pair.fine_shape[1])
    ox, oy = pair.fine_origin
    offsets = np.column_stack((ox + (fx + .5) / 2 - (x + .5),
                               oy + (fy + .5) / 2 - (y + .5)))
    return (np.ascontiguousarray(parent_layout.owner.ravel()[neighbors], dtype=np.uint64),
            np.ascontiguousarray(offsets, dtype=np.float32))


# Required device columns plus optional actuation; observation arms are uploaded
# separately only when enabled. body_id/aperture remain host reduction metadata.
BLOCKED_BOUNDARY_DEVICE_COLUMNS = ("block", "cell", "dir", "q", "fallback", "surface_velocity")


@dataclass
class _BlockedBoundary:
    """One level's stationary BFL data and mutable Jet velocity column."""
    host: object
    tables: dict
    jets: tuple
    actuated: bool
    jet_cache: object = None
    flags: object = None
    device: object = None
    kernel: object = None
    last_raw: object = None
    interval_integral: object = None
    accumulate: object = None
    window_integral: object = None  # optional observation-window FP64 ledger
    spacing: float = None  # level spacing h_l in root units (exchange scaling)
    centers: tuple = ()    # body reference centers in the root edge frame


def _prepare_blocked_boundary(boundary, layout, cfg, jets, *, _shifted_blocked=False):
    if boundary is None or layout is None:
        return None
    fluid_idx, dirs, q, body_id, rx, ry, _, fallback = boundary.curved_host
    tables, _ = blocked_link_tables(layout, boundary.flags, fluid_idx, dirs, q, fallback)
    # blocked_link_tables/map_cut_links already rejects repeated (fluid_idx, dir),
    # independently of body_id, and repeated local writers.
    tables.update(body_id=np.asarray(body_id, np.int32), rx=np.asarray(rx, np.float32),
                  ry=np.asarray(ry, np.float32), aperture=np.zeros(len(dirs), bool))
    actuated = bool(jets) or any(o.surface_actuation is not None for o in boundary.objects)
    jet_cache = None
    if actuated:
        cfg = replace(cfg)
        cfg.validate()
        tables["surface_velocity"], tables["aperture"], shape, jet_id = _surface_links(
            boundary.objects, cfg, boundary.flags, fluid_idx, dirs, body_id, rx, ry,
            jets=jets, hits=boundary.curved_hits, _shifted_blocked=_shifted_blocked)
        if jets:
            jet_cache = (shape, jet_id)
    return _BlockedBoundary(boundary, tables, jets, actuated, jet_cache)