DFT Trait

src/fft/mixed_radix.rs

@@ -0,0 +1,73 @@
+// To perform a mixed radix cooley tuckey fft
+
+use crate::{complex::Complex32, fft::DFT};
+
+pub struct MixedRadixFFT<'a> {
+    input_buffer: &'a [Complex32],
+    output_buffer: &'a mut [Complex32],
+    size: usize,
+
+    p: usize,
+    q: usize,
+}
+
+impl<'a> DFT<'a> for MixedRadixFFT<'a> {
+    fn create(size: usize, input: &'a [Complex32], output: &'a mut [Complex32]) -> Self {
+        let q = decide_radix_factor(size);
+        let p = size / q;
+
+        MixedRadixFFT {
+            input_buffer: input,
+            output_buffer: output,
+            size,
+
+            p,
+            q,
+        }
+    }
+
+    fn execute(&mut self) {}
+}
+
+// This will decide on a good factor to use for the mixed radix fft
+fn decide_radix_factor(n: usize) -> usize {
+    let factors = prime_factors(n);
+    let two_count = factors.iter().take_while(|i| **i == 2).count();
+
+    // If there is a lot of two, we can use them as q factor to be able to use radix2 later on
+    if two_count > 0 {
+        return 1 << two_count;
+    }
+
+    // Otherwise take next big prime
+    return *factors.iter().skip(two_count).next().unwrap();
+}
+
+// Utilities
+fn prime_factors(n: usize) -> Vec<usize> {
+    let mut factors = vec![];
+
+    let mut num = n;
+    // Divide num successively
+    while num != 1 {
+        // Try divisors from 2 up to n (included)
+        for i in 2..n + 1 {
+            // if i divides num, it is a prime factor (if it wasn't, then i would have prime
+            // factors that would divide into num before i)
+            if num % i == 0 {
+                factors.push(i);
+                num /= i;
+            }
+        }
+    }
+
+    // If n = 1 then it does not have any prime factors
+    // The prime factor decomposition theorem states that any integer
+    // greater than TWO has a unique decomposition
+
+    factors
+}
+
+fn is_prime(n: usize) -> bool {
+    prime_factors(n).len() == 1
+}

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;
+}