Compute STFT magnitude squared, returned as flat Vec [n_frames * n_freq].
(samples: &[f32], n_fft: usize, hop_length: usize)
| 50 | |
| 51 | /// Compute STFT magnitude squared, returned as flat Vec [n_frames * n_freq]. |
| 52 | fn compute_stft(samples: &[f32], n_fft: usize, hop_length: usize) -> Vec<f32> { |
| 53 | let n_freq = n_fft / 2 + 1; |
| 54 | let window = hann_window(n_fft); |
| 55 | |
| 56 | let mut planner = FftPlanner::new(); |
| 57 | let fft = planner.plan_fft_forward(n_fft); |
| 58 | |
| 59 | let n_frames = if samples.len() >= n_fft { |
| 60 | (samples.len() - n_fft) / hop_length + 1 |
| 61 | } else { |
| 62 | 0 |
| 63 | }; |
| 64 | |
| 65 | // Pre-allocate result and reusable FFT buffer |
| 66 | let mut result = vec![0.0f32; n_frames * n_freq]; |
| 67 | let mut buffer = vec![Complex::new(0.0f32, 0.0); n_fft]; |
| 68 | |
| 69 | for frame_idx in 0..n_frames { |
| 70 | let start = frame_idx * hop_length; |
| 71 | |
| 72 | // Fill buffer with windowed samples (reuse allocation) |
| 73 | for i in 0..n_fft { |
| 74 | let sample = if start + i < samples.len() { |
| 75 | samples[start + i] |
| 76 | } else { |
| 77 | 0.0 |
| 78 | }; |
| 79 | buffer[i] = Complex::new(sample * window[i], 0.0); |
| 80 | } |
| 81 | |
| 82 | fft.process(&mut buffer); |
| 83 | |
| 84 | // Magnitude squared for first n_freq bins |
| 85 | let out_offset = frame_idx * n_freq; |
| 86 | for (j, item) in buffer.iter().take(n_freq).enumerate() { |
| 87 | result[out_offset + j] = item.norm_sqr(); |
| 88 | } |
| 89 | } |
| 90 | |
| 91 | result |
| 92 | } |
| 93 | |
| 94 | /// Generate a Hann window of the given size. |
| 95 | fn hann_window(size: usize) -> Vec<f32> { |