tests/test_basis.py

"""Basis / Farina formula sanity checks."""

from __future__ import annotations

import math
import sys
from pathlib import Path

import numpy as np

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

from hoa64.basis import (
    MAX_ORDER,
    N_CHANNELS,
    acn_index,
    acn_nm,
    channel_names,
    sh_sn3d,
    unit_vector,
)


def test_layout():
    assert N_CHANNELS == 64
    assert MAX_ORDER == 7
    assert acn_index(0, 0) == 0
    assert acn_index(1, -1) == 1
    assert acn_index(1, 0) == 2
    assert acn_index(1, 1) == 3
    assert acn_index(7, 7) == 63
    for i in range(64):
        n, m = acn_nm(i)
        assert acn_index(n, m) == i
    names = channel_names()
    assert names[0] == "W"
    assert names[1] == "Y"
    assert names[2] == "Z"
    assert names[3] == "X"
    assert len(names) == 64


def test_unit_vector_cardinals():
    # front, left, up
    f = unit_vector(0.0, 0.0)
    l = unit_vector(90.0, 0.0)
    u = unit_vector(0.0, 90.0)
    np.testing.assert_allclose(f, [1, 0, 0], atol=1e-12)
    np.testing.assert_allclose(l, [0, 1, 0], atol=1e-12)
    np.testing.assert_allclose(u, [0, 0, 1], atol=1e-12)


def test_order1_cartesian_identity():
    """Order-1 SN3D: Y=y, Z=z, X=x at unit directions."""
    for az, el, xyz in [
        (0, 0, (1, 0, 0)),
        (90, 0, (0, 1, 0)),
        (0, 90, (0, 0, 1)),
        (-90, 0, (0, -1, 0)),
    ]:
        y = sh_sn3d(az, el)
        assert y.shape == (64,)
        assert abs(y[0] - 1.0) < 1e-12  # W
        x, yy, z = xyz
        assert abs(y[1] - yy) < 1e-12
        assert abs(y[2] - z) < 1e-12
        assert abs(y[3] - x) < 1e-12


def test_front_source_signs():
    y = sh_sn3d(0.0, 0.0)
    # front: X dominant positive among order-1 dipoles
    assert y[3] > 0.9
    assert abs(y[1]) < 1e-12
    assert abs(y[2]) < 1e-12


def test_finite_all_channels():
    y = sh_sn3d(33.0, -12.0)
    assert y.shape == (64,)
    assert np.all(np.isfinite(y))
    # batch
    yb = sh_sn3d([0, 90, 180], [0, 0, 0])
    assert yb.shape == (3, 64)


def test_zenith_only_z_orders():
    y = sh_sn3d(0.0, 90.0)
    # at north pole, only m=0 channels should be nonzero (real SH)
    for n in range(0, 8):
        for m in range(-n, n + 1):
            idx = acn_index(n, m)
            if m != 0:
                assert abs(y[idx]) < 1e-8, f"n={n} m={m} val={y[idx]}"


if __name__ == "__main__":
    for fn in [
        test_layout,
        test_unit_vector_cardinals,
        test_order1_cartesian_identity,
        test_front_source_signs,
        test_finite_all_channels,
        test_zenith_only_z_orders,
    ]:
        fn()
        print("OK", fn.__name__)