From dd17e436a0bb46fd389f58ae28ba758a07d32458 Mon Sep 17 00:00:00 2001 From: yhonag Date: Thu, 30 Jul 2026 17:34:18 +0300 Subject: [PATCH] add spectrally equalized asm --- CITATION.cff | 4 +- README.md | 33 ++++++- benchmarks/README.md | 2 + benchmarks/se_asm_convergence.py | 156 +++++++++++++++++++++++++++++++ examples/binaural_using_asm.py | 3 +- pyproject.toml | 2 +- src/shroom/encoders/asm.py | 89 +++++++++++++++++- tests/test_asm.py | 103 +++++++++++++++++++- 8 files changed, 385 insertions(+), 7 deletions(-) create mode 100644 benchmarks/se_asm_convergence.py diff --git a/CITATION.cff b/CITATION.cff index 3279d74..f940f20 100644 --- a/CITATION.cff +++ b/CITATION.cff @@ -11,8 +11,8 @@ authors: given-names: Yhonatan orcid: "https://orcid.org/0009-0009-1156-9087" affiliation: "Ben-Gurion University of the Negev" -version: 0.2.1 -date-released: "2026-07-12" +version: 0.2.2 +date-released: "2026-07-30" license: MIT repository-code: "https://github.com/Yhonatangayer/shroom" url: "https://github.com/Yhonatangayer/shroom" diff --git a/README.md b/README.md index ae9761e..bdf47f9 100644 --- a/README.md +++ b/README.md @@ -13,7 +13,7 @@ A Python library for simulating room acoustics using Spherical Harmonics (Ambiso * **Spatial Signals**: Unified handling of Time, Frequency, Space, and Spherical Harmonics (SH) domains. * **Processors**: Modular processing chain including: * `ArrayDecoder`: Simulates spherical microphone arrays. - * `ASMEncoder`: Encodes microphone signals to Ambisonics (ASM). + * `ASMEncoder`: Encodes microphone signals to Ambisonics (ASM, optionally spectrally-equalized — SE-ASM). * `BinauralDecoder`: Decodes Ambisonics to Binaural audio using HRTFs. * **Rotation**: Efficient rotation of sound fields and HRTFs using Wigner-D matrices, or via space domain grid rotation. * **Visualization**: 2D and 3D plotting of room geometry, sources, and receiver orientation. @@ -133,6 +133,17 @@ chain = ProcessorChain([ binaural_output = chain.process(room.compute_amb()) ``` +### Spectrally-Equalized ASM (SE-ASM) + +```python +from shroom import ASM + +# Rescales each SH channel so its linear spectral magnitude stays at 0 dB across +# the band, instead of collapsing above the array's spatial-aliasing frequency. +se_asm = ASM(sh_order=1, array=array, fs=fs, duration=duration, spectrally_equalized=True) +cnm = se_asm.cnm # (M, (N+1)^2, F) +``` + ### Optimized Low-Order Rendering (MagLS) ```python @@ -173,6 +184,26 @@ If you use shroom in your research, please cite our paper: ``` ## Changelog +### 0.2.2 + +**New: spectrally-equalized ASM (SE-ASM).** `ASM` gained a `spectrally_equalized` +flag (default `False`, so existing code is unchanged). When enabled, each SH channel +of the ASM solution is rescaled by `1 / xi[nm, f]`, where +`xi[nm, f] = ‖cnm[:, nm, f]^H V[:, :, f]‖ / ‖Y[:, nm]‖` is its linear spectral +magnitude. This keeps every channel at 0 dB across the whole band instead of letting +it collapse above the array's spatial-aliasing frequency, at the cost of a larger +complex MSE. The weights are real and positive, so the phase of the ASM filters is +untouched. + +```python +ASM(sh_order=1, array=array, fs=fs, duration=duration, spectrally_equalized=True) +``` + +Also added: `calculate_se_asm_coefficients` and `linear_spectral_magnitude` in +`shroom.encoders.asm`, and the `benchmarks/se_asm_convergence.py` benchmark comparing +ASM and SE-ASM (per-channel MSE/LSE and binaural magnitude error). No API changes to +existing functions. + ### 0.2.1 Maintenance release — no functional or API changes. Adds the JOSS paper diff --git a/benchmarks/README.md b/benchmarks/README.md index 957a500..5a10966 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -7,6 +7,7 @@ error decays, validating the accuracy of the encoders against high-order referen | Script | What it validates | |--------|-------------------| | `asm_convergence.py` | Ambisonics Signal Matching (ASM) encoder error vs. SH order. | +| `se_asm_convergence.py` | Spectrally-equalized ASM (SE-ASM) vs. plain ASM: per-channel MSE/LSE and binaural magnitude error. | | `bsm_convergence.py` | Binaural Signal Matching (BSM) encoder error vs. SH order (against a MATLAB reference). | | `aa_magls_convergence.py` | Array-aware MagLS binaural magnitude error vs. SH order. | @@ -27,6 +28,7 @@ benchmarks and the examples. ```bash python benchmarks/asm_convergence.py +python benchmarks/se_asm_convergence.py python benchmarks/bsm_convergence.py python benchmarks/aa_magls_convergence.py ``` diff --git a/benchmarks/se_asm_convergence.py b/benchmarks/se_asm_convergence.py new file mode 100644 index 0000000..4dca025 --- /dev/null +++ b/benchmarks/se_asm_convergence.py @@ -0,0 +1,156 @@ +"""Spectrally-equalized ASM (SE-ASM) vs. plain ASM. + +Validates the ``spectrally_equalized`` flag of :class:`shroom.encoders.asm.ASM`: +SE-ASM rescales every SH channel so that its linear spectral magnitude stays at +0 dB across the whole band, where plain ASM collapses towards zero above the +spatial-aliasing frequency of the array. The price is a larger complex MSE. + +Three figures are produced (per-channel complex MSE, per-channel LSE, and +binaural magnitude MSE with a MagLS HRTF), all written to ``benchmarks/figures/``. +""" +import os + +import numpy as np + +from shroom.geometry.sampling import sphereicalGrid +from shroom.paths import DEFAULT_HRTF_PATH +from shroom.utils.file_utils import load_file +from shroom.acoustics.spherical_array import SphericalArray +from shroom.acoustics.hrtf_processing import magls_hrtf +from shroom.utils.grid_utils import from_fibonacci_grid +from shroom.encoders.asm import ASM +from shroom_dev.errors import asm_bin_magnitude_mse_error, asm_mse_error, linear_spectral_error +from shroom_dev.plot import loglog_plot + +FIGURES_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "figures") + +SH_LABELS = ["(0,0)", "(1,-1)", "(1,0)", "(1,1)"] +SH_COLORS = ["#1f77b4", "#ff7f0e", "#2ca02c", "#d62728"] +ENCODER_STYLES = {"ASM": "-", "SE-ASM": "--"} + +SHOW = True + + +def _per_channel_curves(errors_by_encoder): + """Flatten {encoder: (nm, F)} into loglog_plot dicts keyed by label.""" + errors, styles, colors = {}, {}, {} + for encoder, err in errors_by_encoder.items(): + for i, nm in enumerate(SH_LABELS): + label = f"{encoder} {nm}" + errors[label] = err[i, ...] + styles[label] = ENCODER_STYLES[encoder] + colors[label] = SH_COLORS[i] + return errors, styles, colors + + +def main(): + os.makedirs(FIGURES_DIR, exist_ok=True) + + # 1. Setup + fs = 48000 + duration = 512 / 48000 + n_fft = int(duration * fs) + freqs = np.fft.fftfreq(n_fft, 1 / fs) + pos_freqs = np.fft.rfftfreq(n_fft, 1 / fs) + + hrtf = load_file(DEFAULT_HRTF_PATH) + hrtf.resample(fs) + hrtf.zero_pad(n_fft) + + source_grid = from_fibonacci_grid(240) + + hrtf.toFreq() + hrtf.toSH(30) + hrtf.toSpace(source_grid) + space_hrtf = hrtf.copy() + + az = np.deg2rad(np.array([-90, -45, 0, 45, 90])) + co = np.deg2rad(np.array([90, 90 + 18, 90 - 18, 90 + 18, 90])) + mic_grid = sphereicalGrid(az=az, co=co) + + array = SphericalArray( + fs=fs, + duration=duration, + r_sphere=0.08, + r_mics=0.08 * np.ones((mic_grid.n_points,)), + source_grid=source_grid, + mics_grid=mic_grid, + sphere_type="rigid", + sh_order_for_sm_calc=14, + convert_to_time=False, + ) + + Y = array.grid.Y(1) + + # 2. Encoders: plain ASM and its spectrally-equalized counterpart + asm = ASM(sh_order=1, array=array, fs=fs, duration=duration) + se_asm = ASM( + sh_order=1, + array=array, + fs=fs, + duration=duration, + spectrally_equalized=True, + ) + cnm = {"ASM": asm.cnm.data, "SE-ASM": se_asm.cnm.data} + + # 3. Errors + mse = {name: asm_mse_error(c, array.data, Y, freqs) for name, c in cnm.items()} + lse = {name: linear_spectral_error(c, array.data, Y, freqs) for name, c in cnm.items()} + + hrtf_magls = magls_hrtf(hrtf=space_hrtf.copy(), sh_order=1, cutoff_over_freq=1200) + hrtf_magls.toFreq() + bin_mse = { + name: asm_bin_magnitude_mse_error( + hrtf_magls.data, c, array.data, space_hrtf.data, freqs + ) + for name, c in cnm.items() + } + + # 4. Plots + errors, styles, colors = _per_channel_curves(mse) + loglog_plot( + freqs=pos_freqs, + title="ASM vs SE-ASM | Complex MSE per SH Channel", + errors=errors, + styles=styles, + colors=colors, + figsize=(7, 4), + ylim=(-30, 20), + save_path=os.path.join(FIGURES_DIR, "se_asm_mse.png"), + show=SHOW, + ) + + errors, styles, colors = _per_channel_curves(lse) + loglog_plot( + freqs=pos_freqs, + title="ASM vs SE-ASM | Linear Spectral Error per SH Channel", + errors=errors, + styles=styles, + colors=colors, + figsize=(7, 4), + ylim=(-30, 20), + save_path=os.path.join(FIGURES_DIR, "se_asm_lse.png"), + show=SHOW, + ) + + loglog_plot( + freqs=pos_freqs, + title="ASM vs SE-ASM | Binaural Magnitude MSE (MagLS HRTF)", + errors={ + f"{name} {ear}": err[i, :] + for name, err in bin_mse.items() + for i, ear in enumerate(["left", "right"]) + }, + styles={ + f"{name} {ear}": ENCODER_STYLES[name] + for name in bin_mse + for ear in ["left", "right"] + }, + figsize=(7, 4), + save_path=os.path.join(FIGURES_DIR, "se_asm_binaural_magnitude_mse.png"), + show=SHOW, + ) + + +if __name__ == "__main__": + main() diff --git a/examples/binaural_using_asm.py b/examples/binaural_using_asm.py index 6ee556f..b01c107 100644 --- a/examples/binaural_using_asm.py +++ b/examples/binaural_using_asm.py @@ -26,6 +26,7 @@ HEAD_POS = [2.0, 2.0, 1.5] SOURCE_POS = [4.0, 4.0, 1.5] HEAD_ROTATION = [0, 0, 0] +SPECTRALY_EQUALIZATION = True def main(): @@ -65,7 +66,7 @@ def main(): array_time_sh.toSH(1) # 3. Setup ASM - asm = ASM(sh_order=1, array=array, fs=FS, duration=DURATION) + asm = ASM(sh_order=1, array=array, fs=FS, duration=DURATION, spectrally_equalized=SPECTRALY_EQUALIZATION) # 4. Setup HRTF print("Loading HRTF...") diff --git a/pyproject.toml b/pyproject.toml index 5c5f33b..be501ee 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "pyshroom" -version = "0.2.1" +version = "0.2.2" description = "Spherical Harmonics Room (shroom): a Python library for room acoustics simulation, Ambisonics processing, and spherical microphone array modelling." readme = "README.md" license = "MIT" diff --git a/src/shroom/encoders/asm.py b/src/shroom/encoders/asm.py index 7dee7dc..7d0355b 100644 --- a/src/shroom/encoders/asm.py +++ b/src/shroom/encoders/asm.py @@ -3,6 +3,8 @@ import numpy as np from shroom.utils.dsp_utils import convolve_and_sum +_EPS = 1e-12 + def calculate_asm_coefficients( sm: np.ndarray, Y: np.ndarray @@ -50,10 +52,84 @@ def calculate_asm_coefficients( return C.transpose(1, 0, 2) +def linear_spectral_magnitude( + cnm: np.ndarray, sm: np.ndarray, Y: np.ndarray +) -> np.ndarray: + """ + Per-(SH-channel, frequency) linear spectral magnitude of an ASM solution: + + xi[nm, f] = ‖cnm[:, nm, f]^H V[:, :, f]‖_2 / ‖Y[:, nm]‖_2 + + This is the (unsquared) Linear Spectral Error used as the SE-ASM + equalization target: xi = 1 means the encoded SH channel carries the same + spectral energy as the ideal SH pattern. + + Parameters + ---------- + cnm : np.ndarray + Encoder coefficients with shape [M, (N_sh+1)**2, F]. + sm : np.ndarray + Steering matrix (V) with shape [M, Q, F]. + Y : np.ndarray + Spherical harmonic matrix with shape [Q, (N_sh+1)**2]. + + Returns + ------- + np.ndarray + Linear spectral magnitude with shape [(N_sh+1)**2, F], real-valued. + """ + proj = np.einsum("mlf,mqf->lqf", cnm.conj(), sm) # (L, Q, F) + y_norm = np.maximum(np.linalg.norm(Y, axis=0), _EPS) # (L,) + return np.linalg.norm(proj, axis=1) / y_norm[:, np.newaxis] + + +def calculate_se_asm_coefficients(sm: np.ndarray, Y: np.ndarray) -> np.ndarray: + """ + Calculate spectrally-equalized ASM (SE-ASM) coefficients. + + SE-ASM rescales the plain ASM (Tikhonov) solution by a real, per-(SH + channel, frequency) weight derived from the Linear Spectral Error of the + unweighted solution: + + w[nm, f] = 1 / xi[nm, f], + xi[nm, f] = ‖cnm[:, nm, f]^H V[:, :, f]‖_2 / ‖Y[:, nm]‖_2 + + The weight restores unit spectral magnitude in SH channels that plain ASM + attenuates at high frequencies (where the array is spatially aliased and + the regularized solution collapses towards zero), at the cost of a larger + complex MSE. The equalization is applied over the whole spectrum. Because + the weights are real and positive, they leave the phase of the ASM filters + — and hence the SH-domain symmetry of the encoded signal — untouched. + + Parameters + ---------- + sm : np.ndarray + Steering matrix (V) with shape [M, Q, F]. + Y : np.ndarray + Spherical harmonic matrix with shape [Q, (N_sh+1)**2]. + + Returns + ------- + np.ndarray + The SE-ASM filter weights with shape [M, (N_sh+1)**2, F]. + """ + C = calculate_asm_coefficients(sm, Y) # (M, L, F) + + xi = linear_spectral_magnitude(C, sm, Y) # (L, F) + # Bins with a zero solution (DC / Nyquist) carry no energy to equalize. + weights = np.where(xi > _EPS, 1.0 / np.maximum(xi, _EPS), 1.0) # (L, F) + + return C * weights[np.newaxis, :, :] + + class ASM: """ Ambisonics Signal Matching (ASM) encoder. Calculates filters to encode microphone array signals into Ambisonics. + + With ``spectrally_equalized=True`` the encoder instead produces + spectrally-equalized ASM (SE-ASM) filters, see + :func:`calculate_se_asm_coefficients`. """ def __init__( @@ -62,6 +138,7 @@ def __init__( array: SpatialSignal, fs: int = None, duration: float = None, + spectrally_equalized: bool = False, ): """ Initialize ASM encoder. @@ -76,12 +153,17 @@ def __init__( Sampling frequency. Must match array.fs if provided. duration : float, optional Duration of the filters. + spectrally_equalized : bool, optional + If True, equalize the linear spectral magnitude of every SH channel + over the whole spectrum (SE-ASM) instead of returning the plain ASM + solution. Default is False. """ self._validate_inputs(sh_order, array, fs, duration) self.sh_order = sh_order self.array = array self.fs = fs self.duration = duration + self.spectrally_equalized = spectrally_equalized self._cnm = None @@ -96,6 +178,8 @@ def calculate(self) -> SpatialSignal: """ Calculate the ASM coefficients (filters). + Returns SE-ASM coefficients when ``spectrally_equalized`` is True. + Returns ------- SpatialSignal @@ -103,7 +187,10 @@ def calculate(self) -> SpatialSignal: """ sm = self.array.data Y = self.array.grid.Y(N_sp=self.sh_order) - asm_coefficients = calculate_asm_coefficients(sm, Y) + if self.spectrally_equalized: + asm_coefficients = calculate_se_asm_coefficients(sm, Y) + else: + asm_coefficients = calculate_asm_coefficients(sm, Y) self._cnm = SpatialSignal( data=asm_coefficients, fs=self.fs, is_time=False, is_space=False ) diff --git a/tests/test_asm.py b/tests/test_asm.py index a61f41a..43d8971 100644 --- a/tests/test_asm.py +++ b/tests/test_asm.py @@ -1,6 +1,11 @@ import pytest import numpy as np -from shroom.encoders.asm import ASM, calculate_asm_coefficients +from shroom.encoders.asm import ( + ASM, + calculate_asm_coefficients, + calculate_se_asm_coefficients, + linear_spectral_magnitude, +) from shroom.acoustics.spatial_signal import SpatialSignal from shroom.utils.dsp_utils import is_signal_frequency_sh_valid @@ -96,6 +101,102 @@ def test_asm_nyquist_constraint(real_array_signal): assert np.allclose(cnm[:, 0, nyq].imag, 0.0), "ASM Nyquist: (0,0) channel is not real." +def test_se_asm_shape_and_domain(real_array_signal): + """SE-ASM returns the same kind of SpatialSignal as plain ASM.""" + sh_order = 1 + se_asm = ASM( + sh_order=sh_order, + array=real_array_signal, + fs=real_array_signal.fs, + duration=0.1, + spectrally_equalized=True, + ) + cnm_signal = se_asm.calculate() + + assert isinstance(cnm_signal, SpatialSignal) + assert cnm_signal.is_freq and cnm_signal.is_sh + + L = (sh_order + 1) ** 2 + M = real_array_signal.n_channels + F = real_array_signal.data.shape[2] + assert cnm_signal.data.shape == (M, L, F) + + +def test_se_asm_differs_from_asm(real_array_signal): + """The flag must actually change the filters (default stays plain ASM).""" + kwargs = dict(sh_order=1, array=real_array_signal, fs=real_array_signal.fs, duration=0.1) + cnm_asm = ASM(**kwargs).cnm.data + cnm_se = ASM(spectrally_equalized=True, **kwargs).cnm.data + + assert not np.allclose(cnm_asm, cnm_se) + # Default is plain ASM. + assert np.allclose(ASM(spectrally_equalized=False, **kwargs).cnm.data, cnm_asm) + + +def test_se_asm_equalizes_spectral_magnitude(real_array_signal): + """Every equalized SH channel carries unit linear spectral magnitude.""" + sm = real_array_signal.data + Y = real_array_signal.grid.Y(N_sp=1) + + cnm_se = calculate_se_asm_coefficients(sm, Y) + xi = linear_spectral_magnitude(cnm_se, sm, Y) # (L, F) + + F = sm.shape[2] + # DC and Nyquist are constrained to zero by construction and are excluded. + active = np.ones(F, dtype=bool) + active[0] = False + if F % 2 == 0: + active[F // 2] = False + + assert np.allclose(xi[:, active], 1.0, atol=1e-8) + + +def test_se_asm_preserves_asm_phase(real_array_signal): + """SE-ASM only rescales the ASM filters — the weights are real & positive.""" + sm = real_array_signal.data + Y = real_array_signal.grid.Y(N_sp=1) + + cnm_asm = calculate_asm_coefficients(sm, Y) + cnm_se = calculate_se_asm_coefficients(sm, Y) + + nonzero = np.abs(cnm_asm) > 1e-12 + ratio = cnm_se[nonzero] / cnm_asm[nonzero] + assert np.allclose(ratio.imag, 0.0) + assert np.all(ratio.real > 0.0) + + +def test_se_asm_dc_and_nyquist_constraints(real_array_signal): + """Equalization preserves the ASM DC/Nyquist constraints.""" + cnm = ASM( + sh_order=1, + array=real_array_signal, + fs=real_array_signal.fs, + duration=0.1, + spectrally_equalized=True, + ).cnm.data # (M, nm, F) + + assert np.allclose(cnm[:, 1:, 0], 0.0), "SE-ASM DC: higher-order channels are not zero." + assert np.allclose(cnm[:, 0, 0].imag, 0.0), "SE-ASM DC: (0,0) channel is not real." + + F = cnm.shape[-1] + if F % 2 == 0: + nyq = F // 2 + assert np.allclose(cnm[:, 1:, nyq], 0.0), "SE-ASM Nyquist: higher-order channels are not zero." + assert np.allclose(cnm[:, 0, nyq].imag, 0.0), "SE-ASM Nyquist: (0,0) channel is not real." + + +def test_se_asm_filters_are_real_in_time(real_array_signal): + """Real, even weights keep the filter spectrum conjugate-symmetric.""" + cnm = ASM( + sh_order=1, + array=real_array_signal, + fs=real_array_signal.fs, + duration=0.1, + spectrally_equalized=True, + ).cnm.data + assert is_signal_frequency_sh_valid(cnm, freq_axis=-1) + + def test_encode_amb(real_array_signal): """Test encoding microphone signals to Ambisonics.""" sh_order = 1