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 = "../data/simulation/"

# set up data
f_GW = np.array([])
h_0 = np.array([])
h_c = 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
    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)))

# save everything into a dataframe and then a CSV file
df = pd.DataFrame(np.transpose([f_GW, h_0, h_c, labels]),
                  columns=["f_GW [mHz]", "strain_amp", "characteristic_strain_amp", "dco_type"])
df.to_csv("./representative_fiducial_binaries.csv", index=False)