DFT Trait

src/fft/radix2.rs

@@ -0,0 +1,79 @@
+// Cooley-Tukey algorithm
+
+use crate::complex::Complex32;
+use crate::fft::DFT;
+use std::f32::consts::PI;
+
+pub struct Radix2FFT<'a> {
+    output_buffer: &'a mut [Complex32],
+    input_buffer: &'a [Complex32],
+    size: usize,
+    length: usize,
+}
+
+impl<'a> DFT<'a> for Radix2FFT<'a> {
+    // Size as power of two
+    fn create(size: usize, input: &'a [Complex32], output: &'a mut [Complex32]) -> Self {
+        if !is_power_of_two(size) {
+            panic!("Tried to create a Radix2 FFT with a non power of two size.");
+        }
+
+        Radix2FFT {
+            output_buffer: output,
+            input_buffer: input,
+            size: size.ilog2() as usize,
+            length: size,
+        }
+    }
+
+    fn execute(&mut self) {
+        // Reorder samples
+        for (i, x) in self.output_buffer.iter_mut().enumerate() {
+            *x = self.input_buffer[reverse_bits(i, self.size as u32)];
+        }
+
+        for step in 1..self.size {
+            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 = self.output_buffer[s + i];
+                    let b = self.output_buffer[s + i + mid_point];
+                    let angle = 2. * PI * (i as f32) / (pol_length as f32);
+                    let phasor = Complex32::cexp(angle);
+                    self.output_buffer[i + s] = a + phasor * b;
+                    self.output_buffer[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 |= if (num & 1) == 1 { 0 } else { 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;
+}