crabidy/crabidy-server/src/spectrum.rs

//! Turns a window of audio samples into a handful of normalized
//! frequency bars (architecture/spectrum.md). Pure DSP — the tap
//! (`audio-player`) supplies samples, the stream carries the result;
//! this only does the maths.

use realfft::num_complex::Complex;
use realfft::{RealFftPlanner, RealToComplex};
use std::sync::Arc;

/// Number of bars the visualizer shows.
pub const SPECTRUM_BINS: usize = 24;

/// Magnitudes below this (dBFS) map to an empty bar; 0 dB maps to full.
const MIN_DB: f32 = -60.0;

/// A reusable forward-FFT + log-bin folder for a fixed window length.
pub struct SpectrumAnalyzer {
    fft: Arc<dyn RealToComplex<f32>>,
    /// Hann window applied before the transform to cut spectral leakage.
    window: Vec<f32>,
    input: Vec<f32>,
    output: Vec<Complex<f32>>,
    /// `SPECTRUM_BINS + 1` boundaries into the magnitude array, spaced
    /// logarithmically so the bars are roughly musically even.
    edges: Vec<usize>,
}

impl SpectrumAnalyzer {
    pub fn new(window_len: usize) -> Self {
        let fft = RealFftPlanner::<f32>::new().plan_fft_forward(window_len);
        let input = fft.make_input_vec();
        let output = fft.make_output_vec();
        let window = (0..window_len).map(|i| hann(i, window_len)).collect();
        let edges = log_bin_edges(output.len(), SPECTRUM_BINS);
        Self {
            fft,
            window,
            input,
            output,
            edges,
        }
    }

    /// Folds `samples` (exactly the window length) into `SPECTRUM_BINS`
    /// normalized `[0, 1]` magnitudes, low frequency first. A processing
    /// error or a wrong-length input yields all-zero bars rather than a
    /// panic.
    pub fn analyze(&mut self, samples: &[f32]) -> Vec<f32> {
        if samples.len() != self.input.len() {
            return vec![0.0; SPECTRUM_BINS];
        }
        for (dst, (sample, w)) in self
            .input
            .iter_mut()
            .zip(samples.iter().zip(self.window.iter()))
        {
            *dst = sample * w;
        }
        if self.fft.process(&mut self.input, &mut self.output).is_err() {
            return vec![0.0; SPECTRUM_BINS];
        }
        let n = self.window.len() as f32;
        let magnitude = |c: &Complex<f32>| c.norm() / n;
        (0..SPECTRUM_BINS)
            .map(|bar| {
                let lo = self.edges[bar];
                let hi = self.edges[bar + 1].max(lo + 1).min(self.output.len());
                let avg = self.output[lo..hi].iter().map(magnitude).sum::<f32>() / (hi - lo) as f32;
                // Log-compress: dBFS mapped onto [0, 1].
                let db = 20.0 * (avg + 1e-9).log10();
                ((db - MIN_DB) / -MIN_DB).clamp(0.0, 1.0)
            })
            .collect()
    }
}

/// A Hann window coefficient for sample `i` of `len`.
fn hann(i: usize, len: usize) -> f32 {
    let x = std::f32::consts::PI * i as f32 / (len - 1) as f32;
    x.sin().powi(2)
}

/// Geometrically-spaced boundaries into a magnitude array of `mag_len`
/// (skipping the DC bin at 0), giving `bins` groups. Monotonic
/// non-decreasing; callers widen any empty group to at least one bin.
fn log_bin_edges(mag_len: usize, bins: usize) -> Vec<usize> {
    let lo = 1.0_f32;
    let hi = (mag_len - 1).max(2) as f32;
    (0..=bins)
        .map(|b| {
            let t = b as f32 / bins as f32;
            (lo * (hi / lo).powf(t)).round() as usize
        })
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;

    const WINDOW: usize = 2048;

    fn sine(freq: f32, sample_rate: f32, len: usize) -> Vec<f32> {
        (0..len)
            .map(|i| (2.0 * std::f32::consts::PI * freq * i as f32 / sample_rate).sin())
            .collect()
    }

    #[test]
    fn silence_is_all_zero_bars() {
        let mut a = SpectrumAnalyzer::new(WINDOW);
        let bars = a.analyze(&vec![0.0; WINDOW]);
        assert_eq!(bars.len(), SPECTRUM_BINS);
        assert!(bars.iter().all(|&b| b == 0.0), "{bars:?}");
    }

    #[test]
    fn a_tone_lights_one_region_and_stays_normalized() {
        let mut a = SpectrumAnalyzer::new(WINDOW);
        // ~1 kHz at 44.1 kHz lands in the upper-middle of the log bars.
        let bars = a.analyze(&sine(1000.0, 44_100.0, WINDOW));
        assert_eq!(bars.len(), SPECTRUM_BINS);
        assert!(bars.iter().all(|&b| (0.0..=1.0).contains(&b)), "{bars:?}");
        let peak = bars
            .iter()
            .cloned()
            .enumerate()
            .max_by(|a, b| a.1.total_cmp(&b.1))
            .unwrap();
        assert!(
            peak.1 > 0.5,
            "a pure tone should drive its bar high: {bars:?}"
        );
        // Energy is concentrated: most bars stay well below the peak.
        let loud = bars.iter().filter(|&&b| b > peak.1 * 0.5).count();
        assert!(
            loud <= 4,
            "a tone should not light the whole spectrum: {bars:?}"
        );
    }

    #[test]
    fn wrong_length_input_is_zero_not_a_panic() {
        let mut a = SpectrumAnalyzer::new(WINDOW);
        assert!(a.analyze(&[0.1, 0.2, 0.3]).iter().all(|&b| b == 0.0));
    }

    #[test]
    fn bin_edges_are_monotonic() {
        let edges = log_bin_edges(WINDOW / 2 + 1, SPECTRUM_BINS);
        assert_eq!(edges.len(), SPECTRUM_BINS + 1);
        assert!(edges.windows(2).all(|w| w[1] >= w[0]), "{edges:?}");
    }
}