Skip to main content

rill_fft/
spectrum.rs

1// rill-fft/src/spectrum.rs
2//! FFT-based spectrum analyzer.
3//!
4//! Implements `SpectrumAnalyzer` from `rill-core-dsp` using `RealFft`.
5
6use 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
13/// FFT-based spectrum analyzer.
14///
15/// Transforms real input blocks into magnitude spectra via the real FFT.
16/// The window function (Hann by default) is applied before the transform.
17pub 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    /// Create a new spectrum analyzer with FFT size `fft_size` and Hann window.
27    ///
28    /// # Panics
29    ///
30    /// Panics if `fft_size` is not a power of two.
31    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    /// Returns the FFT size.
52    pub fn fft_size(&self) -> usize {
53        self.fft.size()
54    }
55
56    /// Returns the magnitude spectrum (in f32).
57    pub fn spectrum(&self) -> &[f32] {
58        &self.magnitude
59    }
60
61    /// Compute the amplitude at a specific frequency.
62    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}