Implement synchrosqueezing transform for sharp spectrogram display

Co-authored-by: Qwen-Coder <qwen-coder@alibabacloud.com>
This commit is contained in:
2026-04-10 23:27:56 -07:00
parent 5b2d8ddb33
commit 0bc30832be
5 changed files with 341 additions and 10 deletions
+117 -10
View File
@@ -65,6 +65,7 @@ typedef struct {
typedef struct {
FrequencyData* spectrum;
FrequencyData* derivativeSpectrum; // STFT with derivative window (for synchrosqueezing)
int numBins;
int sampleOffset;
int sampleCount;
@@ -116,6 +117,7 @@ typedef struct {
ColormapType colormap;
bool showGrid;
int fftSize; // Current FFT size (128-2048)
bool useSynchrosqueezing; // Enable synchrosqueezing for sharper display
// File browser state
bool showFileBrowser;
@@ -280,6 +282,7 @@ static void ComputeSTFT(AudioSignal* signal, StftResult* result, int fftSize)
int numBins = fftSize / 2 + 1;
float* windowedSamples = (float*)malloc(fftSize * sizeof(float));
float* derivWindowedSamples = (float*)malloc(fftSize * sizeof(float));
float complex *complexInput = (float complex*)malloc(fftSize * sizeof(float complex));
float complex* fftOutput = (float complex*)malloc(fftSize * sizeof(float complex));
@@ -289,18 +292,23 @@ static void ComputeSTFT(AudioSignal* signal, StftResult* result, int fftSize)
if (offset + samplesToCopy > signal->numSamples) {
samplesToCopy = signal->numSamples - offset;
memset(windowedSamples, 0, fftSize * sizeof(float));
memset(derivWindowedSamples, 0, fftSize * sizeof(float));
} else {
memcpy(windowedSamples, signal->samples + offset, fftSize * sizeof(float));
memcpy(derivWindowedSamples, signal->samples + offset, fftSize * sizeof(float));
}
// Apply Hann window: h(t) = 0.5 * (1 - cos(2πt))
// And derivative window: h'(t) = π * sin(2πt)
for (int i = 0; i < fftSize; i++) {
float t = (float)i / (fftSize - 1);
float hann = 0.5f * (1.0f - cosf(2.0f * M_PI * t));
float derivHann = M_PI * sinf(2.0f * M_PI * t);
windowedSamples[i] *= hann;
derivWindowedSamples[i] *= derivHann;
}
// Compute STFT
// Compute normal STFT (V_f)
for (int i = 0; i < fftSize; i++) complexInput[i] = windowedSamples[i] + 0.0f * I;
FFT(complexInput, fftOutput, fftSize, false);
@@ -314,14 +322,26 @@ static void ComputeSTFT(AudioSignal* signal, StftResult* result, int fftSize)
result->segments[seg].spectrum[bin].amplitude = (bin == 0) ? cabsf(fftOutput[bin]) / fftSize : 2.0f * cabsf(fftOutput[bin]) / fftSize;
result->segments[seg].spectrum[bin].phase = cargf(fftOutput[bin]);
}
// Compute derivative-window STFT (V_fd) for synchrosqueezing
result->segments[seg].derivativeSpectrum = (FrequencyData*)malloc(numBins * sizeof(FrequencyData));
for (int i = 0; i < fftSize; i++) complexInput[i] = derivWindowedSamples[i] + 0.0f * I;
FFT(complexInput, fftOutput, fftSize, false);
for (int bin = 0; bin < numBins; bin++) {
result->segments[seg].derivativeSpectrum[bin].frequency = (float)bin * signal->sampleRate / fftSize;
result->segments[seg].derivativeSpectrum[bin].amplitude = cabsf(fftOutput[bin]) / fftSize;
result->segments[seg].derivativeSpectrum[bin].phase = cargf(fftOutput[bin]);
}
}
free(windowedSamples); free(complexInput); free(fftOutput);
free(windowedSamples); free(derivWindowedSamples); free(complexInput); free(fftOutput);
}
static void FreeSTFT(StftResult* result)
{
for (int i = 0; i < result->numSegments; i++) {
free(result->segments[i].spectrum);
if (result->segments[i].derivativeSpectrum) free(result->segments[i].derivativeSpectrum);
}
free(result->segments);
result->segments = NULL;
@@ -433,6 +453,9 @@ static void GenerateSpectrogramTexture(StftResult* stft, Image* image, Texture2D
if (stft->numSegments == 0) return;
int width = stft->numSegments;
int height = stft->segments[0].numBins;
int fftSize = (height - 1) * 2;
float freqPerBin = (float)stft->sampleRate / fftSize;
*image = GenImageColor(width, height, BLACK);
Color* pixels = (Color*)image->data;
@@ -443,16 +466,89 @@ static void GenerateSpectrogramTexture(StftResult* stft, Image* image, Texture2D
if (stft->segments[seg].spectrum[bin].amplitude > maxAmplitude)
maxAmplitude = stft->segments[seg].spectrum[bin].amplitude;
for (int seg = 0; seg < width; seg++) {
for (int bin = 0; bin < height; bin++) {
float amplitude = stft->segments[seg].spectrum[bin].amplitude;
float db = AmplitudeToDecibels(amplitude);
float normalized = (db - app.amplitudeFloorDb) / (app.amplitudeCeilingDb - app.amplitudeFloorDb);
normalized = Clamp(normalized, 0.0f, 1.0f);
int pixelIndex = (height - 1 - bin) * width + seg;
pixels[pixelIndex] = GetColormapColor(normalized, app.colormap);
if (app.useSynchrosqueezing) {
// ===== SYNCHROSQUEEZING =====
// Reassign energy to true frequencies using derivative STFT
// Accumulation buffer for reassigned energy
float* accumBuffer = (float*)calloc(width * height, sizeof(float));
for (int seg = 0; seg < width; seg++) {
for (int bin = 0; bin < height; bin++) {
FrequencyData* V_f = &stft->segments[seg].spectrum[bin];
FrequencyData* V_fd = &stft->segments[seg].derivativeSpectrum[bin];
float amplitude = V_f->amplitude;
if (amplitude < 0.0001f) continue;
// Compute instantaneous frequency using synchrosqueezing formula:
// ω̂ = bin_freq + Re[V_fd / (i * V_f)]
// Complex division: (a+bi)/(c+di) = ((ac+bd) + (bc-ad)i) / (c²+d²)
// We need Re[(a+bi) / (i*(c+di))] = Re[(a+bi) / (-d+ci)] = (ad+bc)/(c²+d²)
float V_f_real = amplitude * cosf(V_f->phase);
float V_f_imag = amplitude * sinf(V_f->phase);
float V_fd_real = V_fd->amplitude * cosf(V_fd->phase);
float V_fd_imag = V_fd->amplitude * sinf(V_fd->phase);
float denom = V_f_real * V_f_real + V_f_imag * V_f_imag;
float trueFreq = V_f->frequency; // Default to bin frequency
if (denom > 1e-10f) {
// Re[V_fd / (i * V_f)] = (V_fd_real * V_f_imag + V_fd_imag * V_f_real) / denom
float correction = (V_fd_real * V_f_imag + V_fd_imag * V_f_real) / denom;
trueFreq = V_f->frequency + correction;
}
// Clamp to valid range
if (trueFreq < 0) trueFreq = 0;
if (trueFreq >= stft->sampleRate / 2.0f) trueFreq = stft->sampleRate / 2.0f - 1;
// Map to bin coordinate
float targetBinF = trueFreq / freqPerBin;
if (targetBinF < 0) targetBinF = 0;
if (targetBinF >= height) targetBinF = height - 0.001f;
// Bilinear splatting to neighboring bins
int bin0 = (int)targetBinF;
int bin1 = bin0 + 1;
if (bin1 >= height) bin1 = height - 1;
float frac = targetBinF - bin0;
int idx0 = (height - 1 - bin0) * width + seg;
int idx1 = (height - 1 - bin1) * width + seg;
accumBuffer[idx0] += amplitude * (1 - frac);
accumBuffer[idx1] += amplitude * frac;
}
}
// Convert accumulation buffer to colors
for (int i = 0; i < width * height; i++) {
if (accumBuffer[i] > 0.0001f) {
float db = AmplitudeToDecibels(accumBuffer[i]);
float normalized = (db - app.amplitudeFloorDb) / (app.amplitudeCeilingDb - app.amplitudeFloorDb);
normalized = Clamp(normalized, 0.0f, 1.0f);
pixels[i] = GetColormapColor(normalized, app.colormap);
}
}
free(accumBuffer);
} else {
// ===== STANDARD STFT (no synchrosqueezing) =====
for (int seg = 0; seg < width; seg++) {
for (int bin = 0; bin < height; bin++) {
float amplitude = stft->segments[seg].spectrum[bin].amplitude;
float db = AmplitudeToDecibels(amplitude);
float normalized = (db - app.amplitudeFloorDb) / (app.amplitudeCeilingDb - app.amplitudeFloorDb);
normalized = Clamp(normalized, 0.0f, 1.0f);
int pixelIndex = (height - 1 - bin) * width + seg;
pixels[pixelIndex] = GetColormapColor(normalized, app.colormap);
}
}
}
if (texture->id != 0) UnloadTexture(*texture);
*texture = LoadTextureFromImage(*image);
SetTextureFilter(*texture, TEXTURE_FILTER_BILINEAR);
@@ -940,6 +1036,16 @@ static void DrawSidebar(void)
DrawRectangleLinesEx(gridCheck, 1, WHITE);
DrawText("Show Grid", x + 25, y + 2, fontSize, LIGHTGRAY); y += 25;
// Synchrosqueezing toggle
Rectangle sqCheck = { x, y, 18, 18 };
if (CheckCollisionPointRec(GetMousePosition(), sqCheck) && IsMouseButtonPressed(MOUSE_LEFT_BUTTON)) {
app.useSynchrosqueezing = !app.useSynchrosqueezing;
needsRegen = true;
}
DrawRectangleRec(sqCheck, app.useSynchrosqueezing ? BLUE : DARKGRAY);
DrawRectangleLinesEx(sqCheck, 1, WHITE);
DrawText("Synchrosqueezing (sharp)", x + 25, y + 2, fontSize, LIGHTGRAY); y += 25;
// File loading
DrawText("File:", x, y, fontSize, LIGHTGRAY); y += 18;
Rectangle fileButton = { x, y, sidebarWidth - 10, 25 };
@@ -1067,6 +1173,7 @@ int main(int argc, char* argv[])
app.fftSize = FFT_SIZE_DEFAULT;
app.isPlaying = false;
app.playbackFinished = false;
app.useSynchrosqueezing = true; // Enabled by default
GenerateColormapTexture();
ScanDirectory(GetWorkingDirectory());