CelerisLab/tests/validation/multigrid/test_refinement_jets.py

"""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._device_boundaries[1].jet_cache
        self.assertGreater(len(np.unique(backend._device_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._device_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._device_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._device_boundaries[1].tables["surface_velocity"]
        b = driven._device_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._device_boundaries[1].tables["surface_velocity"].copy()
        device = backend._device_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._device_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._device_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._device_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()