use crate::{Error, Position3D, Result};
use scirs2_core::ndarray::{Array1, Array2, Axis};
use scirs2_core::Complex32;
use serde::{Deserialize, Serialize};
use std::f32::consts::PI;
pub type AmbisonicsOrder = u32;
pub fn channel_count(order: AmbisonicsOrder) -> usize {
((order + 1) * (order + 1)) as usize
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum NormalizationScheme {
N3D,
SN3D,
FuMa,
MaxN,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ChannelOrdering {
ACN,
FuMa,
SID,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct SphericalCoordinate {
pub azimuth: f32,
pub elevation: f32,
pub distance: f32,
}
impl SphericalCoordinate {
pub fn new(azimuth: f32, elevation: f32, distance: f32) -> Self {
Self {
azimuth,
elevation,
distance,
}
}
pub fn from_cartesian(pos: &Position3D) -> Self {
let distance = (pos.x * pos.x + pos.y * pos.y + pos.z * pos.z).sqrt();
let azimuth = pos.y.atan2(pos.x);
let elevation = (pos.z / distance).asin();
Self {
azimuth,
elevation,
distance,
}
}
pub fn to_cartesian(&self) -> Position3D {
let x = self.distance * self.elevation.cos() * self.azimuth.cos();
let y = self.distance * self.elevation.cos() * self.azimuth.sin();
let z = self.distance * self.elevation.sin();
Position3D::new(x, y, z)
}
}
pub struct SphericalHarmonics {
order: AmbisonicsOrder,
normalization: NormalizationScheme,
}
impl SphericalHarmonics {
pub fn new(order: AmbisonicsOrder, normalization: NormalizationScheme) -> Self {
Self {
order,
normalization,
}
}
pub fn calculate(&self, coord: &SphericalCoordinate) -> Array1<f32> {
let channel_count = channel_count(self.order);
let mut coefficients = Array1::zeros(channel_count);
let cos_el = coord.elevation.cos();
let sin_el = coord.elevation.sin();
let mut idx = 0;
for l in 0..=self.order {
for m in -(l as i32)..=(l as i32) {
let coeff = self.spherical_harmonic_yn(l, m, coord.azimuth, coord.elevation);
coefficients[idx] = self.apply_normalization(coeff, l, m);
idx += 1;
}
}
coefficients
}
fn spherical_harmonic_yn(&self, l: u32, m: i32, azimuth: f32, elevation: f32) -> f32 {
let cos_el = elevation.cos();
let sin_el = elevation.sin();
let legendre = self.associated_legendre_polynomial(l, m.unsigned_abs(), sin_el);
let azimuthal = if m >= 0 {
(m as f32 * azimuth).cos()
} else {
((-m) as f32 * azimuth).sin()
};
legendre * azimuthal
}
fn associated_legendre_polynomial(&self, l: u32, m: u32, x: f32) -> f32 {
if m > l {
return 0.0;
}
let mut result = 1.0;
let one_minus_x2 = 1.0 - x * x;
if m > 0 {
result *= one_minus_x2.powf(m as f32 / 2.0);
for i in 1..=m {
result *= (2 * i - 1) as f32;
}
}
if l > m {
let mut p_prev = result;
let mut p_curr = x * (2 * m + 1) as f32 * p_prev;
for n in (m + 2)..=l {
let p_next = ((2 * n - 1) as f32 * x * p_curr - (n + m - 1) as f32 * p_prev)
/ (n - m) as f32;
p_prev = p_curr;
p_curr = p_next;
}
result = p_curr;
}
result
}
fn apply_normalization(&self, coeff: f32, l: u32, m: i32) -> f32 {
match self.normalization {
NormalizationScheme::N3D => coeff * (2 * l + 1) as f32 / (4.0 * PI),
NormalizationScheme::SN3D => {
let norm_factor = if m == 0 { 1.0 } else { 2.0_f32.sqrt() };
coeff * norm_factor
}
NormalizationScheme::FuMa => {
if l == 0 {
coeff / 2.0_f32.sqrt()
} else {
coeff
}
}
NormalizationScheme::MaxN => coeff / (4.0 * PI).sqrt(),
}
}
}
pub struct AmbisonicsEncoder {
order: AmbisonicsOrder,
normalization: NormalizationScheme,
channel_ordering: ChannelOrdering,
spherical_harmonics: SphericalHarmonics,
}
impl AmbisonicsEncoder {
pub fn new(
order: AmbisonicsOrder,
normalization: NormalizationScheme,
channel_ordering: ChannelOrdering,
) -> Self {
let spherical_harmonics = SphericalHarmonics::new(order, normalization);
Self {
order,
normalization,
channel_ordering,
spherical_harmonics,
}
}
pub fn encode_mono(
&self,
audio_samples: &Array1<f32>,
position: &Position3D,
) -> Result<Array2<f32>> {
let coord = SphericalCoordinate::from_cartesian(position);
let coefficients = self.spherical_harmonics.calculate(&coord);
let channel_count = channel_count(self.order);
let sample_count = audio_samples.len();
let mut encoded = Array2::zeros((channel_count, sample_count));
let distance_gain = if coord.distance > 0.0 {
1.0 / coord.distance.max(0.1)
} else {
1.0
};
for (ch_idx, &coeff) in coefficients.iter().enumerate() {
let channel_idx = self.apply_channel_ordering(ch_idx)?;
for (sample_idx, &sample) in audio_samples.iter().enumerate() {
encoded[[channel_idx, sample_idx]] = sample * coeff * distance_gain;
}
}
Ok(encoded)
}
pub fn encode_multichannel(
&self,
audio_samples: &Array2<f32>, positions: &[Position3D],
) -> Result<Array2<f32>> {
if audio_samples.shape()[0] != positions.len() {
return Err(Error::LegacyProcessing(
"Number of audio channels must match number of positions".to_string(),
));
}
let channel_count = channel_count(self.order);
let sample_count = audio_samples.shape()[1];
let mut encoded = Array2::zeros((channel_count, sample_count));
for (input_ch, position) in positions.iter().enumerate() {
let input_samples = audio_samples.row(input_ch);
let input_array = Array1::from_iter(input_samples.iter().copied());
let channel_encoded = self.encode_mono(&input_array, position)?;
encoded = encoded + channel_encoded;
}
Ok(encoded)
}
fn apply_channel_ordering(&self, acn_index: usize) -> Result<usize> {
match self.channel_ordering {
ChannelOrdering::ACN => Ok(acn_index),
ChannelOrdering::FuMa => {
match acn_index {
0 => Ok(0), 1 => Ok(2), 2 => Ok(3), 3 => Ok(1), _ => Ok(acn_index), }
}
ChannelOrdering::SID => Ok(acn_index), }
}
pub fn get_config(&self) -> (AmbisonicsOrder, NormalizationScheme, ChannelOrdering) {
(self.order, self.normalization, self.channel_ordering)
}
pub fn get_channel_count(&self) -> usize {
channel_count(self.order)
}
}
pub struct AmbisonicsDecoder {
order: AmbisonicsOrder,
normalization: NormalizationScheme,
channel_ordering: ChannelOrdering,
speaker_positions: Vec<SphericalCoordinate>,
decoding_matrix: Array2<f32>,
}
impl AmbisonicsDecoder {
pub fn new(
order: AmbisonicsOrder,
normalization: NormalizationScheme,
channel_ordering: ChannelOrdering,
speaker_positions: Vec<SphericalCoordinate>,
) -> Result<Self> {
let mut decoder = Self {
order,
normalization,
channel_ordering,
speaker_positions,
decoding_matrix: Array2::zeros((1, 1)), };
decoder.calculate_decoding_matrix()?;
Ok(decoder)
}
pub fn for_speaker_config(
order: AmbisonicsOrder,
config: SpeakerConfiguration,
) -> Result<Self> {
let positions = Self::create_speaker_positions(config);
Self::new(
order,
NormalizationScheme::SN3D,
ChannelOrdering::ACN,
positions,
)
}
fn create_speaker_positions(config: SpeakerConfiguration) -> Vec<SphericalCoordinate> {
match config {
SpeakerConfiguration::Stereo => vec![
SphericalCoordinate::new(-PI / 6.0, 0.0, 1.0), SphericalCoordinate::new(PI / 6.0, 0.0, 1.0), ],
SpeakerConfiguration::Quadraphonic => vec![
SphericalCoordinate::new(-PI / 4.0, 0.0, 1.0), SphericalCoordinate::new(PI / 4.0, 0.0, 1.0), SphericalCoordinate::new(-3.0 * PI / 4.0, 0.0, 1.0), SphericalCoordinate::new(3.0 * PI / 4.0, 0.0, 1.0), ],
SpeakerConfiguration::FiveDotOne => vec![
SphericalCoordinate::new(-PI / 6.0, 0.0, 1.0), SphericalCoordinate::new(PI / 6.0, 0.0, 1.0), SphericalCoordinate::new(0.0, 0.0, 1.0), SphericalCoordinate::new(0.0, -PI / 4.0, 1.0), SphericalCoordinate::new(-2.0 * PI / 3.0, 0.0, 1.0), SphericalCoordinate::new(2.0 * PI / 3.0, 0.0, 1.0), ],
SpeakerConfiguration::SevenDotOne => vec![
SphericalCoordinate::new(-PI / 6.0, 0.0, 1.0), SphericalCoordinate::new(PI / 6.0, 0.0, 1.0), SphericalCoordinate::new(0.0, 0.0, 1.0), SphericalCoordinate::new(0.0, -PI / 4.0, 1.0), SphericalCoordinate::new(-PI / 2.0, 0.0, 1.0), SphericalCoordinate::new(PI / 2.0, 0.0, 1.0), SphericalCoordinate::new(-3.0 * PI / 4.0, 0.0, 1.0), SphericalCoordinate::new(3.0 * PI / 4.0, 0.0, 1.0), ],
SpeakerConfiguration::Cube => {
let mut positions = Vec::new();
for &elevation in &[-PI / 4.0, PI / 4.0] {
for i in 0..4 {
let azimuth = i as f32 * PI / 2.0;
positions.push(SphericalCoordinate::new(azimuth, elevation, 1.0));
}
}
positions
}
}
}
fn calculate_decoding_matrix(&mut self) -> Result<()> {
let channel_count = channel_count(self.order);
let speaker_count = self.speaker_positions.len();
let spherical_harmonics = SphericalHarmonics::new(self.order, self.normalization);
let mut matrix = Array2::zeros((speaker_count, channel_count));
for (speaker_idx, position) in self.speaker_positions.iter().enumerate() {
let coefficients = spherical_harmonics.calculate(position);
for (ch_idx, &coeff) in coefficients.iter().enumerate() {
matrix[[speaker_idx, ch_idx]] = coeff;
}
}
self.decoding_matrix = self.pseudo_inverse(&matrix)?;
Ok(())
}
fn pseudo_inverse(&self, matrix: &Array2<f32>) -> Result<Array2<f32>> {
Ok(matrix.clone())
}
pub fn decode(&self, ambisonics_audio: &Array2<f32>) -> Result<Array2<f32>> {
let ambi_channels = ambisonics_audio.shape()[0];
let sample_count = ambisonics_audio.shape()[1];
let speaker_count = self.speaker_positions.len();
if ambi_channels != channel_count(self.order) {
return Err(Error::LegacyProcessing(format!(
"Expected {} ambisonics channels, got {}",
channel_count(self.order),
ambi_channels
)));
}
let mut decoded = Array2::zeros((speaker_count, sample_count));
for sample_idx in 0..sample_count {
let ambi_sample = ambisonics_audio.column(sample_idx);
for speaker_idx in 0..speaker_count {
let mut sum = 0.0;
for ch_idx in 0..ambi_channels {
sum += self.decoding_matrix[[speaker_idx, ch_idx]] * ambi_sample[ch_idx];
}
decoded[[speaker_idx, sample_idx]] = sum;
}
}
Ok(decoded)
}
pub fn get_config(&self) -> (AmbisonicsOrder, NormalizationScheme, ChannelOrdering) {
(self.order, self.normalization, self.channel_ordering)
}
pub fn get_speaker_count(&self) -> usize {
self.speaker_positions.len()
}
pub fn get_speaker_positions(&self) -> &[SphericalCoordinate] {
&self.speaker_positions
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum SpeakerConfiguration {
Stereo,
Quadraphonic,
FiveDotOne,
SevenDotOne,
Cube,
}
pub struct AmbisonicsProcessor {
encoder: AmbisonicsEncoder,
decoder: AmbisonicsDecoder,
}
impl AmbisonicsProcessor {
pub fn new(order: AmbisonicsOrder, speaker_config: SpeakerConfiguration) -> Result<Self> {
let encoder =
AmbisonicsEncoder::new(order, NormalizationScheme::SN3D, ChannelOrdering::ACN);
let decoder = AmbisonicsDecoder::for_speaker_config(order, speaker_config)?;
Ok(Self { encoder, decoder })
}
pub fn process_multichannel(
&self,
audio_samples: &Array2<f32>,
positions: &[Position3D],
) -> Result<Array2<f32>> {
let ambisonics = self.encoder.encode_multichannel(audio_samples, positions)?;
self.decoder.decode(&ambisonics)
}
pub fn process_mono(
&self,
audio_samples: &Array1<f32>,
position: &Position3D,
) -> Result<Array2<f32>> {
let ambisonics = self.encoder.encode_mono(audio_samples, position)?;
self.decoder.decode(&ambisonics)
}
pub fn get_info(&self) -> (AmbisonicsOrder, usize, usize) {
(
self.encoder.order,
self.encoder.get_channel_count(),
self.decoder.get_speaker_count(),
)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_channel_count_calculation() {
assert_eq!(channel_count(0), 1); assert_eq!(channel_count(1), 4); assert_eq!(channel_count(2), 9); assert_eq!(channel_count(3), 16); }
#[test]
fn test_spherical_coordinate_conversion() {
let cartesian = Position3D::new(1.0, 0.0, 0.0);
let spherical = SphericalCoordinate::from_cartesian(&cartesian);
assert!((spherical.azimuth - 0.0).abs() < 1e-6);
assert!((spherical.elevation - 0.0).abs() < 1e-6);
assert!((spherical.distance - 1.0).abs() < 1e-6);
let back_to_cartesian = spherical.to_cartesian();
assert!((back_to_cartesian.x - 1.0).abs() < 1e-6);
assert!((back_to_cartesian.y - 0.0).abs() < 1e-6);
assert!((back_to_cartesian.z - 0.0).abs() < 1e-6);
}
#[test]
fn test_spherical_harmonics_calculation() {
let harmonics = SphericalHarmonics::new(1, NormalizationScheme::SN3D);
let coord = SphericalCoordinate::new(0.0, 0.0, 1.0);
let coefficients = harmonics.calculate(&coord);
assert_eq!(coefficients.len(), 4);
assert!(coefficients[0].abs() > 1e-6);
}
#[test]
fn test_encoder_creation() {
let encoder = AmbisonicsEncoder::new(1, NormalizationScheme::SN3D, ChannelOrdering::ACN);
assert_eq!(encoder.get_channel_count(), 4);
let (order, norm, ordering) = encoder.get_config();
assert_eq!(order, 1);
assert_eq!(norm, NormalizationScheme::SN3D);
assert_eq!(ordering, ChannelOrdering::ACN);
}
#[test]
fn test_mono_encoding() {
let encoder = AmbisonicsEncoder::new(1, NormalizationScheme::SN3D, ChannelOrdering::ACN);
let position = Position3D::new(1.0, 0.0, 0.0); let audio = Array1::from_vec(vec![1.0, 0.5, -0.5, 0.0]);
let encoded = encoder.encode_mono(&audio, &position).unwrap();
assert_eq!(encoded.shape(), [4, 4]);
let sum: f32 = encoded.iter().map(|x| x.abs()).sum();
assert!(sum > 1e-6);
}
#[test]
fn test_decoder_creation() {
let decoder = AmbisonicsDecoder::for_speaker_config(1, SpeakerConfiguration::Stereo);
assert!(decoder.is_ok());
let decoder = decoder.unwrap();
assert_eq!(decoder.get_speaker_count(), 2);
let positions = decoder.get_speaker_positions();
assert_eq!(positions.len(), 2);
}
#[test]
fn test_speaker_configurations() {
let stereo_pos = AmbisonicsDecoder::create_speaker_positions(SpeakerConfiguration::Stereo);
assert_eq!(stereo_pos.len(), 2);
let quad_pos =
AmbisonicsDecoder::create_speaker_positions(SpeakerConfiguration::Quadraphonic);
assert_eq!(quad_pos.len(), 4);
let surround_pos =
AmbisonicsDecoder::create_speaker_positions(SpeakerConfiguration::FiveDotOne);
assert_eq!(surround_pos.len(), 6);
}
#[test]
fn test_combined_processor() {
let processor = AmbisonicsProcessor::new(1, SpeakerConfiguration::Stereo).unwrap();
let (order, ambi_channels, speaker_count) = processor.get_info();
assert_eq!(order, 1);
assert_eq!(ambi_channels, 4);
assert_eq!(speaker_count, 2);
}
#[test]
fn test_full_pipeline() {
let processor = AmbisonicsProcessor::new(1, SpeakerConfiguration::Stereo).unwrap();
let position = Position3D::new(1.0, 0.0, 0.0);
let audio = Array1::from_vec(vec![1.0, 0.5, -0.5, 0.0]);
let output = processor.process_mono(&audio, &position).unwrap();
assert_eq!(output.shape(), [2, 4]);
let sum: f32 = output.iter().map(|x| x.abs()).sum();
assert!(sum > 1e-6);
}
#[test]
fn test_multichannel_processing() {
let processor = AmbisonicsProcessor::new(1, SpeakerConfiguration::Quadraphonic).unwrap();
let positions = vec![
Position3D::new(1.0, 0.0, 0.0), Position3D::new(0.0, 1.0, 0.0), ];
let audio = Array2::from_shape_vec(
(2, 4),
vec![
1.0, 0.5, -0.5, 0.0, 0.8, 0.3, -0.2, 0.1, ],
)
.unwrap();
let output = processor.process_multichannel(&audio, &positions).unwrap();
assert_eq!(output.shape(), [4, 4]);
let sum: f32 = output.iter().map(|x| x.abs()).sum();
assert!(sum > 1e-6);
}
#[test]
fn test_normalization_schemes() {
let harmonics_n3d = SphericalHarmonics::new(1, NormalizationScheme::N3D);
let harmonics_sn3d = SphericalHarmonics::new(1, NormalizationScheme::SN3D);
let harmonics_fuma = SphericalHarmonics::new(1, NormalizationScheme::FuMa);
let coord = SphericalCoordinate::new(0.0, 0.0, 1.0);
let coeff_n3d = harmonics_n3d.calculate(&coord);
let coeff_sn3d = harmonics_sn3d.calculate(&coord);
let coeff_fuma = harmonics_fuma.calculate(&coord);
assert!(coeff_n3d[0] != coeff_sn3d[0]);
assert!(coeff_sn3d[0] != coeff_fuma[0]);
}
}