gwModels kick velocity models

• gwModel_kick_q200 : aligned-spin analytical kick model trained on NR (SXS + RIT, \(q \leq 32\)) and BHPT data (\(q \leq 200\)). Valid for \(1 \leq q \leq 1000\).

• gwModel_kick_q200_GPR : aligned-spin GPR kick model trained on the same dataset. Provides both analytical and GPR predictions with uncertainty.

• gwModel_kick_prec_flow : normalizing-flow model for precessing-spin kick distributions, trained on NR (\(q \leq 32\)) and BHPT (\(q \leq 100\)). Returns samples from the kick distribution marginalized over spin angles.

All models from Islam & Wadekar (2025), https://arxiv.org/abs/2511.11536

Requirements for precessing model: torch, nflows

Requirements for GPR model: scikit-learn

import sys
!{sys.executable} -m pip install -e ../ --no-deps --quiet

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
import warnings
warnings.filterwarnings("ignore", "Wswiglal-redir-stdio")

import gwModels
gwModels.utils.set_rcparams()

lal.MSUN_SI != Msun

1. Aligned-spin kicks : gwModel_kick_q200

q = 2
s1z = 0.6
s2z = -0.7

vk, vk_std = gwModels.remnants.gwModel_kick_q200(q, s1z, s2z, return_std=True)
print(f'Kick velocity: {vk:.2f} +/- {vk_std:.2f} km/s')

Kick velocity: 128.97 +/- 4.29 km/s

q_arr = np.linspace(1, 256, 5000)

plt.figure(figsize=(8, 5))
for chi1z, chi2z, label in [
    (0.9, -0.9, r'$[\chi_1,\chi_2]=[0.9,-0.9]$'),
    (0.0, 0.5,  r'$[\chi_1,\chi_2]=[0.0,0.5]$'),
    (0.5, 0.0,  r'$[\chi_1,\chi_2]=[0.5,0.0]$'),
    (-0.9, 0.5, r'$[\chi_1,\chi_2]=[-0.9,0.5]$'),
    (0.9, -0.3, r'$[\chi_1,\chi_2]=[0.9,-0.3]$'),
]:
    vk = gwModels.remnants.gwModel_kick_q200(q_arr, chi1z, chi2z)
    plt.plot(np.log2(q_arr), vk, lw=2, alpha=0.8, label=label)

plt.xlabel('$q$')
plt.ylabel(r'$v_{\rm kick}$ [km/s]')
plt.xticks([0, 1, 2, 3, 4, 5, 6, 7, 8], [1, 2, 4, 8, 16, 32, 64, 128, 256])
plt.xlim([0, 8])
plt.ylim([0, 500])
plt.legend(frameon=False)
plt.grid(alpha=0.3)
plt.tight_layout()
plt.show()

2. Aligned-spin GPR kicks : gwModel_kick_q200_GPR

gpr_model = gwModels.remnants.gwModel_kick_q200_GPR('../gwModels/data/gwModel_kick_q200_GPR_aligned_spin.pkl')
gpr_model.info()

gwModel_kick_q200_GPR — GPR aligned-spin kick model
  Islam & Wadekar (2025), https://arxiv.org/abs/2511.11536
  Components  : analytical (refitted RIT) + GPR
  GPR features: [log2(q), chi_hat, chi_a]
  Valid range : {'q': [1, 1000], 's1z': [-1, 1], 's2z': [-1, 1]}

/Users/tousifislam/miniforge3/envs/kitp-py310/lib/python3.10/site-packages/sklearn/base.py:442: InconsistentVersionWarning: Trying to unpickle estimator GaussianProcessRegressor from version 1.8.0 when using version 1.7.2. This might lead to breaking code or invalid results. Use at your own risk. For more info please refer to:
https://scikit-learn.org/stable/model_persistence.html#security-maintainability-limitations
  warnings.warn(
/Users/tousifislam/miniforge3/envs/kitp-py310/lib/python3.10/site-packages/sklearn/base.py:442: InconsistentVersionWarning: Trying to unpickle estimator StandardScaler from version 1.8.0 when using version 1.7.2. This might lead to breaking code or invalid results. Use at your own risk. For more info please refer to:
https://scikit-learn.org/stable/model_persistence.html#security-maintainability-limitations
  warnings.warn(

# Point evaluation
vk_a = gpr_model.predict_analytical(q=2, s1z=0.6, s2z=-0.7)
vk_g, vk_g_std = gpr_model.predict_gpr(q=2, s1z=0.6, s2z=-0.7)
print(f'Analytical: {vk_a:.2f} km/s')
print(f'GPR: {vk_g[0]:.2f} +/- {vk_g_std[0]:.2f} km/s')

Analytical: 129.43 km/s
GPR: 132.81 +/- 5.52 km/s

# Comparison: analytical vs GPR across mass ratios
q_plot = np.linspace(1, 256, 2000)

spin_cases = [
    (0.9, -0.9, 'C0'),
    (0.0, 0.5, 'C1'),
    (0.5, 0.0, 'C2'),
    (-0.9, 0.5, 'C3'),
    (0.9, -0.3, 'C4'),
]

plt.figure(figsize=(8, 5))
for chi1z, chi2z, color in spin_cases:
    s1z_arr = np.full_like(q_plot, chi1z)
    s2z_arr = np.full_like(q_plot, chi2z)

    vk_a = gpr_model.predict_analytical(q_plot, s1z_arr, s2z_arr)
    vk_g, _ = gpr_model.predict_gpr(q_plot, s1z_arr, s2z_arr)

    label = f'$[\\chi_1,\\chi_2]=[{chi1z},{chi2z}]$'
    plt.plot(np.log2(q_plot), vk_a, c=color, lw=3, alpha=0.6, label=label)
    plt.plot(np.log2(q_plot), vk_g, c=color, lw=1.5, ls='--')

legend1 = plt.legend(
    handles=[
        Line2D([0], [0], color='gray', lw=3, alpha=0.6, label='Analytical'),
        Line2D([0], [0], color='gray', lw=1.5, ls='--', label='GPR'),
    ],
    frameon=False, loc='upper left'
)
plt.gca().add_artist(legend1)
plt.legend(frameon=False, loc='upper right')

plt.xlabel('$q$')
plt.ylabel(r'$v_{\rm kick}$ [km/s]')
plt.xticks([0, 1, 2, 3, 4, 5, 6, 7, 8], [1, 2, 4, 8, 16, 32, 64, 128, 256])
plt.xlim([0, 8])
plt.ylim(bottom=0)
plt.tight_layout()
plt.show()

# Non-spinning limit with uncertainty bands
q_ns = np.logspace(0, 3, 500)
zeros = np.zeros_like(q_ns)

vk_a, vk_a_std = gpr_model.predict_analytical(q_ns, zeros, zeros, return_std=True)
vk_g, vk_g_std = gpr_model.predict_gpr(q_ns, zeros, zeros)

plt.figure(figsize=(8, 5))
plt.fill_between(np.log2(q_ns), vk_a - 1.96*vk_a_std, vk_a + 1.96*vk_a_std,
                 alpha=0.2, color='C0', label='Analytical 95% CI')
plt.plot(np.log2(q_ns), vk_a, 'C0', lw=2.5, label='Analytical')
plt.fill_between(np.log2(q_ns), vk_g - 1.96*vk_g_std, vk_g + 1.96*vk_g_std,
                 alpha=0.2, color='C1', label='GPR 95% CI')
plt.plot(np.log2(q_ns), vk_g, 'C1--', lw=2, label='GPR')

plt.xlabel('$q$')
plt.ylabel(r'$v_{\rm kick}$ [km/s]')
plt.xticks([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
           [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024])
plt.xlim([0, 10])
plt.ylim(bottom=0)
plt.legend(frameon=False)
plt.title(r'Non-spinning limit ($\chi_{1z}=\chi_{2z}=0$)')
plt.tight_layout()
plt.show()

3. Precessing-spin kicks : gwModel_kick_prec_flow

The flow model samples kick velocity distributions for given \((q, a_1, a_2)\), marginalizing over spin orientation angles.

flow_model = gwModels.remnants.gwModel_kick_prec_flow('../gwModels/data/')
flow_model.info()

gwModel_kick_prec_flow — normalizing-flow kick model for precessing BBH
  Inputs      : q, a1, a2 (spin angles marginalized)
  Flow config : d_in=1, d_hidden=8, d_context=3, n_layers=2

samples = flow_model.sample(q=3, a1=0.5, a2=0.4, num_samples=5000)

plt.figure(figsize=(6, 4))
plt.hist(samples, bins=100, density=True, alpha=0.6)
plt.xlabel(r'$v_{\rm kick}$ [km/s]')
plt.ylabel('Probability')
plt.title(r'Kick distribution: $q=3$, $a_1=0.5$, $a_2=0.4$')
plt.grid(alpha=0.3)
plt.tight_layout()
plt.show()

mass_ratios = [1.0, 5.0, 10.0, 15.0]
a1_test, a2_test = 0.5, 0.4

all_samples = []
for q_test in mass_ratios:
    s = flow_model.sample(q=q_test, a1=a1_test, a2=a2_test, num_samples=5000)
    all_samples.append(s)

plt.figure(figsize=(7, 5))
parts = plt.violinplot(all_samples, positions=range(len(mass_ratios)),
                       widths=0.7, showmedians=True, showextrema=False)
for pc in parts['bodies']:
    pc.set_alpha(0.6)
plt.xticks(range(len(mass_ratios)), [f'$q={q}$' for q in mass_ratios])
plt.xlabel('Mass Ratio')
plt.ylabel('Kick Velocity [km/s]')
plt.title(f'$a_1={a1_test}$, $a_2={a2_test}$')
plt.grid(axis='y', alpha=0.3)
plt.tight_layout()
plt.show()