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 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 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 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 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 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 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}