peerspeak/src/dsp/stft.rs

//! Short-Time Fourier Transform: slice a signal into overlapping windowed
//! frames and take each frame's magnitude spectrum. The result is the
//! time×frequency matrix a spectrogram draws.

use super::fft::real_magnitude_spectrum;
use super::window::{apply, hann};

/// One STFT analysis: a sequence of magnitude frames plus the geometry needed to
/// label axes (sample rate, fft size, hop).
pub struct Spectrogram {
    /// `frames[t][bin]` = linear magnitude of frequency `bin` at time-step `t`.
    /// Each inner vec has `fft_size/2 + 1` bins (DC..Nyquist).
    pub frames: Vec<Vec<f32>>,
    pub sample_rate: u32,
    pub fft_size: usize,
    pub hop: usize,
}

impl Spectrogram {
    /// Number of one-sided frequency bins per frame (`fft_size/2 + 1`).
    pub fn bins(&self) -> usize {
        self.fft_size / 2 + 1
    }

    /// Centre frequency (Hz) of bin index `bin`.
    pub fn bin_hz(&self, bin: usize) -> f32 {
        bin as f32 * self.sample_rate as f32 / self.fft_size as f32
    }

    /// Time (seconds) at the start of frame `t`.
    pub fn frame_time(&self, t: usize) -> f32 {
        (t * self.hop) as f32 / self.sample_rate as f32
    }
}

/// Computes the STFT of `samples` with the given `fft_size` (rounded up to a
/// power of two) and `hop` (frame advance in samples). A Hann window is applied
/// to each frame. The final partial frame is zero-padded so trailing audio is
/// not dropped.
pub fn analyze(samples: &[f32], sample_rate: u32, fft_size: usize, hop: usize) -> Spectrogram {
    let fft_size = fft_size.next_power_of_two().max(2);
    let hop = hop.max(1);
    let window = hann(fft_size);

    let mut frames = Vec::new();
    if !samples.is_empty() {
        let mut start = 0;
        while start < samples.len() {
            let end = (start + fft_size).min(samples.len());
            let mut frame = vec![0.0f32; fft_size];
            frame[..end - start].copy_from_slice(&samples[start..end]);
            let windowed = apply(&frame, &window);
            frames.push(real_magnitude_spectrum(&windowed));
            start += hop;
        }
    }

    Spectrogram {
        frames,
        sample_rate,
        fft_size,
        hop,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::f64::consts::PI;

    #[test]
    fn frame_count_follows_hop() {
        // 1000 samples, hop 250 -> frames start at 0,250,500,750 = 4 frames.
        let sig = vec![0.0f32; 1000];
        let s = analyze(&sig, 48_000, 512, 250);
        assert_eq!(s.frames.len(), 4);
        assert_eq!(s.bins(), 512 / 2 + 1);
    }

    #[test]
    fn tone_lands_in_expected_bin() {
        // A 3 kHz tone at 48 kHz with a 1024-pt FFT -> bin ≈ 3000/(48000/1024) = 64.
        let sr = 48_000;
        let freq = 3000.0;
        let sig: Vec<f32> = (0..4096)
            .map(|i| (2.0 * PI * freq * i as f64 / sr as f64).sin() as f32)
            .collect();
        let s = analyze(&sig, sr, 1024, 512);
        let mid = &s.frames[s.frames.len() / 2];
        let peak_bin = mid
            .iter()
            .enumerate()
            .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
            .unwrap()
            .0;
        let peak_hz = s.bin_hz(peak_bin);
        assert!((peak_hz - freq as f32).abs() < 100.0, "peak at {peak_hz} Hz, want {freq}");
    }

    #[test]
    fn empty_input_yields_no_frames() {
        let s = analyze(&[], 48_000, 512, 256);
        assert!(s.frames.is_empty());
    }
}