2-D FFT — source + log-magnitude spectrum

//! 2-D radix-2 FFT via row-then-column passes.
//!
//! Reuses `fft::fft_in_place`. Layout is row-major interleaved (re, im),
//! `data.len() = 2 · w · h`, both w and h powers of two. V-class L0-closed.
//!
//! Designed to power the "visual FFT" exhibit — a 2-D log-magnitude
//! spectrum next to its source intensity field.

use serde::{Deserialize, Serialize};

use crate::fft::fft_in_place;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Fft2dKernelReport {
    pub w: u32,
    pub h: u32,
    pub bit_ops: u64,
}

/// In-place 2-D FFT. `data` is `2 · w · h` long, interleaved (re, im),
/// row-major. Returns `Err` if dimensions are not powers of two.
pub fn fft2d_in_place(data: &mut [f32], w: u32, h: u32) -> Result<Fft2dKernelReport, &'static str> {
    let w_us = w as usize;
    let h_us = h as usize;
    if !w_us.is_power_of_two() || !h_us.is_power_of_two() || w_us == 0 || h_us == 0 {
        return Err("fft2d: w and h must be powers of two");
    }
    if data.len() != 2 * w_us * h_us {
        return Err("fft2d: buffer size mismatch");
    }

    // ── Row pass: FFT each row in place ──
    let mut row_bits: u64 = 0;
    let row_len = 2 * w_us;
    for y in 0..h_us {
        let s = y * row_len;
        let r = fft_in_place(&mut data[s..s + row_len])?;
        row_bits = row_bits.saturating_add(r.bit_ops);
    }

    // ── Column pass: pack each column, FFT it, write back ──
    let mut col_buf = vec![0.0f32; 2 * h_us];
    let mut col_bits: u64 = 0;
    for x in 0..w_us {
        for y in 0..h_us {
            let src = y * row_len + 2 * x;
            col_buf[2 * y] = data[src];
            col_buf[2 * y + 1] = data[src + 1];
        }
        let r = fft_in_place(&mut col_buf)?;
        col_bits = col_bits.saturating_add(r.bit_ops);
        for y in 0..h_us {
            let dst = y * row_len + 2 * x;
            data[dst] = col_buf[2 * y];
            data[dst + 1] = col_buf[2 * y + 1];
        }
    }

    Ok(Fft2dKernelReport {
        w,
        h,
        bit_ops: row_bits.saturating_add(col_bits),
    })
}

/// Pack `(w · h)` log-magnitude values into `out` (length `w · h`),
/// shifted so DC sits at the centre. Returns the max log-magnitude.
pub fn log_magnitude_shifted(data: &[f32], w: u32, h: u32, out: &mut [f32]) -> f32 {
    let w_us = w as usize;
    let h_us = h as usize;
    debug_assert_eq!(out.len(), w_us * h_us);
    let hx = w_us / 2;
    let hy = h_us / 2;
    let mut maxv = -1e30f32;
    for y in 0..h_us {
        for x in 0..w_us {
            let src_x = (x + hx) % w_us;
            let src_y = (y + hy) % h_us;
            let i = (src_y * w_us + src_x) * 2;
            let re = data[i];
            let im = data[i + 1];
            let lm = (re * re + im * im + 1e-12).ln();
            out[y * w_us + x] = lm;
            if lm > maxv { maxv = lm; }
        }
    }
    maxv
}

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

    fn buf_zero(w: u32, h: u32) -> Vec<f32> {
        vec![0.0; (2 * w * h) as usize]
    }

    #[test]
    fn rejects_non_power_of_two_dims() {
        let mut b = buf_zero(6, 8);
        assert!(fft2d_in_place(&mut b, 6, 8).is_err());
    }

    #[test]
    fn pure_dc_concentrates_at_origin() {
        let w = 16;
        let h = 16;
        let mut b = buf_zero(w, h);
        // Constant 1.0 real signal → DC peak at (0,0).
        for k in 0..(w * h) as usize {
            b[2 * k] = 1.0;
        }
        fft2d_in_place(&mut b, w, h).unwrap();
        let dc_mag = (b[0].powi(2) + b[1].powi(2)).sqrt();
        // First non-DC bin should be ~0.
        let nz = (b[2].powi(2) + b[3].powi(2)).sqrt();
        assert!(dc_mag > 100.0, "DC mag {dc_mag}");
        assert!(nz < 1e-3, "non-DC mag should be ~0, got {nz}");
    }

    #[test]
    fn report_scales_with_n_log_n() {
        let mut a = buf_zero(16, 16);
        let mut b = buf_zero(32, 32);
        let ra = fft2d_in_place(&mut a, 16, 16).unwrap();
        let rb = fft2d_in_place(&mut b, 32, 32).unwrap();
        // Ratio should be ~ (1024 · log2 32) / (256 · log2 16) = 5120/1024 = 5
        let ratio = rb.bit_ops as f64 / ra.bit_ops as f64;
        assert!(ratio > 4.0 && ratio < 6.0, "ratio was {ratio}");
    }
}

This is the exact Rust file compiled into the WASM module the page just loaded. Every line that runs on your device is here. The receipt is a function of this code, not a bespoke benchmark.