hoa64/basis.py

"""Ambix ACN + SN3D real spherical harmonics up to order 7.

Conventions (Angelo Farina / ISO-style Ambix):
  - Cartesian: +X front, +Y left, +Z up (ISO 2631-style).
  - Spherical: azimuth a in XY from +X toward +Y (0=front, +90=left);
               elevation e from horizontal (+90=zenith, -90=nadir).
  - Channel order: ACN  n*(n+1)+m  for m = -n .. +n
  - Normalization: SN3D (Ambix). Order-0 channel W = 1.

Evaluation uses associated Legendre functions (no Condon–Shortley phase)
with Schmidt/SN3D scaling, matching Farina's explicit Ambix formulas for
orders 0–5 and 7. Order 6 is generated by the same recurrence (Farina's
Cartesian transcription for n=6 had inconsistent norms).
"""

from __future__ import annotations

import math
from typing import Dict, Tuple

import numpy as np

MAX_ORDER = 7
N_CHANNELS = (MAX_ORDER + 1) ** 2  # 64


def acn_index(n: int, m: int) -> int:
    """Ambix ACN index for degree n and order m ∈ [-n, n]."""
    if n < 0 or abs(m) > n:
        raise ValueError(f"invalid (n,m)=({n},{m})")
    return n * (n + 1) + m


def acn_nm(acn: int) -> Tuple[int, int]:
    """Inverse of acn_index."""
    if acn < 0:
        raise ValueError(acn)
    n = int(math.floor(math.sqrt(acn)))
    m = acn - n * (n + 1)
    return n, m


_NAME_N3 = {
    (0, 0): "W",
    (1, -1): "Y",
    (1, 0): "Z",
    (1, 1): "X",
    (2, -2): "V",
    (2, -1): "T",
    (2, 0): "R",
    (2, 1): "S",
    (2, 2): "U",
    (3, -3): "Q",
    (3, -2): "O",
    (3, -1): "M",
    (3, 0): "K",
    (3, 1): "L",
    (3, 2): "N",
    (3, 3): "P",
}


def channel_names(max_order: int = MAX_ORDER) -> list[str]:
    names: list[str] = []
    for n in range(max_order + 1):
        for m in range(-n, n + 1):
            names.append(_NAME_N3.get((n, m), f"Y{n}_{m:+d}"))
    return names


def unit_vector(
    azimuth: float | np.ndarray,
    elevation: float | np.ndarray,
    degrees: bool = True,
) -> np.ndarray:
    """Direction → unit vector (x, y, z). Broadcasts over arrays. Shape (..., 3)."""
    a = np.asarray(azimuth, dtype=np.float64)
    e = np.asarray(elevation, dtype=np.float64)
    if degrees:
        a = np.deg2rad(a)
        e = np.deg2rad(e)
    ce = np.cos(e)
    x = np.cos(a) * ce
    y = np.sin(a) * ce
    z = np.sin(e)
    return np.stack([x, y, z], axis=-1)


def az_el_from_unit(vec: np.ndarray, degrees: bool = True) -> Tuple[np.ndarray, np.ndarray]:
    """Unit vector(s) → azimuth, elevation."""
    v = np.asarray(vec, dtype=np.float64)
    x, y, z = v[..., 0], v[..., 1], v[..., 2]
    z = np.clip(z, -1.0, 1.0)
    el = np.arcsin(z)
    az = np.arctan2(y, x)
    if degrees:
        return np.rad2deg(az), np.rad2deg(el)
    return az, el


def _associated_legendre_no_cs(n_max: int, z: float) -> Dict[Tuple[int, int], float]:
    """P_n^m(z) without Condon–Shortley phase, 0 ≤ m ≤ n ≤ n_max."""
    z = float(np.clip(z, -1.0, 1.0))
    st = math.sqrt(max(0.0, 1.0 - z * z))
    P: Dict[Tuple[int, int], float] = {(0, 0): 1.0}
    if n_max >= 1:
        P[(1, 0)] = z
        P[(1, 1)] = st
    for n in range(2, n_max + 1):
        for m in range(0, n + 1):
            if m == n:
                P[(n, n)] = (2 * n - 1) * st * P[(n - 1, n - 1)]
            elif m == n - 1:
                P[(n, n - 1)] = (2 * n - 1) * z * P[(n - 1, n - 1)]
            else:
                P[(n, m)] = (
                    (2 * n - 1) * z * P[(n - 1, m)] - (n + m - 1) * P[(n - 2, m)]
                ) / (n - m)
    return P


def _sn3d_one(x: float, y: float, z: float, max_order: int) -> np.ndarray:
    """SN3D ACN vector at one unit direction."""
    r = math.sqrt(x * x + y * y + z * z)
    if r == 0.0:
        r = 1.0
    x, y, z = x / r, y / r, z / r
    az = math.atan2(y, x)
    P = _associated_legendre_no_cs(max_order, z)
    nch = (max_order + 1) ** 2
    out = np.empty(nch, dtype=np.float64)
    k = 0
    for n in range(max_order + 1):
        for m in range(-n, n + 1):
            if m == 0:
                val = P[(n, 0)]
            else:
                am = abs(m)
                norm = math.sqrt(2.0 * math.factorial(n - am) / math.factorial(n + am))
                if m > 0:
                    val = norm * P[(n, am)] * math.cos(am * az)
                else:
                    val = norm * P[(n, am)] * math.sin(am * az)
            out[k] = val
            k += 1
    return out


def _eval_sn3d_cartesian(x: np.ndarray, y: np.ndarray, z: np.ndarray) -> np.ndarray:
    """Evaluate all 64 SN3D ACN channels at unit direction(s). Output (..., 64)."""
    x = np.asarray(x, dtype=np.float64)
    y = np.asarray(y, dtype=np.float64)
    z = np.asarray(z, dtype=np.float64)
    shape = np.broadcast(x, y, z).shape
    x, y, z = np.broadcast_arrays(x, y, z)
    flat = x.size
    out = np.empty((flat, N_CHANNELS), dtype=np.float64)
    xf, yf, zf = x.reshape(-1), y.reshape(-1), z.reshape(-1)
    for i in range(flat):
        out[i] = _sn3d_one(float(xf[i]), float(yf[i]), float(zf[i]), MAX_ORDER)
    return out.reshape(shape + (N_CHANNELS,))


def sh_sn3d(
    azimuth: float | np.ndarray,
    elevation: float | np.ndarray,
    degrees: bool = True,
    max_order: int = MAX_ORDER,
) -> np.ndarray:
    """Spherical harmonics Y (Ambix SN3D) at direction(s). Shape (..., n_channels)."""
    if max_order < 0 or max_order > MAX_ORDER:
        raise ValueError(f"max_order must be in 0..{MAX_ORDER}")
    uv = unit_vector(azimuth, elevation, degrees=degrees)
    y = _eval_sn3d_cartesian(uv[..., 0], uv[..., 1], uv[..., 2])
    nch = (max_order + 1) ** 2
    return y[..., :nch]


def sh_sn3d_batch(
    directions_xyz: np.ndarray,
    max_order: int = MAX_ORDER,
) -> np.ndarray:
    """Y at unit vectors. directions_xyz: (..., 3) → (..., n_channels)."""
    d = np.asarray(directions_xyz, dtype=np.float64)
    if d.shape[-1] != 3:
        raise ValueError("last dim must be 3")
    y = _eval_sn3d_cartesian(d[..., 0], d[..., 1], d[..., 2])
    nch = (max_order + 1) ** 2
    return y[..., :nch]


def sphere_grid(
    n_azi: int = 72,
    n_el: int = 36,
    degrees: bool = True,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Product azimuth × elevation grid with integration weights ≈ cos(e) de da."""
    if degrees:
        azi = np.linspace(-180.0, 180.0, n_azi, endpoint=False)
        el = np.linspace(-90.0, 90.0, n_el)
        el_r = np.deg2rad(el)
        da = 2.0 * math.pi / n_azi
        de = math.pi / max(n_el - 1, 1)
        w_el = np.cos(el_r) * de
        weights = np.outer(np.full(n_azi, da), np.maximum(w_el, 0.0))
    else:
        azi = np.linspace(-math.pi, math.pi, n_azi, endpoint=False)
        el = np.linspace(-0.5 * math.pi, 0.5 * math.pi, n_el)
        da = 2.0 * math.pi / n_azi
        de = math.pi / max(n_el - 1, 1)
        w_el = np.cos(el) * de
        weights = np.outer(np.full(n_azi, da), np.maximum(w_el, 0.0))
    return azi, el, weights