scope/display/
spectroscope.rs

1use std::collections::VecDeque;
2
3use crossterm::event::{Event, KeyCode};
4use ratatui::{widgets::{Axis, GraphType}, style::Style, text::Span};
5
6use crate::input::Matrix;
7
8use super::{update_value_i, DataSet, Dimension, DisplayMode, GraphConfig};
9
10use rustfft::{FftPlanner, num_complex::Complex};
11
12#[derive(Default)]
13pub struct Spectroscope {
14	pub sampling_rate: u32,
15	pub buffer_size: u32,
16	pub average: u32,
17	pub buf: Vec<VecDeque<Vec<f64>>>,
18	pub window: bool,
19	pub log_y: bool,
20}
21
22fn magnitude(c: Complex<f64>) -> f64 {
23	let squared = (c.re * c.re) + (c.im * c.im);
24	squared.sqrt()
25}
26
27// got this from https://github.com/phip1611/spectrum-analyzer/blob/3c079ec2785b031d304bb381ff5f5fe04e6bcf71/src/windows.rs#L40
28pub fn hann_window(samples: &[f64]) -> Vec<f64> {
29	let mut windowed_samples = Vec::with_capacity(samples.len());
30	let samples_len = samples.len() as f64;
31	for (i, sample) in samples.iter().enumerate() {
32		let two_pi_i = 2.0 * std::f64::consts::PI * i as f64;
33		let idontknowthename = (two_pi_i / samples_len).cos();
34		let multiplier = 0.5 * (1.0 - idontknowthename);
35		windowed_samples.push(sample * multiplier)
36	}
37	windowed_samples
38}
39
40#[cfg(feature = "app")]
41impl From<&crate::cfg::SourceOptions> for Spectroscope {
42	fn from(value: &crate::cfg::SourceOptions) -> Self {
43		Spectroscope {
44			sampling_rate: value.sample_rate,
45			buffer_size: value.buffer,
46			average: 1, buf: Vec::new(),
47			window: false,
48			log_y: true,
49		}
50	}
51}
52
53impl DisplayMode for Spectroscope {
54	fn mode_str(&self) -> &'static str {
55		"spectro"
56	}
57
58	fn channel_name(&self, index: usize) -> String {
59		match index {
60			0 => "L".into(),
61			1 => "R".into(),
62			_ => format!("{}", index),
63		}
64	}
65
66	fn header(&self, _: &GraphConfig) -> String {
67		let window_marker = if self.window { "-|-" } else { "---" };
68		if self.average <= 1 {
69			format!("live  {}  {:.3}Hz bins", window_marker, self.sampling_rate as f64 / self.buffer_size as f64)
70		} else {
71			format!(
72				"{}x avg ({:.1}s)  {}  {:.3}Hz bins",
73				self.average,
74				(self.average * self.buffer_size) as f64 / self.sampling_rate as f64,
75				window_marker,
76				self.sampling_rate as f64 / (self.buffer_size * self.average) as f64,
77			)
78		}
79	}
80
81	fn axis(&self, cfg: &GraphConfig, dimension: Dimension) -> Axis {
82		let (name, bounds) = match dimension {
83			Dimension::X => ("frequency -", [20.0f64.ln(), ((cfg.samples as f64 / cfg.width as f64) * 20000.0).ln()]),
84			Dimension::Y => (
85				if self.log_y { "| level" } else { "| amplitude" },
86				[0.0, cfg.scale * 7.5] // very arbitrary but good default
87			),
88			// TODO super arbitraty! wtf! also ugly inline ifs, get this thing together!
89		};
90		let mut a = Axis::default();
91		if cfg.show_ui { // TODO don't make it necessary to check show_ui inside here
92			a = a.title(Span::styled(name, Style::default().fg(cfg.labels_color)));
93		}
94		a.style(Style::default().fg(cfg.axis_color)).bounds(bounds)
95	}
96
97	fn process(&mut self, cfg: &GraphConfig, data: &Matrix<f64>) -> Vec<DataSet> {
98		if self.average == 0 { self.average = 1 } // otherwise fft breaks
99		if !cfg.pause {
100			for (i, chan) in data.iter().enumerate() {
101				if self.buf.len() <= i {
102					self.buf.push(VecDeque::new());
103				}
104				self.buf[i].push_back(chan.clone());
105				while self.buf[i].len() > self.average as usize {
106					self.buf[i].pop_front();
107				}
108			}
109		}
110
111		let mut out = Vec::new();
112		let mut planner: FftPlanner<f64> = FftPlanner::new();
113		let sample_len = self.buffer_size * self.average;
114		let resolution = self.sampling_rate as f64 / sample_len as f64;
115		let fft = planner.plan_fft_forward(sample_len as usize);
116
117		for (n, chan_queue) in self.buf.iter().enumerate().rev() {
118			let mut chunk = chan_queue.iter().flatten().copied().collect::<Vec<f64>>();
119			if self.window {
120				chunk = hann_window(chunk.as_slice());
121			}
122			let mut max_val = *chunk.iter().max_by(|a, b| a.total_cmp(b)).expect("empty dataset?");
123			if max_val < 1. { max_val = 1.; }
124			let mut tmp : Vec<Complex<f64>> = chunk.iter().map(|x| Complex { re: *x / max_val, im: 0.0 }).collect();
125			fft.process(tmp.as_mut_slice());
126			out.push(DataSet::new(
127				Some(self.channel_name(n)),
128				tmp[..=tmp.len() / 2]
129					.iter()
130					.enumerate()
131					.map(|(i,x)| ((i as f64 * resolution).ln(), if self.log_y { magnitude(*x).ln() } else { magnitude(*x) }))
132					.collect(),
133				cfg.marker_type,
134				if cfg.scatter { GraphType::Scatter } else { GraphType::Line },
135				cfg.palette(n),
136			));
137		}
138
139		out
140	}
141
142	fn handle(&mut self, event: Event) {
143		if let Event::Key(key) = event {
144			match key.code {
145				KeyCode::PageUp   => update_value_i(&mut self.average, true, 1, 1., 1..65535),
146				KeyCode::PageDown => update_value_i(&mut self.average, false, 1, 1., 1..65535),
147				KeyCode::Char('w') => self.window = !self.window,
148				KeyCode::Char('l') => self.log_y = !self.log_y,
149				_ => {}
150			}
151		}
152	}
153
154	fn references(&self, cfg: &GraphConfig) -> Vec<DataSet> {
155		let lower = 0.; // if self.log_y { -(cfg.scale * 5.) } else { 0. };
156		let upper = cfg.scale * 7.5;
157		vec![
158			DataSet::new(None, vec![(0.0, 0.0), ((cfg.samples as f64).ln(), 0.0)], cfg.marker_type, GraphType::Line, cfg.axis_color), 
159
160			// TODO can we auto generate these? lol...
161			DataSet::new(None, vec![(20.0f64.ln(), lower), (20.0f64.ln(), upper)], cfg.marker_type, GraphType::Line, cfg.axis_color),
162			DataSet::new(None, vec![(30.0f64.ln(), lower), (30.0f64.ln(), upper)], cfg.marker_type, GraphType::Line, cfg.axis_color),
163			DataSet::new(None, vec![(40.0f64.ln(), lower), (40.0f64.ln(), upper)], cfg.marker_type, GraphType::Line, cfg.axis_color),
164			DataSet::new(None, vec![(50.0f64.ln(), lower), (50.0f64.ln(), upper)], cfg.marker_type, GraphType::Line, cfg.axis_color),
165			DataSet::new(None, vec![(60.0f64.ln(), lower), (60.0f64.ln(), upper)], cfg.marker_type, GraphType::Line, cfg.axis_color),
166			DataSet::new(None, vec![(70.0f64.ln(), lower), (70.0f64.ln(), upper)], cfg.marker_type, GraphType::Line, cfg.axis_color),
167			DataSet::new(None, vec![(80.0f64.ln(), lower), (80.0f64.ln(), upper)], cfg.marker_type, GraphType::Line, cfg.axis_color),
168			DataSet::new(None, vec![(90.0f64.ln(), lower), (90.0f64.ln(), upper)], cfg.marker_type, GraphType::Line, cfg.axis_color),
169			DataSet::new(None, vec![(100.0f64.ln(), lower), (100.0f64.ln(), upper)], cfg.marker_type, GraphType::Line, cfg.axis_color),
170			DataSet::new(None, vec![(200.0f64.ln(), lower), (200.0f64.ln(), upper)], cfg.marker_type, GraphType::Line, cfg.axis_color),
171			DataSet::new(None, vec![(300.0f64.ln(), lower), (300.0f64.ln(), upper)], cfg.marker_type, GraphType::Line, cfg.axis_color),
172			DataSet::new(None, vec![(400.0f64.ln(), lower), (400.0f64.ln(), upper)], cfg.marker_type, GraphType::Line, cfg.axis_color),
173			DataSet::new(None, vec![(500.0f64.ln(), lower), (500.0f64.ln(), upper)], cfg.marker_type, GraphType::Line, cfg.axis_color),
174			DataSet::new(None, vec![(600.0f64.ln(), lower), (600.0f64.ln(), upper)], cfg.marker_type, GraphType::Line, cfg.axis_color),
175			DataSet::new(None, vec![(700.0f64.ln(), lower), (700.0f64.ln(), upper)], cfg.marker_type, GraphType::Line, cfg.axis_color),
176			DataSet::new(None, vec![(800.0f64.ln(), lower), (800.0f64.ln(), upper)], cfg.marker_type, GraphType::Line, cfg.axis_color),
177			DataSet::new(None, vec![(900.0f64.ln(), lower), (900.0f64.ln(), upper)], cfg.marker_type, GraphType::Line, cfg.axis_color),
178			DataSet::new(None, vec![(1000.0f64.ln(), lower), (1000.0f64.ln(), upper)], cfg.marker_type, GraphType::Line, cfg.axis_color),
179			DataSet::new(None, vec![(2000.0f64.ln(), lower), (2000.0f64.ln(), upper)], cfg.marker_type, GraphType::Line, cfg.axis_color),
180			DataSet::new(None, vec![(3000.0f64.ln(), lower), (3000.0f64.ln(), upper)], cfg.marker_type, GraphType::Line, cfg.axis_color),
181			DataSet::new(None, vec![(4000.0f64.ln(), lower), (4000.0f64.ln(), upper)], cfg.marker_type, GraphType::Line, cfg.axis_color),
182			DataSet::new(None, vec![(5000.0f64.ln(), lower), (5000.0f64.ln(), upper)], cfg.marker_type, GraphType::Line, cfg.axis_color),
183			DataSet::new(None, vec![(6000.0f64.ln(), lower), (6000.0f64.ln(), upper)], cfg.marker_type, GraphType::Line, cfg.axis_color),
184			DataSet::new(None, vec![(7000.0f64.ln(), lower), (7000.0f64.ln(), upper)], cfg.marker_type, GraphType::Line, cfg.axis_color),
185			DataSet::new(None, vec![(8000.0f64.ln(), lower), (8000.0f64.ln(), upper)], cfg.marker_type, GraphType::Line, cfg.axis_color),
186			DataSet::new(None, vec![(9000.0f64.ln(), lower), (9000.0f64.ln(), upper)], cfg.marker_type, GraphType::Line, cfg.axis_color),
187			DataSet::new(None, vec![(10000.0f64.ln(), lower), (10000.0f64.ln(), upper)], cfg.marker_type, GraphType::Line, cfg.axis_color),
188			DataSet::new(None, vec![(20000.0f64.ln(), lower), (20000.0f64.ln(), upper)], cfg.marker_type, GraphType::Line, cfg.axis_color),
189		]
190	}
191}