diff --git a/src/CelerisLab/lbm/refinement_stepper.py b/src/CelerisLab/lbm/refinement_stepper.py index ab4c8cb..053e95d 100644 --- a/src/CelerisLab/lbm/refinement_stepper.py +++ b/src/CelerisLab/lbm/refinement_stepper.py @@ -15,37 +15,14 @@ import numpy as np import pycuda.driver as cuda from ..cuda import compiler_v2 as compiler -from .blocks import BlockLayout2D, blocked_link_tables -from .curved_links import surface_link_velocity -from ..body.jets import jet_link_shapes, jet_link_velocity +from .blocks import BlockLayout2D +from ..body.jets import jet_link_velocity from .block_stepper import MAX_GRID_DIM -from .refinement import (ACTIVE, HALO, OUTER_GHOST, FixedPatch2D, FixedPatchSet2D, +from .refinement import (ACTIVE, HALO, FixedPatch2D, FixedPatchSet2D, nest_patch, nest_patch_set, prepare_fine_boundary) -from .descriptors import SOLID, D2Q9_EX, D2Q9_EY - - -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 +from .refinement_tables import (blocked_base_tables, blocked_fine_tables, # noqa: F401 (re-export) + blocked_kind, _prepare_blocked_boundary) +from .descriptors import SOLID @dataclass @@ -62,36 +39,6 @@ class _Pair: parent_cells: int # parent block window cells (explode donor stride) -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)) - - @dataclass class _Level: shape: tuple @@ -108,35 +55,6 @@ class _Level: boundary: object = None -@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 - - -def _prepare_blocked_boundary(boundary, layout, cfg, jets): - 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) - 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"], _ = surface_link_velocity( - boundary.objects, cfg, boundary.flags, fluid_idx, dirs, body_id, rx, ry, - jets=jets, hits=boundary.curved_hits) - return _BlockedBoundary(boundary, tables, jets, actuated, - jet_link_shapes(jets, boundary.curved_hits) if jets else None) - - class FixedPatchStepper: """D2Q9/SRT/FP32, one unforced internal 2:1 patch in an x-periodic box. diff --git a/src/CelerisLab/lbm/refinement_tables.py b/src/CelerisLab/lbm/refinement_tables.py new file mode 100644 index 0000000..5c7a452 --- /dev/null +++ b/src/CelerisLab/lbm/refinement_tables.py @@ -0,0 +1,95 @@ +# 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)) + + +@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 + + +def _prepare_blocked_boundary(boundary, layout, cfg, jets): + 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) + 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"], _ = surface_link_velocity( + boundary.objects, cfg, boundary.flags, fluid_idx, dirs, body_id, rx, ry, + jets=jets, hits=boundary.curved_hits) + return _BlockedBoundary(boundary, tables, jets, actuated, + jet_link_shapes(jets, boundary.curved_hits) if jets else None)