Skip to main content

embedded_dsp/
psd.rs

1//! Power Spectral Density (PSD) estimation using Welch's Method (Averaged Overlapped Periodogram) and Bartlett/standard periodograms.
2
3#[allow(unused_imports)]
4use crate::math::FloatMath;
5use crate::transform::cfft_f32;
6use crate::types::Status;
7use crate::window::*;
8
9/// Window function choice for spectral estimation.
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum WelchWindow {
12    Rectangular,
13    Hamming,
14    Hanning,
15    Blackman,
16    BlackmanHarris,
17    Bartlett,
18    Welch,
19}
20
21/// Computes the Power Spectral Density (PSD) using Welch's Method (Averaged Overlapped Segment Periodograms).
22///
23/// `src`: continuous input signal time-series.
24/// `dst_psd`: destination slice receiving the one-sided PSD of length `fft_len / 2 + 1` (or `fft_len / 2`).
25/// `fft_len`: FFT segment size (must be power of 2, $\le 512$).
26/// `overlap`: number of overlapping samples between successive FFT segments (must be $< \text{fft\_len}$).
27/// `sample_rate`: sampling frequency in Hz (e.g. 1000.0, 44100.0).
28/// `window`: window function applied to each segment.
29/// `return_db`: if `true`, returns PSD in decibels ($10 \log_{10}(\text{PSD})$). If `false`, returns linear power.
30pub fn welch_psd_f32(
31    src: &[f32],
32    dst_psd: &mut [f32],
33    fft_len: usize,
34    overlap: usize,
35    sample_rate: f32,
36    window: WelchWindow,
37    return_db: bool,
38) -> Status {
39    let out_bins = fft_len / 2;
40    if fft_len < 4 || (fft_len & (fft_len - 1)) != 0 || fft_len > 512 {
41        return Status::ArgumentError;
42    }
43    if overlap >= fft_len || sample_rate <= 0.0 {
44        return Status::ArgumentError;
45    }
46    if src.len() < fft_len || dst_psd.len() < out_bins {
47        return Status::LengthError;
48    }
49
50    let step = fft_len - overlap;
51    let num_segments = (src.len() - fft_len) / step + 1;
52    if num_segments == 0 {
53        return Status::LengthError;
54    }
55
56    // Generate window
57    let mut win = [1.0f32; 512];
58    match window {
59        WelchWindow::Rectangular => win[..fft_len].fill(1.0),
60        WelchWindow::Hamming => hamming_f32(&mut win[..fft_len]),
61        WelchWindow::Hanning => hanning_f32(&mut win[..fft_len]),
62        WelchWindow::Blackman => blackman_f32(&mut win[..fft_len]),
63        WelchWindow::BlackmanHarris => blackman_harris_f32(&mut win[..fft_len]),
64        WelchWindow::Bartlett => bartlett_f32(&mut win[..fft_len]),
65        WelchWindow::Welch => welch_f32(&mut win[..fft_len]),
66    }
67
68    // Window power sum for normalization
69    let mut win_power = 0.0f32;
70    for i in 0..fft_len {
71        win_power += win[i] * win[i];
72    }
73    if win_power == 0.0 {
74        win_power = 1.0;
75    }
76
77    dst_psd[..out_bins].fill(0.0);
78
79    let mut scratch = [0.0f32; 1024];
80
81    for seg in 0..num_segments {
82        let start_idx = seg * step;
83        for i in 0..fft_len {
84            scratch[2 * i] = src[start_idx + i] * win[i];
85            scratch[2 * i + 1] = 0.0;
86        }
87
88        cfft_f32(&mut scratch[..2 * fft_len], fft_len, 0, 1);
89
90        for k in 0..out_bins {
91            let re = scratch[2 * k];
92            let im = scratch[2 * k + 1];
93            let mag_sq = re * re + im * im;
94            dst_psd[k] += mag_sq;
95        }
96    }
97
98    // Normalization factor for one-sided PSD:
99    // 2.0 / (num_segments * sample_rate * win_power)
100    let norm = 2.0f32 / (num_segments as f32 * sample_rate * win_power);
101    for k in 0..out_bins {
102        let linear_psd = dst_psd[k] * norm;
103        if return_db {
104            let clamped = if linear_psd > 1e-14 {
105                linear_psd
106            } else {
107                1e-14
108            };
109            dst_psd[k] = 10.0 * clamped.log10();
110        } else {
111            dst_psd[k] = linear_psd;
112        }
113    }
114
115    Status::Success
116}
117
118/// Computes the single-segment Periodogram Power Spectral Density.
119pub fn periodogram_f32(
120    src: &[f32],
121    dst_psd: &mut [f32],
122    fft_len: usize,
123    sample_rate: f32,
124    window: WelchWindow,
125    return_db: bool,
126) -> Status {
127    welch_psd_f32(src, dst_psd, fft_len, 0, sample_rate, window, return_db)
128}