YrXtals/src/processor.rs

520 lines
17 KiB
Rust

use num_complex::Complex64;
use rustfft::{Fft, FftPlanner};
use std::collections::VecDeque;
use std::f64::consts::PI;
use std::sync::Arc;
use crate::gpu_dsp::GpuFft1D;
use crate::weave;
/// log-spaced spectrum frame plus the real cepstrum used for transient detection.
#[derive(Debug, Clone, Default)]
pub struct Spectrum {
pub freqs: Vec<f32>,
pub db: Vec<f32>,
pub cepstrum: Vec<f32>,
pub activity: Vec<f32>,
}
/// single-channel windowed-fft pipeline with optional gpu offload, log-spaced binning, and cepstral envelope blending.
pub struct Processor {
frame_size: usize,
sample_rate: u32,
device: wgpu::Device,
queue: wgpu::Queue,
cpu_fft: Option<Arc<dyn Fft<f64>>>,
cpu_cep_inv: Option<Arc<dyn Fft<f64>>>,
cpu_planner: FftPlanner<f64>,
gpu_fft: Option<GpuFft1D>,
gpu_blend: f32,
window: Vec<f64>,
buffer: Vec<Complex64>,
sample_freqs: Vec<f64>,
freqs_const: Vec<f64>,
num_bins: usize,
bin_alphas: Vec<f64>,
bin_smoothed: Vec<f64>,
bin_activity: Vec<f64>,
bin_reservoir: Vec<f64>,
bin_prev_delta: Vec<f64>,
bin_initialized: bool,
smoothing_tilt: f64,
smoothing_strength: f64,
smooth_splashes: bool,
history: VecDeque<Vec<f64>>,
smoothing_length: usize,
expand_ratio: f32,
expand_threshold: f32,
hpf_cutoff: f32,
granularity: i32,
detail: i32,
cepstral_strength: f32,
#[allow(dead_code)]
weave_caps: Option<weave::Caps>,
}
impl Processor {
/// allocates fft plans, the analysis window, and a default 26-bin log-spaced layout.
pub fn new(frame_size: usize, sample_rate: u32, device: wgpu::Device, queue: wgpu::Queue) -> Self {
let mut p = Self {
frame_size: 0,
sample_rate,
device,
queue,
cpu_fft: None,
cpu_cep_inv: None,
cpu_planner: FftPlanner::new(),
gpu_fft: None,
gpu_blend: 0.0,
window: Vec::new(),
buffer: Vec::new(),
sample_freqs: Vec::new(),
freqs_const: Vec::new(),
bin_alphas: Vec::new(),
bin_smoothed: Vec::new(),
bin_activity: Vec::new(),
bin_reservoir: Vec::new(),
bin_prev_delta: Vec::new(),
bin_initialized: false,
smoothing_tilt: DEFAULT_SMOOTHING_TILT,
smoothing_strength: DEFAULT_SMOOTHING_STRENGTH,
smooth_splashes: true,
num_bins: 26,
history: VecDeque::new(),
smoothing_length: 3,
expand_ratio: 1.0,
expand_threshold: -60.0,
hpf_cutoff: 0.0,
granularity: 33,
detail: 50,
cepstral_strength: 0.0,
weave_caps: default_caps(),
};
p.set_frame_size(frame_size);
p.set_num_bins(26);
p
}
/// updates the bin-to-hertz sample rate.
pub fn set_sample_rate(&mut self, rate: u32) {
self.sample_rate = rate;
}
/// updates the per-bin smoothing alpha curve's frequency-tilt.
pub fn set_smoothing_tilt(&mut self, tilt: f64) {
self.smoothing_tilt = tilt.clamp(0.0, 1.0);
self.rebuild_smoothing_alphas();
}
/// updates the per-bin smoothing mix between raw FFT and smoothed values.
pub fn set_smoothing_strength(&mut self, strength: f64) {
self.smoothing_strength = strength.clamp(0.0, 1.0);
self.rebuild_smoothing_alphas();
}
/// toggles splash injection between the smoothed-envelope-peak cadence and the raw transient cadence.
pub fn set_smooth_splashes(&mut self, on: bool) {
self.smooth_splashes = on;
}
/// caps the rolling history depth of the per-bin db time-average.
pub fn set_smoothing(&mut self, history_length: usize) {
self.smoothing_length = history_length.max(1);
while self.history.len() > self.smoothing_length {
self.history.pop_front();
}
}
/// configures the upward expander applied below the threshold in db.
pub fn set_expander(&mut self, ratio: f32, threshold_db: f32) {
self.expand_ratio = ratio;
self.expand_threshold = threshold_db;
}
/// sets the corner frequency of the per-bin first-order high-pass roll-off applied during binning.
pub fn set_hpf(&mut self, cutoff_freq: f32) {
self.hpf_cutoff = cutoff_freq;
}
/// dials the bin density, iteration count, and blend weight of the cepstral envelope smoother.
pub fn set_cepstral_params(&mut self, granularity: i32, detail: i32, strength: f32) {
self.granularity = granularity;
self.detail = detail;
self.cepstral_strength = strength;
}
/// crossfades between cpu-fft and gpu-fft magnitudes within zero to one inclusive.
pub fn set_gpu_blend(&mut self, blend: f32) {
self.gpu_blend = blend.clamp(0.0, 1.0);
}
/// rebuilds the 40 hz to 11 khz bin edges and matching center frequencies for the n+1 output columns.
pub fn set_num_bins(&mut self, n: usize) {
self.num_bins = n.max(1);
self.rebuild_bins();
}
/// rebuilds the log-spaced display bin centers from 40Hz to 11kHz with a 10Hz sentinel prepended.
fn rebuild_bins(&mut self) {
self.sample_freqs.clear();
self.freqs_const.clear();
self.history.clear();
let n = self.num_bins.max(1);
let min_freq = 40.0_f64;
let max_freq = 11_000.0_f64;
let mut edges: Vec<f64> = Vec::with_capacity(n + 1);
for i in 0..=n {
let t = i as f64 / n as f64;
edges.push(min_freq * (max_freq / min_freq).powf(t));
}
self.sample_freqs.push(10.0);
self.freqs_const.push(10.0);
for i in 0..n {
let center = (edges[i] + edges[i + 1]) / 2.0;
self.sample_freqs.push(center);
self.freqs_const.push(center);
}
self.rebuild_smoothing_alphas();
let cols = self.freqs_const.len();
self.bin_smoothed = vec![-100.0; cols];
self.bin_activity = vec![0.0; cols];
self.bin_reservoir = vec![0.0; cols];
self.bin_prev_delta = vec![0.0; cols];
self.bin_initialized = false;
}
/// rebuilds per-bin one-pole smoothing alpha values from tilt and strength.
fn rebuild_smoothing_alphas(&mut self) {
let tilt = self.smoothing_tilt.clamp(0.0, 1.0);
let strength = self.smoothing_strength.clamp(0.0, 1.0);
let f_low = 40.0_f64.ln();
let f_high = 11_000.0_f64.ln();
let span = f_high - f_low;
self.bin_alphas.clear();
self.bin_alphas.reserve(self.freqs_const.len());
for &fc in self.freqs_const.iter() {
let f_norm = ((fc.max(1.0).ln() - f_low) / span).clamp(0.0, 1.0);
let curve = 0.5 + 0.45 * tilt * (2.0 * f_norm - 1.0);
let smooth_alpha = curve.clamp(0.05, 0.95);
let effective = 1.0 - strength * (1.0 - smooth_alpha);
self.bin_alphas.push(effective);
}
}
/// zeros the analytic-signal buffer, drops the rolling history, and resets the gpu pipeline.
pub fn clear_state(&mut self) {
for c in self.buffer.iter_mut() {
*c = Complex64::new(0.0, 0.0);
}
self.history.clear();
if let Some(gpu) = self.gpu_fft.as_mut() {
gpu.reset();
}
}
/// rebuilds fft plans, the analysis window, the working buffer, and the bin layout at the requested transform length.
pub fn set_frame_size(&mut self, size: usize) {
if self.frame_size == size {
return;
}
self.frame_size = size;
self.cpu_fft = Some(self.cpu_planner.plan_fft_forward(size));
self.cpu_cep_inv = Some(self.cpu_planner.plan_fft_inverse(size));
self.gpu_fft = Some(GpuFft1D::new(self.device.clone(), self.queue.clone(), size as u32));
self.window = build_window(size);
self.buffer = vec![Complex64::new(0.0, 0.0); size];
self.history.clear();
self.rebuild_bins();
}
/// fills the analytic-signal buffer with the latest frame_size samples of an incoming chunk, evicting the head.
pub fn push_data(&mut self, data: &[Complex64]) {
let n = self.frame_size;
if n == 0 {
return;
}
if data.len() >= n {
let start = data.len() - n;
self.buffer.copy_from_slice(&data[start..]);
} else {
let drop = data.len();
self.buffer.copy_within(drop.., 0);
let tail = n - drop;
self.buffer[tail..].copy_from_slice(data);
}
}
/// produces a spectrum frame through the windowed cpu/gpu fft blend, hpf shaping, cepstral idealisation, expander, and history smoothing.
pub fn get_spectrum(&mut self) -> Spectrum {
let n = self.frame_size;
let blend = self.gpu_blend as f64;
let work_template: Vec<Complex64> = self
.buffer
.iter()
.zip(self.window.iter())
.map(|(c, w)| Complex64::new(c.re * w, c.im * w))
.collect();
let cpu_work = if blend < 1.0 {
let mut w = work_template.clone();
self.cpu_fft
.as_ref()
.expect("cpu fft plan")
.clone()
.process(&mut w);
Some(w)
} else {
None
};
let gpu_work = if blend > 0.0 {
let gpu = self.gpu_fft.as_mut().expect("gpu fft plan");
let mut prev = vec![Complex64::new(0.0, 0.0); n];
let got_prev = gpu.try_collect_into(&mut prev);
gpu.submit_forward(&work_template);
if got_prev {
Some(prev)
} else {
None
}
} else {
let gpu = self.gpu_fft.as_mut().expect("gpu fft plan");
let mut discard = vec![Complex64::new(0.0, 0.0); n];
let _ = gpu.try_collect_into(&mut discard);
None
};
let bins = n / 2 + 1;
let mut freqs_full = vec![0.0_f64; bins];
let mut mag_full = vec![0.0_f64; bins];
for i in 0..bins {
let cpu_m = cpu_work.as_ref().map(|w| {
let re = w[i].re;
let im = w[i].im;
2.0 * (re * re + im * im).sqrt() / n as f64
});
let gpu_m = gpu_work.as_ref().map(|w| {
let re = w[i].re;
let im = w[i].im;
2.0 * (re * re + im * im).sqrt() / n as f64
});
let mut mag = match (cpu_m, gpu_m) {
(Some(c), Some(g)) => c * (1.0 - blend) + g * blend,
(Some(c), None) => c,
(None, Some(g)) => g,
(None, None) => 0.0,
};
let freq = i as f64 * self.sample_rate as f64 / n as f64;
if self.hpf_cutoff > 0.0 && freq > 0.0 {
let ratio = self.hpf_cutoff as f64 / freq;
let r2 = ratio * ratio;
let gain = 1.0 / (1.0 + r2 * r2).sqrt();
mag *= gain;
} else if freq == 0.0 {
mag = 0.0;
}
freqs_full[i] = freq;
mag_full[i] = mag;
}
let mut cep_buf: Vec<Complex64> = Vec::with_capacity(n);
for i in 0..n {
let idx = if i <= n / 2 { i } else { n - i };
let v = mag_full[idx].max(1e-9);
cep_buf.push(Complex64::new(v.ln(), 0.0));
}
self.cpu_cep_inv
.as_ref()
.expect("cpu cep inverse plan")
.clone()
.process(&mut cep_buf);
let cep_scale = 1.0 / n as f64;
let half_n = n / 2;
let cepstrum: Vec<f32> = (0..half_n)
.map(|i| (cep_buf[i].re * cep_scale) as f32)
.collect();
let cols = self.freqs_const.len();
if self.bin_smoothed.len() != cols {
self.bin_smoothed = vec![-100.0; cols];
self.bin_activity = vec![0.0; cols];
self.bin_reservoir = vec![0.0; cols];
self.bin_prev_delta = vec![0.0; cols];
self.bin_initialized = false;
}
let mut current_db = vec![0.0_f64; cols];
for (i, &target) in self.sample_freqs.iter().enumerate() {
let mag = lerp_at(&freqs_full, &mag_full, target);
let raw = 20.0 * mag.max(1e-12).log10();
let smoothed = if !self.bin_initialized {
raw
} else {
let a = self.bin_alphas.get(i).copied().unwrap_or(1.0);
self.bin_smoothed[i] + a * (raw - self.bin_smoothed[i])
};
// suppressed-transient energy as the positive residual between raw and the smoothed display.
let residual = (raw - smoothed).max(0.0);
let norm = (residual / ACTIVITY_REF_DB).clamp(0.0, 1.0);
let delta = smoothed - self.bin_smoothed[i];
if self.smooth_splashes {
self.bin_reservoir[i] += norm;
let peaked = self.bin_prev_delta[i] > 0.0 && delta <= 0.0;
if peaked {
let dump = (self.bin_reservoir[i] * ACTIVITY_RESERVOIR_SCALE).clamp(0.0, 1.0);
self.bin_activity[i] = dump.max(self.bin_activity[i]);
self.bin_reservoir[i] = 0.0;
} else {
self.bin_activity[i] *= ACTIVITY_DECAY;
}
} else {
self.bin_activity[i] = norm.max(self.bin_activity[i] * ACTIVITY_DECAY);
}
self.bin_prev_delta[i] = delta;
self.bin_smoothed[i] = smoothed;
let mut val = smoothed;
if (self.expand_ratio - 1.0).abs() > f32::EPSILON {
let t = self.expand_threshold as f64;
let r = self.expand_ratio as f64;
val = (val - t) * r + t;
}
if val < -100.0 {
val = -100.0;
}
current_db[i] = val;
}
self.bin_initialized = true;
let activity_ret: Vec<f32> = self.bin_activity.iter().map(|&a| a as f32).collect();
self.history.push_back(current_db);
while self.history.len() > self.smoothing_length {
self.history.pop_front();
}
let cols = self.freqs_const.len();
let mut averaged = vec![0.0_f32; cols];
if !self.history.is_empty() {
for v in &self.history {
for (i, &x) in v.iter().enumerate() {
averaged[i] += x as f32;
}
}
let factor = 1.0 / self.history.len() as f32;
for v in &mut averaged {
*v *= factor;
}
}
let freqs_ret: Vec<f32> = self.freqs_const.iter().map(|&f| f as f32).collect();
Spectrum {
freqs: freqs_ret,
db: averaged,
cepstrum,
activity: activity_ret,
}
}
/// returns the active fft transform length.
pub fn frame_size(&self) -> usize {
self.frame_size
}
}
/// linearly interpolates a value from the freqs/values table at the given target frequency with endpoint clamping.
fn lerp_at(freqs: &[f64], values: &[f64], target: f64) -> f64 {
if freqs.is_empty() {
return 0.0;
}
if target <= freqs[0] {
return values[0];
}
if target >= freqs[freqs.len() - 1] {
return values[values.len() - 1];
}
let upper = freqs.partition_point(|&f| f < target);
let lower = upper - 1;
let f0 = freqs[lower];
let f1 = freqs[upper];
let v0 = values[lower];
let v1 = values[upper];
let t = (target - f0) / (f1 - f0);
v0 + t * (v1 - v0)
}
/// Hamming/Blackman-Harris window cutoff in fft samples.
const SMALL_FFT_THRESHOLD: usize = 4096;
/// default per-bin smoothing alpha-curve frequency-tilt.
const DEFAULT_SMOOTHING_TILT: f64 = 0.6;
/// default per-bin smoothing mix between raw FFT and smoothed values.
const DEFAULT_SMOOTHING_STRENGTH: f64 = 0.7;
/// residual dB mapping to full splash activity.
const ACTIVITY_REF_DB: f64 = 12.0;
/// per-frame peak-hold decay for splash activity.
const ACTIVITY_DECAY: f64 = 0.85;
/// reservoir-to-activity scale at a smoothed-envelope peak.
const ACTIVITY_RESERVOIR_SCALE: f64 = 0.25;
/// builds a Hamming or Blackman-Harris analysis window at the requested size.
fn build_window(size: usize) -> Vec<f64> {
if size == 0 {
return Vec::new();
}
let denom = (size - 1).max(1) as f64;
if size <= SMALL_FFT_THRESHOLD {
(0..size)
.map(|i| 0.54 - 0.46 * (2.0 * PI * i as f64 / denom).cos())
.collect()
} else {
(0..size)
.map(|i| {
let a0 = 0.35875;
let a1 = 0.48829;
let a2 = 0.14128;
let a3 = 0.01168;
a0 - a1 * (2.0 * PI * i as f64 / denom).cos()
+ a2 * (4.0 * PI * i as f64 / denom).cos()
- a3 * (6.0 * PI * i as f64 / denom).cos()
})
.collect()
}
}
/// chooses lighter idealisation caps on mobile targets and unlimited caps on desktop.
#[cfg(any(target_os = "ios", target_os = "android"))]
fn default_caps() -> Option<weave::Caps> {
Some(weave::Caps::mobile())
}
#[cfg(not(any(target_os = "ios", target_os = "android")))]
fn default_caps() -> Option<weave::Caps> {
None
}