Frank14f/CelerisLab
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Wire Jets into FixedPatchStepper fine level (dense and blocked)

Browse Source 
FixedPatchStepper(..., jets=()) passes jets (coarse cell-edge frame, as
fine_objects) to prepare_fine_boundary, which runs the shared assignment
checks before CUDA. Dense fine path: assign_surface with jets and the
boundary's true hits. Blocked path: surface_link_velocity with jets/hits;
unit-amplitude shapes cached. set_jet_amplitude (new/ready only) validates
first, then rewrites the dense SoA column or the blocked jet rows and
re-uploads only the velocity buffer after a stream sync. Scheduling,
kernels and launch ABI unchanged.

tests/validation/multigrid/test_refinement_jets.py: mocked-driver CPU
checks and frozen GPU gates G1-G3 (layout identity, jet vs legacy <=1e-6,
amplitude switch).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Claude
2026-10-15 08:42:51 +00:00
co-authored by Claude Opus 5.5
parent f26410f94a
commit aebef52345
2 changed files with 170 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/CelerisLab/lbm/refinement_stepper.py
+36 -7
View File
@@ -16,6 +16,7 @@ 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 .block_stepper import MAX_GRID_DIM
from .refinement import ACTIVE, HALO, OUTER_GHOST, FixedPatch2D, prepare_fine_boundary
from .descriptors import SOLID, D2Q9_EX, D2Q9_EY
@@ -106,9 +107,12 @@ class FixedPatchStepper:
covered cells are explicitly poisoned on every step. Public shapes stay dense.
Blocked mode applies prescribed surface velocity only; its per-link
exchange diagnostic is unavailable (rows keep physical cut-link order).
``jets`` (body.jets.Jet) use the same edge frame as ``fine_objects`` and
drive fine cut links selected at their true hits; ``set_jet_amplitude``
re-uploads the per-link velocity column only.
"""
def __init__(self, cfg, bounds, *, context, coarse_boundary="periodic", fine_objects=(),
fine_block=None, base_block=None):
fine_block=None, base_block=None, jets=()):
if HALO != 4:
raise RuntimeError("refinement.cu requires a four-cell fine halo")
if coarse_boundary not in ("periodic", "slot_reservoir"):
@@ -143,9 +147,10 @@ class FixedPatchStepper:
if (self.fine_layout.nblocks > MAX_GRID_DIM
or self.fine_layout.local_shape[0] > MAX_GRID_DIM):
raise ValueError("fine blocks exceed launch limits")
self._jets = tuple(jets)
self.fine_boundary = prepare_fine_boundary(
cfg, self.patch, fine_objects, wall_aware=self._wall_aware)
self._fine_tables, self._actuated = None, False
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(
@@ -172,12 +177,16 @@ class FixedPatchStepper:
if self._fine_tables is not None:
# Any actuation spec (even zero amplitude) requires the velocity buffer.
b = self.fine_boundary
self._actuated = any(obj.surface_actuation is not None for obj in b.objects)
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)
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()
@@ -257,7 +266,8 @@ class FixedPatchStepper:
# SoA owns its allocations; never also register them in _allocations.
self._fine_curved = CurvedLinkSoA()
self._fine_curved.assign_host(*boundary.curved_host)
self._fine_curved.assign_surface(boundary.objects, fine_cfg, boundary.flags)
self._fine_curved.assign_surface(boundary.objects, fine_cfg, boundary.flags,
jets=self._jets, hits=boundary.curved_hits)
self._fine_curved.upload()
self._curved = self._levels[1].module.get_function("CurvedBoundaryKernel")
except BaseException:
@@ -333,7 +343,8 @@ class FixedPatchStepper:
return
if self._curved is None or cl.fluid_idx_gpu is None:
raise RuntimeError("Fine boundary device buffers are missing")
driven = any(obj.surface_actuation is not None for obj in self.fine_boundary.objects)
driven = bool(self._jets) or any(
obj.surface_actuation is not None for obj in self.fine_boundary.objects)
if driven and (cl.surface_velocity_gpu is None or cl.link_exchange_gpu is None):
raise RuntimeError("Surface boundary device buffers are missing")
self._curved(buffer, cl.fluid_idx_gpu, cl.dir_gpu, cl.q_gpu,
@@ -343,6 +354,24 @@ 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):
"""Replace one jet's amplitude between steps; no geometry/table rebuild."""
with self._lock, self.ctx.activate():
self._require("new", "ready")
jets = list(self._jets)
jets[index] = jets[index].with_amplitude(amplitude) # validates first
if self._fine_curved is not None:
self._fine_curved.set_jet_amplitude(index, amplitude, self._stream)
elif self._jet_cache is not None:
shape, jet_id = self._jet_cache
velocity = self._fine_tables["surface_velocity"].copy()
driven = jet_id >= 0
velocity[driven] = jet_link_velocity(jets, shape, jet_id)[driven]
self._stream.synchronize() # no in-flight step reads the column
cuda.memcpy_htod(self._link_tables["surface_velocity"], velocity)
self._fine_tables["surface_velocity"] = velocity
self._jets = tuple(jets)
def _blocked_fine_boundary(self, buffer, **kw):
if self._fine_tables is None:
return
tests/validation/multigrid/test_refinement_jets.py
+134
View File
@@ -0,0 +1,134 @@
"""Jets on the FixedPatchStepper fine level: CPU wiring checks plus GPU gates.
Jets use the same coarse cell-edge frame as fine_objects. The equivalent jet
of Astra's actuated fixture is SurfaceActuation(0.7, 0.85, U).as_jet((16, 12),
2.15). Frozen GPU gates (8 coarse steps, both coarse boundary modes):
G1 jet run: dense vs 1x1/2x2/4x2 fine blocks bit-identical (values + NaN masks);
G2 jet vs legacy actuation (dense): max |dG| over finite fine ACTIVE fluid
populations <= 1e-6 (velocity rows differ only by FP32 rounding; reported);
G3 amplitude change after 4 steps (+U -> -U), 4 more steps: dense vs blocked
bit-identical and the run differs from the unchanged run.
"""
from contextlib import nullcontext
import unittest
from unittest.mock import Mock, patch
import numpy as np
import pycuda.driver as cuda
from CelerisLab.cuda import compiler_v2 as compiler
from CelerisLab.lbm.descriptors import SOLID
from CelerisLab.lbm.refinement import ACTIVE, FixedPatch2D
from CelerisLab.lbm.refinement_stepper import FixedPatchStepper
from test_refinement_blocked_fine import (BOUNDS, FINE_BLOCKS, NX, NY, STEPS, actuation,
circle, config, fine_circle, raw, state)
SPEED = .002
def jet(speed=SPEED):
return actuation(speed).as_jet((16.0, 12.0), 2.15)
def mocked(objects, jets, fine_block):
"""Real constructor with driver doubles; returns backend and the memcpy mock."""
context = Mock(activate=nullcontext, sm_arch="sm_70")
module = Mock(get_function=Mock(return_value=Mock()))
module.get_global.return_value = (Mock(), 20)
htod = Mock()
with patch.object(cuda, "Stream"), patch.object(compiler, "generate_config"), \
patch.object(compiler, "compile_kernel"), patch.object(compiler, "cleanup_build_dir"), \
patch.object(compiler, "load_module", return_value=module), \
patch.object(cuda, "mem_alloc", side_effect=lambda size: Mock()), \
patch.object(cuda, "memcpy_htod", htod):
backend = FixedPatchStepper(config("periodic"), BOUNDS, context=context,
fine_objects=objects, fine_block=fine_block, jets=jets)
return backend, htod
class JetStepperCPU(unittest.TestCase):
def test_blocked_jet_velocity_matches_legacy_rows(self):
legacy, _ = mocked([circle(actuation(SPEED))], (), FINE_BLOCKS[1])
driven, _ = mocked([circle()], [jet()], FINE_BLOCKS[1])
a = legacy._fine_tables["surface_velocity"]
b = driven._fine_tables["surface_velocity"]
diff = np.abs(a - b).max()
print(f"fine jet vs legacy velocity rows: max |dU| = {diff:.3e} (exact={diff == 0})")
self.assertLessEqual(diff, 4 * np.finfo(np.float32).eps * SPEED)
legacy.close(), driven.close()
def test_blocked_amplitude_update_rewrites_only_velocity_column(self):
backend, htod = mocked([circle()], [jet()], FINE_BLOCKS[1])
before = backend._fine_tables["surface_velocity"].copy()
device = backend._link_tables["surface_velocity"]
with patch.object(cuda, "memcpy_htod", htod):
htod.reset_mock()
backend.set_jet_amplitude(0, -SPEED)
(ptr, sent), _ = htod.call_args
self.assertEqual(htod.call_count, 1)
self.assertIs(ptr, device)
np.testing.assert_allclose(sent, -before, rtol=1e-6, atol=0)
np.testing.assert_array_equal(backend._fine_tables["surface_velocity"], sent)
with self.assertRaises(ValueError):
backend.set_jet_amplitude(0, 1.5)
self.assertEqual(htod.call_count, 1)
self.assertEqual(backend._jets[0].amplitude, -SPEED)
backend.close()
def test_jet_overlapping_legacy_aperture_rejected_before_cuda(self):
with patch.object(cuda, "Stream", side_effect=AssertionError("CUDA touched")):
with self.assertRaisesRegex(ValueError, "SurfaceActuation"):
FixedPatchStepper(config("periodic"), BOUNDS, context=Mock(),
fine_objects=[circle(actuation(SPEED))], jets=[jet()])
def run(mode, fine_block, *, objects, jets, switch=None):
from CelerisLab.cuda.context import CudaContext
cfg = config(mode)
patch_ = FixedPatch2D.from_config(cfg, BOUNDS)
with CudaContext(0) as context, FixedPatchStepper(
cfg, BOUNDS, context=context, coarse_boundary=mode, fine_objects=objects,
fine_block=fine_block, jets=jets) as backend:
backend.initialize(state(NY, NX, 1), state(*patch_.fine_shape, 2))
if switch is None:
backend.step(STEPS)
else:
backend.step(STEPS // 2)
backend.set_jet_amplitude(0, switch)
backend.step(STEPS - STEPS // 2)
return backend.download(), raw(backend)
class JetStepperGPU(unittest.TestCase):
def test_g1_g2_jet_layouts_identical_and_close_to_legacy(self):
for mode in ("periodic", "slot_reservoir"):
patch_, boundary = fine_circle(mode)
fluid = (patch_.fine_kind == ACTIVE) & (
(boundary.flags.reshape(patch_.fine_shape) & SOLID) == 0)
dense = run(mode, None, objects=[circle()], jets=[jet()])
self.assertTrue(np.isfinite(dense[1][1][:, fluid]).all(), mode)
for fine_block in FINE_BLOCKS:
public, raw_arrays = run(mode, fine_block, objects=[circle()], jets=[jet()])
for a, b in zip(public + tuple(raw_arrays), dense[0] + tuple(dense[1])):
np.testing.assert_array_equal(a, b, err_msg=f"G1 {mode} {fine_block}")
legacy = run(mode, None, objects=[circle(actuation(SPEED))], jets=())
diff = np.abs(dense[1][1][:, fluid] - legacy[1][1][:, fluid]).max()
print(f"G2 {mode}: jet vs legacy fine max |dG| = {diff:.3e}", flush=True)
self.assertLessEqual(diff, 1e-6, mode)
def test_g3_amplitude_switch_layouts_identical(self):
for mode in ("periodic", "slot_reservoir"):
dense = run(mode, None, objects=[circle()], jets=[jet()], switch=-SPEED)
steady = run(mode, None, objects=[circle()], jets=[jet()])
self.assertFalse(np.array_equal(dense[1][1], steady[1][1], equal_nan=True), mode)
for fine_block in FINE_BLOCKS:
public, raw_arrays = run(mode, fine_block, objects=[circle()], jets=[jet()],
switch=-SPEED)
for a, b in zip(public + tuple(raw_arrays), dense[0] + tuple(dense[1])):
np.testing.assert_array_equal(a, b, err_msg=f"G3 {mode} {fine_block}")
print(f"G3 {mode}: amplitude switch dense vs {len(FINE_BLOCKS)} layouts identical",
flush=True)
if __name__ == "__main__":
unittest.main()