modewise_stft/
welch.rs

1//! Power spectral density by Welch's method, as scipy.signal.welch (density scaling, one-sided,
2//! mean over segments). One boxcar segment over the whole run is scipy.signal.periodogram.
3
4use std::f64::consts::PI;
5
6use modewise_method::{MethodDescriptor, MethodError};
7use modewise_signals::MultichannelSeries;
8use modewise_spectral::{TimeFrequencyMap, TimeFrequencyTransform};
9use ndarray::Array3;
10use rustfft::FftPlanner;
11use rustfft::num_complex::Complex64;
12
13#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub enum Window {
15    Boxcar,
16    /// Periodic Hann, as scipy's get_window("hann", n).
17    Hann,
18}
19
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21pub enum Detrend {
22    None,
23    /// Subtract each segment's mean.
24    Constant,
25    /// Subtract each segment's least-squares line.
26    Linear,
27}
28
29/// One frame of power spectral density per channel (power units per Hz).
30#[derive(Debug, Clone)]
31pub struct Welch {
32    /// Samples per segment; `None` is the whole signal.
33    pub segment_len: Option<usize>,
34    /// Samples shared by neighbouring segments; `None` is half a segment, as scipy.
35    pub overlap: Option<usize>,
36    pub window: Window,
37    pub detrend: Detrend,
38}
39
40impl Welch {
41    fn window(&self, n: usize) -> Vec<f64> {
42        match self.window {
43            Window::Boxcar => vec![1.0; n],
44            Window::Hann => (0..n)
45                .map(|i| 0.5 - 0.5 * (2.0 * PI * i as f64 / n as f64).cos())
46                .collect(),
47        }
48    }
49}
50
51fn detrend(x: &mut [f64], how: Detrend) {
52    let n = x.len() as f64;
53    match how {
54        Detrend::None => {}
55        Detrend::Constant => {
56            let mean = x.iter().sum::<f64>() / n;
57            x.iter_mut().for_each(|v| *v -= mean);
58        }
59        Detrend::Linear => {
60            let t_mean = (n - 1.0) / 2.0;
61            let y_mean = x.iter().sum::<f64>() / n;
62            let (mut sxy, mut sxx) = (0.0, 0.0);
63            for (i, v) in x.iter().enumerate() {
64                let dt = i as f64 - t_mean;
65                // Σ(t − t̄) = 0, so the mean of y drops out of the covariance.
66                sxy += dt * v;
67                sxx += dt * dt;
68            }
69            let slope = if sxx > 0.0 { sxy / sxx } else { 0.0 };
70            for (i, v) in x.iter_mut().enumerate() {
71                *v -= y_mean + slope * (i as f64 - t_mean);
72            }
73        }
74    }
75}
76
77impl TimeFrequencyTransform for Welch {
78    fn descriptor(&self) -> MethodDescriptor {
79        MethodDescriptor {
80            name: "spectrum",
81            version: env!("CARGO_PKG_VERSION"),
82            multivariate: false,
83        }
84    }
85
86    fn transform(&self, signal: &MultichannelSeries) -> Result<TimeFrequencyMap, MethodError> {
87        let total = signal.n_samples();
88        let n = self.segment_len.unwrap_or(total);
89        if n < 2 {
90            return Err(MethodError::invalid(
91                "segment_len",
92                "must be at least 2 samples",
93            ));
94        }
95        if n > total {
96            return Err(MethodError::SignalTooShort {
97                method: "spectrum",
98                needed: n,
99                actual: total,
100            });
101        }
102        let overlap = self.overlap.unwrap_or(n / 2);
103        if overlap >= n {
104            return Err(MethodError::invalid(
105                "overlap",
106                "must be shorter than a segment",
107            ));
108        }
109        let step = n - overlap;
110        let starts: Vec<usize> = (0..=(total - n) / step).map(|k| k * step).collect();
111        let window = self.window(n);
112        let fs = signal.sampling_rate().hz();
113        let scale = 1.0 / (fs * window.iter().map(|w| w * w).sum::<f64>());
114        let n_bins = n / 2 + 1;
115        let fft = FftPlanner::new().plan_fft_forward(n);
116        let mut power = Array3::zeros((signal.n_channels(), n_bins, 1));
117        for (c, row) in signal.data().rows().into_iter().enumerate() {
118            for &start in &starts {
119                let mut seg: Vec<f64> = row.iter().skip(start).take(n).copied().collect();
120                detrend(&mut seg, self.detrend);
121                let mut buf: Vec<Complex64> = seg
122                    .iter()
123                    .zip(&window)
124                    .map(|(x, w)| Complex64::new(x * w, 0.0))
125                    .collect();
126                fft.process(&mut buf);
127                for (k, z) in buf.iter().take(n_bins).enumerate() {
128                    // One-sided: every bin but DC and an even-length Nyquist counts twice.
129                    let twice = k != 0 && !(n.is_multiple_of(2) && k == n / 2);
130                    power[[c, k, 0]] += z.norm_sqr() * scale * if twice { 2.0 } else { 1.0 };
131                }
132            }
133        }
134        power /= starts.len() as f64;
135        let frequencies = (0..n_bins).map(|k| k as f64 * fs / n as f64).collect();
136        let span = (starts.last().copied().unwrap_or(0) + n) as f64;
137        TimeFrequencyMap::new(
138            power,
139            frequencies,
140            vec![span / (2.0 * fs)],
141            signal.labels().to_vec(),
142        )
143    }
144}
145
146#[cfg(test)]
147mod tests {
148    use modewise_signals::SamplingRate;
149    use ndarray::Array2;
150
151    use super::*;
152
153    fn signal(n: usize) -> MultichannelSeries {
154        let data = Array2::from_shape_fn((1, n), |(_, t)| (t as f64 * 0.7).sin() + 0.1 * t as f64);
155        MultichannelSeries::new(data, SamplingRate::new(2.0).unwrap(), vec!["a".into()]).unwrap()
156    }
157
158    fn welch(segment_len: Option<usize>, overlap: Option<usize>) -> Welch {
159        Welch {
160            segment_len,
161            overlap,
162            window: Window::Hann,
163            detrend: Detrend::Linear,
164        }
165    }
166
167    #[test]
168    fn segments_must_fit_and_the_frame_sits_at_the_centre_of_the_span() {
169        for bad in [Some(0), Some(1)] {
170            assert!(welch(bad, None).transform(&signal(20)).is_err());
171        }
172        assert!(welch(Some(21), None).transform(&signal(20)).is_err());
173        assert!(welch(Some(8), Some(8)).transform(&signal(20)).is_err());
174        assert!(welch(Some(2), None).transform(&signal(20)).is_ok());
175        // Segments of 8 every 4 samples cover 0..20 when 20 samples are given, 10 s at 2 Hz.
176        let map = welch(Some(8), Some(4)).transform(&signal(21)).unwrap();
177        assert_eq!(map.times_s(), [5.0]);
178        let whole = welch(None, None).transform(&signal(21)).unwrap();
179        assert_eq!(whole.times_s(), [5.25]);
180        let d = welch(None, None).descriptor();
181        assert_eq!((d.name, d.multivariate), ("spectrum", false));
182    }
183
184    #[test]
185    fn a_single_sample_segment_has_no_slope() {
186        let mut x = [3.0];
187        detrend(&mut x, Detrend::Linear);
188        assert_eq!(x, [0.0]);
189    }
190}
modewise_stft/welch.rs

modewise_stft/
welch.rs
