Advanced - GW190521 BH1 as 3rd gen BH

pipeline_1g1g = archeo.get_binary_generation_pipeline("2g_aligned_spin")
df_1g1g_binaries, _ = pipeline_1g1g(size=5000, n_workers=1)

pipeline_2g1g = archeo.get_binary_generation_pipeline("ng_aligned_spin")
df_2g1g_binaries, binary_generator = pipeline_2g1g(df_bh1_binaries=df_1g1g_binaries, size=5000, n_workers=1)

# Verify Distribution in Priors

# Here the distributions should be fairly close to each other
#   since we are using the 2g prior samples
#   for the primary black hole in the 2g1g prior.
# However, the sample size is small so there will be some noise.

fig, axes = plt.subplots(1, 2, figsize=(10, 5))

plot_pdf(axes[0], df_1g1g_binaries["m_f"], "blue", name="2g Remnant Mass", unit="M_\odot")
plot_pdf(axes[0], df_2g1g_binaries["m_1"], "lightblue", name="2g1g BH1 Mass", unit="M_\odot")

plot_pdf(axes[1], df_1g1g_binaries["a_f"], "blue", name="2g Remnant Spin")
plot_pdf(axes[1], df_2g1g_binaries["a_1"], "lightblue", name="2g1g BH1 Spin")

for i in range(2):
    axes[i].set_title(
        "{param} Distribution Comparison".format(param="Mass" if i == 0 else "Spin")
    )
    axes[i].set_ylabel("Probability Density")
    axes[i].set_xlabel("Mass (M_\odot)" if i == 0 else "Spin")
    axes[i].grid()
    axes[i].legend()

# Understand the remnant mass and spin distributions of the 2g1g prior
# Here we should at least see the remnant mass distribution shifted to higher masses
#   compared to the 2g prior since the primary black hole is more massive on average.

fig, axes = plt.subplots(1, 2, figsize=(10, 5))

plot_pdf(axes[0], df_1g1g_binaries["m_f"], "blue", name="2g Remnant Mass", unit="M_\odot")
plot_pdf(axes[0], df_2g1g_binaries["m_f"], "lightblue", name="2g1g Remnant Mass", unit="M_\odot")

plot_pdf(axes[1], df_1g1g_binaries["a_f"], "blue", name="2g Remnant Spin")
plot_pdf(axes[1], df_2g1g_binaries["a_f"], "lightblue", name="2g1g Remnant Spin")

for i in range(2):
    axes[i].set_title(
        "{param} Distribution Comparison".format(param="Mass" if i == 0 else "Spin")
    )
    axes[i].set_ylabel("Probability Density")
    axes[i].set_xlabel("Mass (M_\odot)" if i == 0 else "Spin")
    axes[i].grid()
    axes[i].legend()

# This data file is a subset of the LVK posterior samples for GW190521
# See https://github.com/wyhwong/archeo/tree/main/src/tests/test_data/gw190521_lvk_subsampled.json
# For the full dataset, please refer to: https://dcc.ligo.org/P2000158/public

gw190521_pe_samples = pd.read_json("./gw190521_lvk_subsampled.json")
gw190521_pe_samples.head(10)

ancestors_bh1 = archeo.infer_ancestral_posterior_distribution(
    df_binaries=df_2g1g_binaries,
    mass_posterior_samples=gw190521_pe_samples["mass_1_source"].values.tolist(),
    spin_posterior_samples=gw190521_pe_samples["a_1"].values.tolist(),
    n_workers=1,
)
ancestors_bh1.head(10)

table_estimates(
    dfs={
        "Prior (Aligned Spin)": df_2g1g_binaries,
        "GW190521 BH1 Ancestor": ancestors_bh1,
    }
)