FFT spectrum example

spectrum.rs

use std::{f32::consts::PI, sync::Arc};

use realfft::{num_complex::Complex, RealFftPlanner, RealToComplex};

const N: usize = 1024;
const SR: f32 = 44100.;

pub struct Spectrum {
    fft: Arc<dyn RealToComplex<f32>>,
    real_buf: Vec<f32>,
    complex_buf: Vec<Complex<f32>>,
}

impl Spectrum {
    pub fn new() -> Self {
        Spectrum {
            fft: RealFftPlanner::new().plan_fft_forward(N),
            real_buf: vec![0.; N],
            complex_buf: vec![Complex::new(0., 0.); N / 2 + 1],
        }
    }

    pub fn process(&mut self, input: &[f32], output: &mut [f32]) {
        self.real_buf.copy_from_slice(input);

        // Apply Hann window
        for (i, val) in self.real_buf.iter_mut().enumerate() {
            let x = i as f32 / (N - 1) as f32;
            *val *= (PI * x).sin().powi(2);
        }
        self.fft
            .process(&mut self.real_buf, &mut self.complex_buf)
            .unwrap();

        // Collect a log-log plot into `output`
        for (i, res) in output.iter_mut().enumerate() {
            // i in [0, N[
            // x normalized to [0, 1[
            let x = i as f32 / N as f32;
            // We want to map x to frequency in range [40, 20k[, log scale
            // Parameters k, b. f = k*b^x
            // Know: k*b^0 = 40, k*b^1 = 20k
            // Therefore k = 40, b = 500
            let f = 40. * 500f32.powf(x);
            // w in [40, 20k[ / SR * N
            let w = f / SR * N as f32;
            // Closest FFT bin
            let p = (w as isize).clamp(0, (N / 2) as isize);

            // Lanczos interpolation
            let radius = 10;
            // To interpolate in linear space:
            let mut result = Complex::new(0., 0.);
            // To interpolate in dB space:
            // let mut result = 0.;
            for iw in p - radius..=p + radius + 1 {
                if iw < 0 || iw > (N / 2) as isize {
                    continue;
                }
                let delta = w - iw as f32;
                if delta.abs() > radius as f32 {
                    continue;
                }
                let lanczos = if delta == 0. {
                    1.
                } else {
                    radius as f32 * (PI * delta).sin() * (PI * delta / radius as f32).sin()
                        / (PI * delta).powi(2)
                };
                // Linear space
                let value = self.complex_buf[iw as usize];
                // dB space
                // let value = 20. * self.complex_buf[iw as usize].norm().log10();
                result += lanczos * value;
            }
            // If interpolated in linear space, convert to dB now
            *res = 20. * result.norm().log10();
            // Otherwise use result directly
            // *res = result;
            if !(*res).is_finite() {
                *res = -250.;
            }
        }
    }
}