Hann window

2 files changed, 74 insertions, 10 deletions

diff --git a/src/dft.c b/src/dft.c
index b7c7592..7538b95 100644
--- a/src/dft.c
+++ b/src/dft.c
@@ -1,7 +1,9 @@
 #include "dft.h"
 #include "common.h"
+#include "smrt_arena.h"
 #include "vec2.h"
 
+#include <math.h>
 #include <string.h>
 
 dft_data_t discrete_fourier_transform(smrt_arena_t *arena, f64 *samples, u64 sample_count, u64 sample_rate) {
@@ -45,6 +47,17 @@ dft_data_t discrete_fourier_transform(smrt_arena_t *arena, f64 *samples, u64 sam
         return d;
 }
 
+static inline f64 hann_window(u64 i, u64 n) {
+    return 0.5 * (1.0 - cos((f64)i / n * 2.0 * PI));
+}
+
+// Mean value of hann_window over a full period; short_time_fourier_transform scales
+// analysis amplitudes up by 1/HANN_COHERENT_GAIN to compensate for the window's
+// attenuation, so inverse_short_time_fourier_transform must scale back down by
+// HANN_COHERENT_GAIN before resynthesizing, or the reconstructed signal comes out
+// 1/HANN_COHERENT_GAIN times too loud.
+#define HANN_COHERENT_GAIN 0.5
+
 static inline u64 bit_reverse(u64 n, u8 m) {
     u64 out = 0;
     for (u64 i = 0; i < m; i++) {
@@ -83,8 +96,8 @@ dft_data_t fast_fourier_transform(smrt_arena_t *arena, f64 *samples, u64 sample_
             vec2d_t w = VEC2D_FROM(1.0, 0.0);
 
             for (u64 j = 0; j < n; j++) {
-                vec2d_t t = vec2d_cmul(w, vec2d_soa_get(&vs, k + j + n));
                 vec2d_t u = vec2d_soa_get(&vs, k + j);
+                vec2d_t t = vec2d_cmul(w, vec2d_soa_get(&vs, k + j + n));
 
                 vec2d_soa_set(&vs, k + j    , vec2d_add(u, t));
                 vec2d_soa_set(&vs, k + j + n, vec2d_sub(u, t));
@@ -165,6 +178,14 @@ stft_data_t short_time_fourier_transform(smrt_arena_t *arena, u64 window_size, u
 
         f64 *seg_samples = samples + start;
 
+        smrta_temp_t scratch = smrta_scratch_start(NULL, 0);
+
+        f64 *seg_hann = SMRTA_ALLOC_ARRAY(scratch.arena, f64, window_size);
+
+        for (u64 h = 0; h < window_size; h++) {
+            seg_hann[h] = seg_samples[h] * hann_window(h, window_size);
+        }
+
         stft_segment_t seg = {
             .start_index = start,
             .sample_count = window_size,
@@ -172,12 +193,18 @@ stft_data_t short_time_fourier_transform(smrt_arena_t *arena, u64 window_size, u
 
         seg.data = fast_fourier_transform(
             arena,
-            seg_samples,
+            seg_hann,
             window_size,
             sample_rate,
             NULL, 0
         );
 
+        for (u64 a = 0; a < seg.data.freq_count; a++) {
+            seg.data.amplitudes[a] /= HANN_COHERENT_GAIN;
+        }
+
+        smrta_scratch_end(scratch);
+
         segments[i] = seg;
     }
 
@@ -260,8 +287,16 @@ f64 *inverse_short_time_fourier_transform(smrt_arena_t *arena, stft_data_t stft,
     for (u64 k = 0;  k < stft.segment_count; k++) {
         smrt_arena_mark(scratch.arena);
 
+        // Undo short_time_fourier_transform's coherent-gain compensation (a copy, so
+        // the caller's stft.segments data is left untouched for repeated/analysis use).
+        dft_data_t seg = frames[k];
+        seg.amplitudes = SMRTA_ALLOC_ARRAY(scratch.arena, f64, seg.freq_count);
+        for (u64 a = 0; a < seg.freq_count; a++) {
+            seg.amplitudes[a] = frames[k].amplitudes[a] * HANN_COHERENT_GAIN;
+        }
+
         f64 *spec_real, *spec_imag;
-        reconstruct_spectrum(scratch.arena, &frames[k], window_size, &spec_real, &spec_imag);
+        reconstruct_spectrum(scratch.arena, &seg, window_size, &spec_real, &spec_imag);
 
         f64 *real_o;
         inverse_fast_fourier_transform(scratch.arena, spec_real, spec_imag, window_size, &real_o, NULL, NULL, 0);
@@ -269,8 +304,9 @@ f64 *inverse_short_time_fourier_transform(smrt_arena_t *arena, stft_data_t stft,
         u64 start = k * hop_size;
 
         for (u64 j = 0; j < window_size; j++) {
-            output[start + j] += real_o[j];
-            weights[start + j] += 1.0;
+            f64 w = hann_window(j, window_size);
+            output[start + j]  += real_o[j] * w;
+            weights[start + j] += w * w;
         }
 
         smrt_arena_pop_to_mark(scratch.arena);
diff --git a/tests/dft.c b/tests/dft.c
index 3f3d7f7..d4e88cb 100644
--- a/tests/dft.c
+++ b/tests/dft.c
@@ -8,6 +8,7 @@
 #include <criterion/criterion.h>
 #include <criterion/internal/assert.h>
 #include <criterion/internal/test.h>
+#include <criterion/redirect.h>
 
 f64 *mock_amplitude_data(smrt_arena_t *arena, f64 *freqs, f64 *amps, u64 n, f64 duration, u64 sample_count) {
     f64 *samples = SMRTA_ALLOC_ARRAY(arena, f64, sample_count);
@@ -197,21 +198,37 @@ Test(dft, stft_segments_match_direct_transform) {
 
     u64 expected_starts[3] = { 0, 4, 8 };
 
+    // short_time_fourier_transform hann-windows each segment before transforming it, then
+    // scales amplitudes by 1/HANN_COHERENT_GAIN (0.5) to compensate for the window's average
+    // attenuation. Mirror both steps here to build a comparable direct transform of the same
+    // windowed data, rather than comparing against a direct transform of the raw slice.
     for (u64 s = 0; s < stft.segment_count; s++) {
         stft_segment_t seg = stft.segments[s];
         cr_expect_eq(seg.start_index, expected_starts[s]);
         cr_expect_eq(seg.sample_count, window_size);
 
-        dft_data_t direct = discrete_fourier_transform(arena, samples + seg.start_index, window_size, sample_rate);
+        f64 windowed[8];
+        for (u64 h = 0; h < window_size; h++) {
+            f64 hann = 0.5 * (1.0 - cos((f64)h / window_size * 2.0 * PI));
+            windowed[h] = samples[seg.start_index + h] * hann;
+        }
+
+        dft_data_t direct = discrete_fourier_transform(arena, windowed, window_size, sample_rate);
 
         cr_assert_eq(seg.data.freq_count, direct.freq_count);
         for (u64 i = 0; i < direct.freq_count; i++) {
+            f64 expected_amplitude = direct.amplitudes[i] / 0.5; // undo HANN_COHERENT_GAIN
+
             cr_expect(F64_EQ(seg.data.frequencies[i], direct.frequencies[i], 1e-9));
-            cr_expect(F64_EQ(seg.data.amplitudes[i],  direct.amplitudes[i],  1e-9));
+            cr_expect(F64_EQ(seg.data.amplitudes[i],  expected_amplitude,    1e-9));
 
-            // phase is only meaningful at bins with non-negligible amplitude
+            // phase is only meaningful at bins with non-negligible amplitude. Compare
+            // angles modulo 2*PI: atan2 can return either +PI or -PI for the same angle,
+            // which a plain subtraction would wrongly see as a ~2*PI difference.
             if (direct.amplitudes[i] > 1e-6) {
-                cr_expect(F64_EQ(seg.data.phases[i], direct.phases[i], 1e-9));
+                f64 phase_diff = seg.data.phases[i] - direct.phases[i];
+                phase_diff -= 2.0 * PI * round(phase_diff / (2.0 * PI));
+                cr_expect(F64_EQ(phase_diff, 0.0, 1e-9));
             }
         }
     }
@@ -376,6 +393,8 @@ Test(dft, data_to_wav_clamps_out_of_range_amplitude) {
 }
 
 Test(dft, stft_rejects_zero_hop_size, .signal = SIGABRT) {
+    cr_redirect_stderr(); // assert() prints to stderr before aborting; that's expected
+
     smrt_arena_t *arena = smrt_arena_create(KiB(16), KiB(4), false);
 
     f64 samples[8] = { 0 };
@@ -383,6 +402,8 @@ Test(dft, stft_rejects_zero_hop_size, .signal = SIGABRT) {
 }
 
 Test(dft, stft_rejects_zero_sample_count, .signal = SIGABRT) {
+    cr_redirect_stderr(); // assert() prints to stderr before aborting; that's expected
+
     smrt_arena_t *arena = smrt_arena_create(KiB(16), KiB(4), false);
 
     f64 samples[1] = { 0 };
@@ -390,6 +411,8 @@ Test(dft, stft_rejects_zero_sample_count, .signal = SIGABRT) {
 }
 
 Test(dft, reconstruct_spectrum_rejects_mismatched_freq_count, .signal = SIGABRT) {
+    cr_redirect_stderr(); // assert() prints to stderr before aborting; that's expected
+
     smrt_arena_t *arena = smrt_arena_create(KiB(16), KiB(4), false);
 
     // sample_count=8 expects freq_count == 5; give it 3 instead
@@ -415,7 +438,12 @@ Test(dft, istft_round_trip_recovers_samples) {
 
     f64 *output = inverse_short_time_fourier_transform(arena, stft, NULL, 0);
 
-    for (u64 i = 0; i < sample_count; i++) {
+    // Sample 0 is inherently unrecoverable: periodic hann_window(0, window_size) == 0,
+    // so the analysis window zeroes out samples[0]'s contribution before it ever reaches
+    // the FFT, and no other segment covers index 0 to make up for it. Every other sample
+    // is covered by a segment with a nonzero window weight, which weighted overlap-add
+    // recovers exactly.
+    for (u64 i = 1; i < sample_count; i++) {
         cr_expect(F64_EQ(output[i], samples[i], 1e-9));
     }