Skip to main content

sva_samples/
fft.rs

1// Concern: the radix-2 complex transform, its inverse from half the bins | Non-concern: what a bin means (stft.rs, measure/) | IO: (&mut [f64], &mut [f64]) -> ()
2
3use std::f64::consts::TAU;
4
5/// One twiddle table serves every stage.
6pub 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
64/// The caller fills `0..=n/2`; this mirrors the rest.
65pub 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}