79 lines
2.1 KiB
Rust
79 lines
2.1 KiB
Rust
// Cooley-Tukey algorithm
|
|
|
|
use crate::complex::Complex32;
|
|
use crate::fft::{DFT, FFTDirection};
|
|
use std::f32::consts::PI;
|
|
|
|
pub struct Radix2FFT {
|
|
direction: FFTDirection,
|
|
size: usize,
|
|
length: usize,
|
|
}
|
|
|
|
impl DFT for Radix2FFT {
|
|
// Size as power of two
|
|
fn create(size: usize, direction: FFTDirection) -> Self {
|
|
if !is_power_of_two(size) {
|
|
panic!("Tried to create a Radix2 FFT with a non power of two size.");
|
|
}
|
|
|
|
Radix2FFT {
|
|
size: size.ilog2() as usize,
|
|
direction,
|
|
length: size,
|
|
}
|
|
}
|
|
|
|
fn execute(&mut self, input: &[Complex32], output: &mut [Complex32], window: fn(f32) -> f32) {
|
|
// Reorder samples
|
|
for (i, x) in output.iter_mut().enumerate() {
|
|
let k = reverse_bits(i, self.size as u32);
|
|
*x = input[k] * window(k as f32 / self.size as f32);
|
|
}
|
|
|
|
for step in 1..(self.size + 1) {
|
|
let pol_length = 2usize.pow(step as u32);
|
|
let mid_point = pol_length / 2;
|
|
for s in (0..(self.length / pol_length)).map(|i| i * pol_length) {
|
|
for i in 0..mid_point {
|
|
// Compute current polynomial at each unit root
|
|
let a = output[s + i];
|
|
let b = output[s + i + mid_point];
|
|
let angle = -2. * self.direction.sign() * PI * (i as f32) / (pol_length as f32);
|
|
let phasor = Complex32::cexp(angle);
|
|
output[i + s] = a + phasor * b;
|
|
output[i + s + mid_point] = a - phasor * b;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Utilities
|
|
pub fn reverse_bits(n: usize, bit_count: u32) -> usize {
|
|
let mut num = n;
|
|
let mut output = 0usize;
|
|
for _ in 0..bit_count {
|
|
output <<= 1;
|
|
output |= num & 1;
|
|
num >>= 1;
|
|
}
|
|
|
|
output
|
|
}
|
|
|
|
fn is_power_of_two(n: usize) -> bool {
|
|
if n == 0 {
|
|
return false;
|
|
}
|
|
|
|
let mut num = n;
|
|
while num != 1 {
|
|
if num % 2 != 0 {
|
|
return false;
|
|
}
|
|
num /= 2;
|
|
}
|
|
return true;
|
|
}
|