rustfft/
twiddles.rs

1use crate::{common::FftNum, FftDirection};
2use num_complex::Complex;
3
4use strength_reduce::{StrengthReducedU128, StrengthReducedU64};
5
6pub fn compute_twiddle<T: FftNum>(
7    index: usize,
8    fft_len: usize,
9    direction: FftDirection,
10) -> Complex<T> {
11    let constant = -2f64 * std::f64::consts::PI / fft_len as f64;
12    let angle = constant * index as f64;
13
14    let result = Complex {
15        re: T::from_f64(angle.cos()).unwrap(),
16        im: T::from_f64(angle.sin()).unwrap(),
17    };
18
19    match direction {
20        FftDirection::Forward => result,
21        FftDirection::Inverse => result.conj(),
22    }
23}
24
25pub fn fill_bluesteins_twiddles<T: FftNum>(
26    destination: &mut [Complex<T>],
27    direction: FftDirection,
28) {
29    let twice_len = destination.len() * 2;
30
31    // Standard bluestein's twiddle computation requires us to square the index before usingit to compute a twiddle factor
32    // And since twiddle factors are cyclic, we can improve precision once the squared index gets converted to floating point by taking a modulo
33    // Modulo is expensive, so we're going to use strength-reduction to keep it manageable
34
35    // Strength-reduced u128s are very heavy, so we only want to use them if we need them - and we only need them if
36    // len * len doesn't fit in a u64, AKA if len doesn't fit in a u32
37    if destination.len() < std::u32::MAX as usize {
38        let twice_len_reduced = StrengthReducedU64::new(twice_len as u64);
39
40        for (i, e) in destination.iter_mut().enumerate() {
41            let i_squared = i as u64 * i as u64;
42            let i_mod = i_squared % twice_len_reduced;
43            *e = compute_twiddle(i_mod as usize, twice_len, direction);
44        }
45    } else {
46        // Sadly, the len doesn't fit in a u64, so we have to crank it up to u128 arithmetic
47        let twice_len_reduced = StrengthReducedU128::new(twice_len as u128);
48
49        for (i, e) in destination.iter_mut().enumerate() {
50            // Standard bluestein's twiddle computation requires us to square the index before usingit to compute a twiddle factor
51            // And since twiddle factors are cyclic, we can improve precision once the squared index gets converted to floating point by taking a modulo
52            let i_squared = i as u128 * i as u128;
53            let i_mod = i_squared % twice_len_reduced;
54            *e = compute_twiddle(i_mod as usize, twice_len, direction);
55        }
56    }
57}
58
59pub fn rotate_90<T: FftNum>(value: Complex<T>, direction: FftDirection) -> Complex<T> {
60    match direction {
61        FftDirection::Forward => Complex {
62            re: value.im,
63            im: -value.re,
64        },
65        FftDirection::Inverse => Complex {
66            re: -value.im,
67            im: value.re,
68        },
69    }
70}
71
72#[cfg(test)]
73mod unit_tests {
74    use super::*;
75
76    #[test]
77    fn test_rotate() {
78        // Verify that the rotate90 function does the same thing as multiplying by twiddle(1,4), in the forward direction
79        let value = Complex { re: 9.1, im: 2.2 };
80        let rotated_forward = rotate_90(value, FftDirection::Forward);
81        let twiddled_forward = value * compute_twiddle(1, 4, FftDirection::Forward);
82
83        assert_eq!(value.re, -rotated_forward.im);
84        assert_eq!(value.im, rotated_forward.re);
85
86        assert!(value.re + twiddled_forward.im < 0.0001);
87        assert!(value.im - rotated_forward.re < 0.0001);
88
89        // Verify that the rotate90 function does the same thing as multiplying by twiddle(1,4), in the inverse direction
90        let rotated_forward = rotate_90(value, FftDirection::Inverse);
91        let twiddled_forward = value * compute_twiddle(1, 4, FftDirection::Inverse);
92
93        assert_eq!(value.re, rotated_forward.im);
94        assert_eq!(value.im, -rotated_forward.re);
95
96        assert!(value.re - twiddled_forward.im < 0.0001);
97        assert!(value.im + rotated_forward.re < 0.0001);
98    }
99}