1#[allow(unused_imports)]
4use crate::math::FloatMath;
5use crate::transform::cfft_f32;
6use crate::types::Status;
7use crate::window::*;
8
9#[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
21pub 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 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 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 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
118pub 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}