Spectral waterfall 3D

Python Code

"""Signal Processing — Spectral waterfall 3D (LIGO GW150914 chirp)."""


import numpy as np
import plotly.graph_objects as go


LIGO_FILE = Path(__file__).parents[4] / "plotly-studio/qa/super_prompts/data/H-H1_GWOSC_16KHZ_R1-1126259447-32.hdf5"

FS = 16384         # Hz sampling rate of the GWOSC 16 kHz strain file
START_GPS = 1126259447
EVENT_GPS = 1126259462.4
EVENT_T   = EVENT_GPS - START_GPS   # ≈ 15.4 s into the segment

N_LINES = 48       # time slices, each drawn as one stacked spectrum line

# Cool noise floor → warm peaks (orange reserved for the energetic ridge)
WATERFALL_SCALE = [
    [0.0, TEAL],
    [0.5, GREEN],
    [1.0, ORANGE],
]


def generate():
    import h5py
    from scipy import signal
    from scipy.signal import resample_poly

    print(f"  loading {LIGO_FILE.name}")
    with h5py.File(LIGO_FILE, "r") as f:
        strain = f["strain"]["Strain"][:]
    print(f"  strain shape: {strain.shape}, fs={FS} Hz")
    strain = strain.astype(float)

    # ── Standard GW pre-processing (same as the chart-139 spectrogram) ─────────
    # Whiten by dividing by the noise amplitude spectral density, then band-pass
    # to LIGO's sensitive band. A raw spectrum is dominated by coloured detector
    # noise and the chirp is invisible; this is what exposes the signal.
    freqs_full = np.fft.rfftfreq(len(strain), 1 / FS)
    strain_fft = np.fft.rfft(strain)
    f_w, psd = signal.welch(strain, fs=FS, nperseg=4 * FS)
    psd_interp = np.interp(freqs_full, f_w, psd)
    psd_interp[psd_interp <= 0] = 1e-40
    whitened = np.fft.irfft(strain_fft / np.sqrt(psd_interp), n=len(strain))

    b, a = signal.butter(4, [20, 350], btype="bandpass", fs=FS)
    whitened = signal.filtfilt(b, a, whitened)

    # Down-sample to 2048 Hz — the chirp lives below 350 Hz.
    fs2 = 2048
    wd = resample_poly(whitened, fs2, FS)
    t = np.arange(len(wd)) / fs2

    # ── Morlet wavelet scalogram (constant-Q) ──────────────────────────────────
    # A fixed-window STFT can't resolve the fast frequency sweep; a wavelet
    # transform uses a window that shrinks with frequency, which is how LIGO's
    # own Q-scans make the chirp pop. Compute on a padded window so edge
    # artefacts fall outside the crop.
    pad = 0.6
    sel = (t >= EVENT_T - pad) & (t <= EVENT_T + 0.35)
    seg, t_seg = wd[sel], t[sel]

    w0 = 8.0
    f_spec = np.linspace(20, 350, 220)
    widths = w0 * fs2 / (2 * np.pi * f_spec)
    power = np.abs(signal.cwt(seg, signal.morlet2, widths, w=w0)) ** 2  # (freq, time)

    # Crop to ~0.35 s around the merger.
    d0, d1 = EVENT_T - 0.28, EVENT_T + 0.07
    dm = (t_seg >= d0) & (t_seg <= d1)
    t_crop, power = t_seg[dm], power[:, dm]

    # Thin the time axis down to N_LINES evenly-spaced spectrum slices.
    idx = np.linspace(0, power.shape[1] - 1, N_LINES).round().astype(int)
    t_lines = t_crop[idx] - EVENT_T            # seconds relative to merger
    spectra = power[:, idx].T                  # (N_LINES, n_freq)
    spectra = spectra / spectra.max()          # normalise to peak energy
    print(f"  spectra: {spectra.shape}, freqs: {len(f_spec)}")

    # Each time slice is its own spectrum line, stacked into the page along the
    # time axis → a true 3D waterfall. As the binary spirals in, the peak climbs
    # from ~35 Hz toward ~250 Hz, so every line differs: the warm ridge sweeps up
    # in frequency and grows toward the merger. Far lines are added first so
    # nearer ones render in front.
    fig = go.Figure()
    for i in range(N_LINES - 1, -1, -1):
        fig.add_trace(go.Scatter3d(
            x=f_spec,
            y=np.full_like(f_spec, t_lines[i]),
            z=spectra[i],
            mode="lines",
            line=dict(
                width=3,
                color=spectra[i],
                colorscale=WATERFALL_SCALE,
                cmin=0.0,
                cmax=1.0,
                showscale=(i == 0),
                colorbar=dict(title="Normalized<br>energy", thickness=14, len=0.6),
            ),
            hoverinfo="skip",
            showlegend=False,
        ))

    fig.update_layout(
        scene=dict(
            xaxis_title="Frequency (Hz)",
            yaxis_title="Time from merger (s)",
            zaxis_title="Energy",
            camera=dict(eye=dict(x=1.6, y=1.55, z=0.65)),
            aspectratio=dict(x=1.5, y=1.5, z=0.55),
        ),
        margin=dict(t=20, b=20, l=20, r=20),
        height=620,
    )
    return fig


if __name__ == "__main__":
    generate()

Made with Plotly