From 1646dfb9159e66c30821a5659e8f0a5eb3d52f05 Mon Sep 17 00:00:00 2001 From: steven-na Date: Wed, 5 Aug 2026 21:23:31 -0700 Subject: Formatting, bug fix, test coverage --- src/dft.c | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) (limited to 'src/dft.c') diff --git a/src/dft.c b/src/dft.c index 7f90318..b7c7592 100644 --- a/src/dft.c +++ b/src/dft.c @@ -139,6 +139,7 @@ wav_data_t dft_data_to_wav(smrt_arena_t *arena, dft_data_t dft, u64 sample_rate, f64 p = dft.phases[freq_index]; amp += cos((t * f * 2.0 * PI) + p) * a; } + amp = amp < -1.0 ? -1.0 : amp > 1.0 ? 1.0 : amp; data[s_num] = (u8)((amp+1.0) * (UINT8_MAX/2)); } @@ -148,9 +149,12 @@ wav_data_t dft_data_to_wav(smrt_arena_t *arena, dft_data_t dft, u64 sample_rate, stft_data_t short_time_fourier_transform(smrt_arena_t *arena, u64 window_size, u64 hop_size, f64 *samples, u64 sample_count, u64 sample_rate) { assert(F64_EQ(round(log2(window_size)), log2(window_size), 1e-9) && "STFT input window_size must be a power of 2"); + assert(hop_size != 0 && "STFT input hop_size must be nonzero"); sample_count = ALIGN_UP_POW2(sample_count * sizeof(f64), STFT_SAMPLE_ALIGN_BYTES(window_size)) / sizeof(f64); + assert(sample_count >= window_size && "STFT input sample_count must be at least window_size"); + u64 segment_count = ((sample_count - window_size) / hop_size) + 1; stft_segment_t *segments = SMRTA_ALLOC_ARRAY(arena, stft_segment_t, segment_count); @@ -179,8 +183,11 @@ stft_data_t short_time_fourier_transform(smrt_arena_t *arena, u64 window_size, u return (stft_data_t){ .sample_rate=sample_rate, + .window_size=window_size, + .hop_size=hop_size, .segment_count=segment_count, .segments=segments, + .total_samples=sample_count, }; } @@ -233,10 +240,13 @@ void inverse_fast_fourier_transform(smrt_arena_t *arena, f64 const *real, f64 co smrta_scratch_end(scratch); } -f64 *inverse_short_time_fourier_transform(smrt_arena_t *arena, stft_data_t stft, u64 window_size, u64 hop_size, smrt_arena_t **conflicts, u64 num_conflicts) { +f64 *inverse_short_time_fourier_transform(smrt_arena_t *arena, stft_data_t stft, smrt_arena_t **conflicts, u64 num_conflicts) { + u64 window_size = stft.window_size; + u64 hop_size = stft.hop_size; + smrta_temp_t scratch = smrta_scratch_start(conflicts, num_conflicts); - u64 output_len = (stft.segment_count - 1) * hop_size + window_size; + u64 output_len = stft.total_samples; dft_data_t *frames = SMRTA_ALLOC_ARRAY(scratch.arena, dft_data_t, stft.segment_count); @@ -277,6 +287,9 @@ f64 *inverse_short_time_fourier_transform(smrt_arena_t *arena, stft_data_t stft, } void reconstruct_spectrum(smrt_arena_t *arena, dft_data_t const *data, u64 sample_count, f64 **real_o, f64 **imag_o) { + assert(data->freq_count == (sample_count / 2) + 1 && + "reconstruct_spectrum: data->freq_count must equal (sample_count / 2) + 1"); + f64 *real = *real_o = SMRTA_ALLOC_ARRAY(arena, f64, sample_count); f64 *imag = *imag_o = SMRTA_ALLOC_ARRAY(arena, f64, sample_count); -- cgit v1.2.3