diff --git a/src/CelerisLab/lbm/refinement_stepper.py b/src/CelerisLab/lbm/refinement_stepper.py index a63be11..387f383 100644 --- a/src/CelerisLab/lbm/refinement_stepper.py +++ b/src/CelerisLab/lbm/refinement_stepper.py @@ -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 diff --git a/tests/validation/multigrid/test_refinement_jets.py b/tests/validation/multigrid/test_refinement_jets.py new file mode 100644 index 0000000..f28fda6 --- /dev/null +++ b/tests/validation/multigrid/test_refinement_jets.py @@ -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()