sound_to_patch/dsp/
stft.rs

1//! The short-time Fourier transform with centred frames.
2
3use std::error::Error;
4use std::fmt;
5use std::sync::Arc;
6
7use realfft::num_complex::Complex64;
8use realfft::{RealFftPlanner, RealToComplex};
9
10use super::Window;
11use crate::audio::SampleRate;
12use crate::cast;
13use crate::units::{Hz, Seconds};
14
15/// How a centred frame continues past either end of the signal.
16#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
17pub enum Padding {
18    /// Zeros. The convention for descriptors.
19    Zero,
20    /// The signal mirrored at its end sample, which is not repeated
21    /// (`d c b | a b c d | c b a`).
22    Reflect,
23}
24
25/// The frame layout of an STFT: frame length `N`, hop, window, padding and
26/// FFT oversampling.
27///
28/// Frames are centred: frame `m` covers samples `m hop - N/2` to
29/// `m hop - N/2 + N - 1`, so a signal of `len` samples gives
30/// `1 + floor((len + 2 floor(N/2) - N) / hop)` frames, which is
31/// `1 + floor(len / hop)` for an even `N`. The default hop is `N/4`.
32///
33/// Oversampling by a factor `p` zero-pads each windowed frame to `p N`
34/// samples, so the bins sit `fs / (p N)` apart. It refines peak positions
35/// and changes no frequency resolution.
36#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
37pub struct StftParams {
38    frame_len: usize,
39    hop: usize,
40    window: Window,
41    padding: Padding,
42    oversample: usize,
43}
44
45impl StftParams {
46    /// A periodic Hann window of `frame_len` samples, hop `frame_len / 4`,
47    /// zero padding and no oversampling.
48    ///
49    /// # Errors
50    ///
51    /// Returns an [`StftError`] for a frame shorter than 4 samples.
52    pub fn new(frame_len: usize) -> Result<Self, StftError> {
53        if frame_len < 4 {
54            return Err(StftError::FrameTooShort(frame_len));
55        }
56        Ok(Self {
57            frame_len,
58            hop: frame_len / 4,
59            window: Window::Hann,
60            padding: Padding::Zero,
61            oversample: 1,
62        })
63    }
64
65    /// The defaults of [`StftParams::new`] for a frame length that the code
66    /// fixes as a constant of at least 4 samples.
67    pub(crate) const fn fixed(frame_len: usize) -> Self {
68        Self {
69            frame_len,
70            hop: frame_len / 4,
71            window: Window::Hann,
72            padding: Padding::Zero,
73            oversample: 1,
74        }
75    }
76
77    /// Sets the hop in samples.
78    ///
79    /// # Errors
80    ///
81    /// Returns an [`StftError`] for a hop of zero.
82    pub fn with_hop(self, hop: usize) -> Result<Self, StftError> {
83        if hop == 0 {
84            return Err(StftError::ZeroHop);
85        }
86        Ok(Self { hop, ..self })
87    }
88
89    /// Sets the window.
90    #[must_use]
91    pub fn with_window(self, window: Window) -> Self {
92        Self { window, ..self }
93    }
94
95    /// Sets the padding at both ends of the signal.
96    #[must_use]
97    pub fn with_padding(self, padding: Padding) -> Self {
98        Self { padding, ..self }
99    }
100
101    /// Sets the FFT oversampling factor.
102    ///
103    /// # Errors
104    ///
105    /// Returns an [`StftError`] for a factor of zero.
106    pub fn with_oversample(self, factor: usize) -> Result<Self, StftError> {
107        if factor == 0 {
108            return Err(StftError::ZeroOversample);
109        }
110        Ok(Self {
111            oversample: factor,
112            ..self
113        })
114    }
115
116    /// Returns the frame length `N` in samples.
117    #[must_use]
118    pub fn frame_len(self) -> usize {
119        self.frame_len
120    }
121
122    /// Returns the hop in samples.
123    #[must_use]
124    pub fn hop(self) -> usize {
125        self.hop
126    }
127
128    /// Returns the window.
129    #[must_use]
130    pub fn window(self) -> Window {
131        self.window
132    }
133
134    /// Returns the padding.
135    #[must_use]
136    pub fn padding(self) -> Padding {
137        self.padding
138    }
139
140    /// Returns the FFT length, `N` times the oversampling factor.
141    #[must_use]
142    pub fn fft_len(self) -> usize {
143        self.frame_len * self.oversample
144    }
145
146    /// Returns the number of frequency bins, `fft_len / 2 + 1`.
147    #[must_use]
148    pub fn bins(self) -> usize {
149        self.fft_len() / 2 + 1
150    }
151
152    /// Returns the number of centred frames for a signal of `len` samples.
153    #[must_use]
154    pub fn frame_count(self, len: usize) -> usize {
155        let padded = len + 2 * (self.frame_len / 2);
156        1 + padded.saturating_sub(self.frame_len) / self.hop
157    }
158
159    /// Returns the frequency of bin `k`, `k fs / fft_len`.
160    #[must_use]
161    pub fn bin_hz(self, bin: usize, rate: SampleRate) -> Hz {
162        Hz::clamped(cast::usize_to_f64(bin) * rate.hz().get() / cast::usize_to_f64(self.fft_len()))
163    }
164
165    /// Returns the bin spacing, `fs / fft_len`.
166    #[must_use]
167    pub fn bin_width(self, rate: SampleRate) -> Hz {
168        self.bin_hz(1, rate)
169    }
170
171    /// Returns the time of the centre of frame `m`, `m hop / fs`.
172    #[must_use]
173    pub fn frame_time(self, frame: usize, rate: SampleRate) -> Seconds {
174        rate.seconds(frame * self.hop)
175    }
176
177    /// Returns the first sample of frame `m`; negative indices lie in the
178    /// padding before the signal.
179    pub(crate) fn frame_start(self, frame: usize) -> isize {
180        // A frame index times a hop stays far below isize::MAX for any
181        // signal that fits in memory.
182        let centre = isize::try_from(frame * self.hop).unwrap_or(isize::MAX);
183        let half = isize::try_from(self.frame_len / 2).unwrap_or(isize::MAX);
184        centre.saturating_sub(half)
185    }
186}
187
188/// Reads sample `index` of `samples` under a padding rule, for an index
189/// that may lie before or after the signal.
190pub(crate) fn padded_sample(samples: &[f64], index: isize, padding: Padding) -> f64 {
191    if let Ok(i) = usize::try_from(index)
192        && let Some(&sample) = samples.get(i)
193    {
194        return sample;
195    }
196    match padding {
197        Padding::Zero => 0.0,
198        Padding::Reflect => {
199            let len = samples.len();
200            if len < 2 {
201                return samples.first().copied().unwrap_or(0.0);
202            }
203            let period = 2 * (len - 1);
204            // The mirror is symmetric about sample 0, so x[-i] = x[i].
205            let wrapped = index.unsigned_abs() % period;
206            samples[if wrapped >= len {
207                period - wrapped
208            } else {
209                wrapped
210            }]
211        }
212    }
213}
214
215/// A planned STFT. It computes one frame at a time into a buffer it owns,
216/// so a long stem never holds its whole spectrogram in memory.
217///
218/// The transform is not normalised: a frame holds
219/// `X[m, k] = sum_n w[n] x[m hop - N/2 + n] exp(-j 2 pi k n / (p N))`.
220///
221/// ```
222/// use sound_to_patch::audio::SampleRate;
223/// use sound_to_patch::dsp::{Stft, StftParams};
224///
225/// let rate = SampleRate::new(8_000)?;
226/// let params = StftParams::new(256)?;
227/// let mut stft = Stft::new(params);
228/// let signal: Vec<f64> = (0..1_000).map(|n| (0.1 * f64::from(n)).sin()).collect();
229/// assert_eq!(params.frame_count(signal.len()), 16);
230/// let peak = stft
231///     .frame(&signal, 8)
232///     .iter()
233///     .enumerate()
234///     .max_by(|a, b| a.1.norm().total_cmp(&b.1.norm()))
235///     .map(|(bin, _)| bin);
236/// // 0.1 rad per sample is 127.3 Hz at 8 kHz: bin 4 of 31.25 Hz.
237/// assert_eq!(peak, Some(4));
238/// # Ok::<(), Box<dyn std::error::Error>>(())
239/// ```
240pub struct Stft {
241    params: StftParams,
242    window: Vec<f64>,
243    fft: Arc<dyn RealToComplex<f64>>,
244    input: Vec<f64>,
245    output: Vec<Complex64>,
246    scratch: Vec<Complex64>,
247}
248
249impl Stft {
250    /// Plans the FFT and the window for these parameters.
251    #[must_use]
252    pub fn new(params: StftParams) -> Self {
253        let fft = RealFftPlanner::<f64>::new().plan_fft_forward(params.fft_len());
254        Self {
255            window: params.window().samples(params.frame_len()),
256            input: fft.make_input_vec(),
257            output: fft.make_output_vec(),
258            scratch: fft.make_scratch_vec(),
259            fft,
260            params,
261        }
262    }
263
264    /// Returns the parameters.
265    #[must_use]
266    pub fn params(&self) -> StftParams {
267        self.params
268    }
269
270    /// Computes frame `m` of `samples` and returns its `bins()` complex
271    /// values. A frame past the end holds only padding.
272    pub fn frame(&mut self, samples: &[f64], frame: usize) -> &[Complex64] {
273        let start = self.params.frame_start(frame);
274        let padding = self.params.padding();
275        self.input.fill(0.0);
276        for (offset, (slot, weight)) in self.input.iter_mut().zip(&self.window).enumerate() {
277            let index = start.saturating_add(isize::try_from(offset).unwrap_or(isize::MAX));
278            *slot = weight * padded_sample(samples, index, padding);
279        }
280        let result =
281            self.fft
282                .process_with_scratch(&mut self.input, &mut self.output, &mut self.scratch);
283        debug_assert!(result.is_ok(), "buffers come from the plan: {result:?}");
284        &self.output
285    }
286
287    /// Computes every frame of `samples`. Holds the whole spectrogram, so
288    /// use it for short signals; iterate with [`Stft::frame`] otherwise.
289    pub fn frames(&mut self, samples: &[f64]) -> Vec<Vec<Complex64>> {
290        (0..self.params.frame_count(samples.len()))
291            .map(|frame| self.frame(samples, frame).to_vec())
292            .collect()
293    }
294}
295
296impl fmt::Debug for Stft {
297    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
298        f.debug_struct("Stft")
299            .field("params", &self.params)
300            .finish_non_exhaustive()
301    }
302}
303
304/// STFT parameters that cannot work.
305#[derive(Clone, Copy, Debug, PartialEq, Eq)]
306pub enum StftError {
307    /// A frame shorter than 4 samples.
308    FrameTooShort(usize),
309    /// A hop of zero samples.
310    ZeroHop,
311    /// An oversampling factor of zero.
312    ZeroOversample,
313}
314
315impl fmt::Display for StftError {
316    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
317        match self {
318            Self::FrameTooShort(len) => {
319                write!(
320                    f,
321                    "an STFT frame of {len} samples is too short; use 4 or more"
322                )
323            }
324            Self::ZeroHop => write!(f, "an STFT hop of 0 samples never advances; use 1 or more"),
325            Self::ZeroOversample => write!(
326                f,
327                "an oversampling factor of 0 leaves no FFT; use 1 or more"
328            ),
329        }
330    }
331}
332
333impl Error for StftError {}
334
335#[cfg(test)]
336mod tests {
337    use super::*;
338
339    #[test]
340    fn frame_count_follows_the_centred_layout() {
341        let params = StftParams::new(8).unwrap();
342        assert_eq!(params.hop(), 2);
343        assert_eq!(params.frame_count(0), 1);
344        assert_eq!(params.frame_count(9), 5);
345        assert_eq!(params.frame_count(10), 6);
346        let odd = StftParams::new(9).unwrap();
347        // padded 10 + 8 = 18, (18 - 9) / 2 = 4
348        assert_eq!(odd.frame_count(10), 5);
349    }
350
351    #[test]
352    fn bins_and_times_use_the_fft_length_and_the_hop() {
353        let rate = SampleRate::new(8_000).unwrap();
354        let params = StftParams::new(256).unwrap().with_oversample(4).unwrap();
355        assert_eq!(params.fft_len(), 1024);
356        assert_eq!(params.bins(), 513);
357        assert_eq!(params.bin_width(rate).get(), 7.8125);
358        assert_eq!(params.frame_time(2, rate).get(), 0.016);
359    }
360
361    #[test]
362    fn reflect_padding_mirrors_without_repeating_the_end() {
363        let x = [1.0, 2.0, 3.0, 4.0];
364        let read = |i| padded_sample(&x, i, Padding::Reflect);
365        assert_eq!([read(-3), read(-2), read(-1)], [4.0, 3.0, 2.0]);
366        assert_eq!([read(4), read(5), read(6)], [3.0, 2.0, 1.0]);
367        assert_eq!(read(7), 2.0);
368        assert_eq!(padded_sample(&x, -1, Padding::Zero), 0.0);
369        assert_eq!(padded_sample(&[5.0], 3, Padding::Reflect), 5.0);
370        assert_eq!(padded_sample(&[], 3, Padding::Reflect), 0.0);
371    }
372
373    #[test]
374    fn a_constant_signal_puts_the_window_sum_in_bin_zero() {
375        let params = StftParams::new(16).unwrap();
376        let mut stft = Stft::new(params);
377        let frame = stft.frame(&[1.0; 64], 4).to_vec();
378        let window_sum: f64 = Window::Hann.samples(16).iter().sum();
379        assert!((frame[0].re - window_sum).abs() < 1e-12);
380        // A periodic Hann leaks into bin 1 only.
381        assert!(frame[2..].iter().all(|bin| bin.norm() < 1e-12));
382    }
383
384    #[test]
385    fn bad_parameters_fail_with_a_fix() {
386        assert_eq!(StftParams::new(2).unwrap_err(), StftError::FrameTooShort(2));
387        let params = StftParams::new(8).unwrap();
388        assert_eq!(
389            params.with_hop(0).unwrap_err().to_string(),
390            "an STFT hop of 0 samples never advances; use 1 or more"
391        );
392        assert!(params.with_oversample(0).is_err());
393    }
394}
sound_to_patch/dsp/stft.rs

sound_to_patch/dsp/
stft.rs
