ClaudeandClaude Opus 5.5
0f8a428f07
MovingCircleLinks computes the cut links of each translating circle from an
FP64 {cx, cy, r} table over its swept band (same q/hit/clamp/fallback order
as build_cut_links, explicit _rn operations) and applies the unchanged
apply_bouzidi_link with spin wall velocity. obs keeps its contract; a
separate per-body diagnostics buffer records CME, the Wen term u_w*dm,
link_dm and the link count. Moving tables are built at initialize and
attached to the stepper, which then launches the kernel before the static
curved kernel; until then direct and run steps stay blocked. Displacement
is disabled: set_body offset/vx/vy on a translating circle raises before
mutation. Contract section 10a documents S2a and the S3 pending rule.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
225 lines
8.6 KiB
Python
225 lines
8.6 KiB
Python
# CelerisLab/lbm/stepper.py
|
|
"""
|
|
LBMStepper — time-advance driver for LBM kernels.
|
|
|
|
Owns kernel function handles and manages the double-buffer swap.
|
|
Does NOT return observations by default — callers use field.get_macroscopic()
|
|
or body manager methods to pull data when they need it.
|
|
"""
|
|
|
|
import numpy as np
|
|
import pycuda.driver as cuda
|
|
|
|
|
|
class LBMStepper:
|
|
"""Drive the LBM kernel forward in time."""
|
|
|
|
def __init__(self, field, module: cuda.Module, cfg):
|
|
self.field = field
|
|
self.module = module
|
|
self.cfg = cfg
|
|
|
|
self._esopull = (cfg.streaming == "esopull")
|
|
# Set by the owner when a step would be wrong (translating circles
|
|
# excluded from the packed links before the moving kernels exist).
|
|
self.blocked: str | None = None
|
|
self.moving = None # MovingTables attached by the owner (translating circles)
|
|
|
|
if self._esopull:
|
|
self.step_fn = module.get_function("EsoPullStep")
|
|
self.init_fn = module.get_function("InitEsoPull")
|
|
self.prepare_boundaries_fn = module.get_function("PrepareEsoPullBoundaries")
|
|
self.boundary_values_gpu = cuda.mem_alloc((6 * cfg.nx + 2 * cfg.ny * cfg.nq) * 4)
|
|
else:
|
|
self.step_fn = module.get_function("OneStep")
|
|
self.init_fn = module.get_function("InitTubeFlow")
|
|
|
|
self.curved_fn = module.get_function("CurvedBoundaryKernel")
|
|
self.prepare_esopull_curved_fn = module.get_function("PrepareEsoPullCurvedKernel")
|
|
self.apply_esopull_curved_fn = module.get_function("ApplyEsoPullCurvedKernel")
|
|
self.sensor_fn = module.get_function("SensorKernel")
|
|
self.force_region_fn = module.get_function("ForceRegionKernel")
|
|
self.moving_fn = None if self._esopull else module.get_function("MovingCircleLinks")
|
|
|
|
tpb = cfg.threads_per_block
|
|
self.block = (tpb, 1, 1)
|
|
if cfg.is_d2q9:
|
|
self.grid = ((cfg.nx + tpb - 1) // tpb, cfg.ny, 1)
|
|
else:
|
|
self.grid = ((cfg.nx + tpb - 1) // tpb, cfg.ny, cfg.nz)
|
|
|
|
self.boundary_grid = ((max(cfg.nx, cfg.ny) + tpb - 1) // tpb, 2, 1)
|
|
self._step_count = 0
|
|
|
|
def initialize(self):
|
|
"""Run the init kernel to set up channel flow + flags."""
|
|
f = self.field
|
|
self.init_fn(
|
|
f.flag_gpu, f.ddf_gpu,
|
|
block=self.block, grid=self.grid,
|
|
)
|
|
if not self._esopull:
|
|
_gpu_bytes = f.n * f.nq * f.store_bytes
|
|
cuda.memcpy_dtod(f.temp_gpu, f.ddf_gpu, _gpu_bytes)
|
|
cuda.memcpy_dtoh(f.flag, f.flag_gpu)
|
|
# Initialization publishes a new t=0 state, including EsoPull parity.
|
|
self._step_count = 0
|
|
f.download_ddf(step_id=0, force=True)
|
|
|
|
def step(
|
|
self,
|
|
n: int = 1,
|
|
*,
|
|
action_gpu: cuda.DeviceAllocation,
|
|
obs_gpu: cuda.DeviceAllocation,
|
|
stream: cuda.Stream | None = None,
|
|
):
|
|
"""Advance *n* time steps."""
|
|
if n > 0 and self.blocked is not None:
|
|
raise RuntimeError(self.blocked)
|
|
f = self.field
|
|
launch_kw = {}
|
|
if stream is not None:
|
|
launch_kw["stream"] = stream
|
|
|
|
for _ in range(n):
|
|
if self._esopull:
|
|
self.prepare_boundaries_fn(
|
|
f.ddf_gpu, self.boundary_values_gpu, np.uint64(self._step_count),
|
|
block=self.block, grid=self.boundary_grid, **launch_kw,
|
|
)
|
|
if self.moving is not None:
|
|
self._launch_moving(action_gpu, obs_gpu, **launch_kw)
|
|
self._launch_curved(action_gpu, obs_gpu, **launch_kw)
|
|
|
|
if self._esopull:
|
|
self.step_fn(
|
|
f.ddf_gpu, f.flag_gpu,
|
|
action_gpu, obs_gpu, self.boundary_values_gpu,
|
|
np.uint64(self._step_count),
|
|
block=self.block, grid=self.grid,
|
|
**launch_kw,
|
|
)
|
|
else:
|
|
self.step_fn(
|
|
f.flag_gpu, f.ddf_gpu, f.temp_gpu,
|
|
block=self.block, grid=self.grid,
|
|
**launch_kw,
|
|
)
|
|
f.ddf_gpu, f.temp_gpu = f.temp_gpu, f.ddf_gpu
|
|
|
|
self._launch_force_region(action_gpu, **launch_kw)
|
|
self._launch_sensor(obs_gpu, **launch_kw)
|
|
self._step_count += 1
|
|
if n > 0:
|
|
# Any completed GPU step invalidates host DDF mirror.
|
|
f.invalidate_host_ddf_cache()
|
|
|
|
def _launch_moving(self, action_gpu, obs_gpu, **launch_kw):
|
|
m, tpb = self.moving, self.cfg.threads_per_block
|
|
self.moving_fn(
|
|
self.field.ddf_gpu, self.field.flag_gpu, m.band_gpu, m.geom_gpu,
|
|
np.uint32(m.n_bodies), np.uint32(m.n_slots), action_gpu, obs_gpu,
|
|
m.diag_gpu, m.slot_debug_gpu if m.slot_debug_gpu is not None else np.uintp(0),
|
|
block=(tpb, 1, 1), grid=((m.n_slots + tpb - 1) // tpb, 1, 1), **launch_kw,
|
|
)
|
|
|
|
def _launch_curved(self, action_gpu, obs_gpu, **launch_kw):
|
|
f = self.field
|
|
if f.n_curved == 0:
|
|
return
|
|
tpb = self.cfg.threads_per_block
|
|
grid_c = ((f.n_curved + tpb - 1) // tpb, 1, 1)
|
|
cl = f.curved
|
|
assert cl.fluid_idx_gpu is not None
|
|
assert cl.dir_gpu is not None
|
|
assert cl.q_gpu is not None
|
|
assert cl.rx_gpu is not None
|
|
assert cl.ry_gpu is not None
|
|
assert cl.rz_gpu is not None
|
|
assert cl.fallback_class_gpu is not None
|
|
assert cl.body_id_gpu is not None
|
|
|
|
if self._esopull:
|
|
eb = f.esopull_curved
|
|
assert eb.corr_value_gpu is not None
|
|
assert eb.target_cell_gpu is not None
|
|
assert eb.target_slot_gpu is not None
|
|
self.prepare_esopull_curved_fn(
|
|
f.ddf_gpu,
|
|
cl.fluid_idx_gpu, cl.dir_gpu, cl.q_gpu,
|
|
cl.rx_gpu, cl.ry_gpu, cl.fallback_class_gpu, cl.body_id_gpu,
|
|
action_gpu, obs_gpu,
|
|
eb.corr_value_gpu, eb.target_cell_gpu, eb.target_slot_gpu,
|
|
np.uint32(f.n_curved),
|
|
np.uint64(self._step_count),
|
|
block=(tpb, 1, 1), grid=grid_c,
|
|
**launch_kw,
|
|
)
|
|
self.apply_esopull_curved_fn(
|
|
f.ddf_gpu,
|
|
eb.corr_value_gpu, eb.target_cell_gpu, eb.target_slot_gpu,
|
|
np.uint32(f.n_curved),
|
|
block=(tpb, 1, 1), grid=grid_c,
|
|
**launch_kw,
|
|
)
|
|
else:
|
|
assert cl.fluid_idx_gpu is not None
|
|
if cl.surface_velocity is not None and (
|
|
cl.surface_velocity_gpu is None or cl.link_exchange_gpu is None):
|
|
raise RuntimeError("Surface boundary device buffers are missing")
|
|
self.curved_fn(
|
|
f.ddf_gpu,
|
|
cl.fluid_idx_gpu, cl.dir_gpu, cl.q_gpu,
|
|
cl.rx_gpu, cl.ry_gpu, cl.rz_gpu, cl.fallback_class_gpu, cl.body_id_gpu,
|
|
action_gpu, obs_gpu,
|
|
np.uint32(f.n_curved),
|
|
cl.surface_velocity_gpu if cl.surface_velocity_gpu is not None else np.uintp(0),
|
|
cl.link_exchange_gpu if cl.link_exchange_gpu is not None else np.uintp(0),
|
|
block=(tpb, 1, 1), grid=grid_c,
|
|
**launch_kw,
|
|
)
|
|
|
|
def _launch_sensor(self, obs_gpu, **launch_kw):
|
|
f = self.field
|
|
if f.n_sensor == 0:
|
|
return
|
|
tpb = self.cfg.threads_per_block
|
|
grid_s = ((f.n_sensor + tpb - 1) // tpb, 1, 1)
|
|
fi_in = f.ddf_gpu
|
|
se = f.sensors
|
|
assert se.cells_gpu is not None and se.obj_id_gpu is not None
|
|
t = np.uint64(self._step_count + 1 if self._esopull else 0)
|
|
self.sensor_fn(
|
|
fi_in, f.flag_gpu,
|
|
se.cells_gpu, se.obj_id_gpu,
|
|
obs_gpu,
|
|
np.uint32(f.n_sensor),
|
|
t,
|
|
block=(tpb, 1, 1), grid=grid_s,
|
|
**launch_kw,
|
|
)
|
|
|
|
def _launch_force_region(self, action_gpu, **launch_kw):
|
|
f = self.field
|
|
if f.n_force_region == 0:
|
|
return
|
|
tpb = self.cfg.threads_per_block
|
|
grid_f = ((f.n_force_region + tpb - 1) // tpb, 1, 1)
|
|
fr = f.force_regions
|
|
assert fr.cells_gpu is not None and fr.obj_id_gpu is not None
|
|
t = np.uint64(self._step_count + 1 if self._esopull else 0)
|
|
self.force_region_fn(
|
|
f.ddf_gpu, f.flag_gpu,
|
|
fr.cells_gpu, fr.obj_id_gpu,
|
|
action_gpu,
|
|
np.uint32(f.n_force_region),
|
|
t,
|
|
block=(tpb, 1, 1), grid=grid_f,
|
|
**launch_kw,
|
|
)
|
|
|
|
@property
|
|
def step_count(self) -> int:
|
|
return self._step_count
|