Ridge detection#

track dominant frequencies over time of Zxx#

Find the most prominent modes (like chirping MHD modes) in the plasma.

We compute the absolute amplitude of Zxx, then for each time slice (Zxx[:, t]), we find the index of the maximum amplitude, and map that back to a frequency using f. The x value in FFT shows a time slice.
max_ridges determines the number of max value of ridges to find. So max_ridges=2, picks the top 2 frequencies for each time bin by sorting Zxx amplitudes.

def plot_stft_ridge(
    t,
    f,
    Zxx,
    shot_id=None,
    threshold=1e-3,
    max_ridges=1,
    fmax_kHz=50,
    smooth=True,
    smooth_window=3,
    return_ridges=False,
    min_bin_distance=2
):
    """
    Extract and plot top-N ridges in STFT with better control.
    Prevents overlap and handles low-power regions.
    """
    fmax = fmax_kHz * 1000
    valid_freq_idx = f <= fmax
    f = f[valid_freq_idx]
    Zxx = Zxx[valid_freq_idx, :]

    magnitude = np.abs(Zxx)
    time_bins = magnitude.shape[1]
    freq_bins = magnitude.shape[0]

    # Set low amplitudes to zero
    magnitude[magnitude < threshold] = 0

    ridge_list = []

    used_mask = np.zeros_like(magnitude, dtype=bool)

    plt.figure(figsize=(10, 4))

    for r in range(max_ridges):
        ridge_freq = np.full(time_bins, np.nan)

        for t_idx in range(time_bins):
            # Get amplitudes for this time bin
            col = magnitude[:, t_idx].copy()

            # Mask already used indices (avoid overlapping ridges)
            col[used_mask[:, t_idx]] = 0

            if col.max() >= threshold:
                i = np.argmax(col)
                ridge_freq[t_idx] = f[i] / 1000  # to kHz

                # Mark this bin and neighboring bins as used
                i_start = max(0, i - min_bin_distance)
                i_end = min(freq_bins, i + min_bin_distance + 1)
                used_mask[i_start:i_end, t_idx] = True

        if smooth:
            filled = np.nan_to_num(ridge_freq, nan=0.0)
            smoothed = uniform_filter1d(filled, size=smooth_window)
            ridge_freq = np.where(np.isnan(ridge_freq), np.nan, smoothed)

        plt.plot(t, ridge_freq, label=f'Ridge {r+1}', lw=2)
        ridge_list.append(ridge_freq)

    plt.xlabel("Time [sec]")
    plt.ylabel("Frequency [kHz]")
    title = "STFT Ridge Detection"
    if shot_id is not None:
        title += f" – Shot {shot_id}"
    plt.title(title)
    plt.grid(True)
    plt.legend()
    plt.tight_layout()

    if return_ridges:
        return ridge_list

[plot_stft_ridge(
    t_list[i],
    f_list[i],
    Zxx_list[i],
    shot_id=shot_ids[i],
    threshold=1e-3,
    max_ridges=1,
    fmax_kHz=50,
    smooth=True
) for i in range(len(ds_list))]

[None, None, None, None]

../../_images/b77ae9c7e68f1d5a2885bb0469c6b5b3990c369cbd2fe1b7f26c77523f7047a9.png

../../_images/0487ef9fb88df0244b2be89d786231a0d2fb80f5c5606d736d2d31528e085db3.png

../../_images/5c771d5c7c525d9b01e1f063a4a7e8b9d498f87ecefe3ef38c57d789d464eff2.png

../../_images/86eaa809360a079bcb4c1c9365a8371b3504f17067f956f527510c91abdf737a.png

[plot_stft_ridge(
    t_list[i],
    f_list[i],
    Zxx_list[i],
    shot_id=shot_ids[i],
    threshold=1e-3,
    max_ridges=2,
    fmax_kHz=50,
    smooth=True
) for i in range(len(ds_list))]

[None, None, None, None]

../../_images/261a6dd872b58312a0b6a25f9b53c51cc94f397f27705fc4e1144f6dc81e0b38.png

../../_images/9a9ad1cf1ab741061667f7cc558a89e2a650bf7cb2fed95bf272bab82e36642f.png

../../_images/e62e63c3bc2ef2969d9c30537a2254b89f96bab8bfd8bbac9952e709991b162b.png

../../_images/69822d5928f49e7c53e082f3439fcb408a51c689d56cb71eac8da20c28dd4a94.png

def plot_spectrogram_with_ridges(t, f, Zxx, ridge_freqs_kHz, shot_id=None, vmin=1e-5, fmax_kHz=50):
    """
    Plot the STFT spectrogram with overlayed ridge tracks.
    """
    import matplotlib.pyplot as plt
    from matplotlib.colors import LogNorm

    fig, ax = plt.subplots(figsize=(12, 5))
    f_kHz = f / 1000

    cax = ax.pcolormesh(t, f_kHz, np.abs(Zxx), shading='nearest',
                        norm=LogNorm(vmin=vmin), cmap='jet')

    # Custom color cycle
    ridge_colors = ['black', 'c', 'magenta', 'lime', 'orange']

    for i, ridge in enumerate(ridge_freqs_kHz):
        color = ridge_colors[i % len(ridge_colors)]
        ax.plot(t, ridge, label=f'Ridge {i+1}', lw=2, color=color)

    ax.set_ylim(0, fmax_kHz)
    ax.set_xlim(0.1, t[-1])
    ax.set_xlabel('Time [sec]')
    ax.set_ylabel('Frequency [kHz]')
    title = "STFT Spectrogram with Ridge Overlay"
    if shot_id is not None:
        title += f" – Shot {shot_id}"
    ax.set_title(title)
    plt.colorbar(cax, ax=ax, label="Amplitude")
    ax.legend()
    ax.grid(True)
    plt.tight_layout()

ridges = plot_stft_ridge(
    t_list[0],
    f_list[0],
    Zxx_list[0],
    shot_id=shot_ids[0],
    threshold=1e-3,
    max_ridges=2,
    fmax_kHz=50,
    smooth=True,
    return_ridges=True
)

# Then overlay them:
plot_spectrogram_with_ridges(t_list[0], f_list[0], Zxx_list[0], ridge_freqs_kHz=ridges, shot_id=shot_ids[0])

/tmp/ipykernel_1087196/2863097040.py:32: UserWarning: Creating legend with loc="best" can be slow with large amounts of data.
  plt.tight_layout()

../../_images/261a6dd872b58312a0b6a25f9b53c51cc94f397f27705fc4e1144f6dc81e0b38.png

../../_images/788339521bd66dfc75e40781e6a79281f0a6cfe05364955a4a2aa4d051d3a0b6.png

ridges = plot_stft_ridge(
    t_list[0],
    f_list[0],
    Zxx_list[0],
    shot_id=shot_ids[0],
    threshold=1e-3,
    max_ridges=3,
    fmax_kHz=50,
    smooth=True,
    return_ridges=True
)

# Then overlay them:
plot_spectrogram_with_ridges(t_list[0], f_list[0], Zxx_list[0], ridge_freqs_kHz=ridges, shot_id=shot_ids[0])

../../_images/d3f2158c86f4ebd55613ffa1717932b2e0fd83c4be6729e7c0e1d5c99c7417ab.png

../../_images/64c2ee46c2564c137e896a9f0a62526450bee61faae26b9c27a4e2124ebdb8b5.png