1use std::f64::consts::TAU;
4
5pub fn fft(re: &mut [f64], im: &mut [f64]) {
7 let n = re.len();
8 assert_eq!(n, im.len(), "one real and one imaginary plane");
9 assert!(n.is_power_of_two(), "radix-2 needs a power-of-two length");
10 let mut j = 0usize;
11 for i in 1..n {
12 let mut bit = n >> 1;
13 while j & bit != 0 {
14 j ^= bit;
15 bit >>= 1;
16 }
17 j |= bit;
18 if i < j {
19 re.swap(i, j);
20 im.swap(i, j);
21 }
22 }
23
24 let twiddle: Vec<(f64, f64)> = (0..n / 2)
25 .map(|k| {
26 let angle = -TAU * k as f64 / n as f64;
27 (angle.cos(), angle.sin())
28 })
29 .collect();
30
31 let mut len = 2;
32 while len <= n {
33 let stride = n / len;
34 for start in (0..n).step_by(len) {
35 for k in 0..len / 2 {
36 let (wr, wi) = twiddle[k * stride];
37 let (i0, i1) = (start + k, start + k + len / 2);
38 let tr = re[i1] * wr - im[i1] * wi;
39 let ti = re[i1] * wi + im[i1] * wr;
40 re[i1] = re[i0] - tr;
41 im[i1] = im[i0] - ti;
42 re[i0] += tr;
43 im[i0] += ti;
44 }
45 }
46 len <<= 1;
47 }
48}
49
50pub fn ifft(re: &mut [f64], im: &mut [f64]) {
51 for x in im.iter_mut() {
52 *x = -*x;
53 }
54 fft(re, im);
55 let scale = 1.0 / re.len() as f64;
56 for x in re.iter_mut() {
57 *x *= scale;
58 }
59 for x in im.iter_mut() {
60 *x *= -scale;
61 }
62}
63
64pub fn irfft(bins_re: &[f64], bins_im: &[f64], n: usize) -> Vec<f64> {
66 assert_eq!(bins_re.len(), n / 2 + 1);
67 assert_eq!(bins_im.len(), n / 2 + 1);
68 let mut re = vec![0.0; n];
69 let mut im = vec![0.0; n];
70 for k in 0..=n / 2 {
71 re[k] = bins_re[k];
72 im[k] = bins_im[k];
73 if k > 0 && k < n - k {
74 re[n - k] = bins_re[k];
75 im[n - k] = -bins_im[k];
76 }
77 }
78 ifft(&mut re, &mut im);
79 re
80}