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: 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()