"""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