1use crate::{
4 audio::{
5 AudioFeatureLayout, AudioFrontendConfig, MelFilterConfig, PadMode, StftConfig, WindowKind,
6 },
7 error::{Error, Result},
8 utility::checked_element_count,
9};
10
11pub fn stft_power(input: &[f32], config: StftConfig) -> Result<Vec<f32>> {
12 Ok(power_from_r2c(&stft_r2c(input, config)?))
13}
14
15pub fn stft_r2c(input: &[f32], config: StftConfig) -> Result<Vec<f32>> {
16 validate_stft_input(input.len(), config)?;
17 let window = stft_window(config.window, config.win_length);
18 let (dft_real, dft_imag) = dft_tables(config.n_fft, config.frequency_bin_count())?;
19 let frames = config.frame_count(input.len())?;
20 let mut out = vec![0.0; frames * config.frequency_bin_count() * 2];
21 let pad = if config.center { config.n_fft / 2 } else { 0 };
22
23 for frame in 0..frames {
24 for frequency in 0..config.frequency_bin_count() {
25 let mut real = 0.0f32;
26 let mut imag = 0.0f32;
27 for (sample, window_value) in window.iter().copied().enumerate().take(config.n_fft) {
28 let raw_index = frame * config.hop_length + sample;
29 let value =
30 padded_sample(input, raw_index as isize - pad as isize, config.pad_mode)?
31 * window_value;
32 let table_offset = frequency * config.n_fft + sample;
33 real += value * dft_real[table_offset];
34 imag += value * dft_imag[table_offset];
35 }
36 let output_offset = (frame * config.frequency_bin_count() + frequency) * 2;
37 out[output_offset] = real;
38 out[output_offset + 1] = imag;
39 }
40 }
41
42 Ok(out)
43}
44
45pub fn padded_sample(input: &[f32], index: isize, pad_mode: PadMode) -> Result<f32> {
46 match pad_mode {
47 PadMode::Reflect => Ok(input[reflect_index(index, input.len())?]),
48 PadMode::Zero => {
49 if index < 0 || index >= input.len() as isize {
50 Ok(0.0)
51 } else {
52 Ok(input[index as usize])
53 }
54 }
55 }
56}
57
58pub fn power_from_r2c(spectrum: &[f32]) -> Vec<f32> {
59 spectrum
60 .chunks_exact(2)
61 .map(|complex| complex[0] * complex[0] + complex[1] * complex[1])
62 .collect()
63}
64
65pub fn log_mel(power: &[f32], frames: usize, config: AudioFrontendConfig) -> Result<Vec<f32>> {
66 config.validate()?;
67 let expected_power = checked_element_count(frames, config.mel.frequency_bin_count)?;
68 if power.len() < expected_power || config.log_mel.dynamic_range.is_none() {
69 return Err(Error::InvalidLength);
70 }
71 let filters = slaney_mel_filters(config.mel)?;
72 let dynamic_range = config.log_mel.dynamic_range.ok_or(Error::InvalidLength)?;
73 let mut raw = vec![0.0; frames * config.mel.mel_bin_count];
74
75 for frame in 0..frames {
76 for mel in 0..config.mel.mel_bin_count {
77 let mut sum = 0.0f32;
78 for frequency in 0..config.mel.frequency_bin_count {
79 sum += power[frame * config.mel.frequency_bin_count + frequency]
80 * filters[mel * config.mel.frequency_bin_count + frequency];
81 }
82 raw[frame * config.mel.mel_bin_count + mel] = sum.max(config.log_mel.floor).log10();
83 }
84 }
85
86 let reference_max = config
87 .log_mel
88 .reference_max
89 .or_else(|| raw.iter().copied().reduce(f32::max))
90 .ok_or(Error::InvalidLength)?;
91 let lower = reference_max - dynamic_range;
92 let mut out = vec![0.0; frames * config.mel.mel_bin_count];
93
94 for frame in 0..frames {
95 for mel in 0..config.mel.mel_bin_count {
96 let mut value = raw[frame * config.mel.mel_bin_count + mel];
97 value = value.max(lower);
98 value = (value + config.log_mel.offset) / config.log_mel.scale;
99 let output_offset = match config.layout {
100 AudioFeatureLayout::FramesFirst => frame * config.mel.mel_bin_count + mel,
101 AudioFeatureLayout::MelsFirst => mel * frames + frame,
102 };
103 out[output_offset] = value;
104 }
105 }
106
107 Ok(out)
108}
109
110pub fn sparse_log_mel_range(
111 full: &[f32],
112 config: AudioFrontendConfig,
113 first_frame: usize,
114 frames: usize,
115 total_frames: usize,
116) -> Vec<f32> {
117 let mut out = vec![0.0; full.len()];
118 for local_frame in 0..frames {
119 let frame = first_frame + local_frame;
120 for mel in 0..config.mel.mel_bin_count {
121 let offset = match config.layout {
122 AudioFeatureLayout::FramesFirst => frame * config.mel.mel_bin_count + mel,
123 AudioFeatureLayout::MelsFirst => mel * total_frames + frame,
124 };
125 out[offset] = full[offset];
126 }
127 }
128 out
129}
130
131pub fn convert_log_mel_layout(
132 input: &[f32],
133 batch: usize,
134 frames: usize,
135 mel_bins: usize,
136 input_layout: AudioFeatureLayout,
137 output_layout: AudioFeatureLayout,
138) -> Vec<f32> {
139 let mut output = vec![0.0; input.len()];
140 for batch_index in 0..batch {
141 for frame in 0..frames {
142 for mel in 0..mel_bins {
143 let input_offset =
144 log_mel_offset(batch_index, frame, mel, frames, mel_bins, input_layout);
145 let output_offset =
146 log_mel_offset(batch_index, frame, mel, frames, mel_bins, output_layout);
147 output[output_offset] = input[input_offset];
148 }
149 }
150 }
151 output
152}
153
154pub fn log_mel_offset(
155 batch_index: usize,
156 frame: usize,
157 mel: usize,
158 frames: usize,
159 mel_bins: usize,
160 layout: AudioFeatureLayout,
161) -> usize {
162 let feature_len = frames * mel_bins;
163 match layout {
164 AudioFeatureLayout::FramesFirst => batch_index * feature_len + frame * mel_bins + mel,
165 AudioFeatureLayout::MelsFirst => batch_index * feature_len + mel * frames + frame,
166 }
167}
168
169pub fn reflect_index(index: isize, len: usize) -> Result<usize> {
170 if len == 0 {
171 return Err(Error::InvalidLength);
172 }
173 if len == 1 {
174 return if index == 0 {
175 Ok(0)
176 } else {
177 Err(Error::InvalidLength)
178 };
179 }
180 let reflected = if index < 0 {
181 (-index) as usize
182 } else if index >= len as isize {
183 (2 * len - 2) - index as usize
184 } else {
185 index as usize
186 };
187 if reflected >= len {
188 return Err(Error::InvalidLength);
189 }
190 Ok(reflected)
191}
192
193fn validate_stft_input(input_len: usize, config: StftConfig) -> Result<()> {
194 validate_stft_config(config)?;
195 if config.pad_mode != PadMode::Reflect {
196 return Ok(());
197 }
198 if input_len == 0 {
199 return Err(Error::InvalidLength);
200 }
201 if config.center && input_len <= config.n_fft / 2 {
202 return Err(Error::InvalidLength);
203 }
204 Ok(())
205}
206
207fn validate_stft_config(config: StftConfig) -> Result<()> {
208 if config.n_fft == 0 || config.hop_length == 0 || config.win_length == 0 {
209 return Err(Error::InvalidLength);
210 }
211 if config.win_length > config.n_fft {
212 return Err(Error::InvalidLength);
213 }
214 if !config.n_fft.is_multiple_of(2) {
215 return Err(Error::UnsupportedParameter {
216 op: "stft".into(),
217 parameter: "n_fft".into(),
218 value: config.n_fft,
219 });
220 }
221 Ok(())
222}
223
224fn validate_mel_filter_config(config: MelFilterConfig) -> Result<()> {
225 if config.sample_rate == 0 || config.frequency_bin_count == 0 || config.mel_bin_count == 0 {
226 return Err(Error::InvalidLength);
227 }
228 let nyquist = config.sample_rate as f32 / 2.0;
229 if !config.min_frequency.is_finite()
230 || !config.max_frequency.is_finite()
231 || config.min_frequency < 0.0
232 || config.max_frequency <= config.min_frequency
233 || config.max_frequency > nyquist
234 {
235 return Err(Error::InvalidLength);
236 }
237 Ok(())
238}
239
240fn stft_window(kind: WindowKind, len: usize) -> Vec<f32> {
241 match kind {
242 WindowKind::Rectangular => vec![1.0; len],
243 WindowKind::PeriodicHann => (0..len)
244 .map(|index| {
245 let angle = 2.0 * std::f32::consts::PI * index as f32 / len as f32;
246 0.5 * (1.0 - angle.cos())
247 })
248 .collect(),
249 WindowKind::PeriodicHamming => {
250 const ALPHA: f32 = 0.54;
251 const BETA: f32 = 1.0 - ALPHA;
252 (0..len)
253 .map(|index| {
254 let angle = 2.0 * std::f32::consts::PI * index as f32 / len as f32;
255 ALPHA - BETA * angle.cos()
256 })
257 .collect()
258 }
259 }
260}
261
262fn dft_tables(n_fft: usize, frequency_bin_count: usize) -> Result<(Vec<f32>, Vec<f32>)> {
263 let len = checked_element_count(n_fft, frequency_bin_count)?;
264 let mut real = Vec::with_capacity(len);
265 let mut imag = Vec::with_capacity(len);
266 for frequency in 0..frequency_bin_count {
267 for sample in 0..n_fft {
268 let angle =
269 2.0 * std::f32::consts::PI * frequency as f32 * sample as f32 / n_fft as f32;
270 real.push(angle.cos());
271 imag.push(-angle.sin());
272 }
273 }
274 Ok((real, imag))
275}
276
277fn slaney_mel_filters(config: MelFilterConfig) -> Result<Vec<f32>> {
278 validate_mel_filter_config(config)?;
279 let len = checked_element_count(config.mel_bin_count, config.frequency_bin_count)?;
280 let mut filters = vec![0.0; len];
281
282 let min_mel = hertz_to_slaney_mel(config.min_frequency);
283 let max_mel = hertz_to_slaney_mel(config.max_frequency);
284 let step = (max_mel - min_mel) / (config.mel_bin_count + 1) as f32;
285
286 let mut mel_edges = Vec::with_capacity(config.mel_bin_count + 2);
287 for index in 0..config.mel_bin_count + 2 {
288 mel_edges.push(slaney_mel_to_hertz(min_mel + step * index as f32));
289 }
290
291 let mut fft_frequencies = Vec::with_capacity(config.frequency_bin_count);
292 for bin in 0..config.frequency_bin_count {
293 let frequency =
294 bin as f32 * config.sample_rate as f32 / (2 * (config.frequency_bin_count - 1)) as f32;
295 fft_frequencies.push(frequency);
296 }
297
298 for mel in 0..config.mel_bin_count {
299 let lower = mel_edges[mel];
300 let center = mel_edges[mel + 1];
301 let upper = mel_edges[mel + 2];
302 let enorm = 2.0 / (upper - lower);
303
304 for (frequency_bin, frequency) in fft_frequencies.iter().copied().enumerate() {
305 let lower_slope = (frequency - lower) / (center - lower);
306 let upper_slope = (upper - frequency) / (upper - center);
307 let weight = lower_slope.min(upper_slope).max(0.0) * enorm;
308 filters[mel * config.frequency_bin_count + frequency_bin] = weight;
309 }
310 }
311
312 Ok(filters)
313}
314
315fn hertz_to_slaney_mel(frequency: f32) -> f32 {
316 const MIN_LOG_HZ: f32 = 1000.0;
317 const MIN_LOG_MEL: f32 = 15.0;
318 const LOG_STEP: f32 = 0.06875178;
319
320 if frequency < MIN_LOG_HZ {
321 frequency * 0.015
322 } else {
323 MIN_LOG_MEL + (frequency / MIN_LOG_HZ).ln() / LOG_STEP
324 }
325}
326
327fn slaney_mel_to_hertz(mel: f32) -> f32 {
328 const MIN_LOG_HZ: f32 = 1000.0;
329 const MIN_LOG_MEL: f32 = 15.0;
330 const LOG_STEP: f32 = 0.06875178;
331
332 if mel < MIN_LOG_MEL {
333 mel / 0.015
334 } else {
335 MIN_LOG_HZ * (LOG_STEP * (mel - MIN_LOG_MEL)).exp()
336 }
337}
338
339#[cfg(test)]
340mod tests {
341 use super::*;
342 use crate::audio::{LogMelConfig, SpectralNormalization, SpectrumKind};
343
344 #[test]
345 fn padded_sample_rejects_empty_reflect_input() {
346 assert!(matches!(
347 padded_sample(&[], 0, PadMode::Reflect),
348 Err(Error::InvalidLength)
349 ));
350 }
351
352 #[test]
353 fn stft_power_rejects_too_short_centered_reflect_input() {
354 let config = test_stft_config(PadMode::Reflect);
355
356 assert!(matches!(
357 stft_power(&[0.0, 1.0], config),
358 Err(Error::InvalidLength)
359 ));
360 }
361
362 #[test]
363 fn stft_power_allows_empty_centered_zero_padded_input() -> Result<()> {
364 let config = test_stft_config(PadMode::Zero);
365
366 let actual = stft_power(&[], config)?;
367
368 assert_eq!(
369 actual.len(),
370 config.frame_count(0)? * config.frequency_bin_count()
371 );
372 assert!(actual.iter().all(|value| *value == 0.0));
373 Ok(())
374 }
375
376 #[test]
377 fn log_mel_rejects_empty_dynamic_reference_input() {
378 let mut config = test_frontend_config();
379 config.log_mel.reference_max = None;
380
381 assert!(matches!(log_mel(&[], 0, config), Err(Error::InvalidLength)));
382 }
383
384 fn test_stft_config(pad_mode: PadMode) -> StftConfig {
385 StftConfig {
386 n_fft: 8,
387 hop_length: 4,
388 win_length: 8,
389 center: true,
390 pad_mode,
391 window: WindowKind::Rectangular,
392 spectrum: SpectrumKind::OneSide,
393 normalization: SpectralNormalization::None,
394 drop_last_frame: false,
395 }
396 }
397
398 fn test_frontend_config() -> AudioFrontendConfig {
399 AudioFrontendConfig {
400 stft: test_stft_config(PadMode::Zero),
401 mel: MelFilterConfig {
402 sample_rate: 16_000,
403 frequency_bin_count: 5,
404 mel_bin_count: 2,
405 min_frequency: 0.0,
406 max_frequency: 8_000.0,
407 },
408 log_mel: LogMelConfig {
409 floor: 1e-10,
410 reference_max: Some(0.0),
411 dynamic_range: Some(8.0),
412 offset: 4.0,
413 scale: 4.0,
414 },
415 layout: AudioFeatureLayout::FramesFirst,
416 }
417 }
418}