diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/dft.c | 139 | ||||
| -rw-r--r-- | src/dft.h | 48 |
2 files changed, 187 insertions, 0 deletions
diff --git a/src/dft.c b/src/dft.c new file mode 100644 index 0000000..430f5e8 --- /dev/null +++ b/src/dft.c @@ -0,0 +1,139 @@ +#include "dft.h" +#include "vec2.h" + +#include <string.h> + +dft_data_t discrete_fourier_transform(smrt_arena_t *arena, f64 *samples, u64 sample_count, u64 sample_rate) { + u64 freq_count = (sample_count / 2) + 1; + f64 freq_step = (f64)sample_rate / (f64)sample_count; + + dft_data_t d = {.freq_count = freq_count}; + d.frequencies = SMRTA_ALLOC_ARRAY(arena, f64, freq_count); + d.amplitudes = SMRTA_ALLOC_ARRAY(arena, f64, freq_count); + d.phases = SMRTA_ALLOC_ARRAY(arena, f64, freq_count); + + for (u64 freq_index = 0; freq_index < freq_count; freq_index++) { + smrta_temp_t scratch = smrta_scratch_start(NULL, 0); + + vec2d_soa_t vs; + vs.xs = SMRTA_ALLOC_ARRAY(scratch.arena, f64, sample_count); + vs.ys = SMRTA_ALLOC_ARRAY(scratch.arena, f64, sample_count); + vs.size = sample_count; + + if (!vs.xs || !vs.ys) { smrta_scratch_end(scratch); return (dft_data_t){ 0 }; } + + for (u64 i = 0; i < sample_count; i++) { + f64 angle = i / (f64)sample_count * PI * 2.0 * freq_index; + + vs.xs[i] = cos(angle) * samples[i]; + vs.ys[i] = sin(angle) * samples[i]; + } + + vec2d_t average_pos = vec2d_soa_average(&vs); + smrta_scratch_end(scratch); + + b8 is_zero_hz = freq_index == 0; + b8 is_nyquist = freq_index == freq_count - 1 && sample_count % 2 == 0; + f64 amp_coeff = is_zero_hz || is_nyquist ? 1.0 : 2.0; + + d.frequencies[freq_index] = freq_index * freq_step; + d.amplitudes[freq_index] = vec2d_length(average_pos) * amp_coeff; + d.phases[freq_index] = -atan2(average_pos.y, average_pos.x); + + } + return d; +} + +wav_data_t dft_data_to_wav(smrt_arena_t *arena, dft_data_t dft, u64 sample_rate, f64 duration) { + u64 sample_count = (u64)(sample_rate * duration); + + u8 *data = SMRTA_ALLOC_ARRAY(arena, u8, sample_count); + if (!data) return (wav_data_t){ 0 }; + + wav_data_t d = { + .sample_count=sample_count, + .samples = data, + }; + + for (u64 s_num = 0; s_num < sample_count; s_num++) { + f64 amp = 0.0; + f64 t = ((f64)s_num / sample_count) * duration; + for (u64 freq_index = 0; freq_index < dft.freq_count; freq_index++) { + f64 f = dft.frequencies[freq_index]; + f64 a = dft.amplitudes[freq_index]; + f64 p = dft.phases[freq_index]; + amp += cos((t * f * 2.0 * PI) + p) * a; + } + data[s_num] = (u8)((amp+1.0) * (UINT8_MAX/2)); + } + + return d; +} + +stft_data_t short_time_fourier_transform(smrt_arena_t *arena, u64 samples_per_segment, f64 *samples, u64 sample_count, u64 sample_rate) { + u64 segment_count = (u64)ceil((f64)sample_count / samples_per_segment); + + stft_segment_t *segments = SMRTA_ALLOC_ARRAY(arena, stft_segment_t, segment_count); + if (!segments) return (stft_data_t){ 0 }; + + for (u64 seg_num = 0; seg_num < segment_count; seg_num++) { + + u64 samples_in_this_seg = samples_per_segment; + if (seg_num == segment_count - 1 && !(sample_count % samples_per_segment == 0)) + samples_in_this_seg = sample_count % samples_per_segment; + + stft_segment_t seg = { + .start_index = samples_per_segment * seg_num, + .sample_count = samples_in_this_seg, + }; + + seg.data = discrete_fourier_transform( + arena, + samples + seg.start_index, + samples_in_this_seg, + sample_rate + ); + + segments[seg_num] = seg; + } + + stft_data_t data = { + .segment_count = segment_count, + .sample_rate = sample_rate, + .segments = segments + }; + + return data; +} + +wav_data_t stft_data_to_wav(smrt_arena_t *arena, stft_data_t stft, u64 sample_count) { + u8 *data = SMRTA_ALLOC_ARRAY(arena, u8, sample_count); + if (!data) return (wav_data_t){ 0 }; + + wav_data_t d = { + .sample_count=sample_count, + .samples = data, + }; + + u64 offset = 0; + for (u64 seg_num = 0; seg_num < stft.segment_count; seg_num++) { + smrta_temp_t scratch = smrta_scratch_start(NULL, 0); + + stft_segment_t seg = stft.segments[seg_num]; + + wav_data_t wav_data = dft_data_to_wav( + scratch.arena, + seg.data, + stft.sample_rate, + (f64)seg.sample_count / stft.sample_rate + ); + + memcpy(data + offset, wav_data.samples, wav_data.sample_count); + + offset += wav_data.sample_count * sizeof(u8); + + smrta_scratch_end(scratch); + } + + return d; +} diff --git a/src/dft.h b/src/dft.h new file mode 100644 index 0000000..982cc6f --- /dev/null +++ b/src/dft.h @@ -0,0 +1,48 @@ +#pragma once + +#include "common.h" +#include "smrt_arena.h" +#include "wav.h" + +typedef struct { + f64 *frequencies; + f64 * amplitudes; + f64 * phases; + u64 freq_count; +} dft_data_t; + +/// Run DFT algorithm on samples +dft_data_t discrete_fourier_transform(smrt_arena_t * arena , + f64 * samples , + u64 sample_count , + u64 sample_rate); + +/// Convert DFT frequency data to unsigned 8-bit PCM WAV amplitude data +wav_data_t dft_data_to_wav(smrt_arena_t * arena , + dft_data_t dft , + u64 sample_rate , + f64 duration); + +typedef struct { + dft_data_t data; + u64 sample_count; + u64 start_index; +} stft_segment_t; + +typedef struct { + stft_segment_t * segments; + u64 segment_count; + u64 sample_rate; +} stft_data_t; + +/// Run STFT algorithm on samples. +stft_data_t short_time_fourier_transform(smrt_arena_t * arena , + u64 samples_per_segment , + f64 * samples , + u64 sample_count , + u64 sample_rate); + +/// Convert STFT frequency data to unsigned 8-bit PCM WAV amplitude data +wav_data_t stft_data_to_wav(smrt_arena_t * arena , + stft_data_t stft , + u64 sample_count); |