Files
CelerisLab/src/CelerisLab/lbm/refinement_tables.py
T

108 lines
4.9 KiB
Python

# 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 jet_link_shapes, surface_link_velocity
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_base_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))
# 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
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)
if actuated:
cfg = replace(cfg)
cfg.validate()
tables["surface_velocity"], tables["aperture"] = surface_link_velocity(
boundary.objects, cfg, boundary.flags, fluid_idx, dirs, body_id, rx, ry,
jets=jets, hits=boundary.curved_hits, _shifted_blocked=_shifted_blocked)
return _BlockedBoundary(boundary, tables, jets, actuated,
jet_link_shapes(jets, boundary.curved_hits) if jets else None)