1use font8x8::UnicodeFonts as _;
10
11use crate::audio::SampleRate;
12use crate::cast;
13use crate::dsp::{ActiveFrames, Stft, StftParams, Window};
14use crate::units::{Decibels, Hz, Seconds};
15
16pub const WIDTH: usize = 1500;
19pub const HEIGHT: usize = 512;
21pub const PASSAGE_S: Seconds = Seconds::from_const(10.0);
23pub const RANGE_DB: Decibels = Decibels::from_const(80.0);
25
26#[derive(Clone, Copy, Debug, PartialEq)]
28pub struct PanelSpec {
29 pub frame_len: usize,
31 pub lo: Hz,
33 pub hi: Hz,
35}
36
37pub const PANELS: [PanelSpec; 2] = [
39 PanelSpec {
40 frame_len: 4096,
41 lo: Hz::from_const(20.0),
42 hi: Hz::from_const(20_000.0),
43 },
44 PanelSpec {
45 frame_len: 16384,
46 lo: Hz::from_const(20.0),
47 hi: Hz::from_const(500.0),
48 },
49];
50
51const LEFT: usize = 48;
52const RIGHT: usize = 60;
53const TITLE: usize = 12;
54const PANEL_HEIGHT: usize = 232;
55pub const PLOT_WIDTH: usize = WIDTH - LEFT - RIGHT;
57pub const PLOT_HEIGHT: usize = PANEL_HEIGHT;
59const AXIS_TOP: usize = 2 * (TITLE + PANEL_HEIGHT);
60const BACKGROUND: [u8; 3] = [0, 0, 0];
61const INK: [u8; 3] = [255, 255, 255];
62const FRAME: [u8; 3] = [128, 128, 128];
63
64#[derive(Clone, Debug, PartialEq)]
67pub struct PanelDb {
68 spec: PanelSpec,
69 db: Vec<f64>,
70}
71
72impl PanelDb {
73 #[must_use]
75 pub fn spec(&self) -> PanelSpec {
76 self.spec
77 }
78
79 #[must_use]
82 pub fn db(&self) -> &[f64] {
83 &self.db
84 }
85}
86
87#[derive(Clone, Debug, PartialEq)]
89pub struct Spectrogram {
90 start: Seconds,
91 span: Seconds,
92 panels: Vec<PanelDb>,
93}
94
95impl Spectrogram {
96 #[must_use]
104 pub fn of(samples: &[f64], rate: SampleRate) -> Self {
105 let (first, len) = passage(samples, rate);
106 let panels = PANELS
107 .iter()
108 .map(|&spec| PanelDb {
109 spec,
110 db: panel_db(samples, rate, spec, first, len),
111 })
112 .collect();
113 Self {
114 start: rate.seconds(first),
115 span: rate.seconds(len),
116 panels,
117 }
118 }
119
120 #[must_use]
122 pub fn start(&self) -> Seconds {
123 self.start
124 }
125
126 #[must_use]
128 pub fn span(&self) -> Seconds {
129 self.span
130 }
131
132 #[must_use]
134 pub fn panels(&self) -> &[PanelDb] {
135 &self.panels
136 }
137
138 #[must_use]
140 pub fn pixels(&self) -> Vec<u8> {
141 let mut canvas = Canvas::new();
142 for (index, panel) in self.panels.iter().enumerate() {
143 let top = index * (TITLE + PANEL_HEIGHT);
144 let spec = panel.spec;
145 canvas.text(
146 LEFT,
147 top + 2,
148 &format!(
149 "N = {}, {} Hz to {} Hz, dB below the stem peak",
150 spec.frame_len,
151 hz_label(spec.lo.get()),
152 hz_label(spec.hi.get())
153 ),
154 );
155 let plot_top = top + TITLE;
156 for (row, levels) in panel.db.chunks(PLOT_WIDTH).enumerate() {
157 for (column, &db) in levels.iter().enumerate() {
158 canvas.set(LEFT + column, plot_top + row, colour(db));
159 }
160 }
161 canvas.outline(LEFT - 1, plot_top - 1, PLOT_WIDTH + 2, PANEL_HEIGHT + 2);
162 for &tick in frequency_ticks(spec) {
163 let row = frequency_row(spec, tick);
164 let y = plot_top + row;
165 canvas.hline(LEFT - 4, y, 3);
166 let label = hz_label(tick);
167 let x = LEFT - 5 - 8 * label.len();
168 canvas.text(x, y.saturating_sub(3), &label);
169 }
170 }
171 canvas.text(0, 2, "Hz");
172 self.time_axis(&mut canvas);
173 colour_bar(&mut canvas);
174 canvas.pixels
175 }
176
177 fn time_axis(&self, canvas: &mut Canvas) {
178 let span = self.span.get();
179 let start = self.start.get();
180 let step = if span >= 2.0 { 1.0 } else { 0.1 };
181 let first = (start / step).ceil();
182 let mut index = 0.0;
183 while (first + index) * step <= start + span + 1e-9 {
184 let time = (first + index) * step;
185 index += 1.0;
186 let offset = if span > 0.0 {
187 (time - start) / span
188 } else {
189 0.0
190 };
191 let column = cast::round_to_usize(offset * cast::usize_to_f64(PLOT_WIDTH - 1))
192 .unwrap_or(0)
193 .min(PLOT_WIDTH - 1);
194 let x = LEFT + column;
195 canvas.vline(x, AXIS_TOP, 4);
196 let label = if step < 1.0 {
197 format!("{time:.1}")
198 } else {
199 format!("{time:.0}")
200 };
201 let width = 8 * label.len();
202 canvas.text(
203 x.saturating_sub(width / 2).min(WIDTH - RIGHT - width),
204 AXIS_TOP + 6,
205 &label,
206 );
207 }
208 canvas.text(WIDTH - RIGHT + 8, AXIS_TOP + 6, "time s");
209 }
210
211 pub fn to_png(&self) -> Result<Vec<u8>, png::EncodingError> {
218 let mut bytes = Vec::new();
219 let width = u32::try_from(WIDTH).unwrap_or(u32::MAX);
220 let height = u32::try_from(HEIGHT).unwrap_or(u32::MAX);
221 let mut encoder = png::Encoder::new(&mut bytes, width, height);
222 encoder.set_color(png::ColorType::Rgb);
223 encoder.set_depth(png::BitDepth::Eight);
224 let mut writer = encoder.write_header()?;
225 writer.write_image_data(&self.pixels())?;
226 writer.finish()?;
227 Ok(bytes)
228 }
229}
230
231fn passage(samples: &[f64], rate: SampleRate) -> (usize, usize) {
233 let window = cast::round_to_usize(PASSAGE_S.get() * rate.hz().get()).unwrap_or(usize::MAX);
234 if samples.len() <= window {
235 return (0, samples.len());
236 }
237 let params = StftParams::fixed(PANELS[0].frame_len);
238 let hop = params.hop();
239 let active = ActiveFrames::of(samples, params);
240 let mut prefix = vec![0_usize; active.len() + 1];
242 for frame in 0..active.len() {
243 prefix[frame + 1] = prefix[frame] + usize::from(active.is_active(frame));
244 }
245 let span = window / hop;
246 let last_start = (samples.len() - window) / hop;
247 let (mut best, mut best_count) = (0, 0);
248 for start in 0..=last_start {
249 let end = (start + span).min(active.len());
250 let count = prefix[end] - prefix[start.min(end)];
251 if count > best_count {
252 (best, best_count) = (start, count);
253 }
254 }
255 (best * hop, window)
256}
257
258fn panel_db(
260 samples: &[f64],
261 rate: SampleRate,
262 spec: PanelSpec,
263 first: usize,
264 len: usize,
265) -> Vec<f64> {
266 let params = StftParams::fixed(spec.frame_len);
267 let window_sum: f64 = Window::Hann.samples(spec.frame_len).iter().sum();
268 let bin_hz = params.bin_width(rate).get();
269 let rows = row_bins(spec, bin_hz, params.bins());
270 let mut stft = Stft::new(params);
271 let mut peak = 0.0_f64;
272 for frame in 0..params.frame_count(samples.len()) {
273 for value in stft.frame(samples, frame).iter().skip(1) {
275 peak = peak.max(value.norm() / window_sum);
276 }
277 }
278 let mut columns = Stft::new(params.with_hop(1).unwrap_or(params));
280 let mut db = vec![-RANGE_DB.get(); PLOT_WIDTH * PLOT_HEIGHT];
281 for column in 0..PLOT_WIDTH {
282 let centre = first
283 + cast::floor_to_usize(
284 (cast::usize_to_f64(column) + 0.5) * cast::usize_to_f64(len)
285 / cast::usize_to_f64(PLOT_WIDTH),
286 )
287 .unwrap_or(0);
288 let magnitudes: Vec<f64> = columns
290 .frame(samples, centre)
291 .iter()
292 .map(|value| value.norm() / window_sum)
293 .collect();
294 for (row, bins) in rows.iter().enumerate() {
295 db[row * PLOT_WIDTH + column] = level_db(bins.magnitude(&magnitudes), peak);
296 }
297 }
298 db
299}
300
301fn level_db(magnitude: f64, peak: f64) -> f64 {
304 let floor = -RANGE_DB.get();
305 if magnitude > 0.0 && peak > 0.0 {
306 (20.0 * (magnitude / peak).log10()).clamp(floor, 0.0)
307 } else {
308 floor
309 }
310}
311
312#[derive(Clone, Copy, Debug)]
314enum RowBins {
315 Max(usize, usize),
317 Between(usize, f64),
320 Empty,
322}
323
324impl RowBins {
325 fn magnitude(self, magnitudes: &[f64]) -> f64 {
326 match self {
327 Self::Max(first, last) => magnitudes
328 .get(first..=last)
329 .unwrap_or(&[])
330 .iter()
331 .copied()
332 .fold(0.0, f64::max),
333 Self::Between(k, fraction) => match (magnitudes.get(k), magnitudes.get(k + 1)) {
334 (Some(a), Some(b)) => a + fraction * (b - a),
335 _ => 0.0,
336 },
337 Self::Empty => 0.0,
338 }
339 }
340}
341
342fn row_bins(spec: PanelSpec, bin_hz: f64, bins: usize) -> Vec<RowBins> {
345 let height = cast::usize_to_f64(PLOT_HEIGHT);
346 let edge = |position: f64| {
347 spec.lo.get() * (spec.hi.get() / spec.lo.get()).powf(1.0 - position / height)
348 };
349 (0..PLOT_HEIGHT)
350 .map(|row| {
351 let r = cast::usize_to_f64(row);
352 let (low, high) = (edge(r + 1.0), edge(r));
353 let first = cast::ceil_to_usize(low / bin_hz).unwrap_or(0);
354 let last = cast::ceil_to_usize(high / bin_hz)
355 .unwrap_or(0)
356 .saturating_sub(1);
357 if first >= bins {
358 RowBins::Empty
359 } else if first <= last {
360 RowBins::Max(first, last.min(bins - 1))
361 } else {
362 let centre = (low * high).sqrt() / bin_hz;
363 let k = cast::floor_to_usize(centre).unwrap_or(0);
364 RowBins::Between(k, centre - cast::usize_to_f64(k))
365 }
366 })
367 .collect()
368}
369
370fn frequency_row(spec: PanelSpec, frequency: f64) -> usize {
372 let position = (frequency / spec.lo.get()).ln() / (spec.hi.get() / spec.lo.get()).ln();
373 cast::round_to_usize((1.0 - position) * cast::usize_to_f64(PLOT_HEIGHT))
374 .unwrap_or(0)
375 .min(PLOT_HEIGHT - 1)
376}
377
378fn frequency_ticks(spec: PanelSpec) -> &'static [f64] {
379 if spec.hi.get() > 1000.0 {
380 &[
381 20.0, 50.0, 100.0, 200.0, 500.0, 1000.0, 2000.0, 5000.0, 10_000.0, 20_000.0,
382 ]
383 } else {
384 &[20.0, 30.0, 50.0, 100.0, 200.0, 300.0, 500.0]
385 }
386}
387
388fn hz_label(frequency: f64) -> String {
389 if frequency >= 1000.0 {
390 format!("{}k", frequency / 1000.0)
391 } else {
392 format!("{frequency}")
393 }
394}
395
396fn colour(db: f64) -> [u8; 3] {
398 let colour = colorous::VIRIDIS.eval_continuous((db / RANGE_DB.get() + 1.0).clamp(0.0, 1.0));
399 [colour.r, colour.g, colour.b]
400}
401
402fn colour_bar(canvas: &mut Canvas) {
404 let x = WIDTH - RIGHT + 8;
405 let (top, bottom) = (TITLE, AXIS_TOP - 1);
406 let height = bottom - top;
407 for y in top..=bottom {
408 let db = -RANGE_DB.get() * cast::usize_to_f64(y - top) / cast::usize_to_f64(height);
409 for dx in 0..16 {
410 canvas.set(x + dx, y, colour(db));
411 }
412 }
413 canvas.outline(x - 1, top - 1, 18, height + 3);
414 canvas.text(x + 20, top, "0 dB");
415 canvas.text(x + 20, top + height / 2 - 4, "-40");
416 canvas.text(x + 20, bottom - 7, "-80");
417}
418
419struct Canvas {
420 pixels: Vec<u8>,
421}
422
423impl Canvas {
424 fn new() -> Self {
425 Self {
426 pixels: BACKGROUND.repeat(WIDTH * HEIGHT),
427 }
428 }
429
430 fn set(&mut self, x: usize, y: usize, colour: [u8; 3]) {
431 if x < WIDTH && y < HEIGHT {
432 let index = 3 * (y * WIDTH + x);
433 self.pixels[index..index + 3].copy_from_slice(&colour);
434 }
435 }
436
437 fn hline(&mut self, x: usize, y: usize, len: usize) {
438 for dx in 0..len {
439 self.set(x + dx, y, INK);
440 }
441 }
442
443 fn vline(&mut self, x: usize, y: usize, len: usize) {
444 for dy in 0..len {
445 self.set(x, y + dy, INK);
446 }
447 }
448
449 fn outline(&mut self, x: usize, y: usize, width: usize, height: usize) {
451 for dx in 0..width {
452 self.set(x + dx, y, FRAME);
453 self.set(x + dx, y + height - 1, FRAME);
454 }
455 for dy in 0..height {
456 self.set(x, y + dy, FRAME);
457 self.set(x + width - 1, y + dy, FRAME);
458 }
459 }
460
461 fn text(&mut self, x: usize, y: usize, text: &str) {
463 for (index, character) in text.chars().enumerate() {
464 let glyph = font8x8::BASIC_FONTS.get(character).unwrap_or([0; 8]);
465 for (dy, bits) in glyph.iter().enumerate() {
466 for dx in 0..8 {
467 if bits >> dx & 1 == 1 {
468 self.set(x + 8 * index + dx, y + dy, INK);
469 }
470 }
471 }
472 }
473 }
474}