Support stationary bodies and geometry Jets on nested L2

This commit is contained in:
TRON committed 2026-10-16 08:49:08 -07:00
1 parent 64665319a4
commit 20fb16ded5
5 files changed
+400 -67

No files matched your search

+19 -8
View File
@@ -349,12 +349,16 @@ def nest_patch(cfg, parent, child_bounds):
return child, _readonly(composed, np.uint8)
def prepare_fine_boundary(cfg, patch, objects, *, wall_aware=False, jets=()):
def prepare_fine_boundary(cfg, patch, objects, *, wall_aware=False, jets=(),
parent_origin=(0., 0.), parent_spacing=1.):
"""Project fixed circles and reject boundary/transfer overlap before CUDA.
Input centers/radii are in coarse physical coordinates (dx=dt=1).
Input centers/radii and Jets are in the root edge frame. ``parent_origin``
and ``parent_spacing`` map this patch's parent grid into that frame; the
default preserves the L0->L1 convention. cfg must describe the parent grid.
SurfaceActuation retains its lattice-velocity meaning: input coarse units
are converted into independent fine-local objects with dx=dt=.5.
are converted into independent fine-local objects with
dx=dt=parent_spacing/2.
Original objects are not retained or mutated. No runtime geometry updates.
``patch`` may be a FixedPatchSet2D: objects are then projected once onto
the level-wide fine grid (fine_origin (0, 0)); each circle must lie inside
@@ -364,6 +368,10 @@ def prepare_fine_boundary(cfg, patch, objects, *, wall_aware=False, jets=()):
selection is checked against true fine hits here. Jet amplitudes need no
rescaling (lattice velocity scales by dt/dx = 1 under acoustic 2:1).
"""
origin = np.asarray(parent_origin, dtype=np.float64)
if (origin.shape != (2,) or not np.isfinite(origin).all()
or not np.isfinite(parent_spacing) or parent_spacing <= 0):
raise ValueError("Parent frame must have a finite origin and positive spacing")
if not objects:
if jets:
raise ValueError("Fine jets need fine boundary objects")
@@ -386,7 +394,8 @@ def prepare_fine_boundary(cfg, patch, objects, *, wall_aware=False, jets=()):
local, links = [], []
rects = patch.bounds if isinstance(patch, FixedPatchSet2D) else (patch.bounds,)
covered = patch.covered
dx = dt = .5
dx = dt = .5 * parent_spacing
fine_origin = origin + parent_spacing * np.asarray(patch.fine_origin)
for bid, obj in enumerate(objects):
if (type(obj.geometry) is not CircleGeometry or obj.is_sensor or obj.is_force_region
or any(getattr(obj.state, name) != 0 for name in ("vx", "vy", "omega"))
@@ -394,11 +403,13 @@ def prepare_fine_boundary(cfg, patch, objects, *, wall_aware=False, jets=()):
raise ValueError("Fine boundaries require stationary circles without other roles")
cx, cy = obj.geometry.center
radius = obj.geometry._radius
px, py = (np.asarray((cx, cy)) - origin) / parent_spacing
pr = radius / parent_spacing
if (not np.isfinite([cx, cy, radius]).all() or radius <= 0
or not any(x0 <= cx-radius and cx+radius <= x1 and y0 <= cy-radius
and cy+radius <= y1 for x0, y0, x1, y1 in rects)):
or not any(x0 <= px-pr and px+pr <= x1 and y0 <= py-pr
and py+pr <= y1 for x0, y0, x1, y1 in rects)):
raise ValueError("Entire fine circle must lie inside one physical patch")
center = tuple((v-o)/dx-.5 for v, o in zip((cx, cy), patch.fine_origin))
center = tuple((v-o)/dx-.5 for v, o in zip((cx, cy), fine_origin))
spec = obj.surface_actuation
if spec is not None:
spec = replace(spec, normal_velocity=spec.normal_velocity * dt / dx)
@@ -465,7 +476,7 @@ def prepare_fine_boundary(cfg, patch, objects, *, wall_aware=False, jets=()):
raise ValueError("Transfer footprint touches the coarse wall reconstruction stencil")
from ..body.jets import fine_hits_to_edge_frame, surface_link_velocity
hits = fine_hits_to_edge_frame([(link.hit_x, link.hit_y) for link in links],
patch.fine_origin, dx)
fine_origin, dx)
packed = pack_cut_links_to_soa(links)
if jets:
# Same assignment checks as the stepper: empty/overlapping jets and
+115 -51
View File
@@ -93,6 +93,36 @@ class _Level:
halo_copy: object = None
halo_dst: object = None
halo_src: object = None
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:
@@ -140,9 +170,18 @@ class FixedPatchStepper:
uses one level-wide fine grid (2*ny, 2*nx): public fine arrays have that
shape, NaN outside the patches' ACTIVE cells. The step sequence is the
blocked one above, with every patch's transfers in one launch each.
Three-level single chain: ``child`` bounds use L1 window cells; all three
levels must be blocked. ``child_objects`` and ``child_jets`` use the L0
edge frame, with L2 dx=dt=.25 and unchanged lattice velocity. L1 bodies
and Jets remain unsupported in a chain. Every whole circle and its BFL
stencil must fit strictly away from the L2 transfer footprint. L2 runs
halo -> BFL -> step four times per base step. Runtime child Jet changes
require ``set_jet_amplitude(index, amplitude, level=2)``.
"""
def __init__(self, cfg, bounds, *, context, coarse_boundary="periodic", fine_objects=(),
fine_block=None, base_block=None, jets=(), child=None, child_block=None):
fine_block=None, base_block=None, jets=(), child=None, child_block=None,
child_objects=(), child_jets=()):
if HALO != 4:
raise RuntimeError("refinement.cu requires a four-cell fine halo")
if coarse_boundary not in ("periodic", "slot_reservoir"):
@@ -164,7 +203,10 @@ class FixedPatchStepper:
rects = self.patch.bounds if multipatch else (self.patch.bounds,)
if any(r[1] < 4 or r[3] > cfg.ny - 4 for r in rects):
raise ValueError("slot_reservoir requires y0 >= 4 and y1 <= ny-4")
fine_objects = tuple(fine_objects)
fine_objects, jets = tuple(fine_objects), tuple(jets)
child_objects, child_jets = tuple(child_objects), tuple(child_jets)
if child is None and (child_objects or child_jets):
raise ValueError("child_objects/child_jets require a nested child")
if base_block is not None and fine_block is None:
raise ValueError("base_block requires fine_block")
self.base_layout = None
@@ -193,7 +235,7 @@ class FixedPatchStepper:
raise ValueError("a nested child needs one patch plus fine_block, base_block "
"and child_block")
if tuple(fine_objects) or tuple(jets):
raise ValueError("a nested child does not support fine_objects or jets yet")
raise ValueError("a nested child supports bodies/Jets only on L2; use child_objects/child_jets")
if len(child_block) != 2:
raise ValueError("child_block must be (by, bx)")
self.child_patch, self._middle_kind = nest_patch(cfg, self.patch, child)
@@ -205,15 +247,33 @@ class FixedPatchStepper:
self._jets = tuple(jets)
self.fine_boundary = prepare_fine_boundary(
cfg, self.patch, fine_objects, wall_aware=self._wall_aware, jets=self._jets)
self._fine_tables, self._actuated, self._jet_cache = None, False, None
if self.fine_layout is not None and self.fine_boundary is not None:
fluid_idx, dirs, q, _, _, _, _, fallback = self.fine_boundary.curved_host
self._fine_tables, _ = blocked_link_tables(
self.fine_layout, self.fine_boundary.flags, fluid_idx, dirs, q, fallback)
fine_cfg = replace(cfg, nx=self.patch.fine_shape[1], ny=self.patch.fine_shape[0],
viscosity=2 * cfg.viscosity)
if self._wall_aware:
fine_cfg = replace(fine_cfg, y_wall_bc="bounce_back", top_wall_bc="velocity",
wall_slot_width=0, wall_normal_velocity=0.0, velocity=0.0)
self.child_boundary = (None if self.child_patch is None else prepare_fine_boundary(
fine_cfg, self.child_patch, child_objects, jets=child_jets,
parent_origin=self.patch.fine_origin, parent_spacing=.5))
child_cfg = (None if self.child_patch is None else replace(
fine_cfg, nx=self.child_patch.fine_shape[1], ny=self.child_patch.fine_shape[0],
viscosity=4 * cfg.viscosity))
self._boundaries = {
1: _prepare_blocked_boundary(self.fine_boundary, self.fine_layout, fine_cfg, self._jets),
2: _prepare_blocked_boundary(self.child_boundary, self.child_layout, child_cfg, child_jets),
}
# Compatibility aliases for the existing L1 inspection API.
b = self._boundaries[1]
self._fine_tables = None if b is None else b.tables
self._actuated = False if b is None else b.actuated
self._jet_cache = None if b is None else b.jet_cache
self._fine_leaf = (self.patch.fine_kind if self._middle_kind is None
else self._middle_kind) == ACTIVE
self._child_leaf = (None if self.child_patch is None
else self.child_patch.fine_kind == ACTIVE)
if self.child_boundary is not None:
self._child_leaf &= (self.child_boundary.flags.reshape(
self.child_patch.fine_shape) & SOLID) == 0
self._child = None
if self.fine_boundary is not None:
self._fine_leaf &= (self.fine_boundary.flags.reshape(self.patch.fine_shape) & SOLID) == 0
@@ -228,31 +288,13 @@ class FixedPatchStepper:
self._coalesce = None
self._link_tables = {}
self._block = (cfg.threads_per_block, 1, 1)
fine_cfg = replace(cfg, nx=self.patch.fine_shape[1], ny=self.patch.fine_shape[0],
viscosity=2 * cfg.viscosity)
if self._wall_aware:
fine_cfg = replace(fine_cfg, y_wall_bc="bounce_back", top_wall_bc="velocity",
wall_slot_width=0, wall_normal_velocity=0.0, velocity=0.0)
if self._fine_tables is not None:
# Any actuation spec (even zero amplitude) requires the velocity buffer.
b = self.fine_boundary
self._actuated = bool(self._jets) or any(
obj.surface_actuation is not None for obj in b.objects)
if self._actuated:
fine_cfg.validate() # derives nq etc.; the unique-writer key needs them
fluid_idx, dirs, _, body_id, rx, ry, _, _ = b.curved_host
self._fine_tables["surface_velocity"], _ = surface_link_velocity(
b.objects, fine_cfg, b.flags, fluid_idx, dirs, body_id, rx, ry,
jets=self._jets, hits=b.curved_hits)
if self._jets:
self._jet_cache = jet_link_shapes(self._jets, b.curved_hits)
with self.ctx.activate():
try:
self._stream = cuda.Stream()
specs = [(replace(cfg), self.base_layout, self.patch.tau_coarse, "PatchCoarseStep"),
(fine_cfg, self.fine_layout, self.patch.tau_fine, "PatchFineStep")]
if self.child_patch is not None:
specs.append((replace(fine_cfg, viscosity=4 * cfg.viscosity), self.child_layout,
specs.append((child_cfg, self.child_layout,
self.child_patch.tau_fine, "PatchFineStep"))
for index, (config, layout, tau, name) in enumerate(specs):
source = "refinement.cu"
@@ -284,6 +326,7 @@ class FixedPatchStepper:
1 if layout is None else layout.nblocks)
level = _Level(shape, module, self._allocate(size), self._allocate(size),
module.get_function(name), grid, layout)
level.boundary = self._boundaries.get(index)
self._levels.append(level)
if layout is not None:
level.halo_copy = module.get_function("BlockHaloCopy")
@@ -330,11 +373,15 @@ class FixedPatchStepper:
self._upload(ghosts), self._upload(donors.astype(np.uint64)),
self._upload(targets), self._upload(children),
len(ghosts), len(targets), int(np.prod(f.layout.local_shape)))
if self._fine_tables is not None:
self._fine_flags = self._upload(self._fine_tables["flags"])
self._link_tables = {k: self._upload(v) for k, v in self._fine_tables.items()
if k != "flags"}
self._curved = self._levels[1].module.get_function("BlockCurvedBoundary")
for level in self._levels:
b = level.boundary
if b is not None:
b.flags = self._upload(b.tables["flags"])
b.device = {k: self._upload(v) for k, v in b.tables.items() if k != "flags"}
b.kernel = level.module.get_function("BlockCurvedBoundary")
if self._boundaries[1] is not None:
b = self._boundaries[1]
self._fine_flags, self._link_tables, self._curved = b.flags, b.device, b.kernel
elif self.fine_boundary is not None:
from .curved_links import CurvedLinkSoA
boundary = self.fine_boundary
@@ -431,58 +478,72 @@ class FixedPatchStepper:
cl.link_exchange_gpu if cl.link_exchange_gpu is not None else np.uintp(0),
grid=((cl.count + self._block[0] - 1) // self._block[0], 1), **kw)
def set_jet_amplitude(self, index, amplitude):
def set_jet_amplitude(self, index, amplitude, *, level=1):
"""Replace one jet's amplitude between steps; no geometry/table rebuild.
``level=1`` preserves the L1 default; use ``level=2`` for child_jets.
Invalid input raises before any change. A device/sync error marks the
instance failed (as step/initialize do); host state commits only after
a successful upload.
"""
with self._lock, self.ctx.activate():
self._require("new", "ready")
jets = list(self._jets)
if (isinstance(level, bool) or not isinstance(level, Integral)
or level not in (1, 2) or level >= len(self._levels)):
raise ValueError("Jet level must be 1 or an existing child level 2")
boundary = self._boundaries[level]
jets = list(self._jets if level == 1 else (() if boundary is None else boundary.jets))
jets[index] = jets[index].with_amplitude(amplitude) # validates first
velocity = None
if self._fine_curved is None and self._jet_cache is not None:
shape, jet_id = self._jet_cache
velocity = self._fine_tables["surface_velocity"].copy()
if boundary is not None and boundary.jet_cache is not None:
shape, jet_id = boundary.jet_cache
velocity = boundary.tables["surface_velocity"].copy()
driven = jet_id >= 0
velocity[driven] = jet_link_velocity(jets, shape, jet_id)[driven]
try:
if self._fine_curved is not None:
if level == 1 and self._fine_curved is not None:
self._fine_curved.set_jet_amplitude(index, amplitude, self._stream)
elif velocity is not None:
self._stream.synchronize() # no in-flight step reads the column
cuda.memcpy_htod(self._link_tables["surface_velocity"], velocity)
cuda.memcpy_htod(boundary.device["surface_velocity"], velocity)
except BaseException:
self._status = "failed"
raise
if velocity is not None:
self._fine_tables["surface_velocity"] = velocity
self._jets = tuple(jets)
boundary.tables["surface_velocity"] = velocity
if boundary is not None:
boundary.jets = tuple(jets)
if level == 1:
self._jets = tuple(jets)
def _blocked_fine_boundary(self, buffer, **kw):
if self._fine_tables is None:
self._blocked_boundary(self._levels[1], buffer, **kw)
def _blocked_boundary(self, level, buffer, **kw):
b = level.boundary
if b is None:
return
if self._curved is None or self._fine_flags is None:
raise RuntimeError("Fine boundary device buffers are missing")
t, count = self._link_tables, len(self._fine_tables["cell"])
if b.kernel is None or b.flags is None:
raise RuntimeError("Boundary device buffers are missing")
t, count = b.device, len(b.tables["cell"])
velocity = t.get("surface_velocity")
if self._actuated and velocity is None:
if b.actuated and velocity is None:
raise RuntimeError("Surface velocity device buffer is missing")
self._curved(buffer, t["block"], t["cell"], t["dir"], t["q"], t["fallback"],
np.uint32(count), velocity if velocity is not None else np.uintp(0),
grid=((count + self._block[0] - 1) // self._block[0], 1), **kw)
b.kernel(buffer, t["block"], t["cell"], t["dir"], t["q"], t["fallback"],
np.uint32(count), velocity if velocity is not None else np.uintp(0),
grid=((count + self._block[0] - 1) // self._block[0], 1), **kw)
def _advance_child(self, parent_buffer, **kw):
"""Explode from the given L1 buffer, then two L2 hops (halo -> step)."""
"""Explode from the given L1 buffer, then two L2 hops (halo -> BFL -> step)."""
pair, g = self._child, self._levels[2]
pair.explode(parent_buffer, g.current, pair.ghosts, pair.donors, np.uint32(pair.n_ghosts),
np.uint64(pair.parent_cells),
grid=((pair.n_ghosts + self._block[0] - 1) // self._block[0], 1), **kw)
for src, dst in ((g.current, g.spare), (g.spare, g.current)):
self._gather_halo(g, src, **kw)
g.kernel(src, dst, np.uint32(g.layout.halo), np.uintp(0), g.kind, grid=g.grid, **kw)
self._blocked_boundary(g, src, **kw)
flags = np.uintp(0) if g.boundary is None else g.boundary.flags
g.kernel(src, dst, np.uint32(g.layout.halo), flags, g.kind, grid=g.grid, **kw)
def _gather_halo(self, level, buffer, **kw):
pairs = len(level.layout.halo_dst)
@@ -589,6 +650,9 @@ class FixedPatchStepper:
self._stream.synchronize()
except Exception as exc:
error = exc
for b in getattr(self, "_boundaries", {}).values():
if b is not None:
b.flags = b.device = b.kernel = None
self._levels.clear()
self._explode = self._coalesce = self._child = None
if getattr(self, "_fine_curved", None) is not None: