acoustic_ofdm/
spectrogram.rs

1// Copyright (c) 2026 Elias S. G. Carotti
2
3use std::error::Error;
4use std::path::Path;
5
6use plotters::prelude::*;
7use rustfft::{num_complex::Complex32, FftPlanner};
8
9/// Analysis window used for STFT-based spectrogram generation.
10#[derive(Clone, Copy, Debug, PartialEq, Eq)]
11pub enum SpectrogramWindow {
12    /// Hann window, good general-purpose sidelobe suppression.
13    ///
14    /// Reference: [Wikipedia: Hann window](https://en.wikipedia.org/wiki/Window_function#Hann_window)
15    Hann,
16    /// Hamming window, slightly narrower main lobe than Hann.
17    ///
18    /// Reference: [Wikipedia: Hamming window](https://en.wikipedia.org/wiki/Window_function#Hamming_window)
19    Hamming,
20    /// Blackman window, stronger sidelobe suppression at the cost of width.
21    ///
22    /// Reference: [Wikipedia: Blackman window](https://en.wikipedia.org/wiki/Window_function#Blackman_window)
23    Blackman,
24    /// Rectangular window, no tapering.
25    ///
26    /// Reference: [Wikipedia: Rectangular window](https://en.wikipedia.org/wiki/Window_function#Rectangular_window)
27    Rect,
28}
29
30/// Configuration for spectrogram rendering.
31#[derive(Clone, Copy, Debug)]
32pub struct SpectrogramOptions {
33    /// FFT size used for each STFT frame.
34    pub nfft: usize,
35    /// Hop size, in samples, between adjacent STFT frames.
36    pub hop: usize,
37    /// Analysis window applied before each FFT.
38    pub window: SpectrogramWindow,
39}
40
41impl Default for SpectrogramOptions {
42    fn default() -> Self {
43        Self {
44            nfft: 512,
45            hop: 128,
46            window: SpectrogramWindow::Hann,
47        }
48    }
49}
50
51fn window_samples(n: usize, kind: SpectrogramWindow) -> Vec<f32> {
52    if n <= 1 {
53        return vec![1.0; n.max(1)];
54    }
55    match kind {
56        SpectrogramWindow::Rect => vec![1.0; n],
57        SpectrogramWindow::Hann => (0..n)
58            .map(|i| 0.5 - 0.5 * (2.0 * std::f32::consts::PI * (i as f32) / ((n - 1) as f32)).cos())
59            .collect(),
60        SpectrogramWindow::Hamming => (0..n)
61            .map(|i| {
62                0.54 - 0.46 * (2.0 * std::f32::consts::PI * (i as f32) / ((n - 1) as f32)).cos()
63            })
64            .collect(),
65        SpectrogramWindow::Blackman => (0..n)
66            .map(|i| {
67                let phi = 2.0 * std::f32::consts::PI * (i as f32) / ((n - 1) as f32);
68                0.42 - 0.5 * phi.cos() + 0.08 * (2.0 * phi).cos()
69            })
70            .collect(),
71    }
72}
73
74fn jet_like(t: f32) -> RGBColor {
75    let t = t.clamp(0.0, 1.0);
76    let anchors = [
77        (0.0, (0.0, 0.0, 131.0)),
78        (0.125, (0.0, 60.0, 170.0)),
79        (0.375, (5.0, 255.0, 255.0)),
80        (0.625, (255.0, 255.0, 0.0)),
81        (0.875, (250.0, 0.0, 0.0)),
82        (1.0, (128.0, 0.0, 0.0)),
83    ];
84    for w in anchors.windows(2) {
85        let (t0, c0) = w[0];
86        let (t1, c1) = w[1];
87        if t <= t1 {
88            let a = ((t - t0) / (t1 - t0)).clamp(0.0, 1.0);
89            let lerp = |x0: f32, x1: f32| (x0 + a * (x1 - x0)).round() as u8;
90            return RGBColor(lerp(c0.0, c1.0), lerp(c0.1, c1.1), lerp(c0.2, c1.2));
91        }
92    }
93    RGBColor(128, 0, 0)
94}
95
96/// Saves a human-readable spectrogram PNG for a mono waveform.
97///
98/// Parameters:
99/// - `path`: output PNG path.
100/// - `x`: input mono waveform.
101/// - `fs`: sample rate in Hz.
102/// Returns:
103/// - `Result<(), Box<dyn Error>>`: `Ok(())` when the PNG is written.
104pub fn save_spectrogram_png(path: &Path, x: &[f32], fs: f32) -> Result<(), Box<dyn Error>> {
105    save_spectrogram_png_with_options(path, x, fs, SpectrogramOptions::default())
106}
107
108pub fn save_spectrogram_png_with_options(
109    path: &Path,
110    x: &[f32],
111    fs: f32,
112    opts: SpectrogramOptions,
113) -> Result<(), Box<dyn Error>> {
114    if x.is_empty() {
115        return Ok(());
116    }
117    let nfft = opts.nfft.max(16);
118    let hop = opts.hop.max(1).min(nfft);
119    if x.len() < nfft {
120        return Ok(());
121    }
122
123    let window = window_samples(nfft, opts.window);
124    let mut planner = FftPlanner::<f32>::new();
125    let fft = planner.plan_fft_forward(nfft);
126    let n_frames = 1 + (x.len() - nfft) / hop;
127    let n_bins = nfft / 2 + 1;
128    let t_max = ((n_frames - 1) * hop + nfft) as f32 / fs;
129    let f_max_khz = fs / 2000.0;
130
131    let mut spec = vec![0.0f32; n_frames * n_bins];
132    let mut frame = vec![Complex32::new(0.0, 0.0); nfft];
133    let mut max_db = -120.0f32;
134    let mut min_db = 0.0f32;
135    for t in 0..n_frames {
136        let s0 = t * hop;
137        for i in 0..nfft {
138            frame[i] = Complex32::new(x[s0 + i] * window[i], 0.0);
139        }
140        fft.process(&mut frame);
141        for k in 0..n_bins {
142            let p = frame[k].norm_sqr() / (nfft as f32);
143            let db = 10.0 * p.max(1e-12).log10();
144            spec[t * n_bins + k] = db;
145            max_db = max_db.max(db);
146            min_db = min_db.min(db);
147        }
148    }
149    let floor_db = (max_db - 70.0).max(min_db);
150
151    let root = BitMapBackend::new(path, (1280, 720)).into_drawing_area();
152    root.fill(&RGBColor(245, 245, 240))?;
153    let (main_area, legend_area) = root.split_horizontally(1120);
154
155    let mut chart = ChartBuilder::on(&main_area)
156        .caption("Spectrogram", ("sans-serif", 28).into_font())
157        .margin(20)
158        .x_label_area_size(45)
159        .y_label_area_size(60)
160        .build_cartesian_2d(0f32..t_max, 0f32..f_max_khz)?;
161
162    chart
163        .configure_mesh()
164        .x_desc("Time [s]")
165        .y_desc("Frequency [kHz]")
166        .axis_desc_style(("sans-serif", 20))
167        .label_style(("sans-serif", 16))
168        .light_line_style(RGBAColor(0, 0, 0, 0.08))
169        .bold_line_style(RGBAColor(0, 0, 0, 0.18))
170        .draw()?;
171
172    let dt = hop as f32 / fs;
173    let df_khz = fs / (nfft as f32) / 1000.0;
174    let mut cells = Vec::with_capacity(n_frames * n_bins);
175    for t in 0..n_frames {
176        for k in 0..n_bins {
177            let db = spec[t * n_bins + k];
178            let norm = ((db - floor_db) / (max_db - floor_db).max(1e-6)).clamp(0.0, 1.0);
179            let color = jet_like(norm).filled();
180            let x0 = t as f32 * dt;
181            let x1 = (t as f32 + 1.0) * dt;
182            let y0 = k as f32 * df_khz;
183            let y1 = (k as f32 + 1.0) * df_khz;
184            cells.push(Rectangle::new([(x0, y0), (x1, y1)], color));
185        }
186    }
187    chart.draw_series(cells)?;
188
189    legend_area.fill(&RGBColor(245, 245, 240))?;
190    let (_lw, lh) = legend_area.dim_in_pixel();
191    let lh = lh as i32;
192    let bar_left = 28i32;
193    let bar_right = 52i32;
194    let bar_top = 36i32;
195    let bar_bottom = (lh - 42).max(bar_top + 1);
196    let n_steps = 256usize;
197    for i in 0..n_steps {
198        let y0 = bar_bottom - ((i as i32) * (bar_bottom - bar_top) / (n_steps as i32));
199        let y1 = bar_bottom - (((i + 1) as i32) * (bar_bottom - bar_top) / (n_steps as i32));
200        legend_area.draw(&Rectangle::new(
201            [(bar_left, y1), (bar_right, y0.max(y1 + 1))],
202            jet_like((i as f32) / ((n_steps - 1) as f32)).filled(),
203        ))?;
204    }
205    legend_area.draw(&Rectangle::new(
206        [(bar_left, bar_top), (bar_right, bar_bottom)],
207        ShapeStyle::from(&BLACK).stroke_width(1),
208    ))?;
209
210    let tick_vals = [
211        max_db,
212        max_db - 20.0,
213        max_db - 40.0,
214        max_db - 60.0,
215        floor_db,
216    ];
217    for &db in &tick_vals {
218        let a = ((db - floor_db) / (max_db - floor_db).max(1e-6)).clamp(0.0, 1.0);
219        let y = bar_bottom - ((a * ((bar_bottom - bar_top) as f32)).round() as i32);
220        legend_area.draw(&PathElement::new(
221            vec![(bar_right + 2, y), (bar_right + 10, y)],
222            BLACK,
223        ))?;
224        legend_area.draw(&Text::new(
225            format!("{db:.0}"),
226            (bar_right + 14, y + 5),
227            ("sans-serif", 15).into_font(),
228        ))?;
229    }
230    legend_area.draw(&Text::new(
231        "Power [dB]",
232        (bar_right + 42, lh / 2),
233        ("sans-serif", 18)
234            .into_font()
235            .transform(FontTransform::Rotate90),
236    ))?;
237
238    root.present()?;
239    Ok(())
240}
241
242// vim: set ts=4 sw=4 et:
acoustic_ofdm/spectrogram.rs

acoustic_ofdm/
spectrogram.rs
