1use crate::signal::{Spectrogram, ndarray};
2use eframe::egui;
3use egui_plot::{GridInput, GridMark, Plot, PlotImage, PlotPoint};
4
5use crate::plot::native_options_any_thread;
6use crate::plot::save::SavePlotState;
7
8const TURBO_PALETTE: &[egui::Color32] = &[
9 egui::Color32::from_rgb(48, 18, 59),
10 egui::Color32::from_rgb(72, 42, 121),
11 egui::Color32::from_rgb(55, 87, 176),
12 egui::Color32::from_rgb(35, 135, 207),
13 egui::Color32::from_rgb(30, 178, 185),
14 egui::Color32::from_rgb(45, 209, 145),
15 egui::Color32::from_rgb(88, 223, 98),
16 egui::Color32::from_rgb(164, 227, 39),
17 egui::Color32::from_rgb(228, 207, 14),
18 egui::Color32::from_rgb(250, 170, 10),
19 egui::Color32::from_rgb(246, 117, 6),
20 egui::Color32::from_rgb(220, 62, 2),
21 egui::Color32::from_rgb(122, 4, 2),
22];
23
24#[derive(Clone, Copy, Debug, PartialEq, Eq)]
25enum SpectrogramScale {
26 Linear,
27 Log,
28}
29
30#[derive(Clone, Copy, Debug, PartialEq)]
31pub struct SpectrogramPlotOptions {
32 pub log_freq: bool,
33 pub db_floor: f64,
34}
35
36impl Default for SpectrogramPlotOptions {
37 fn default() -> Self {
38 Self {
39 log_freq: false,
40 db_floor: -80.0,
41 }
42 }
43}
44
45fn apply_colormap(t: f32, palette: &[egui::Color32]) -> egui::Color32 {
46 if palette.len() < 2 {
47 return palette.first().copied().unwrap_or(egui::Color32::BLACK);
48 }
49 let v = t.clamp(0.0, 1.0);
50 let idx = v * (palette.len() - 1) as f32;
51 let lo = idx.floor() as usize;
52 let hi = (lo + 1).min(palette.len() - 1);
53 let t = idx - lo as f32;
54 let c0 = palette[lo];
55 let c1 = palette[hi];
56 egui::Color32::from_rgb(
57 (c0.r() as f32 * (1.0 - t) + c1.r() as f32 * t).round() as u8,
58 (c0.g() as f32 * (1.0 - t) + c1.g() as f32 * t).round() as u8,
59 (c0.b() as f32 * (1.0 - t) + c1.b() as f32 * t).round() as u8,
60 )
61}
62
63fn draw_colorbar(ui: &mut egui::Ui, db_min: f64, db_max: f64) {
64 const BAR_WIDTH: f32 = 16.0;
65 const TICK_LEN: f32 = 5.0;
66 const LABEL_GAP: f32 = 4.0;
67 const MARGIN_Y: f32 = 10.0;
68
69 let available = ui.available_size();
70 let (response, painter) = ui.allocate_painter(available, egui::Sense::hover());
71
72 let bar_rect = egui::Rect::from_min_max(
73 egui::pos2(response.rect.min.x, response.rect.min.y + MARGIN_Y),
74 egui::pos2(
75 response.rect.min.x + BAR_WIDTH,
76 response.rect.max.y - MARGIN_Y,
77 ),
78 );
79 let bar_height = bar_rect.height();
80
81 for i in 0..128usize {
82 let t = 1.0 - (i as f32 + 0.5) / 128.0;
83 let y_top = bar_rect.min.y + i as f32 * bar_height / 128.0;
84 let y_bot = bar_rect.min.y + (i + 1) as f32 * bar_height / 128.0;
85 let strip = egui::Rect::from_min_max(
86 egui::pos2(bar_rect.min.x, y_top),
87 egui::pos2(bar_rect.max.x, y_bot),
88 );
89 painter.rect_filled(strip, 0.0, apply_colormap(t, TURBO_PALETTE));
90 }
91
92 let stroke = egui::Stroke::new(1.0, ui.visuals().text_color());
93 painter.rect_stroke(bar_rect, 0.0, stroke, egui::StrokeKind::Outside);
94
95 let font_id = egui::TextStyle::Small.resolve(ui.style());
96 let text_color = ui.visuals().text_color();
97 for tick in 0..=4 {
98 let t = tick as f32 / 4.0;
99 let db = db_min + t as f64 * (db_max - db_min);
100 let y = bar_rect.min.y + (1.0 - t) * bar_height;
101 painter.line_segment(
102 [
103 egui::pos2(bar_rect.max.x, y),
104 egui::pos2(bar_rect.max.x + TICK_LEN, y),
105 ],
106 egui::Stroke::new(1.0, text_color),
107 );
108 painter.text(
109 egui::pos2(bar_rect.max.x + TICK_LEN + LABEL_GAP, y),
110 egui::Align2::LEFT_CENTER,
111 format!("{db:.0} dB"),
112 font_id.clone(),
113 text_color,
114 );
115 }
116}
117
118fn format_hz(freq_hz: f64) -> String {
119 if freq_hz >= 1000.0 {
120 format!("{:.2} kHz", freq_hz / 1000.0)
121 } else {
122 format!("{freq_hz:.0} Hz")
123 }
124}
125
126fn format_hz_tick(freq_hz: f64) -> String {
127 if freq_hz >= 1000.0 {
128 format!("{:.0} kHz", freq_hz / 1000.0)
129 } else {
130 format!("{freq_hz:.0} Hz")
131 }
132}
133
134fn spectrum_grid_spacer(input: GridInput) -> Vec<GridMark> {
135 let (lo, hi) = (input.bounds.0, input.bounds.1);
136 let mut marks = Vec::new();
137
138 let start_exp = lo.floor() as i32 - 1;
139 let end_exp = hi.ceil() as i32 + 1;
140
141 for exp in start_exp..=end_exp {
142 let decade = 10f64.powi(exp);
143 for mult in 1..10 {
144 let hz = decade * mult as f64;
145 let log_pos = hz.log10();
146 if log_pos < lo || log_pos > hi {
147 continue;
148 }
149 let step_size = match mult {
150 1 => 1.0,
151 2 | 5 => 0.3,
152 _ => 0.1,
153 };
154 marks.push(GridMark {
155 value: log_pos,
156 step_size,
157 });
158 }
159 }
160
161 marks
162}
163
164fn positive_freq_bins(spec: &Spectrogram) -> Vec<f64> {
165 spec.freq_bins()
166 .iter()
167 .copied()
168 .filter(|&freq| freq > 0.0)
169 .collect()
170}
171
172fn make_log_freq_grid(freq_bins: &[f64]) -> Vec<f64> {
173 if freq_bins.is_empty() {
174 return Vec::new();
175 }
176 if freq_bins.len() == 1 {
177 return vec![freq_bins[0]];
178 }
179
180 let start = freq_bins[0].log10();
181 let end = freq_bins[freq_bins.len() - 1].log10();
182 (0..freq_bins.len())
183 .map(|idx| {
184 let t = idx as f64 / (freq_bins.len() - 1) as f64;
185 10.0_f64.powf(start + t * (end - start))
186 })
187 .collect()
188}
189
190#[derive(Clone, Copy, Debug, PartialEq)]
191struct InterpolationStep {
192 lower: usize,
193 upper: usize,
194 weight: f64,
195}
196
197fn make_interpolation_plan(source_freqs: &[f64], target_freqs: &[f64]) -> Vec<InterpolationStep> {
198 target_freqs
199 .iter()
200 .map(|&target_freq| {
201 match source_freqs.binary_search_by(|probe| probe.total_cmp(&target_freq)) {
202 Ok(index) => InterpolationStep {
203 lower: index,
204 upper: index,
205 weight: 0.0,
206 },
207 Err(0) => InterpolationStep {
208 lower: 0,
209 upper: 0,
210 weight: 0.0,
211 },
212 Err(index) if index >= source_freqs.len() => {
213 let last = source_freqs.len() - 1;
214 InterpolationStep {
215 lower: last,
216 upper: last,
217 weight: 0.0,
218 }
219 }
220 Err(upper) => {
221 let lower = upper - 1;
222 let f0 = source_freqs[lower];
223 let f1 = source_freqs[upper];
224 let weight = if f1 > f0 {
225 (target_freq - f0) / (f1 - f0)
226 } else {
227 0.0
228 };
229 InterpolationStep {
230 lower,
231 upper,
232 weight,
233 }
234 }
235 }
236 })
237 .collect()
238}
239
240#[cfg(test)]
241fn interpolate_row(source_freqs: &[f64], source_values: &[f64], target_freq: f64) -> f64 {
242 let step = make_interpolation_plan(source_freqs, &[target_freq])[0];
243 let v0 = source_values[step.lower];
244 let v1 = source_values[step.upper];
245 v0 + step.weight * (v1 - v0)
246}
247
248fn make_linear_values(
249 spec: &Spectrogram,
250 mag_db: &ndarray::Array3<f64>,
251) -> (Vec<Vec<f64>>, Vec<f64>) {
252 let num_channels = spec.num_channels();
253 let num_frames = spec.num_frames();
254 let num_freq_bins = spec.num_freq_bins();
255
256 let values = (0..num_channels)
257 .map(|ch| {
258 let mut v = Vec::with_capacity(num_frames * num_freq_bins);
259 for freq_bin in 0..num_freq_bins {
260 for frame in 0..num_frames {
261 v.push(mag_db[[ch, frame, freq_bin]]);
262 }
263 }
264 v
265 })
266 .collect();
267
268 (values, spec.freq_bins().to_vec())
269}
270
271fn make_log_values(spec: &Spectrogram, mag_db: &ndarray::Array3<f64>) -> (Vec<Vec<f64>>, Vec<f64>) {
272 let source_freqs = positive_freq_bins(spec);
273 if source_freqs.len() < 2 {
274 return make_linear_values(spec, mag_db);
275 }
276
277 let target_freqs = make_log_freq_grid(&source_freqs);
278 let interpolation_plan = make_interpolation_plan(&source_freqs, &target_freqs);
279 let num_channels = spec.num_channels();
280 let num_frames = spec.num_frames();
281
282 let values = (0..num_channels)
283 .map(|ch| {
284 let mut v = Vec::with_capacity(num_frames * target_freqs.len());
285 for step in &interpolation_plan {
286 for frame in 0..num_frames {
287 let v0 = mag_db[[ch, frame, step.lower + 1]];
288 let v1 = mag_db[[ch, frame, step.upper + 1]];
289 v.push(v0 + step.weight * (v1 - v0));
290 }
291 }
292 v
293 })
294 .collect();
295
296 (values, target_freqs)
297}
298
299fn values_to_images(
300 values: &[Vec<f64>],
301 num_frames: usize,
302 num_freq_bins: usize,
303 db_floor: f64,
304 db_peak: f64,
305) -> Vec<egui::ColorImage> {
306 let db_span = (db_peak - db_floor).max(f64::EPSILON);
307 values
308 .iter()
309 .map(|channel_values| {
310 let mut pixels = Vec::with_capacity(num_frames * num_freq_bins);
311 for freq_bin in (0..num_freq_bins).rev() {
312 let row_start = freq_bin * num_frames;
313 for frame in 0..num_frames {
314 let db = channel_values[row_start + frame];
315 let t = ((db - db_floor) / db_span).clamp(0.0, 1.0) as f32;
316 pixels.push(apply_colormap(t, TURBO_PALETTE));
317 }
318 }
319 egui::ColorImage::new([num_frames, num_freq_bins], pixels)
320 })
321 .collect()
322}
323
324pub(crate) struct SpectrogramPlot {
325 values: Vec<Vec<f64>>,
326 images: Vec<egui::ColorImage>,
327 textures: Vec<Option<egui::TextureHandle>>,
328 num_frames: usize,
329 num_freq_bins: usize,
330 db_floor: f64,
331 db_peak: f64,
332 pos: PlotPoint,
333 tile_size: (f32, f32),
334 image_center: PlotPoint,
335 image_size: egui::Vec2,
336 num_channels: usize,
337 current_channel: usize,
338 channel_labels: Vec<String>,
339 save: SavePlotState,
340 scale: SpectrogramScale,
341}
342
343impl SpectrogramPlot {
344 pub(crate) fn new(spec: &Spectrogram, title: &str, options: SpectrogramPlotOptions) -> Self {
345 let scale = if options.log_freq {
346 SpectrogramScale::Log
347 } else {
348 SpectrogramScale::Linear
349 };
350 let db_floor = options.db_floor;
351 let num_channels = spec.num_channels();
352 let num_frames = spec.num_frames();
353 let frame_times = spec.frame_times();
354 let mag_db = spec.amplitude_spectrum_db(db_floor);
355 let db_peak = mag_db.iter().copied().fold(f64::NEG_INFINITY, f64::max);
356
357 let (values, plotted_freqs_hz) = match scale {
358 SpectrogramScale::Linear => make_linear_values(spec, &mag_db),
359 SpectrogramScale::Log => make_log_values(spec, &mag_db),
360 };
361 let num_freq_bins = plotted_freqs_hz.len();
362 let images = values_to_images(&values, num_frames, num_freq_bins, db_floor, db_peak);
363
364 let dt = if num_frames > 1 {
365 (frame_times[num_frames - 1] - frame_times[0]) / (num_frames - 1) as f64
366 } else {
367 spec.hop_size() as f64 / spec.sample_rate()
368 };
369
370 let plotted_y: Vec<f64> = match scale {
371 SpectrogramScale::Linear => plotted_freqs_hz.clone(),
372 SpectrogramScale::Log => plotted_freqs_hz.iter().map(|freq| freq.log10()).collect(),
373 };
374 let dy = if num_freq_bins > 1 {
375 (plotted_y[num_freq_bins - 1] - plotted_y[0]) / (num_freq_bins - 1) as f64
376 } else if scale == SpectrogramScale::Log {
377 0.1
378 } else {
379 spec.sample_rate() / spec.window_size() as f64
380 };
381
382 let pos = PlotPoint {
383 x: frame_times.first().copied().unwrap_or(0.0) - dt / 2.0,
384 y: plotted_y.first().copied().unwrap_or(0.0) - dy / 2.0,
385 };
386 let image_size = egui::vec2(
387 (num_frames as f64 * dt) as f32,
388 (num_freq_bins as f64 * dy) as f32,
389 );
390 let image_center = PlotPoint::new(
391 pos.x + f64::from(image_size.x) / 2.0,
392 pos.y + f64::from(image_size.y) / 2.0,
393 );
394
395 let channel_labels: Vec<String> = (0..num_channels)
396 .map(|ch| match spec.channel_label(ch) {
397 Some(label) => crate::plot::legend::clip_label(label),
398 None => format!("Channel {ch}"),
399 })
400 .collect();
401
402 Self {
403 values,
404 images,
405 textures: vec![None; num_channels],
406 num_frames,
407 num_freq_bins,
408 db_floor,
409 db_peak,
410 pos,
411 tile_size: (dt as f32, dy as f32),
412 image_center,
413 image_size,
414 num_channels,
415 current_channel: 0,
416 channel_labels,
417 save: SavePlotState::new(title),
418 scale,
419 }
420 }
421
422 fn freq_from_plot_y(&self, y: f64) -> f64 {
423 match self.scale {
424 SpectrogramScale::Linear => y,
425 SpectrogramScale::Log => 10.0_f64.powf(y),
426 }
427 }
428
429 fn axis_label(&self) -> &'static str {
430 match self.scale {
431 SpectrogramScale::Linear => "Frequency (Hz)",
432 SpectrogramScale::Log => "Frequency (Hz, log)",
433 }
434 }
435}
436
437impl eframe::App for SpectrogramPlot {
438 fn ui(&mut self, ui: &mut egui::Ui, _frame: &mut eframe::Frame) {
439 egui::Panel::top("controls").show(ui, |ui| {
440 self.save.show_panel(ui);
441 if self.num_channels > 1 {
442 ui.horizontal(|ui| {
443 for ch in 0..self.num_channels {
444 ui.selectable_value(
445 &mut self.current_channel,
446 ch,
447 &self.channel_labels[ch],
448 );
449 }
450 });
451 }
452 });
453
454 egui::Panel::right("colorbar")
455 .exact_size(80.0)
456 .resizable(false)
457 .show(ui, |ui| {
458 draw_colorbar(ui, self.db_floor, self.db_peak);
459 });
460
461 egui::CentralPanel::default().show(ui, |ui| {
462 let texture = self.textures[self.current_channel].get_or_insert_with(|| {
463 ui.ctx().load_texture(
464 format!(
465 "spectrogram-{:?}-channel-{}",
466 self.scale, self.current_channel
467 ),
468 self.images[self.current_channel].clone(),
469 egui::TextureOptions::NEAREST,
470 )
471 });
472 let image = PlotImage::new(
473 self.channel_labels[self.current_channel].clone(),
474 texture.id(),
475 self.image_center,
476 self.image_size,
477 );
478
479 let scale = self.scale;
480 let mut plot = Plot::new("spectrogram")
481 .x_axis_label("Time (s)")
482 .y_axis_label(self.axis_label())
483 .y_axis_formatter(move |mark, _range| {
484 let freq_hz = match scale {
485 SpectrogramScale::Linear => mark.value,
486 SpectrogramScale::Log => 10.0_f64.powf(mark.value),
487 };
488 format_hz_tick(freq_hz)
489 })
490 .set_margin_fraction(egui::vec2(0.0, 0.0));
491 if self.scale == SpectrogramScale::Log {
492 plot = plot.y_grid_spacer(spectrum_grid_spacer);
493 }
494
495 let mut hovered_coord = None;
496 let inner = plot.show(ui, |plot_ui| {
497 if plot_ui.response().hovered() {
498 hovered_coord = plot_ui.pointer_coordinate();
499 }
500 plot_ui.image(image);
501 });
502
503 if let Some(coord) = hovered_coord {
504 let freq_hz = self.freq_from_plot_y(coord.y);
505 let frame = ((coord.x - self.pos.x) / self.tile_size.0 as f64).floor() as isize;
506 let freq_bin = ((coord.y - self.pos.y) / self.tile_size.1 as f64).floor() as isize;
507 let db_value = if frame >= 0 && (frame as usize) < self.num_frames && freq_bin >= 0
508 {
509 if (freq_bin as usize) < self.num_freq_bins {
510 let idx = freq_bin as usize * self.num_frames + frame as usize;
511 Some(self.values[self.current_channel][idx])
512 } else {
513 None
514 }
515 } else {
516 None
517 };
518
519 let text = match db_value {
520 Some(db) => format!(
521 "t = {:.3} s\nf = {}\n{:.1} dB",
522 coord.x,
523 format_hz(freq_hz),
524 db
525 ),
526 None => format!("t = {:.3} s\nf = {}", coord.x, format_hz(freq_hz)),
527 };
528 let mut tooltip = egui::Tooltip::always_open(
529 ui.ctx().clone(),
530 inner.response.layer_id,
531 inner.response.id.with("hover_tooltip"),
532 egui::PopupAnchor::Pointer,
533 );
534 tooltip.popup = tooltip.popup.gap(12.0).width(f32::INFINITY);
535 tooltip.show(|ui: &mut egui::Ui| {
536 ui.add(egui::Label::new(text).wrap_mode(egui::TextWrapMode::Extend));
537 });
538 }
539 });
540
541 self.save.handle_screenshot(ui.ctx());
542 }
543}
544
545pub fn show_spectrogram(
546 title: &str,
547 spectrogram: &Spectrogram,
548 options: SpectrogramPlotOptions,
549) -> Result<(), crate::plot::Error> {
550 Ok(eframe::run_native(
551 title,
552 native_options_any_thread(),
553 Box::new(|_cc| Ok(Box::new(SpectrogramPlot::new(spectrogram, title, options)))),
554 )?)
555}
556
557#[cfg(test)]
558mod tests {
559 use crate::dsp::WindowFn;
560 use crate::signal::{Spectrogram, ndarray::Array3};
561 use num::complex::Complex64;
562
563 use super::*;
564
565 fn make_spectrogram() -> Spectrogram {
566 let num_channels = 2;
567 let num_frames = 50;
568 let window_size = 32;
569 let num_freq_bins = window_size / 2 + 1;
570 let sample_rate = 16_000.0;
571 let data =
572 Array3::from_shape_fn((num_channels, num_frames, num_freq_bins), |(ch, f, b)| {
573 Complex64::new((ch as f64 + 1.0) * (f + b) as f64, 0.0)
574 });
575 Spectrogram::new(
576 data,
577 sample_rate,
578 window_size,
579 window_size / 2,
580 WindowFn::Rectangular,
581 )
582 .unwrap()
583 }
584
585 #[test]
586 fn log_grid_uses_positive_frequencies() {
587 let freqs = vec![62.5, 125.0, 250.0, 500.0, 1000.0];
588 let grid = make_log_freq_grid(&freqs);
589
590 assert_eq!(grid.len(), freqs.len());
591 assert!((grid[0] - 62.5).abs() < 1e-12);
592 assert!((grid[grid.len() - 1] - 1000.0).abs() < 1e-12);
593 assert!(grid.windows(2).all(|window| window[1] > window[0]));
594 }
595
596 #[test]
597 fn interpolation_hits_exact_bin_values() {
598 let freqs = [100.0, 200.0, 400.0];
599 let values = [-60.0, -20.0, -10.0];
600
601 assert!((interpolate_row(&freqs, &values, 200.0) + 20.0).abs() < 1e-12);
602 assert!((interpolate_row(&freqs, &values, 300.0) + 15.0).abs() < 1e-12);
603 }
604
605 #[test]
606 fn interpolation_plan_reuses_bin_mapping() {
607 let freqs = [100.0, 200.0, 400.0];
608 let plan = make_interpolation_plan(&freqs, &[100.0, 300.0, 500.0]);
609
610 assert_eq!(
611 plan,
612 vec![
613 InterpolationStep {
614 lower: 0,
615 upper: 0,
616 weight: 0.0
617 },
618 InterpolationStep {
619 lower: 1,
620 upper: 2,
621 weight: 0.5
622 },
623 InterpolationStep {
624 lower: 2,
625 upper: 2,
626 weight: 0.0
627 }
628 ]
629 );
630 }
631
632 #[test]
633 #[ignore]
634 fn test_show_spectrogram() {
635 show_spectrogram(
636 "Spectrogram Test",
637 &make_spectrogram(),
638 SpectrogramPlotOptions::default(),
639 )
640 .unwrap();
641 }
642}