Live Code - LIGO
import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import hilbert
from scipy.linalg import hankel, svd
from gwpy.timeseries import TimeSeries
# ----------------------------------------------------------------------
# 1. Core ESPRIT Function (Your Optimized Implementation)
# ----------------------------------------------------------------------
def esprit_optimized(x, fs, order=2, L=None):
"""
ESPRIT for real damped sinusoids.
Optimized for long signals with a limited Hankel matrix size.
"""
xa = hilbert(x)
N = len(xa)
if L is None:
L = min(2000, N // 2) # <-- optimization: cap L
if L < order:
return np.nan, np.nan
M = N - L + 1
H = np.zeros((L, M), dtype=complex)
for i in range(L):
H[i, :] = xa[i:i+M]
try:
U, s, Vh = svd(H, full_matrices=False)
except:
return np.nan, np.nan
if len(s) < order:
return np.nan, np.nan
U_s = U[:, :order]
J1 = np.eye(L-1, L, dtype=int)[:, :L-1]
J2 = np.eye(L-1, L, dtype=int)[:, 1:]
A = J1 @ U_s
B = J2 @ U_s
Psi = np.linalg.lstsq(A, B, rcond=None)[0]
eigvals = np.linalg.eigvals(Psi)
freqs = np.abs(np.angle(eigvals)) * fs / (2 * np.pi)
dampings = -np.log(np.abs(eigvals)) * fs
pos = freqs > 0
if not np.any(pos):
return np.nan, np.nan
return freqs[pos][0], dampings[pos][0]
# ----------------------------------------------------------------------
# 2. Download and Preprocess LIGO Data (GW150914 as example)
# ----------------------------------------------------------------------
# Step 2.1: Download the data from GWOSC
print("Downloading LIGO data from GWOSC...")
# Use the known GPS times for the GW150914 event (centered at gps 1126259462.4)
start_gps = 1126259446 # 16 seconds before the event
end_gps = 1126259478 # 16 seconds after the event
strain = TimeSeries.fetch_open_data('H1', start_gps, end_gps)
# Step 2.2: Preprocess the data
fs = strain.sample_rate.value
print(f"Sampling frequency: {fs:.0f} Hz")
print(f"Data range: {strain.t0.value} to {strain.t0.value + strain.duration.value} seconds GPS")
# Apply a bandpass filter to focus on the frequency range of interest (your invariant)
low_freq = 35
high_freq = 400
filtered_strain = strain.bandpass(low_freq, high_freq)
# Whiten the data to flatten the noise spectrum
whitened_strain = filtered_strain.whiten(fftlength=4)
# Convert the whitened GWpy TimeSeries object to a NumPy array for ESPRIT
signal_array = whitened_strain.value
print(f"Signal length: {len(signal_array)} samples")
# ----------------------------------------------------------------------
# 3. Run ESPRIT Analysis on the LIGO Data
# ----------------------------------------------------------------------
print("\nRunning ESPRIT on LIGO strain data...")
f_est, d_est = esprit_optimized(signal_array, fs, order=2)
if not np.isnan(f_est):
print(f"ESPRIT Results on LIGO Data:")
print(f" Frequency: {f_est:.4f} Hz")
print(f" Damping: {d_est:.4f} s⁻¹")
else:
print("ESPRIT could not detect a clear dominant frequency in this segment.")
# ----------------------------------------------------------------------
# 4. Visualization: Compare the PSD to the ESPRIT Result
# ----------------------------------------------------------------------
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 8))
# Top plot: Spectrum of the filtered data
spec = filtered_strain.spectrum(fftlength=4)
ax1.loglog(spec.frequencies.value, spec.value, label='Filtered LIGO Strain (H1)')
ax1.set_xlim(low_freq, high_freq)
ax1.set_ylabel('Amplitude Spectral Density (1/√Hz)')
ax1.set_title('LIGO Strain Data Spectrum (GW150914)')
ax1.legend()
if not np.isnan(f_est):
ax1.axvline(f_est, color='r', linestyle='--', label=f'ESPRIT Frequency = {f_est:.2f} Hz')
ax1.legend()
ax1.grid(True, alpha=0.3)
# Bottom plot: The whitened time series around the event
time_axis = strain.times.value - strain.t0.value # seconds relative to segment start
event_time_offset = 1126259462.4 - strain.t0.value # Center of GW150914
t_window = 1.0 # seconds
t_mask = np.abs(time_axis - event_time_offset) < t_window
ax2.plot(time_axis[t_mask], whitened_strain[t_mask], label='Whitened LIGO Strain')
ax2.set_xlabel('Time (seconds from segment start)')
ax2.set_ylabel('Whitened Amplitude')
ax2.set_title('Whitened LIGO Strain around GW150914')
ax2.legend()
ax2.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()