audio_processor_analysis/
fft_processor.rs

1// Augmented Audio: Audio libraries and applications
2// Copyright (c) 2022 Pedro Tacla Yamada
3//
4// The MIT License (MIT)
5//
6// Permission is hereby granted, free of charge, to any person obtaining a copy
7// of this software and associated documentation files (the "Software"), to deal
8// in the Software without restriction, including without limitation the rights
9// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
10// copies of the Software, and to permit persons to whom the Software is
11// furnished to do so, subject to the following conditions:
12//
13// The above copyright notice and this permission notice shall be included in
14// all copies or substantial portions of the Software.
15//
16// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
17// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
18// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
19// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
20// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
21// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
22// THE SOFTWARE.
23
24//! FFT processor implementation with windowing & overlap, wraps `rustfft`.
25//!
26//! `rustfft` audio-processor, forwards or backwards, real-time safe, FFT.
27//!
28//! Applies a Hann window by default. Several window functions are exported by [`audio_processor_analysis::window_functions`].
29//!
30//! ![](https://raw.githubusercontent.com/yamadapc/augmented-audio/master/crates/augmented/audio/audio-processor-analysis/src/window_functions/windows--HannWindow.png)
31//!
32//! Then performs FFT with N bins.
33//!
34//! ![](https://raw.githubusercontent.com/yamadapc/augmented-audio/master/crates/augmented/audio/audio-processor-analysis/src/fft_processor.png--FFT_sine_440Hz.png)
35//!
36//! Overlap is configurable
37//!
38//! ![](https://raw.githubusercontent.com/yamadapc/augmented-audio/master/crates/augmented/audio/audio-processor-analysis/screen.png)
39
40use std::sync::Arc;
41
42use rustfft::num_complex::Complex;
43pub use rustfft::FftDirection;
44use rustfft::{Fft, FftNum, FftPlanner};
45
46use audio_processor_traits::num::traits::FloatConst;
47use audio_processor_traits::simple_processor::MonoAudioProcessor;
48use audio_processor_traits::{AudioContext, Float};
49
50use crate::window_functions::{make_window_vec, WindowFunctionType};
51
52pub struct FftProcessorOptions {
53    pub size: usize,
54    pub direction: FftDirection,
55    pub overlap_ratio: f32,
56    pub window_function: WindowFunctionType,
57}
58
59impl Default for FftProcessorOptions {
60    fn default() -> Self {
61        Self {
62            size: 8192,
63            direction: FftDirection::Forward,
64            overlap_ratio: 0.0,
65            window_function: WindowFunctionType::Hann,
66        }
67    }
68}
69
70/// Default f32 FFT processor
71pub type FftProcessor = FftProcessorImpl<f32>;
72
73/// An FFT processor with overlap and windowing.
74///
75/// This processor will collect samples onto a circular buffer and perform FFTs whenever hop size is
76/// reached.
77pub struct FftProcessorImpl<ST> {
78    input_buffer: Vec<ST>,
79    fft_buffer: Vec<Complex<ST>>,
80    scratch: Vec<Complex<ST>>,
81    cursor: usize,
82    window: Vec<ST>,
83    step_len: usize,
84    size: usize,
85    fft: Arc<dyn Fft<ST>>,
86    has_changed: bool,
87}
88
89impl<ST: FftNum + std::iter::Sum + Float + FloatConst> Default for FftProcessorImpl<ST> {
90    fn default() -> Self {
91        Self::new(Default::default())
92    }
93}
94
95impl<ST: FftNum + std::iter::Sum + Float + FloatConst> FftProcessorImpl<ST> {
96    /// Constructs a new `FftProcessor`
97    ///
98    /// * size: Size of the FFT
99    /// * direction: Direction of the FFT
100    /// * overlap_ratio: 0.0 will do no overlap, 0.5 will do half a window of overlap and 0.75 will
101    ///   do 3/4 window overlap
102    /// * window_function: The window function to use
103    pub fn new(options: FftProcessorOptions) -> Self {
104        let FftProcessorOptions {
105            size,
106            direction,
107            overlap_ratio,
108            window_function,
109        } = options;
110        let mut planner = FftPlanner::new();
111        let fft = planner.plan_fft(size, direction);
112
113        let mut input_buffer = Vec::with_capacity(size);
114        input_buffer.resize(size, ST::zero());
115        let mut fft_buffer = Vec::with_capacity(size);
116        fft_buffer.resize(size, ST::zero().into());
117
118        let scratch_size = fft.get_inplace_scratch_len();
119        let mut scratch = Vec::with_capacity(scratch_size);
120        scratch.resize(scratch_size, ST::zero().into());
121
122        let window = make_window_vec(size, window_function);
123        let step_len = Self::calculate_hop_size(size, overlap_ratio);
124
125        Self {
126            input_buffer,
127            fft_buffer,
128            window,
129            scratch,
130            size,
131            step_len,
132            cursor: 0,
133            fft,
134            has_changed: false,
135        }
136    }
137
138    fn calculate_hop_size(size: usize, overlap_ratio: f32) -> usize {
139        (size as f32 * (1.0 - overlap_ratio)) as usize
140    }
141
142    /// The number of frequency bins this FFT processor operates with
143    pub fn size(&self) -> usize {
144        self.size
145    }
146
147    /// Get a reference to the FFT bins buffer
148    pub fn buffer(&self) -> &Vec<Complex<ST>> {
149        &self.fft_buffer
150    }
151
152    /// Get a reference to the rustfft instance
153    pub fn fft(&self) -> &Arc<dyn Fft<ST>> {
154        &self.fft
155    }
156
157    /// Get a mutable reference to the FFT bins buffer
158    pub fn buffer_mut(&mut self) -> &mut Vec<Complex<ST>> {
159        &mut self.fft_buffer
160    }
161
162    pub fn input_mut(&mut self) -> &mut [ST] {
163        &mut self.input_buffer
164    }
165
166    /// Get a mutable reference to the scratch buffer
167    pub fn scratch_mut(&mut self) -> &mut Vec<Complex<ST>> {
168        &mut self.scratch
169    }
170
171    /// Get the hop size of this processor. This is the number of samples between each FFT.
172    pub fn step_len(&self) -> usize {
173        self.step_len
174    }
175
176    /// Manually process an external FFT buffer in-place.
177    pub fn process_fft_buffer(&mut self, samples: &mut [Complex<ST>]) {
178        self.fft.process_with_scratch(samples, &mut self.scratch);
179    }
180
181    /// Returns true if an FFT has just been performed on the last call to `s_process`
182    pub fn has_changed(&self) -> bool {
183        self.has_changed
184    }
185
186    /// Returns the sum of the power of the current input buffer window.
187    pub fn input_buffer_sum(&self) -> ST {
188        self.input_buffer.iter().map(|f| f.abs()).sum()
189    }
190
191    /// Manually perform an FFT; offset the input buffer by a certain index.
192    #[inline]
193    pub fn perform_fft(&mut self, start_idx: usize) {
194        for i in 0..self.size {
195            let index = (start_idx + i) % self.size;
196            let sample = self.input_buffer[index];
197
198            let magnitude = sample * self.window[i];
199            assert!(!magnitude.is_nan());
200            let complex = Complex::new(magnitude, ST::zero());
201            assert!(!complex.re.is_nan());
202            assert!(!complex.im.is_nan());
203
204            self.fft_buffer[i] = complex;
205        }
206
207        self.fft
208            .process_with_scratch(&mut self.fft_buffer, &mut self.scratch);
209    }
210}
211
212impl<ST: FftNum + Float + std::iter::Sum + FloatConst> MonoAudioProcessor for FftProcessorImpl<ST> {
213    type SampleType = ST;
214
215    #[inline]
216    fn m_process(
217        &mut self,
218        _context: &mut AudioContext,
219        sample: Self::SampleType,
220    ) -> Self::SampleType {
221        self.has_changed = false;
222        self.input_buffer[self.cursor] = sample;
223        self.cursor = self.cursor + 1;
224        if self.cursor >= self.size {
225            self.cursor = 0;
226        }
227
228        if self.cursor % self.step_len == 0 {
229            // Offset FFT so it's reading from the input buffer at the start of this window
230            let start_idx = (self.cursor as i32 - self.size as i32) as usize % self.size;
231            self.perform_fft(start_idx);
232            self.has_changed = true;
233        }
234
235        sample
236    }
237}
238
239#[cfg(test)]
240mod test {
241    use std::cmp::Ordering;
242    use std::time::Duration;
243
244    use audio_processor_testing_helpers::{
245        charts::draw_vec_chart, oscillator_buffer, relative_path, sine_generator,
246    };
247
248    use audio_processor_traits::simple_processor::process_buffer;
249    use audio_processor_traits::{AudioBuffer, AudioProcessorSettings};
250
251    use super::*;
252
253    #[test]
254    fn test_hop_size_is_correct() {
255        let hop_size = FftProcessor::calculate_hop_size(2048, 0.75);
256        assert_eq!(hop_size, 512);
257        let hop_size = FftProcessor::calculate_hop_size(2048, 0.875);
258        assert_eq!(hop_size, 256);
259    }
260
261    #[test]
262    fn test_create_fft() {
263        let mut fft_processor = FftProcessor::new(FftProcessorOptions::default());
264        assert_eq!(fft_processor.size(), FftProcessorOptions::default().size);
265        assert_eq!(
266            fft_processor.buffer().len(),
267            FftProcessorOptions::default().size
268        );
269        assert_eq!(
270            fft_processor.scratch_mut().len(),
271            FftProcessorOptions::default().size
272        );
273        // By default there is no overlap
274        assert_eq!(
275            fft_processor.step_len(),
276            FftProcessorOptions::default().size
277        );
278    }
279
280    #[test]
281    fn test_create_fft_with_4x_overlap() {
282        let mut fft_processor = FftProcessor::new(FftProcessorOptions {
283            overlap_ratio: 0.75,
284            ..FftProcessorOptions::default()
285        });
286        assert_eq!(fft_processor.size(), FftProcessorOptions::default().size);
287        assert_eq!(
288            fft_processor.buffer().len(),
289            FftProcessorOptions::default().size
290        );
291        assert_eq!(
292            fft_processor.scratch_mut().len(),
293            FftProcessorOptions::default().size
294        );
295        assert_eq!(
296            fft_processor.step_len(),
297            FftProcessorOptions::default().size / 4
298        );
299    }
300
301    #[test]
302    fn test_perform_fft_over_silence() {
303        let mut fft_processor = FftProcessor::default();
304        let input_buffer = vec![0.0; fft_processor.size()];
305        fft_processor.input_mut().copy_from_slice(&input_buffer);
306        fft_processor.perform_fft(0);
307
308        let result = fft_processor.buffer();
309        assert_eq!(result.len(), fft_processor.size());
310        assert_eq!(result.iter().map(|c| c.norm()).sum::<f32>(), 0.0);
311    }
312
313    #[test]
314    fn test_perform_fft_over_sine() {
315        let mut fft_processor = FftProcessor::default();
316        let mut input_buffer = vec![0.0; fft_processor.size()];
317        let frequency = 441.0;
318        let sample_rate = 44100.0;
319        for i in 0..input_buffer.len() {
320            let sine = (2.0 * std::f32::consts::PI * frequency * i as f32 / sample_rate).sin();
321            input_buffer[i] = sine;
322        }
323        fft_processor.input_mut().copy_from_slice(&input_buffer);
324        fft_processor.perform_fft(0);
325
326        let result = fft_processor.buffer();
327        assert_eq!(result.len(), fft_processor.size());
328        let (maximum_index, _) = result
329            .iter()
330            .take(fft_processor.size() / 2)
331            .enumerate()
332            .max_by(|(_, c1), (_, c2)| c1.norm().partial_cmp(&c2.norm()).unwrap_or(Ordering::Equal))
333            .unwrap();
334        assert_eq!(
335            maximum_index,
336            (440.0 / sample_rate * fft_processor.size() as f32) as usize + 1
337        );
338    }
339
340    #[test]
341    fn test_draw_fft() {
342        println!("Generating signal");
343        let signal = oscillator_buffer(44100.0, 440.0, Duration::from_millis(1000), sine_generator);
344        let mut context = AudioContext::from(AudioProcessorSettings::new(44100.0, 1, 1, 512));
345        let mut signal = AudioBuffer::from_interleaved(1, &signal);
346
347        println!("Processing");
348        let mut fft_processor = FftProcessor::default();
349        process_buffer(&mut context, &mut fft_processor, &mut signal);
350
351        println!("Drawing chart");
352        let mut output: Vec<f32> = fft_processor
353            .buffer()
354            .iter()
355            .map(|c| 20.0 * (c.norm() / 10.0).log10())
356            .collect();
357        output.reverse();
358        let output: Vec<f32> = output.iter().take(1000).copied().collect();
359
360        draw_vec_chart(
361            &relative_path!("src/fft_processor.png"),
362            "FFT_sine_440Hz",
363            output,
364        );
365    }
366
367    #[test]
368    fn test_usize_cast() {
369        let i = -1;
370        let i = i as usize % 2;
371        assert_eq!(i, 1)
372    }
373}
audio_processor_analysis/fft_processor.rs

audio_processor_analysis/
fft_processor.rs
