Implement synchrosqueezing transform for sharp spectrogram display

Co-authored-by: Qwen-Coder <qwen-coder@alibabacloud.com>

src/spectrogram.c

@@ -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++) {
+    if (app.useSynchrosqueezing) {
-        for (int bin = 0; bin < height; bin++) {
+        // ===== SYNCHROSQUEEZING =====
-            float amplitude = stft->segments[seg].spectrum[bin].amplitude;
+        // Reassign energy to true frequencies using derivative STFT
-            float db = AmplitudeToDecibels(amplitude);
+        
-            float normalized = (db - app.amplitudeFloorDb) / (app.amplitudeCeilingDb - app.amplitudeFloorDb);
+        // Accumulation buffer for reassigned energy
-            normalized = Clamp(normalized, 0.0f, 1.0f);
+        float* accumBuffer = (float*)calloc(width * height, sizeof(float));
-            int pixelIndex = (height - 1 - bin) * width + seg;
+        
-            pixels[pixelIndex] = GetColormapColor(normalized, app.colormap);
+        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());

synchrosqueezing_for_programmers.aux

@@ -0,0 +1,5 @@
+\relax 
+\providecommand\hyper@newdestlabel[2]{}
+\providecommand\HyField@AuxAddToFields[1]{}
+\providecommand\HyField@AuxAddToCoFields[2]{}
+\gdef \@abspage@last{5}

synchrosqueezing_for_programmers.tex

@@ -0,0 +1,219 @@
+\documentclass[12pt]{article}
+\usepackage[margin=1in]{geometry}
+\usepackage{amsmath}
+\usepackage{tikz}
+\usepackage{color}
+\usepackage{hyperref}
+
+\title{\textbf{Why Spectrograms Look Smudged (And How to Fix Them)}\\
+\small\textit{A beginner-friendly guide to synchrosqueezing}}
+\author{}
+\date{}
+
+\begin{document}
+\maketitle
+
+\section*{What You Already Know}
+
+You know what an FFT does. You feed it a chunk of audio, it tells you which frequencies are present. You slide a window across your audio, stack up the FFTs, and you get a spectrogram.
+
+But your spectrogram looks \textbf{smudged}. A pure tone doesn't look like a clean line --- it looks like someone took a finger and smeared it vertically. You're wondering: \textit{``Is this just how FFTs work, or am I doing something wrong?''}
+
+You're not doing anything wrong. This is a known limitation of the standard approach. And there's a technique called \textbf{synchrosqueezing} that fixes it.
+
+\section*{Why the Smudge Happens}
+
+Your FFT has bins at fixed frequencies. If your FFT size is 512 and your sample rate is 48\,kHz, the bins are spaced 93.75\,Hz apart:
+
+\begin{center}
+\begin{tabular}{|c|c|c|c|c|c|c|}
+\hline
+Bin 0 & Bin 1 & Bin 2 & Bin 3 & Bin 4 & Bin 5 & Bin 6 \\
+0 Hz & 94 Hz & 188 Hz & 281 Hz & 375 Hz & 469 Hz & 563 Hz \\
+\hline
+\end{tabular}
+\end{center}
+
+Now imagine your signal has a pure 500\,Hz tone. Where does it go?
+
+It doesn't match any bin exactly. It falls \textbf{between} bin 5 (469\,Hz) and bin 6 (563\,Hz). The FFT doesn't know what to do with it, so it spreads the energy across both bins.
+
+\textbf{Result:} Your clean 500\,Hz tone now shows up as energy at both 469\,Hz and 563\,Hz. That's the smudge.
+
+\textbf{This is called \textit{spectral leakage}, and it's a fundamental property of the FFT.} It's not a bug --- it's just how the math works.
+
+\section*{The Hidden Information: Phase}
+
+Every FFT output value is a \textbf{complex number}. That means it has two parts:
+
+\begin{itemize}
+    \item \textbf{Magnitude} --- how strong is this frequency? (This is what you use to draw the spectrogram.)
+    \item \textbf{Phase} --- where are we in the wave cycle? (This is what you've been ignoring.)
+\end{itemize}
+
+Here's the thing: \textbf{the phase contains information about the true frequency.}
+
+\subsection*{Think of It Like This: Engine Timing}
+
+Imagine you're tuning an engine. There's a timing mark painted on the crankshaft pulley, and a sensor that detects when the mark passes by.
+
+You \textbf{expect} the engine to run at 3000 RPM. At that speed, the timing mark should pass the sensor at exactly the same point in each revolution.
+
+But the engine is actually running at 3100 RPM --- slightly faster than expected. What happens?
+
+\begin{itemize}
+    \item Revolution 1: the mark passes right where you expect it
+    \item Revolution 2: the mark arrives a little \textbf{early} (the engine completed the rotation faster than expected)
+    \item Revolution 3: the mark arrives even earlier
+    \item Revolution 4: the mark is now noticeably ahead of where it should be
+\end{itemize}
+
+The mark is \textbf{drifting forward} relative to your expectation. By measuring \textbf{how much it drifts per revolution}, you can calculate the \textbf{actual RPM}.
+
+\textbf{This is exactly what happens with FFT bins.}
+
+Each FFT bin is like expecting the engine to run at a specific RPM (the bin center frequency). The phase tells you where the ``timing mark'' is. If the phase drifts forward from one window to the next, the true frequency is higher than the bin center. If it drifts backward, the true frequency is lower.
+
+\textbf{Measure the drift rate, and you know the true frequency.}
+
+\section*{The Formula}
+
+The instantaneous frequency estimate is:
+
+\begin{equation}
+\hat{\omega}(a,b) = \text{Re}\left[ \frac{\partial_b V_f(a,b)}{i \cdot V_f(a,b)} \right]
+\end{equation}
+
+Let's decode this:
+
+\begin{itemize}
+    \item $V_f(a,b)$ = the complex FFT value for bin $a$ at time window $b$
+    \item $\partial_b V_f$ = how much that value changed from the previous time window
+    \item $i$ = the imaginary unit (a math trick that converts phase rotation into Hz)
+    \item $\text{Re}[\ldots]$ = take the real part of the result
+\end{itemize}
+
+\textbf{In plain English:}
+
+\begin{center}
+\fbox{\parbox{0.7\textwidth}{\centering
+``How fast is the phase changing from one window to the next?'' \\
+$=$ \\
+``What is the true frequency at this point?''
+}}
+\end{center}
+
+\section*{How to Actually Compute This}
+
+``Wait, what's that $\partial_b$ symbol? That's a derivative!''
+
+Good news: you don't need to compute derivatives numerically. There's a neat trick.
+
+You compute \textbf{two FFTs} per window:
+
+\begin{enumerate}
+    \item \textbf{Normal STFT} using your Hann window $h(t)$ $\rightarrow$ gives you $V_f$
+    \item \textbf{Second STFT} using the \textbf{derivative of the Hann window} $h'(t)$ $\rightarrow$ gives you $V_{fd}$
+\end{enumerate}
+
+The Hann window is:
+$$h(t) = 0.5 \times \bigl(1 - \cos(2\pi t)\bigr)$$
+
+Its derivative is:
+$$h'(t) = \pi \times \sin(2\pi t)$$
+
+\textbf{That's all.} You already know how to window and FFT. Just use a different window function for the second FFT.
+
+Then the formula becomes:
+
+\begin{equation}
+\hat{\omega} = \text{bin\_frequency} + \text{Re}\left[ \frac{V_{fd}}{i \cdot V_f} \right]
+\end{equation}
+
+\section*{What You Do With This Number}
+
+For every FFT bin at every time window, you now know the \textbf{true frequency}, not just the bin center.
+
+So instead of drawing the energy at the bin center (469\,Hz), you draw it at the true frequency (500\,Hz).
+
+\textbf{This is synchrosqueezing:} you're taking the smeared energy and ``squeezing'' it back to where it actually belongs.
+
+\section*{The Algorithm}
+
+\begin{verbatim}
+for each time window:
+    # Step 1: Normal STFT with Hann window
+    windowed = samples * hann_window
+    V_f = FFT(windowed)
+    
+    # Step 2: STFT with derivative-of-Hann window
+    deriv_windowed = samples * hann_derivative
+    V_fd = FFT(deriv_windowed)
+    
+    # Step 3: For each bin, compute true frequency
+    for each bin:
+        # Complex division
+        correction = Re(V_fd / (i * V_f))
+        true_freq = bin_frequency + correction
+        
+        # Step 4: Put the energy at the true frequency
+        target_bin = round(true_freq / freq_per_bin)
+        output[target_bin] += magnitude(V_f)
+\end{verbatim}
+
+\section*{What Changes Visually}
+
+\begin{center}
+\begin{tabular}{|p{0.45\textwidth}|p{0.45\textwidth}|}
+\hline
+\textbf{Standard STFT} & \textbf{Synchrosqueezed} \\
+\hline
+Energy spread across 2--4 bins & Energy concentrated in 1 bin \\
+Tone looks like a fuzzy bar & Tone looks like a sharp line \\
+Hard to distinguish close tones & Close tones are clearly separate \\
+\hline
+\end{tabular}
+\end{center}
+
+For an FSK signal (like your FSK4 at 50 baud):
+\begin{itemize}
+    \item \textbf{Before:} each symbol looks like a blurry blob
+    \item \textbf{After:} each symbol looks like a clean rectangle
+\end{itemize}
+
+\section*{The Trade-offs}
+
+\begin{itemize}
+    \item \textbf{Computation:} 2 FFTs per window instead of 1 (roughly 1.5--2$\times$ slower)
+    \item \textbf{Memory:} need to store both $V_f$ and $V_{fd}$
+    \item \textbf{Quality:} significantly sharper time-frequency representation
+\end{itemize}
+
+For a 2-second audio file at FFT size 512, we're talking about 0.1 seconds vs.\ 0.15 seconds. Not a big deal for offline analysis.
+
+\section*{Why Not Just Use a Bigger FFT?}
+
+You could increase your FFT size to get narrower bins. But:
+\begin{itemize}
+    \item Bigger FFT = wider time window = worse \textbf{time} resolution
+    \item You trade frequency smearing for time smearing
+    \item Synchrosqueezing gives you sharp frequency \textbf{and} sharp time
+\end{itemize}
+
+It's not a replacement for choosing the right FFT size. It's a way to get the most out of whatever FFT size you choose.
+
+\section*{Further Reading}
+
+If you want to go deeper:
+\begin{itemize}
+    \item Thakur et al., \textit{The Synchrosqueezing algorithm for time-varying spectral analysis} (2011) --- the standard reference
+    \item \texttt{librosa.reassigned\_spectrogram} (Python) --- a working implementation
+    \item \texttt{scipy.signal.spectrogram} with phase-based reassignment --- another approach
+\end{itemize}
+
+\vspace{1cm}
+\begin{center}
+\textit{The smudge isn't your fault. But now you know how to fix it.}
+\end{center}
+
+\end{document}