peerspeak/src/dsp/fft.rs

//! A small, dependency-free radix-2 Cooley-Tukey FFT.
//!
//! We hand-roll this rather than pull in `rustfft`/`realfft` because (a) the
//! whole DSP toolkit is a developer/measurement aid, not a hot real-time path,
//! and (b) it keeps the supply-chain surface at zero new crates. The transform
//! is the textbook iterative in-place algorithm (bit-reversal permutation +
//! log2(N) butterfly stages), computed in `f64` for headroom even though the
//! audio it analyses is `f32`.
//!
//! Only power-of-two lengths are supported; the STFT layer always pads frames
//! up to a power of two before calling in.

use std::f64::consts::PI;

/// A minimal complex number for the transform. Kept local (rather than pulling a
/// `num-complex` dependency) since the FFT is the only thing that needs it.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Complex {
    pub re: f64,
    pub im: f64,
}

impl Complex {
    pub const ZERO: Complex = Complex { re: 0.0, im: 0.0 };

    pub fn new(re: f64, im: f64) -> Self {
        Complex { re, im }
    }

    /// Magnitude `sqrt(re^2 + im^2)`.
    pub fn magnitude(self) -> f64 {
        self.re.hypot(self.im)
    }

    fn add(self, o: Complex) -> Complex {
        Complex::new(self.re + o.re, self.im + o.im)
    }

    fn sub(self, o: Complex) -> Complex {
        Complex::new(self.re - o.re, self.im - o.im)
    }

    fn mul(self, o: Complex) -> Complex {
        Complex::new(
            self.re * o.re - self.im * o.im,
            self.re * o.im + self.im * o.re,
        )
    }
}

/// In-place forward FFT. `buf.len()` must be a power of two.
///
/// Uses the standard sign convention `X[k] = sum_n x[n] * exp(-2πi·kn/N)`.
pub fn fft(buf: &mut [Complex]) {
    transform(buf, false);
}

/// In-place inverse FFT (normalized by `1/N`), the exact inverse of [`fft`].
pub fn ifft(buf: &mut [Complex]) {
    transform(buf, true);
    let n = buf.len() as f64;
    for c in buf.iter_mut() {
        c.re /= n;
        c.im /= n;
    }
}

fn transform(buf: &mut [Complex], inverse: bool) {
    let n = buf.len();
    assert!(n.is_power_of_two(), "FFT length {n} must be a power of two");
    if n <= 1 {
        return;
    }

    // Bit-reversal permutation: reorder so the iterative butterflies can run
    // bottom-up in place.
    let mut j = 0usize;
    for i in 1..n {
        let mut bit = n >> 1;
        while j & bit != 0 {
            j ^= bit;
            bit >>= 1;
        }
        j ^= bit;
        if i < j {
            buf.swap(i, j);
        }
    }

    // Butterfly stages: combine length-`len` DFTs from length-`len/2` halves,
    // doubling `len` each pass.
    let sign = if inverse { 1.0 } else { -1.0 };
    let mut len = 2;
    while len <= n {
        let ang = sign * 2.0 * PI / len as f64;
        let wlen = Complex::new(ang.cos(), ang.sin());
        let mut i = 0;
        while i < n {
            let mut w = Complex::new(1.0, 0.0);
            for k in 0..len / 2 {
                let u = buf[i + k];
                let v = buf[i + k + len / 2].mul(w);
                buf[i + k] = u.add(v);
                buf[i + k + len / 2] = u.sub(v);
                w = w.mul(wlen);
            }
            i += len;
        }
        len <<= 1;
    }
}

/// Forward FFT of a real signal, returning the **one-sided** magnitude spectrum:
/// bins `0..=N/2` (DC through Nyquist), where `N` is the next power of two ≥
/// `samples.len()`. The input is zero-padded up to `N`.
///
/// Magnitudes are raw (un-normalized) linear amplitudes; callers convert to dB
/// or normalize as needed.
pub fn real_magnitude_spectrum(samples: &[f32]) -> Vec<f32> {
    let n = samples.len().next_power_of_two().max(2);
    let mut buf = vec![Complex::ZERO; n];
    for (i, &s) in samples.iter().enumerate() {
        buf[i].re = s as f64;
    }
    fft(&mut buf);
    buf[..=n / 2].iter().map(|c| c.magnitude() as f32).collect()
}

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

    fn approx(a: f64, b: f64, eps: f64) -> bool {
        (a - b).abs() <= eps
    }

    #[test]
    fn impulse_transforms_to_flat_spectrum() {
        // FFT of a unit impulse at n=0 is all-ones (flat spectrum).
        let mut buf = vec![Complex::ZERO; 8];
        buf[0] = Complex::new(1.0, 0.0);
        fft(&mut buf);
        for c in &buf {
            assert!(
                approx(c.magnitude(), 1.0, 1e-9),
                "expected flat 1.0, got {c:?}"
            );
        }
    }

    #[test]
    fn single_bin_sine_peaks_in_that_bin() {
        // A cosine at exactly bin k=2 over N=16 should put all energy in bin 2
        // (and its mirror N-2). Check the one-sided spectrum peaks at bin 2.
        let n = 16;
        let k = 2;
        let samples: Vec<f32> = (0..n)
            .map(|i| (2.0 * PI * k as f64 * i as f64 / n as f64).cos() as f32)
            .collect();
        let mag = real_magnitude_spectrum(&samples);
        let peak_bin = mag
            .iter()
            .enumerate()
            .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
            .unwrap()
            .0;
        assert_eq!(peak_bin, k, "energy should land in bin {k}, got {peak_bin}");
    }

    #[test]
    fn ifft_inverts_fft() {
        let original: Vec<Complex> = (0..32)
            .map(|i| Complex::new((i as f64 * 0.3).sin(), (i as f64 * 0.1).cos()))
            .collect();
        let mut buf = original.clone();
        fft(&mut buf);
        ifft(&mut buf);
        for (a, b) in original.iter().zip(&buf) {
            assert!(approx(a.re, b.re, 1e-9) && approx(a.im, b.im, 1e-9));
        }
    }

    #[test]
    #[should_panic(expected = "power of two")]
    fn non_power_of_two_panics() {
        let mut buf = vec![Complex::ZERO; 6];
        fft(&mut buf);
    }
}