Skip to main content

scope/display/
spectroscope.rs

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