Skip to main content

fft_convolver/
fft.rs

1use realfft::num_complex::Complex;
2use realfft::{ComplexToReal, FftError, FftNum, RealFftPlanner, RealToComplex};
3use rtsan_standalone::nonblocking;
4use std::{fmt, sync::Arc};
5
6#[derive(Clone)]
7pub struct Fft<F: FftNum> {
8    fft_forward: Arc<dyn RealToComplex<F>>,
9    scratch_forward: Vec<Complex<F>>,
10    fft_inverse: Arc<dyn ComplexToReal<F>>,
11    scratch_inverse: Vec<Complex<F>>,
12}
13
14impl<F: FftNum + fmt::Debug> fmt::Debug for Fft<F> {
15    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
16        f.debug_struct("Fft")
17            .field("scratch_forward", &self.scratch_forward)
18            .field("scratch_inverse", &self.scratch_inverse)
19            .finish_non_exhaustive()
20    }
21}
22
23impl<F: FftNum> Default for Fft<F> {
24    fn default() -> Self {
25        let mut planner = RealFftPlanner::new();
26        Self {
27            fft_forward: planner.plan_fft_forward(0),
28            fft_inverse: planner.plan_fft_inverse(0),
29            scratch_forward: Vec::new(),
30            scratch_inverse: Vec::new(),
31        }
32    }
33}
34
35impl<F: FftNum> Fft<F> {
36    pub fn init(&mut self, length: usize) {
37        let mut planner = RealFftPlanner::new();
38        self.fft_forward = planner.plan_fft_forward(length);
39        self.fft_inverse = planner.plan_fft_inverse(length);
40        self.scratch_forward = self.fft_forward.make_scratch_vec();
41        self.scratch_inverse = self.fft_inverse.make_scratch_vec();
42    }
43
44    #[nonblocking]
45    pub fn forward(&mut self, input: &mut [F], output: &mut [Complex<F>]) -> Result<(), FftError> {
46        self.fft_forward
47            .process_with_scratch(input, output, &mut self.scratch_forward)?;
48        Ok(())
49    }
50
51    #[nonblocking]
52    pub fn inverse(&mut self, input: &mut [Complex<F>], output: &mut [F]) -> Result<(), FftError> {
53        self.fft_inverse
54            .process_with_scratch(input, output, &mut self.scratch_inverse)?;
55
56        // FFT Normalization
57        let len = output.len();
58        output.iter_mut().for_each(|bin| {
59            *bin = *bin / F::from_usize(len).expect("usize can be converted to FftNum");
60        });
61
62        Ok(())
63    }
64}
65
66#[cfg(test)]
67mod tests {
68    use realfft::num_traits::Zero;
69
70    use super::*;
71
72    #[test]
73    fn test_fft_roundtrip() {
74        let mut fft = Fft::<f32>::default();
75        let size = 128;
76        fft.init(size);
77
78        // Create a simple test signal
79        let mut input = vec![0.0; size];
80        for i in 0..size {
81            input[i] = ((i * 7 + 13) % 50) as f32 / 25.0 - 1.0;
82        }
83        let original = input.clone();
84        let mut freq = vec![Complex::<f32>::zero(); size / 2 + 1];
85
86        // Forward then inverse should give us back the original
87        fft.forward(&mut input, &mut freq).unwrap();
88        fft.inverse(&mut freq, &mut input).unwrap();
89
90        for (&result, &expected) in input.iter().zip(original.iter()) {
91            assert!((result - expected).abs() < 1e-5);
92        }
93    }
94
95    #[test]
96    fn test_fft_sine_wave_frequency() {
97        let mut fft = Fft::<f32>::default();
98        let size = 128;
99        fft.init(size);
100
101        // Create a pure sine wave at bin 10
102        let freq_bin = 10;
103        let mut input = vec![0.0; size];
104        for i in 0..size {
105            input[i] =
106                (2.0 * std::f32::consts::PI * freq_bin as f32 * i as f32 / size as f32).sin();
107        }
108
109        let mut freq = vec![Complex::<f32>::zero(); size / 2 + 1];
110        fft.forward(&mut input, &mut freq).unwrap();
111
112        // The energy should be concentrated at bin 10
113        // For a real sine wave, the FFT magnitude at the frequency bin should be size/2
114        let magnitude = freq[freq_bin].norm();
115        let expected_magnitude = size as f32 / 2.0;
116
117        assert!(
118            (magnitude - expected_magnitude).abs() < 0.1,
119            "Expected magnitude ~{} at bin {}, got {}",
120            expected_magnitude,
121            freq_bin,
122            magnitude
123        );
124
125        // Other bins should have much smaller magnitudes (near zero)
126        for (i, f) in freq.iter().enumerate() {
127            if i != freq_bin {
128                assert!(
129                    f.norm() < 1.0,
130                    "Bin {} should be near zero, got magnitude {}",
131                    i,
132                    f.norm()
133                );
134            }
135        }
136    }
137}