108 lines
4.4 KiBLFS
Python
108 lines
4.4 KiBLFS
Python
#!/usr/bin/env python3
|
|
"""
|
|
Gravitational Wave Detection Solution
|
|
|
|
Detect gravitational wave signals using grid search over mass parameters
|
|
and different waveform approximants with PyCBC.
|
|
"""
|
|
|
|
import pandas as pd
|
|
from pycbc.filter import highpass, matched_filter, resample_to_delta_t
|
|
from pycbc.frame import read_frame
|
|
from pycbc.psd import interpolate, inverse_spectrum_truncation
|
|
from pycbc.waveform import get_td_waveform
|
|
|
|
|
|
def main():
|
|
# Load data from frame file
|
|
fname = "/root/data/PyCBC_T2_2.gwf"
|
|
strain = read_frame(fname, "H1:TEST-STRAIN")
|
|
|
|
# Condition the data: highpass and downsample
|
|
strain = resample_to_delta_t(highpass(strain, 15.0), 1.0 / 4096)
|
|
|
|
# Remove filter wraparound - crop 2 seconds from both ends
|
|
conditioned = strain.crop(2, 2)
|
|
|
|
# Calculate power spectral density
|
|
psd = conditioned.psd(4)
|
|
psd = interpolate(psd, conditioned.delta_f)
|
|
psd = inverse_spectrum_truncation(psd, int(4 * conditioned.sample_rate), low_frequency_cutoff=15)
|
|
|
|
# Different waveform approximants to test
|
|
approximants = ["SEOBNRv4_opt", "IMRPhenomD", "TaylorT4"]
|
|
|
|
# Grid search over mass parameter space and approximants
|
|
# Collect all results
|
|
all_results = []
|
|
|
|
# Loop over different approximants
|
|
for approx in approximants:
|
|
# Loop over different masses (10-40 solar masses, integer steps)
|
|
# Use convention m1 >= m2 (primary mass >= secondary mass) to avoid redundant combinations
|
|
for m1 in range(10, 41):
|
|
for m2 in range(10, m1 + 1):
|
|
try:
|
|
# Generate template waveform
|
|
hp, hc = get_td_waveform(approximant=approx, mass1=m1, mass2=m2, delta_t=conditioned.delta_t, f_lower=20)
|
|
|
|
# Resize the waveform to match the data length
|
|
hp.resize(len(conditioned))
|
|
|
|
# Shift template so merger is at the start
|
|
# By convention waveforms from get_td_waveform have merger at time zero
|
|
# This cyclic shift aligns the merger properly for matched filtering
|
|
template = hp.cyclic_time_shift(hp.start_time)
|
|
|
|
# Calculate signal-to-noise ratio time series
|
|
snr = matched_filter(template, conditioned, psd=psd, low_frequency_cutoff=20)
|
|
|
|
# Remove time corrupted by template filter and PSD filter
|
|
# Remove 4 seconds at beginning and end for PSD filtering
|
|
# Remove 4 additional seconds at beginning for template length
|
|
snr = snr.crop(4 + 4, 4)
|
|
|
|
# Find peak SNR
|
|
# matched_filter returns complex SNR; abs() maximizes over phase
|
|
peak_idx = abs(snr).numpy().argmax()
|
|
snr_peak = abs(snr[peak_idx])
|
|
|
|
# Calculate total mass: M_total = m1 + m2
|
|
total_mass = m1 + m2
|
|
|
|
# Store result
|
|
all_results.append({"approximant": approx, "m1": m1, "m2": m2, "snr": float(snr_peak), "total_mass": float(total_mass)})
|
|
except Exception:
|
|
# Some approximants may fail for certain mass combinations
|
|
# Skip and continue
|
|
continue
|
|
|
|
# Find the best result for each approximant
|
|
best_per_approximant = {}
|
|
for result in all_results:
|
|
approx = result["approximant"]
|
|
if approx not in best_per_approximant or result["snr"] > best_per_approximant[approx]["snr"]:
|
|
best_per_approximant[approx] = result
|
|
|
|
# Create results DataFrame with best result for each approximant
|
|
# Sort by approximant name for consistent output
|
|
if best_per_approximant:
|
|
results_list = [best_per_approximant[approx] for approx in approximants if approx in best_per_approximant]
|
|
results_df = pd.DataFrame(results_list)
|
|
# Only output approximant, snr, and total_mass
|
|
results = results_df[["approximant", "snr", "total_mass"]]
|
|
else:
|
|
# No signal found
|
|
results = pd.DataFrame({"approximant": [None], "snr": [0.0], "total_mass": [0.0]})
|
|
|
|
results.to_csv("/root/detection_results.csv", index=False)
|
|
print("Solution completed. Best signal found for each approximant.")
|
|
if len(results) > 0:
|
|
for _, row in results.iterrows():
|
|
if row["snr"] > 0:
|
|
print(f"Best {row['approximant']}: SNR={row['snr']:.2f}, total_mass={row['total_mass']:.1f}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|