1use rill_core::traits::algorithm::{Algorithm, AlgorithmCategory, AlgorithmMetadata};
7use rill_core::traits::ProcessResult;
8use rill_core::Transcendental;
9use rill_core_dsp::analyzer::{Analyzer, SpectrumAnalyzer};
10
11use crate::real_fft::RealFft;
12
13pub struct FftSpectrumAnalyzer<T: Transcendental> {
18 fft: RealFft<T>,
19 window: Vec<T>,
20 scratch: Vec<num_complex::Complex<T>>,
21 magnitude: Vec<f32>,
22 block_buf: Vec<T>,
23}
24
25impl<T: Transcendental> FftSpectrumAnalyzer<T> {
26 pub fn new(fft_size: usize) -> Self {
32 let fft = RealFft::new(fft_size);
33 let half_plus_one = fft_size / 2 + 1;
34
35 let window = (0..fft_size)
36 .map(|i| {
37 let phase = T::from_f64(2.0 * std::f64::consts::PI * i as f64 / fft_size as f64);
38 T::ONE - phase.cos()
39 })
40 .collect();
41
42 Self {
43 fft,
44 window,
45 scratch: vec![num_complex::Complex::new(T::ZERO, T::ZERO); half_plus_one],
46 magnitude: vec![0.0f32; half_plus_one],
47 block_buf: vec![T::ZERO; fft_size],
48 }
49 }
50
51 pub fn fft_size(&self) -> usize {
53 self.fft.size()
54 }
55
56 pub fn spectrum(&self) -> &[f32] {
58 &self.magnitude
59 }
60
61 pub fn amplitude_at(&self, freq: f32, sample_rate: f32) -> f32 {
63 let bin = (freq * self.fft.size() as f32 / sample_rate) as usize;
64 self.magnitude.get(bin).copied().unwrap_or(0.0)
65 }
66}
67
68impl<T: Transcendental> Algorithm<T> for FftSpectrumAnalyzer<T> {
69 fn process(&mut self, input: Option<&[T]>, output: &mut [T]) -> ProcessResult<()> {
70 match input {
71 Some(samples) => {
72 let len = samples.len().min(self.fft.size());
73 self.block_buf.fill(T::ZERO);
74 self.block_buf[..len].copy_from_slice(&samples[..len]);
75
76 for i in 0..len {
77 self.block_buf[i] = self.block_buf[i] * self.window[i];
78 }
79
80 self.fft.forward(&self.block_buf, &mut self.scratch);
81
82 for (i, c) in self.scratch.iter().enumerate() {
83 let mag_sq = c.re * c.re + c.im * c.im;
84 self.magnitude[i] = mag_sq.to_f64().sqrt() as f32;
85 }
86
87 output.fill(T::ZERO);
88 Ok(())
89 }
90 None => {
91 output.fill(T::ZERO);
92 Ok(())
93 }
94 }
95 }
96
97 fn reset(&mut self) {
98 self.block_buf.fill(T::ZERO);
99 self.scratch
100 .fill(num_complex::Complex::new(T::ZERO, T::ZERO));
101 self.magnitude.fill(0.0);
102 }
103
104 fn metadata(&self) -> AlgorithmMetadata {
105 AlgorithmMetadata {
106 name: "FftSpectrumAnalyzer",
107 category: AlgorithmCategory::Analyzer,
108 description: "FFT-based spectrum analyzer",
109 author: "Rill",
110 version: env!("CARGO_PKG_VERSION"),
111 }
112 }
113}
114
115impl<T: Transcendental> Analyzer<T> for FftSpectrumAnalyzer<T> {
116 type Output = Vec<f32>;
117
118 fn result(&self) -> &Self::Output {
119 &self.magnitude
120 }
121
122 fn reset_analysis(&mut self) {
123 self.magnitude.fill(0.0);
124 }
125}
126
127impl<T: Transcendental> SpectrumAnalyzer<T> for FftSpectrumAnalyzer<T> {
128 fn fft_size(&self) -> usize {
129 self.fft.size()
130 }
131}
132
133#[cfg(test)]
134mod tests {
135 use super::*;
136
137 #[test]
138 fn test_dc_input_gives_dc_bin() {
139 let mut analyzer = FftSpectrumAnalyzer::<f32>::new(16);
140 let input = [2.5f32; 16];
141 let mut output = [0.0f32; 16];
142 analyzer.process(Some(&input), &mut output).unwrap();
143
144 let dc_level = analyzer.spectrum()[0];
145 assert!(dc_level > 5.0, "DC bin too low: {dc_level}");
146
147 for i in 3..analyzer.spectrum().len() {
148 assert!(
149 analyzer.spectrum()[i] < 0.5,
150 "bin {i} should be near zero, got {}",
151 analyzer.spectrum()[i]
152 );
153 }
154 }
155
156 #[test]
157 fn test_sine_input_gives_peak() {
158 let fft_size = 128;
159 let mut analyzer = FftSpectrumAnalyzer::<f32>::new(fft_size);
160 let freq = 1000.0;
161 let sr = 44100.0;
162
163 let input: Vec<f32> = (0..fft_size)
164 .map(|i| {
165 let t = i as f32 / sr;
166 (2.0 * std::f32::consts::PI * freq * t).sin()
167 })
168 .collect();
169
170 let mut output = vec![0.0f32; fft_size];
171 analyzer.process(Some(&input), &mut output).unwrap();
172
173 let expected_bin = (freq * fft_size as f32 / sr) as usize;
174 let peak = analyzer.spectrum()[expected_bin];
175 let nearby = if expected_bin > 0 {
176 analyzer.spectrum()[expected_bin - 1]
177 } else {
178 0.0
179 };
180
181 assert!(
182 peak > nearby * 2.0,
183 "peak at bin {expected_bin} should dominate"
184 );
185 assert!(peak > 0.5, "peak magnitude too low: {peak}");
186 }
187
188 #[test]
189 fn test_amplitude_at() {
190 let fft_size = 256;
191 let mut analyzer = FftSpectrumAnalyzer::<f32>::new(fft_size);
192 let sr = fft_size as f32;
193 let freq = 2.0;
194
195 let input: Vec<f32> = (0..fft_size)
196 .map(|i| {
197 let t = i as f32 / sr;
198 (2.0 * std::f32::consts::PI * freq * t).sin()
199 })
200 .collect();
201
202 let mut output = vec![0.0f32; fft_size];
203 analyzer.process(Some(&input), &mut output).unwrap();
204
205 let amp = analyzer.amplitude_at(freq, sr);
206 assert!(amp > 0.5, "amplitude at {freq} Hz too low: {amp}");
207
208 let amp_far = analyzer.amplitude_at(freq * 5.0, sr);
209 assert!(
210 amp_far < 0.15,
211 "amplitude at 5x freq should be low: {amp_far}"
212 );
213 }
214
215 #[test]
216 fn test_reset_clears_spectrum() {
217 let mut analyzer = FftSpectrumAnalyzer::<f32>::new(16);
218 let input = [1.0f32; 16];
219 let mut output = [0.0f32; 16];
220 analyzer.process(Some(&input), &mut output).unwrap();
221
222 analyzer.reset();
223 assert!(analyzer.spectrum().iter().all(|&v| v == 0.0));
224 }
225}