Quick script to get a representative population of DCOs from my LISA simulations and extract their frequencies/GW amplitudes

get_representative_fiducial_binaries.py

import h5py as h5
import numpy as np
import legwork
import astropy.units as u
import pandas as pd

# constants
dco_types = ["BHBH", "BHNS", "NSNS"]
det_rates = [74, 42, 8]

# CHANGE THIS FOR WHERE YOU PUT YOUR FILES
path_to_simulation_files = "simulation/"

# set up data
f_GW = np.array([])
h_0 = np.array([])
h_c = np.array([])
m1 = np.array([])
m2 = np.array([])
ecc = np.array([])
R = np.array([])
z = np.array([])
theta = np.array([])
snr = np.array([])

labels = np.concatenate([[dco_type] * det_rate for dco_type, det_rate in zip(dco_types, det_rates)])

for dco_type, detections in zip(dco_types, det_rates):
    # get full dataset
    with h5.File(path_to_simulation_files + f"{dco_type}_fiducial_all.h5") as f:
        dco = f["simulation"][...].squeeze()

    # draw a random subset
    subset = np.random.choice(np.arange(len(dco)), size=detections, p=dco["weight"] / dco["weight"].sum())
    dco = dco[subset]

    # set up a LEGWORK source class
    sources = legwork.source.Source(m_1=dco["m_1"] * u.Msun, m_2=dco["m_2"] * u.Msun, a=dco["a_LISA"] * u.AU,
                                    ecc=dco["e_LISA"], dist=dco["dist"] * u.kpc,
                                    interpolate_g=False, gw_lum_tol=1e-3)

    # calculate SNR, get frequency of max SNR harmonic and sum up the strain amplitudes in first 10,000 harmonics
    snr = np.concatenate((snr, sources.get_snr()))
    f_GW = np.concatenate((f_GW, (sources.f_orb * sources.max_snr_harmonic).to(u.mHz).value))
    h_0 = np.concatenate((h_0, sources.get_h_0_n(np.arange(1, 10_000)).sum(axis=1)))
    h_c = np.concatenate((h_c, sources.get_h_c_n(np.arange(1, 10_000)).sum(axis=1)))
    m1 = np.concatenate((m1, sources.m_1))
    m2 = np.concatenate((m2, sources.m_2))
    ecc = np.concatenate((ecc, sources.ecc))
    R = np.concatenate((R, dco["R"]))
    z = np.concatenate((z, dco["z"]))
    theta = np.concatenate((theta, dco["theta"]))

print((f_GW > 1e-2).all())

# save everything into a dataframe and then a CSV file
df = pd.DataFrame(np.transpose([m1, m2, f_GW, ecc, R, z, theta, h_0, h_c, snr, labels]),
                  columns=["m1", "m2", "f_GW [mHz]", "ecc", "galactocentric_radius [kpc]", "galactocentric_height [kpc]", "azimuthal_angle [rad]", "strain_amp", "characteristic_strain_amp", "snr", "dco_type"])
df.to_csv("./representative_fiducial_binaries.csv", index=False)