modewise_stft/
lib.rs

1//! Short-time Fourier transform with a periodic Hann window, and Welch power spectra.
2
3mod welch;
4
5pub use welch::{Detrend, Welch, Window};
6
7use std::f64::consts::PI;
8
9use modewise_method::{MethodDescriptor, MethodError};
10use modewise_signals::MultichannelSeries;
11use modewise_spectral::{TimeFrequencyMap, TimeFrequencyTransform};
12use ndarray::Array3;
13use rustfft::FftPlanner;
14use rustfft::num_complex::Complex64;
15
16/// Frames start at sample 0 and step by `hop`. Frames that would run past the
17/// end are dropped. Power is the unscaled squared magnitude, |X|².
18#[derive(Debug, Clone)]
19pub struct Stft {
20    window_len: usize,
21    hop: usize,
22}
23
24impl Stft {
25    pub fn new(window_len: usize, hop: usize) -> Result<Self, MethodError> {
26        if window_len < 2 {
27            return Err(MethodError::invalid("window_len", "must be at least 2"));
28        }
29        if hop == 0 {
30            return Err(MethodError::invalid("hop", "must be at least 1"));
31        }
32        Ok(Self { window_len, hop })
33    }
34
35    fn window(&self) -> Vec<f64> {
36        let n = self.window_len as f64;
37        (0..self.window_len)
38            .map(|i| 0.5 - 0.5 * (2.0 * PI * i as f64 / n).cos())
39            .collect()
40    }
41}
42
43impl TimeFrequencyTransform for Stft {
44    fn descriptor(&self) -> MethodDescriptor {
45        MethodDescriptor {
46            name: "stft",
47            version: env!("CARGO_PKG_VERSION"),
48            multivariate: false,
49        }
50    }
51
52    fn transform(&self, signal: &MultichannelSeries) -> Result<TimeFrequencyMap, MethodError> {
53        let (n, hop) = (self.window_len, self.hop);
54        if signal.n_samples() < n {
55            return Err(MethodError::SignalTooShort {
56                method: "stft",
57                needed: n,
58                actual: signal.n_samples(),
59            });
60        }
61        let n_frames = (signal.n_samples() - n) / hop + 1;
62        let n_bins = n / 2 + 1;
63        let window = self.window();
64        let fft = FftPlanner::new().plan_fft_forward(n);
65
66        let mut power = Array3::zeros((signal.n_channels(), n_bins, n_frames));
67        for (c, channel) in signal.data().rows().into_iter().enumerate() {
68            for frame in 0..n_frames {
69                let start = frame * hop;
70                let mut buffer: Vec<Complex64> = (0..n)
71                    .map(|i| Complex64::new(channel[start + i] * window[i], 0.0))
72                    .collect();
73                fft.process(&mut buffer);
74                for (bin, value) in buffer.iter().take(n_bins).enumerate() {
75                    power[[c, bin, frame]] = value.norm_sqr();
76                }
77            }
78        }
79
80        let fs = signal.sampling_rate().hz();
81        let frequencies_hz = (0..n_bins).map(|k| k as f64 * fs / n as f64).collect();
82        let times_s = (0..n_frames)
83            .map(|f| (f * hop) as f64 / fs + n as f64 / (2.0 * fs))
84            .collect();
85        TimeFrequencyMap::new(power, frequencies_hz, times_s, signal.labels().to_vec())
86    }
87}
modewise_stft/lib.rs

modewise_stft/
lib.rs
