CelerisLab/tests/unit/test_jets.py

"""Body-independent Jet: true-hit spatial selection, fields, profiles, overlap, amplitude."""
import math
from types import SimpleNamespace

import numpy as np
import pytest

from CelerisLab.body.jets import (Jet, Parallel, Radial, Sector, Tube,
                                  fine_hits_to_physical, jet_link_shapes)
from CelerisLab.body.objects import SurfaceActuation
from CelerisLab.lbm.curved_links import CurvedLinkSoA, surface_link_velocity
from .test_state_regressions import device, make_field, manager
from .test_surface_actuation import circle, projected

CENTER, RADIUS = (6., 5.), 2.999


def scene(spec=None):
    """Single circle links plus their true FP64 hits, row-aligned with the SoA."""
    obj = circle(spec)
    cl, cfg, flags = projected(obj)
    links = obj.geometry.build_cut_links(cfg.nx, cfg.ny, domain_flags=flags)
    hits = np.array([(link.hit_x, link.hit_y) for link in links])
    assert len(hits) == len(cl.dir)
    return obj, cl, cfg, flags, hits


def arc(hits, angle, half_width):
    d = hits - CENTER
    delta = np.angle(np.exp(1j*(np.arctan2(d[:, 1], d[:, 0]) - angle)))
    return np.abs(delta) < half_width


def test_concentric_sector_gives_inward_radial_and_tube_parallel_on_same_links():
    *_, hits = scene()
    half = math.radians(10)
    expected = arc(hits, 0., half)
    assert expected.sum() >= 2
    sector = Jet(Sector(CENTER, 0., half, RADIUS + .1), Radial(CENTER), "uniform", 1e-3)
    shape, jet_id = jet_link_shapes([sector], hits)
    assert np.array_equal(jet_id >= 0, expected)
    d = hits[expected] - CENTER
    np.testing.assert_allclose(shape[expected], -d/np.hypot(*d.T)[:, None], atol=1e-15)
    # Tube along +x through the same arc: lateral |y-5| < R sin(10 deg).
    tube = Jet(Tube((CENTER[0], CENTER[1]), 0., RADIUS + .1, RADIUS*math.sin(half)),
               Parallel(math.pi), "uniform", 1e-3)
    shape_t, id_t = jet_link_shapes([tube], hits)
    assert np.array_equal(id_t >= 0, expected)
    np.testing.assert_allclose(shape_t[expected], [[-1., 0.]]*expected.sum(), atol=1e-15)


def test_off_centre_radial_points_to_field_origin_not_body_centre():
    *_, hits = scene()
    origin = (12., 9.)
    jet = Jet(Sector(CENTER, .1, math.pi, 10.), Radial(origin, inward=False))
    shape, jet_id = jet_link_shapes([jet], hits)
    assert np.all(jet_id == 0)
    d = hits - origin
    np.testing.assert_allclose(shape, d/np.hypot(*d.T)[:, None], atol=1e-15)
    assert not np.allclose(shape, (hits - CENTER)/RADIUS, atol=1e-3)


@pytest.mark.parametrize("profile,expected", [("uniform", 1.), ("parabolic", .75), ("quartic", .5625)])
def test_profiles_on_lateral_coordinate(profile, expected):
    hits = np.array([[1., .5]])
    jet = Jet(Tube((0., 0.), 0., 2., 1.), Parallel(0.), profile)
    shape, _ = jet_link_shapes([jet], hits)
    np.testing.assert_allclose(shape[0], [expected, 0.], rtol=1e-15)


def test_selection_uses_true_hit_not_clipped_q():
    # Same link as test_aperture_uses_true_hit_not_fluid_node_or_clamped_q:
    # receiver (9,5), direction (-1,+1); true t < .01 but packed q == .01.
    t = (6-math.sqrt(36-8*(9-2.999**2)))/4
    angle = math.atan2(t, 3-t)
    _, cl, _, _, hits = scene()
    row = np.flatnonzero((cl.fluid_idx == 5*12 + 9) & (cl.dir == 8))
    assert len(row) == 1 and np.isclose(cl.q[row[0]], .01)
    clipped = np.array([9., 5.]) + .01*np.array([-1., 1.])
    jet = Jet(Sector(CENTER, angle, 1e-4, RADIUS + .1), Radial(CENTER))
    _, jet_id = jet_link_shapes([jet], hits)
    assert jet_id[row[0]] == 0
    with pytest.raises(ValueError, match="selects no"):
        jet_link_shapes([jet], clipped[None, :])


def test_fine_half_cell_frame():
    np.testing.assert_allclose(fine_hits_to_physical([[0., 0.], [3., 1.5]], (4., 6.), .5),
                               [[4.25, 6.25], [5.75, 7.]], rtol=0, atol=0)


def test_links_are_distinct_by_receiver_and_direction_and_overlap_rejected():
    # Two physical links (different receiver/direction) sharing one hit point
    # are two rows, both driven; two jets on one link are rejected.
    hits = np.array([[1., 1.], [1., 1.]])
    a = Jet(Tube((0., 1.), 0., 2., .5), Parallel(0.))
    shape, jet_id = jet_link_shapes([a], hits)
    assert jet_id.tolist() == [0, 0]
    b = Jet(Sector((1., 0.), math.pi/2, .1, 2.), Parallel(0.))
    with pytest.raises(ValueError, match="overlaps"):
        jet_link_shapes([a, b], hits)
    obj, cl, cfg, flags, real = scene(SurfaceActuation(0., .3, .002))
    with pytest.raises(ValueError, match="SurfaceActuation"):
        surface_link_velocity([obj], cfg, flags, cl.fluid_idx, cl.dir, cl.body_id, cl.rx, cl.ry,
                              jets=[Jet(Sector(CENTER, 0., .1, 4.), Radial(CENTER))], hits=real)
    with pytest.raises(ValueError, match="row-aligned"):
        surface_link_velocity([], cfg, flags, cl.fluid_idx, cl.dir, cl.body_id, cl.rx, cl.ry,
                              jets=[a], hits=real[:-1])


def test_amplitude_update_reuses_cached_shapes():
    obj, cl, cfg, flags, hits = scene()
    jet = Jet(Sector(CENTER, 0., .5, 4.), Radial(CENTER), "parabolic", 1e-3)
    cl.assign_surface([obj], cfg, flags, jets=[jet], hits=hits)
    shape, ids = cl.jet_shape, cl.jet_id
    before = cl.surface_velocity.copy()
    cl.set_jet_amplitude(0, -2e-3)
    assert cl.jet_shape is shape and cl.jet_id is ids
    np.testing.assert_allclose(cl.surface_velocity, -2*before, rtol=1e-6, atol=0)
    assert np.array_equal(cl.aperture, ids >= 0)
    with pytest.raises(ValueError):
        cl.set_jet_amplitude(0, 1.5)
    assert cl.jets[0].amplitude == -2e-3 and jet.amplitude == 1e-3  # registered Jet untouched
    np.testing.assert_allclose(cl.surface_velocity, -2*before, rtol=1e-6, atol=0)


def test_jet_definition_is_immutable_and_amplitude_validated():
    jet = Jet(Sector(CENTER, 0., .5, 4.), Radial(CENTER), "parabolic", 1e-3)
    for name, value in (("field", Parallel(0.)), ("region", Tube(CENTER, 0., 1., 1.)),
                        ("profile", "uniform"), ("amplitude", float("nan"))):
        with pytest.raises(AttributeError):
            setattr(jet, name, value)
    for bad in (float("nan"), float("inf"), 1.0):
        with pytest.raises(ValueError):
            jet.with_amplitude(bad)
    assert jet.with_amplitude(-1e-3).field is jet.field and jet.amplitude == 1e-3
    with pytest.raises(ValueError, match="finite"):
        jet_link_shapes([jet], [[np.nan, 5.]])


@pytest.mark.parametrize("spec", [SurfaceActuation(0., .6, .002), SurfaceActuation(2., 1.1, -.003)])
def test_legacy_adapter_measured_equivalence(spec):
    obj, legacy, cfg, flags, hits = scene(spec)
    cl = CurvedLinkSoA()
    cl.assign_host(legacy.fluid_idx, legacy.dir, legacy.q, legacy.body_id,
                   legacy.rx, legacy.ry, legacy.rz, legacy.fallback_class)
    cl.assign_surface([circle()], cfg, flags, jets=[spec.as_jet(CENTER, RADIUS)], hits=hits)
    assert np.array_equal(cl.aperture, legacy.aperture)
    diff = np.abs(cl.surface_velocity - legacy.surface_velocity).max()
    print(f"legacy adapter max |dU| = {diff:.3e} (exact={diff == 0})")
    assert diff <= 4*np.finfo(np.float32).eps*abs(spec.normal_velocity)


def jet_case(device, jets):
    field = make_field(device)
    bodies = manager(field, [circle(radius=2.)])
    bodies.jets.extend(jets)
    bodies.sync_to_gpu(field)
    return field, bodies


def test_manager_passes_true_hits_and_reuploads_only_velocity(device):
    # CPU memory double only: checks host/device bookkeeping, not GPU numerics.
    jet = Jet(Sector(CENTER, 0., .6, 3.), Radial(CENTER, inward=False), "quartic", 2e-3)
    field, bodies = jet_case(device, [jet])
    cl = field.curved
    assert cl.jet_id is not None and (cl.jet_id == 0).sum() >= 2
    velocity_gpu = cl.surface_velocity_gpu
    np.testing.assert_array_equal(np.frombuffer(velocity_gpu.data, np.float32).reshape(-1, 2),
                                  cl.surface_velocity)
    geometry = {name: bytes(getattr(cl, name + "_gpu").data)
                for name in ("fluid_idx", "dir", "q", "body_id", "rx", "ry")}
    writes, before = device.writes, cl.surface_velocity.copy()
    cl.set_jet_amplitude(0, -1e-3, SimpleNamespace(synchronize=lambda: None))
    assert device.writes == writes + 1 and cl.surface_velocity_gpu is velocity_gpu
    np.testing.assert_array_equal(np.frombuffer(velocity_gpu.data, np.float32).reshape(-1, 2),
                                  cl.surface_velocity)
    assert all(bytes(getattr(cl, name + "_gpu").data) == data for name, data in geometry.items())
    np.testing.assert_allclose(cl.surface_velocity, -.5*before, rtol=1e-6, atol=0)
    np.testing.assert_array_equal(cl.surface_velocity[cl.jet_id < 0], 0.)


def test_manager_rejects_overlapping_jets_and_moving_bodies(device):
    a = Jet(Sector(CENTER, 0., .6, 3.), Radial(CENTER))
    b = Jet(Tube(CENTER, 0., 3., .5), Parallel(0.))
    with pytest.raises(ValueError, match="overlaps"):
        jet_case(device, [a, b])
    field = make_field(device)
    moving = circle(radius=2.)
    moving.state.omega = 1e-3
    bodies = manager(field, [moving])
    bodies.jets.append(a)
    with pytest.raises(ValueError, match="stationary"):
        bodies.sync_to_gpu(field)