spectrum

src/engine/audio.rs

@@ -2,6 +2,7 @@ use cpal::traits::{DeviceTrait, HostTrait, StreamTrait};
 use cpal::Stream;
 use crossbeam_channel::Receiver;
 use doux::{Engine, EngineMetrics};
+use rustfft::{num_complex::Complex, FftPlanner};
 use std::sync::atomic::{AtomicU32, Ordering};
 use std::sync::Arc;
 
@@ -50,6 +51,106 @@ impl ScopeBuffer {
     }
 }
 
+pub struct SpectrumBuffer {
+    pub bands: [AtomicU32; 32],
+}
+
+impl SpectrumBuffer {
+    pub fn new() -> Self {
+        Self {
+            bands: std::array::from_fn(|_| AtomicU32::new(0)),
+        }
+    }
+
+    pub fn write(&self, data: &[f32; 32]) {
+        for (atom, &val) in self.bands.iter().zip(data.iter()) {
+            atom.store(val.to_bits(), Ordering::Relaxed);
+        }
+    }
+
+    pub fn read(&self) -> [f32; 32] {
+        std::array::from_fn(|i| f32::from_bits(self.bands[i].load(Ordering::Relaxed)))
+    }
+}
+
+const FFT_SIZE: usize = 512;
+const NUM_BANDS: usize = 32;
+
+struct SpectrumAnalyzer {
+    ring: Vec<f32>,
+    pos: usize,
+    fft: Arc<dyn rustfft::Fft<f32>>,
+    window: [f32; FFT_SIZE],
+    scratch: Vec<Complex<f32>>,
+    band_edges: [usize; NUM_BANDS + 1],
+}
+
+impl SpectrumAnalyzer {
+    fn new(sample_rate: f32) -> Self {
+        let mut planner = FftPlanner::new();
+        let fft = planner.plan_fft_forward(FFT_SIZE);
+        let scratch_len = fft.get_inplace_scratch_len();
+
+        let window: [f32; FFT_SIZE] = std::array::from_fn(|i| {
+            0.5 * (1.0 - (2.0 * std::f32::consts::PI * i as f32 / (FFT_SIZE - 1) as f32).cos())
+        });
+
+        let nyquist = sample_rate / 2.0;
+        let min_freq: f32 = 20.0;
+        let log_min = min_freq.ln();
+        let log_max = nyquist.ln();
+        let band_edges: [usize; NUM_BANDS + 1] = std::array::from_fn(|i| {
+            let freq = (log_min + (log_max - log_min) * i as f32 / NUM_BANDS as f32).exp();
+            let bin = (freq * FFT_SIZE as f32 / sample_rate).round() as usize;
+            bin.min(FFT_SIZE / 2)
+        });
+
+        Self {
+            ring: vec![0.0; FFT_SIZE],
+            pos: 0,
+            fft,
+            window,
+            scratch: vec![Complex::default(); scratch_len],
+            band_edges,
+        }
+    }
+
+    fn feed(&mut self, samples: &[f32], output: &SpectrumBuffer) {
+        for &s in samples {
+            self.ring[self.pos] = s;
+            self.pos += 1;
+            if self.pos >= FFT_SIZE {
+                self.pos = 0;
+                self.run_fft(output);
+            }
+        }
+    }
+
+    fn run_fft(&mut self, output: &SpectrumBuffer) {
+        let mut buf: Vec<Complex<f32>> = (0..FFT_SIZE)
+            .map(|i| {
+                let idx = (self.pos + i) % FFT_SIZE;
+                Complex::new(self.ring[idx] * self.window[i], 0.0)
+            })
+            .collect();
+
+        self.fft.process_with_scratch(&mut buf, &mut self.scratch);
+
+        let mut bands = [0.0f32; NUM_BANDS];
+        for (band, mag) in bands.iter_mut().enumerate() {
+            let lo = self.band_edges[band];
+            let hi = self.band_edges[band + 1].max(lo + 1);
+            let sum: f32 = buf[lo..hi].iter().map(|c| c.norm()).sum();
+            let avg = sum / (hi - lo) as f32;
+            let amplitude = avg / (FFT_SIZE as f32 / 2.0);
+            let db = 20.0 * amplitude.max(1e-10).log10();
+            *mag = ((db + 60.0) / 60.0).clamp(0.0, 1.0);
+        }
+
+        output.write(&bands);
+    }
+}
+
 pub struct AudioStreamConfig {
     pub output_device: Option<String>,
     pub channels: u16,
@@ -60,6 +161,7 @@ pub fn build_stream(
     config: &AudioStreamConfig,
     audio_rx: Receiver<AudioCommand>,
     scope_buffer: Arc<ScopeBuffer>,
+    spectrum_buffer: Arc<SpectrumBuffer>,
     metrics: Arc<EngineMetrics>,
     initial_samples: Vec<doux::sample::SampleEntry>,
 ) -> Result<(Stream, f32), String> {
@@ -95,6 +197,8 @@ pub fn build_stream(
     let mut engine = Engine::new_with_metrics(sample_rate, channels, Arc::clone(&metrics));
     engine.sample_index = initial_samples;
 
+    let mut analyzer = SpectrumAnalyzer::new(sample_rate);
+
     let stream = device
         .build_output_stream(
             &stream_config,
@@ -128,6 +232,12 @@ pub fn build_stream(
                 engine.metrics.load.set_buffer_time(buffer_time_ns);
                 engine.process_block(data, &[], &[]);
                 scope_buffer.write(&engine.output);
+
+                // Feed mono mix to spectrum analyzer
+                let mono: Vec<f32> = engine.output.chunks(channels)
+                    .map(|ch| ch.iter().sum::<f32>() / channels as f32)
+                    .collect();
+                analyzer.feed(&mono, &spectrum_buffer);
             },
             |err| eprintln!("stream error: {err}"),
             None,