use microfft::{real::*, Complex32};
macro_rules! impl_ff {
( $name:ident, $n:tt, $fn:ident) => {
#[allow(dead_code)]
pub fn $name<const N: usize>(arr: &mut [f32; $n], result: &mut [Complex32; 1 + N / 2]) {
let out = $fn(arr);
result[0] = Complex32::new(out[0].re, 0.0);
(1..($n / 2)).for_each(|i| {
result[i] = out[i];
});
result[$n / 2] = Complex32::new(out[0].im, 0.0);
}
};
}
impl_ff!(fft_4, 4, rfft_4);
impl_ff!(fft_8, 8, rfft_8);
impl_ff!(fft_16, 16, rfft_16);
impl_ff!(fft_32, 32, rfft_32);
impl_ff!(fft_64, 64, rfft_64);
impl_ff!(fft_128, 128, rfft_128);
impl_ff!(fft_256, 256, rfft_256);
pub unsafe fn fft<const N: usize>(arr: &mut [f32], result: &mut [Complex32; 1 + N / 2]) {
match N {
4 => {
let arr = arr.as_ptr() as *mut [f32; 4];
fft_4(&mut unsafe { *arr }, result)
}
8 => {
let arr = arr.as_ptr() as *mut [f32; 8];
fft_8(&mut unsafe { *arr }, result)
}
16 => {
let arr = arr.as_ptr() as *mut [f32; 16];
fft_16(&mut unsafe { *arr }, result)
}
32 => {
let arr = arr.as_ptr() as *mut [f32; 32];
fft_32(&mut unsafe { *arr }, result)
}
64 => {
let arr = arr.as_ptr() as *mut [f32; 64];
fft_64(&mut unsafe { *arr }, result)
}
128 => {
let arr = arr.as_ptr() as *mut [f32; 128];
fft_128(&mut unsafe { *arr }, result)
}
256 => {
let arr = arr.as_ptr() as *mut [f32; 256];
fft_256(&mut unsafe { *arr }, result)
}
_ => unreachable!(),
}
}
#[cfg(test)]
mod test {
use super::*;
#[test]
#[allow(clippy::excessive_precision)]
fn test_fft_4() {
let mut input = [1.0f32, 2.19328005, 4.81047738, 10.55072407];
let expected_scipy = [
Complex32::new(18.55448151, -0.),
Complex32::new(-3.8104773, 8.35744402),
Complex32::new(-6.933527, -0.),
];
let mut out: [Complex32; 3] = Default::default();
fft_4::<4>(&mut input, &mut out);
assert_eq!(out, expected_scipy);
}
#[test]
#[allow(clippy::excessive_precision)]
fn test_fft_8() {
let mut input = [
1.0f32,
2.19328005,
4.81047738,
10.55072407,
23.14069263,
50.75401951,
111.31777849,
244.15106285,
];
let expected_scipy = [
Complex32::new(447.91803499, -0.),
Complex32::new(108.70207, 306.02533),
Complex32::new(-91.98756324, 201.75448737),
Complex32::new(-152.9834481, 93.0107107),
Complex32::new(-167.38013, -0.),
];
let mut out: [Complex32; 5] = Default::default();
fft_8::<8>(&mut input, &mut out);
assert_eq!(out, expected_scipy);
}
#[test]
#[allow(clippy::excessive_precision)]
fn test_fft_16() {
let mut input = [
0.65600173, 0.63951888, 0.62348176, 0.60798841, 0.59066761, 0.57568833, 0.56063477,
0.54556932, 0.53043256, 0.51324547, 0.4984534, 0.48411299, 0.46861073, 0.45568773,
0.4428301, 0.42829952,
];
let expected_scipy = [
Complex32::new(8.62122329, -0.),
Complex32::new(0.14047813, -0.61227777),
Complex32::new(0.12938741, -0.2879077),
Complex32::new(0.12032691, -0.18132925),
Complex32::new(0.12031269, -0.11817002),
Complex32::new(0.120595664, -0.082025215),
Complex32::new(0.12492439, -0.050967306),
Complex32::new(0.12087595, -0.02474621),
Complex32::new(0.1210022, -0.),
];
let mut out: [Complex32; 9] = Default::default();
fft_16::<16>(&mut input, &mut out);
assert_eq!(out, expected_scipy);
}
}