The repo never enforced rustfmt, so formatting had drifted broadly. This is a single mechanical `cargo fmt` pass over the whole crate (no behavioral change; lib suite green, 493 passed). Going forward fmt should be enforced (planned CI fmt --check step). Part of the 0.6.1 hygiene pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
109 lines
3.4 KiB
Rust
109 lines
3.4 KiB
Rust
//! 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());
|
||
}
|
||
}
|