"""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: 1x1/2x2/4x2 fine blocks (blocked (12, 8) root) bit-identical to the FINE_BLOCKS[0] reference (values + NaN masks); G2 jet vs legacy actuation (blocked reference): 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: blocked layouts bit-identical to the FINE_BLOCKS[0] reference 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 from CelerisLab.lbm.descriptors import SOLID from CelerisLab.lbm.refinement import ACTIVE, FixedPatch2D, prepare_fine_boundary 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 root_and(fine_block): """Blocked levels with a (12, 8) root tile.""" return [(12, 8), fine_block] 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(cuda, "mem_get_info", return_value=(1 << 40, 1 << 40)), 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"), context=context, refinements=[BOUNDS], level_blocks=root_and(fine_block), level_objects=[(), tuple(objects)], level_jets=[(), tuple(jets)]) return backend, htod class JetStepperCPU(unittest.TestCase): def test_one_body_across_blocks_accepted(self): # Fine circle (19.5, 15.5) R4.3 crosses the 2x2 block seams; one body, one jet. import math from CelerisLab.body.jets import Jet, Radial, Sector C = (16.0, 12.0) whole = Jet(Sector(C, .1, math.pi, 3.), Radial(C, inward=False), amplitude=0.) backend, _ = mocked([circle()], [whole], FINE_BLOCKS[1]) _, jet_id = backend._boundaries[1].jet_cache self.assertGreater(len(np.unique(backend._boundaries[1].tables["block"][jet_id == 0])), 1) backend.close() def test_fine_ownership_uses_remapped_physical_owner(self): # Caller ids (99, 7) are remapped to fine owners (0, 1); block ids never count. import math from CelerisLab.body.geometry.circle import CircleGeometry from CelerisLab.body.jets import Jet, Parallel, Radial, Sector, Tube from CelerisLab.body.objects import SimObject C = (16.0, 12.0) source = SimObject(99, geometry=CircleGeometry(*C, 2.15), center=C, radius=2.15) whole = Jet(Sector(C, .1, math.pi, 3.), Radial(C, inward=False), amplitude=0.) backend, _ = mocked([source], [whole], FINE_BLOCKS[1]) _, jet_id = backend._boundaries[1].jet_cache owners = backend.boundaries[1].curved_host[3][jet_id == 0] self.assertEqual(np.unique(owners).tolist(), [0]) self.assertGreater(len(np.unique(backend._boundaries[1].tables["block"][jet_id == 0])), 1) self.assertEqual(source.obj_id, 99) backend.close() cfg = config("periodic") patch_ = FixedPatch2D.from_config(cfg, BOUNDS) A, B = (12.5, 12.0), (19.5, 12.0) pair = [SimObject(i, geometry=CircleGeometry(*c, 1.5), center=c, radius=1.5) for i, c in ((99, A), (7, B))] for amplitude in (0., .002): with self.assertRaisesRegex(ValueError, "one jet may drive only one body"): prepare_fine_boundary(cfg, patch_, pair, jets=[Jet(Tube((9., 12.), 0., 14., .5), Parallel(0.), amplitude=amplitude)]) self.assertEqual([o.obj_id for o in pair], [99, 7]) 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._boundaries[1].tables["surface_velocity"] b = driven._boundaries[1].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._boundaries[1].tables["surface_velocity"].copy() device = backend._boundaries[1].device["surface_velocity"] with patch.object(cuda, "memcpy_htod", htod): htod.reset_mock() backend.set_jet_amplitude(0, -SPEED, level=1) (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._boundaries[1].tables["surface_velocity"], sent) with self.assertRaises(ValueError): backend.set_jet_amplitude(0, 1.5, level=1) self.assertEqual(htod.call_count, 1) self.assertEqual(backend.jets[1][0].amplitude, -SPEED) backend.close() def test_copy_failure_marks_failed_and_keeps_host_amplitude(self): for fine_block in (FINE_BLOCKS[1],): # blocked tables backend, _ = mocked([circle()], [jet()], fine_block) status = backend._status for bad, error in ((1.5, ValueError), (float("nan"), ValueError)): with self.assertRaises(error): backend.set_jet_amplitude(0, bad, level=1) self.assertEqual(backend._status, status) # invalid input stays usable with self.assertRaises(IndexError): backend.set_jet_amplitude(1, -SPEED, level=1) self.assertEqual(backend._status, status) before = backend._boundaries[1].tables["surface_velocity"].copy() boom = Mock(side_effect=RuntimeError("copy failed")) with patch.object(cuda, "memcpy_htod", boom), patch.object(cuda, "memcpy_htod_async", boom), \ patch.object(cuda, "pagelocked_empty", side_effect=lambda shape, dtype: np.empty(shape, dtype)): with self.assertRaisesRegex(RuntimeError, "copy failed"): backend.set_jet_amplitude(0, -SPEED, level=1) self.assertEqual(backend._status, "failed", fine_block) self.assertEqual(backend.jets[1][0].amplitude, SPEED) np.testing.assert_array_equal(backend._boundaries[1].tables["surface_velocity"], before) with self.assertRaisesRegex(RuntimeError, "failed"): backend.step(1) with self.assertRaisesRegex(RuntimeError, "failed"): backend.set_jet_amplitude(0, -SPEED, level=1) 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"), context=Mock(), refinements=[BOUNDS], level_blocks=root_and(FINE_BLOCKS[1]), level_objects=[(), (circle(actuation(SPEED)),)], level_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, context=context, coarse_boundary=mode, refinements=[BOUNDS], level_blocks=root_and(fine_block), level_objects=[(), tuple(objects)], level_jets=[(), tuple(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, level=1) 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) reference = run(mode, FINE_BLOCKS[0], objects=[circle()], jets=[jet()]) self.assertTrue(np.isfinite(reference[1][1][:, fluid]).all(), mode) for fine_block in FINE_BLOCKS[1:]: public, raw_arrays = run(mode, fine_block, objects=[circle()], jets=[jet()]) for a, b in zip(public + tuple(raw_arrays), reference[0] + tuple(reference[1])): np.testing.assert_array_equal(a, b, err_msg=f"G1 {mode} {fine_block}") legacy = run(mode, FINE_BLOCKS[0], objects=[circle(actuation(SPEED))], jets=()) diff = np.abs(reference[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"): reference = run(mode, FINE_BLOCKS[0], objects=[circle()], jets=[jet()], switch=-SPEED) steady = run(mode, FINE_BLOCKS[0], objects=[circle()], jets=[jet()]) self.assertFalse(np.array_equal(reference[1][1], steady[1][1], equal_nan=True), mode) for fine_block in FINE_BLOCKS[1:]: public, raw_arrays = run(mode, fine_block, objects=[circle()], jets=[jet()], switch=-SPEED) for a, b in zip(public + tuple(raw_arrays), reference[0] + tuple(reference[1])): np.testing.assert_array_equal(a, b, err_msg=f"G3 {mode} {fine_block}") print(f"G3 {mode}: amplitude switch, {len(FINE_BLOCKS) - 1} layouts " f"identical to {FINE_BLOCKS[0]}", flush=True) if __name__ == "__main__": unittest.main()