tests/test_encode_decode.py

"""Encode / decode / DOA hypothesis tests."""

from __future__ import annotations

import sys
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from hoa64.analysis import (
    angular_error_deg,
    doa_from_intensity,
    field_energy,
    peak_direction,
)
from hoa64.decode import beamform, decode_directions
from hoa64.encode import encode_plane_waves, encode_points, mix
from hoa64.basis import sh_sn3d


def test_encode_matches_basis():
    a = encode_points([45.0], [10.0], [1.0])
    y = sh_sn3d(45.0, 10.0)
    np.testing.assert_allclose(a, y, atol=1e-12)


def test_beamform_peaks_at_source():
    a = encode_points([30.0], [-15.0], [1.0])
    on = float(beamform(a, 30.0, -15.0))
    off = float(beamform(a, -150.0, 40.0))
    assert on > off
    assert on > 0.5  # SN3D self-inner-product roughly order-dependent


def test_intensity_doa_cardinal():
    for az, el in [(0, 0), (90, 0), (-90, 0), (0, 45)]:
        a = encode_points([az], [el], [1.0])
        az_h, el_h = doa_from_intensity(a)
        err = angular_error_deg(az, el, az_h, el_h)
        assert err < 2.0, f"src=({az},{el}) hat=({az_h},{el_h}) err={err}"


def test_peak_direction_near_source():
    a = encode_points([120.0], [25.0], [1.0])
    paz, pel, pval = peak_direction(a, n_azi=120, n_el=60)
    err = angular_error_deg(120.0, 25.0, paz, pel)
    assert err < 5.0, f"peak=({paz},{pel}) err={err}"
    assert pval > 0


def test_superposition_linearity():
    a1 = encode_points([0], [0], [1.0])
    a2 = encode_points([90], [0], [0.5])
    m = mix(a1, a2)
    np.testing.assert_allclose(m, a1 + a2)
    assert field_energy(m) > field_energy(a1)


def test_plane_wave_time_series():
    t = np.linspace(0, 1, 100, endpoint=False)
    sig = np.sin(2 * np.pi * 5 * t)[None, :]  # (1, T)
    a = encode_plane_waves([0.0], [0.0], sig)
    assert a.shape == (64, 100)
    # W channel tracks the signal
    np.testing.assert_allclose(a[0], sig[0], atol=1e-12)
    # X (front) also tracks for front source
    np.testing.assert_allclose(a[3], sig[0], atol=1e-12)


def test_two_source_separation_energy():
    a = mix(
        encode_points([0], [0], [1.0]),
        encode_points([180], [0], [1.0]),
    )
    front = float(beamform(a, 0, 0))
    back = float(beamform(a, 180, 0))
    side = float(beamform(a, 90, 0))
    assert front > side
    assert back > side


if __name__ == "__main__":
    for fn in [
        test_encode_matches_basis,
        test_beamform_peaks_at_source,
        test_intensity_doa_cardinal,
        test_peak_direction_near_source,
        test_superposition_linearity,
        test_plane_wave_time_series,
        test_two_source_separation_energy,
    ]:
        fn()
        print("OK", fn.__name__)