Add tests and make API breaking changes

- Add a test to make sure inverse twiddle factors multiplied by forward
  twiddle factors always gives 1.0

- `fft` now takes the real/imaginary input and the `Direction` of the `fft`

- Pre-built planners can be used with the `fft_with_opts_and_plan`
  function

Author: smu160

examples/benchmark.rs

@@ -1,19 +1,19 @@
 use std::env;
 use std::str::FromStr;
 
-use phastft::fft;
-use phastft::planner::{Direction, Planner};
 use utilities::gen_random_signal;
 
+use phastft::fft;
+use phastft::planner::Direction;
+
 fn benchmark_fft(n: usize) {
     let big_n = 1 << n;
     let mut reals = vec![0.0; big_n];
     let mut imags = vec![0.0; big_n];
     gen_random_signal(&mut reals, &mut imags);
 
     let now = std::time::Instant::now();
-    let mut planner = Planner::new(big_n, Direction::Forward);
-    fft(&mut reals, &mut imags, &mut planner);
+    fft(&mut reals, &mut imags, Direction::Forward);
     let elapsed = now.elapsed().as_micros();
     println!("{elapsed}");
 }

examples/profile.rs

@@ -2,14 +2,13 @@ use std::env;
 use std::str::FromStr;
 
 use phastft::fft;
-use phastft::planner::Planner;
+use phastft::planner::Direction;
 
 fn benchmark_fft(num_qubits: usize) {
     let n = 1 << num_qubits;
     let mut reals: Vec<f64> = (1..=n).map(|i| i as f64).collect();
     let mut imags: Vec<f64> = (1..=n).map(|i| i as f64).collect();
-    let mut planner = Planner::new(n);
-    fft(&mut reals, &mut imags, &mut planner);
+    fft(&mut reals, &mut imags, Direction::Forward);
 }
 
 fn main() {

pybindings/src/lib.rs

@@ -1,14 +1,20 @@
 use numpy::PyReadwriteArray1;
-use phastft::{fft as fft_rs, planner::Planner};
+use phastft::{fft as fft_rs, planner::Direction};
 use pyo3::prelude::*;
 
 #[pyfunction]
-fn fft(mut reals: PyReadwriteArray1<f64>, mut imags: PyReadwriteArray1<f64>) {
-    let mut planner = Planner::new(reals.len());
+fn fft(mut reals: PyReadwriteArray1<f64>, mut imags: PyReadwriteArray1<f64>, direction: char) {
+    assert!(direction == 'f' || direction == 'r');
+    let dir = if direction == 'f' {
+        Direction::Forward
+    } else {
+        Direction::Reverse
+    };
+
     fft_rs(
         reals.as_slice_mut().unwrap(),
         imags.as_slice_mut().unwrap(),
-        &mut planner,
+        dir,
     );
 }
 

src/lib.rs

@@ -3,7 +3,7 @@
 use crate::cobra::cobra_apply;
 use crate::kernels::{fft_chunk_2, fft_chunk_4, fft_chunk_n, fft_chunk_n_simd, Float};
 use crate::options::Options;
-use crate::planner::Planner;
+use crate::planner::{Direction, Planner};
 use crate::twiddles::filter_twiddles;
 
 mod cobra;
@@ -22,20 +22,36 @@ mod twiddles;
 /// Panics if `reals.len() != imags.len()`
 ///
 /// [1] https://inst.eecs.berkeley.edu/~ee123/sp15/Notes/Lecture08_FFT_and_SpectAnalysis.key.pdf
-pub fn fft(reals: &mut [Float], imags: &mut [Float], planner: &mut Planner) {
+pub fn fft(reals: &mut [Float], imags: &mut [Float], direction: Direction) {
+    assert_eq!(
+        reals.len(),
+        imags.len(),
+        "real and imaginary inputs must be of equal size, but got: {} {}",
+        reals.len(),
+        imags.len()
+    );
+
+    let mut planner = Planner::new(reals.len(), direction);
+    assert!(planner.num_twiddles().is_power_of_two() && planner.num_twiddles() == reals.len() / 2);
+
     let opts = Options::guess_options(reals.len());
-    fft_with_opts(reals, imags, &opts, planner);
+    fft_with_opts_and_plan(reals, imags, &opts, &mut planner);
 }
 
-/// Same as [fft], but also accepts [`Options`] that control optimization strategies.
+/// Same as [fft], but also accepts [`Options`] that control optimization strategies, as well as
+/// a [`Planner`] in the case that this FFT will need to be run multiple times.
 ///
-/// `fft_dif` automatically guesses the best strategy for a given input,
+/// `fft` automatically guesses the best strategy for a given input,
 /// so you only need to call this if you are tuning performance for a specific hardware platform.
 ///
+/// In addition, `fft` automatically creates a planner to be used. In the case that you plan
+/// on running an FFT many times on inputs of the same size, use this function with the pre-built
+/// [`Planner`].
+///
 /// # Panics
 ///
-/// Panics if `reals.len() != imags.len()`
-pub fn fft_with_opts(
+/// Panics if `reals.len() != imags.len()`, or if the input length is *not* a power of two.
+pub fn fft_with_opts_and_plan(
     reals: &mut [Float],
     imags: &mut [Float],
     opts: &Options,
@@ -87,6 +103,8 @@ mod tests {
         rustfft::{num_complex::Complex64, FftPlanner},
     };
 
+    use crate::planner::Direction;
+
     use super::*;
 
     #[should_panic]
@@ -95,11 +113,14 @@ mod tests {
         let num_points = 5;
 
         // this test will actually always fail at this stage
-        let mut planner = Planner::new(num_points);
+        let mut planner = Planner::new(num_points, Direction::Forward);
 
         let mut reals = vec![0.0; num_points];
         let mut imags = vec![0.0; num_points];
-        fft(&mut reals, &mut imags, &mut planner);
+        let opts = Options::guess_options(reals.len());
+
+        // but this call should panic as well
+        fft_with_opts_and_plan(&mut reals, &mut imags, &opts, &mut planner);
     }
 
     #[test]
@@ -111,8 +132,7 @@ mod tests {
 
             let mut reals: Vec<Float> = (1..=n).map(|i| i as f64).collect();
             let mut imags: Vec<Float> = (1..=n).map(|i| i as f64).collect();
-            let mut planner = Planner::new(n);
-            fft(&mut reals, &mut imags, &mut planner);
+            fft(&mut reals, &mut imags, Direction::Forward);
 
             let mut buffer: Vec<Complex64> = (1..=n)
                 .map(|i| Complex64::new(i as f64, i as f64))

src/planner.rs

@@ -10,21 +10,17 @@ pub struct Planner {
     pub twiddles_im: Vec<f64>,
 }
 
-// TODO(saveliy yusufov): Add a parameter to `new` that will take into consideration whether we do inverse FFT (IFFT)
-// In this case, the twiddle factors should be pre-computed as follows:
-//
-//                  FFT Twiddle Factor: e^{i2π*k/N}
-//                  IFFT Twiddle Factor: e^{-i2π*k/N}
-//
-// source: https://dsp.stackexchange.com/q/73367
 impl Planner {
-    /// Create a `Planner` for an FFT of size `num_points`
+    /// Create a `Planner` for an FFT of size `num_points`.
+    /// The twiddle factors are pre-computed based on the provided [`Direction`].
+    /// For `Forward`, use [`Direction::Forward`].
+    /// For `Reverse`, use [`Direction::Reverse`].
     ///
     /// # Panics
     ///
     /// Panics if `num_points` is less than 1
     pub fn new(num_points: usize, direction: Direction) -> Self {
-        assert!(num_points > 0);
+        assert!(num_points > 0 && num_points.is_power_of_two());
         if num_points <= 4 {
             return Self {
                 twiddles_re: vec![],
@@ -46,17 +42,53 @@ impl Planner {
             twiddles_im,
         }
     }
+
+    pub(crate) fn num_twiddles(&self) -> usize {
+        assert_eq!(self.twiddles_re.len(), self.twiddles_im.len());
+        self.twiddles_re.len()
+    }
 }
 
 #[cfg(test)]
 mod tests {
+    use utilities::assert_f64_closeness;
+
     use crate::planner::{Direction, Planner};
 
     #[test]
     fn no_twiddles() {
-        for num_points in 2..=4 {
+        for num_points in [2, 4] {
             let planner = Planner::new(num_points, Direction::Forward);
             assert!(planner.twiddles_im.is_empty() && planner.twiddles_re.is_empty());
         }
     }
+
+    #[test]
+    fn forward_mul_inverse_eq_identity() {
+        for i in 3..25 {
+            let num_points = 1 << i;
+            let planner_forward = Planner::new(num_points, Direction::Forward);
+            let planner_reverse = Planner::new(num_points, Direction::Reverse);
+
+            assert_eq!(
+                planner_reverse.num_twiddles(),
+                planner_forward.num_twiddles()
+            );
+
+            // (a + ib) (c + id) = ac + iad + ibc - bd
+            //                   = ac - bd + i(ad + bc)
+            planner_forward
+                .twiddles_re
+                .iter()
+                .zip(planner_forward.twiddles_im.iter())
+                .zip(planner_reverse.twiddles_re.iter())
+                .zip(planner_reverse.twiddles_im)
+                .for_each(|(((a, b), c), d)| {
+                    let temp_re = a * c - b * d;
+                    let temp_im = a * d + b * c;
+                    assert_f64_closeness(temp_re, 1.0, 1e-6);
+                    assert_f64_closeness(temp_im, 0.0, 1e-6);
+                });
+        }
+    }
 }

src/twiddles.rs

@@ -219,8 +219,8 @@ mod tests {
         for n in 4..28 {
             let dist = 1 << n;
 
-            let (twiddles_re_ref, twiddles_im_ref) = generate_twiddles(dist);
-            let (twiddles_re, twiddles_im) = generate_twiddles_simd(dist);
+            let (twiddles_re_ref, twiddles_im_ref) = generate_twiddles(dist, Direction::Forward);
+            let (twiddles_re, twiddles_im) = generate_twiddles_simd(dist, Direction::Forward);
 
             twiddles_re
                 .iter()
@@ -245,7 +245,7 @@ mod tests {
         let dist = 1 << (n - 1);
         let mut twiddles_iter = Twiddles::new(dist);
 
-        let (mut twiddles_re, mut twiddles_im) = generate_twiddles(dist);
+        let (mut twiddles_re, mut twiddles_im) = generate_twiddles(dist, Direction::Forward);
 
         for i in 0..dist {
             let (tw_re, tw_im) = twiddles_iter.next().unwrap();