ferraus_denoise/
stft.rs

1//! STFT framing shared by the profile learner, the gate engine, and the metrics.
2//!
3//! Frames use a periodic Hann window with 75% overlap. Magnitudes are scaled by
4//! `2 / sum(window)` so a full-scale sine reads 0 dB at its bin; profiles store levels on
5//! this scale (`magnitude_scale: "dbfs_sine"`).
6
7pub use realfft::num_complex::Complex32;
8use realfft::{ComplexToReal, RealFftPlanner, RealToComplex};
9use std::sync::Arc;
10
11/// Floor for levels in dB (digital silence reads this instead of -inf).
12pub const FLOOR_DB: f32 = -160.0;
13/// Floor for power values (10^(FLOOR_DB/10)).
14pub const FLOOR_POW: f32 = 1e-16;
15
16/// FFT size for a sample rate: 2048 at 44.1/48 kHz, scaled by powers of two with the rate
17/// (1024 at 22.05/24 kHz, 4096 at 88.2/96 kHz, 512 at 16 kHz).
18pub fn fft_size_for(sample_rate: u32) -> usize {
19    let ratio = sample_rate.max(1) as f64 / 46_050.0; // geometric mean of 44.1k and 48k
20    let e = ratio.log2().round() as i32;
21    (2048f64 * 2f64.powi(e)).round().clamp(256.0, 16384.0) as usize
22}
23
24/// Hop for an FFT size (75% overlap).
25pub fn hop_for(n_fft: usize) -> usize {
26    n_fft / 4
27}
28
29/// Periodic Hann window (constant overlap-add at hop n/4 for both w and w²).
30pub fn hann(n: usize) -> Vec<f32> {
31    (0..n)
32        .map(|i| {
33            let x = std::f64::consts::PI * 2.0 * i as f64 / n as f64;
34            (0.5 - 0.5 * x.cos()) as f32
35        })
36        .collect()
37}
38
39/// A planned forward/inverse real FFT pair with scratch buffers.
40pub struct Fft {
41    pub n: usize,
42    pub window: Vec<f32>,
43    /// Scale that maps |X| to the dBFS-sine magnitude scale.
44    pub mag_scale: f32,
45    fwd: Arc<dyn RealToComplex<f32>>,
46    inv: Arc<dyn ComplexToReal<f32>>,
47    time: Vec<f32>,
48    fwd_scratch: Vec<Complex32>,
49    inv_scratch: Vec<Complex32>,
50}
51
52impl Fft {
53    pub fn new(n: usize) -> Fft {
54        let mut planner = RealFftPlanner::<f32>::new();
55        let fwd = planner.plan_fft_forward(n);
56        let inv = planner.plan_fft_inverse(n);
57        let window = hann(n);
58        let wsum: f32 = window.iter().sum();
59        Fft {
60            n,
61            mag_scale: 2.0 / wsum,
62            fwd_scratch: fwd.make_scratch_vec(),
63            inv_scratch: inv.make_scratch_vec(),
64            fwd,
65            inv,
66            window,
67            time: vec![0.0; n],
68        }
69    }
70
71    pub fn bins(&self) -> usize {
72        self.n / 2 + 1
73    }
74
75    /// Window `frame` (length n) and transform it into `spec` (length n/2+1).
76    pub fn forward(&mut self, frame: &[f32], spec: &mut [Complex32]) {
77        for ((t, x), w) in self.time.iter_mut().zip(frame).zip(&self.window) {
78            *t = x * w;
79        }
80        self.fwd.process_with_scratch(&mut self.time, spec, &mut self.fwd_scratch).expect("fft sizes match");
81    }
82
83    /// Forward transform of a frame given as two pieces (ring-buffer order).
84    pub fn forward_split(&mut self, a: &[f32], b: &[f32], spec: &mut [Complex32]) {
85        for ((t, x), w) in self.time.iter_mut().zip(a.iter().chain(b)).zip(&self.window) {
86            *t = x * w;
87        }
88        self.fwd.process_with_scratch(&mut self.time, spec, &mut self.fwd_scratch).expect("fft sizes match");
89    }
90
91    /// Inverse transform `spec` (modified in place as scratch) and return the windowed,
92    /// overlap-add normalized time frame. Summing these at hop n/4 reconstructs the input.
93    pub fn inverse_windowed(&mut self, spec: &mut [Complex32]) -> &[f32] {
94        // DC and Nyquist must be real for the c2r transform
95        spec[0].im = 0.0;
96        let last = spec.len() - 1;
97        spec[last].im = 0.0;
98        self.inv.process_with_scratch(spec, &mut self.time, &mut self.inv_scratch).expect("fft sizes match");
99        // periodic Hann: sum of w² at hop n/4 is exactly 1.5
100        let norm = 1.0 / (self.n as f32 * 1.5);
101        for (t, w) in self.time.iter_mut().zip(&self.window) {
102            *t *= w * norm;
103        }
104        &self.time
105    }
106}
107
108/// Convert a power value on the dBFS-sine scale to dB.
109#[inline]
110pub fn pow_db(p: f32) -> f32 {
111    10.0 * p.max(FLOOR_POW).log10()
112}
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117
118    #[test]
119    fn sizes_scale_with_rate() {
120        assert_eq!(fft_size_for(44_100), 2048);
121        assert_eq!(fft_size_for(48_000), 2048);
122        assert_eq!(fft_size_for(96_000), 4096);
123        assert_eq!(fft_size_for(22_050), 1024);
124        assert_eq!(fft_size_for(16_000), 512);
125        assert_eq!(fft_size_for(8_000), 256);
126    }
127
128    #[test]
129    fn hann_squared_is_cola() {
130        let n = 2048;
131        let w = hann(n);
132        let hop = hop_for(n);
133        for i in 0..hop {
134            let s: f32 = (0..4).map(|m| w[i + m * hop].powi(2)).sum();
135            assert!((s - 1.5).abs() < 1e-5, "{s}");
136        }
137    }
138
139    #[test]
140    fn sine_reads_zero_db() {
141        let n = 2048;
142        let mut f = Fft::new(n);
143        let sr = 48_000.0;
144        let k = 100.0;
145        let freq = k * sr / n as f64;
146        let x: Vec<f32> = (0..n).map(|i| (2.0 * std::f64::consts::PI * freq * i as f64 / sr).sin() as f32).collect();
147        let mut spec = vec![Complex32::new(0.0, 0.0); f.bins()];
148        f.forward(&x, &mut spec);
149        let db = 20.0 * (spec[100].norm() * f.mag_scale).log10();
150        assert!(db.abs() < 0.01, "{db}");
151    }
152}
ferraus_denoise/stft.rs

ferraus_denoise/
stft.rs
