sage.data.waveform.approximants.IMRPhenomPv2

Filename : IMRPhenomPv2.py Description : Short description of the file

Created on 2026-01-21 05:26:04

__author__ = Narenraju Nagarajan __copyright__ = Copyright 2026, ProjectName __license__ = MIT Licence __version__ = 0.0.1 __maintainer__ = Narenraju Nagarajan __affiliation__ = N/A __email__ = N/A __status__ = [‘inProgress’, ‘Archived’, ‘inUsage’, ‘Debugging’]

GitHub Repository: NULL

Documentation: NULL

Classes

IMRPhenomPv2

GPU-native batched IMRPhenomPv2 precessing-spin waveform generator.

Module Contents

class IMRPhenomPv2(param_sampler=None, waveform_project=None, augment=None, append_per_det_targets=False, extra_batch=0)[source]

Bases: sage.data.waveform.approximants.IMRPhenomD.IMRPhenomD, torch.nn.Module

GPU-native batched IMRPhenomPv2 precessing-spin waveform generator.

Extends IMRPhenomD with precessing-spin corrections to the polarisations (hp, hc), then projects through ConstantProjection to produce detector-frame strain. Optional SNR rescaling and data augmentation are applied before the final output.

GRAPH_READY = True indicates that the entire forward pass is compatible with torch.compile(fullgraph=True).

Parameters:

  • param_sampler (callable or None) – Waveform parameter sampler; if None, a default DistributionSampler is used.

  • waveform_project (callable or None) – Detector projection module; if None, defaults to ConstantProjection.

  • augment (callable or None) – Optional augmentation callable applied to the projected strain.

  • state (Initialize internal Module)

  • ScriptModule. (shared by both nn.Module and)

GRAPH_READY = True[source]
cfg = None[source]
data_cfg = None[source]
signal_batch_size[source]
f[source]
df[source]
sample_length_in_s[source]
f_numel[source]
f_ref[source]
B[source]
n_pad[source]
hp_buffer[source]
hc_buffer[source]
param_sampler = None[source]
waveform_project = None[source]
augment = None[source]
append_per_det_targets = False[source]
mc_pe_idx[source]
param_names = ['mass1', 'mass2', 'spin1x', 'spin1y', 'spin1z', 'spin2x', 'spin2y', 'spin2z', 'distance', 'tc',...[source]
dist_col[source]
tc_col[source]
forward(return_theta=False)[source]
get_hphc(theta, reproduce_lal=False)[source]

Compute frequency-domain plus and cross polarisations with precessing-spin corrections.

Calls compute_derived_parameters(), convert_spins(), and PhenomPCoreTwistUp() to apply IMRPhenomPv2 precessing corrections on top of the aligned-spin IMRPhenomD backbone.

Parameters:

  • theta (torch.Tensor, shape (B, 15)) – Waveform parameters: mass1, mass2, spin1x, spin1y, spin1z, spin2x, spin2y, spin2z, distance, tc, coa_phase, inclination, polarization, ra, dec.

  • reproduce_lal (bool, optional) – If True, skip tapering, time-shifting, and df normalisation so output matches the raw LAL convention. Default is False.

Returns:

  • hp (torch.Tensor, shape (B, n_freq)) – Plus polarisation (complex64).

  • hc (torch.Tensor, shape (B, n_freq)) – Cross polarisation (complex64).

apply_tc(hp, hc, tc)[source]

Apply a time-of-coalescence phase shift to hp and hc.

Converts tc from duration-space into a frequency-domain phase ramp and applies it to the plus and cross polarisations in polar form.

Parameters:
Returns:

  • hp (torch.Tensor, shape (B, n_freq)) – Phase-shifted plus polarisation.

  • hc (torch.Tensor, shape (B, n_freq)) – Phase-shifted cross polarisation.

pad_missing_frequencies(hp, hc)[source]

Zero-pad hp and hc from DC to the low-frequency cutoff.

The waveform is only computed above f_min; this method prefixes the required number of zero bins so the output spans [0, f_max] with uniform df spacing.

Parameters:

  • hp (torch.Tensor, shape (B, n_active)) – Plus polarisation on the active frequency grid.

  • hc (torch.Tensor, shape (B, n_active)) – Cross polarisation on the active frequency grid.

Returns:

  • hp_pad (torch.Tensor, shape (B, n_pad + n_active)) – Zero-padded plus polarisation.

  • hc_pad (torch.Tensor, shape (B, n_pad + n_active)) – Zero-padded cross polarisation.

compute_derived_parameters(theta)[source]

Compute PhenomPv2-specific derived parameters from raw masses.

Overrides the IMRPhenomD base method. Internally, mass ordering is m1 ≤ m2 (Pv2 convention), whereas PhenomD expects m1 ≥ m2 and the swap is applied before calling any PhenomD helper.

Parameters:

theta (torch.Tensor, shape (B, 2+)) – Columns 0 and 1 are mass1 and mass2 in solar masses.

Returns:

derived – Columns: M (total mass, M☉), eta (symmetric mass ratio), q = m1/m2 ≥ 1 (Pv2 convention), M_s (M in seconds).

Return type:

torch.Tensor, shape (B, 4)

compute_pv2_coeffs(theta, derived, converted_spins)[source]

Compute NNLO precession-angle coefficients and reference-frequency offsets.

Evaluates the five alpha (precession) and five epsilon (rotation) PN coefficients via ComputeNNLOanglecoeffs() and integrates them at f_ref to obtain the reference-frame offsets used in PhenomPCoreTwistUp().

Parameters:
Returns:

  • angcoeffs (torch.Tensor, shape (B, 10)) – Stacked alpha (cols 0–4) and epsilon (cols 5–9) PN coefficients.

  • alphaNNLOoffset (torch.Tensor, shape (B, 1)) – Precession angle at f_ref for reference-frame subtraction.

  • epsilonNNLOoffset (torch.Tensor, shape (B, 1)) – Rotation angle at f_ref for reference-frame subtraction.

compute_spin_weighted_Y(converted_spins)[source]

Evaluate the five l=2 spin-weight-(-2) spherical harmonics at thetaJN.

Parameters:

converted_spins (torch.Tensor, shape (B, 7)) – Output of convert_spins(); column 3 is thetaJN.

Returns:

Y2 – Columns: Y₂₋₂, Y₂₋₁, Y₂₀, Y₂₁, Y₂₂ (complex, s=-2).

Return type:

torch.Tensor, shape (B, 5)

get_derived_freqs(theta_swapped, derived, phd_derived, coeffs, converted_spins)[source]

Compute all dimensionful frequency quantities needed for PhenomPv2.

Calls phP_get_transition_frequencies() with the mass-swapped parameters and collects the full set of f × M_s scale products.

Parameters:
Returns:

  • f_Ms (torch.Tensor, shape (B, n_freq)) – Frequency grid scaled by M_s.

  • fx_Ms (torch.Tensor, shape (B, 8)) – Special frequency scale products: fref, f1, f2, f3, f4, fRD, fdamp, fmid.

  • fcut_true (torch.Tensor, shape (B, 1)) – Physical frequency cutoff in Hz.

  • trans_fs (torch.Tensor, shape (B, 6)) – Transition frequencies: f1, f2, f3, f4, fRD, fdamp.

correct_time_and_phase(hp, hc, theta_swapped, derived, phd_derived, trans_fs, fx_Ms, coeffs, fcut_true)[source]

Apply time-shift and phase corrections so the PhenomPv2 waveform coalesces at t=0.

Evaluates the PhenomD phase on a fixed frequency grid near the ringdown frequency, estimates d(phase)/df via central difference, and calls apply_time_shift_phase_correction().

Parameters:

  • hp (torch.Tensor, shape (B, n_freq)) – Plus polarisation before correction.

  • hc (torch.Tensor, shape (B, n_freq)) – Cross polarisation before correction.

  • theta_swapped (torch.Tensor, shape (B, 6)) – Mass-swapped reduced parameters.

  • derived (torch.Tensor, shape (B, 4)) – PhenomPv2 derived parameters.

  • phd_derived (torch.Tensor, shape (B, 4+)) – PhenomD derived parameters.

  • trans_fs (torch.Tensor, shape (B, 6)) – Transition frequencies (f1, f2, f3, f4, fRD, fdamp).

  • fx_Ms (torch.Tensor, shape (B, 8)) – Special frequency scale products.

  • coeffs (torch.Tensor, shape (B, 7+)) – PhenomD coefficients.

  • fcut_true (torch.Tensor, shape (B, 1)) – Physical frequency cutoff in Hz.

Returns:

  • hp (torch.Tensor, shape (B, n_freq)) – Corrected plus polarisation.

  • hc (torch.Tensor, shape (B, n_freq)) – Corrected cross polarisation.

convert_spins(theta, derived)[source]

Convert Cartesian spin components to the PhenomPv2 spin parameterisation.

Maps (spin1x/y/z, spin2x/y/z) plus masses and inclination into the seven quantities used throughout PhenomPv2: aligned spins (chi1_l, chi2_l), precessing-plane spin magnitude (chip), tilt of J w.r.t. line-of-sight (thetaJN), precession reference angle (alpha0), aligned-frame orbital phase offset (phi_aligned), and polarisation rotation (zeta_polariz).

Parameters:

  • theta (torch.Tensor, shape (B, 12+)) – Columns: mass1[0], mass2[1], spin1x[2], spin1y[3], spin1z[4], spin2x[5], spin2y[6], spin2z[7], distance[8], tc[9], coa_phase[10], inclination[11]. Sage convention: mass1 >= mass2 (enforced by mass_order constraint), so spin1 belongs to the LARGER body and spin2 to the SMALLER body. Therefore theta[:,7] (spin2z) = aligned spin of the SMALLER body and theta[:,4] (spin1z) = aligned spin of the LARGER body, which is why chi1_l reads from index 7 and chi2_l from index 4 below.

  • derived (torch.Tensor, shape (B, 4)) – Output of compute_derived_parameters().

Returns:

converted – Columns: chi1_l, chi2_l, chip, thetaJN, alpha0, phi_aligned, zeta_polariz.

Return type:

torch.Tensor, shape (B, 7)

L2PNR(v, eta)[source]

Compute the 2PN orbital angular momentum magnitude L (reduced units).

Parameters:
Returns:

L – 2PN orbital angular momentum in units of M² (G=c=1).

Return type:

torch.Tensor

static ROTATEZ(angle, x, y, z)[source]

Rotate vector (x, y, z) about the z-axis by angle radians.

static ROTATEY(angle, x, y, z)[source]

Rotate vector (x, y, z) about the y-axis by angle radians.

ComputeNNLOanglecoeffs(q, chil, chip)[source]

Compute the ten NNLO PN precession-angle coefficients.

Returns the five alpha (precession angle) and five epsilon (rotation angle) post-Newtonian coefficients as stacked columns. See Appendix A of arXiv:1408.1810 (Hannam et al.) for the analytic expressions.

Parameters:

  • q (torch.Tensor, shape (B, 1)) – Mass ratio q = m1/m2 ≥ 1.

  • chil (torch.Tensor, shape (B, 1)) – Effective aligned spin χ_eff weighted by (1+q)/q.

  • chip (torch.Tensor, shape (B, 1)) – In-plane spin magnitude parameter.

Returns:

angcoeffs – Columns 0–4: alphacoeff1…5; columns 5–9: epsiloncoeff1…5.

Return type:

torch.Tensor, shape (B, 10)

SpinWeightedY(theta, phi, s, l, m)[source]

Evaluate a spin-weighted spherical harmonic Y^s_{lm}(theta, phi).

Currently supports only s=-2, l=2, m in {-2,-1,0,1,2} (dominant GW modes). Ported from SphericalHarmonics.c in LALSuite.

Parameters:

  • theta (torch.Tensor, shape (B, 1)) – Polar angle in radians.

  • phi (float or torch.Tensor) – Azimuthal angle in radians; typically 0 in the J-frame.

  • s (int) – Spin weight; must be -2.

  • l (int) – Degree; must be 2.

  • m (int) – Order; must satisfy |m| ≤ l.

Returns:

Y – Complex spin-weighted spherical harmonic value.

Return type:

torch.Tensor, shape (B, 1)

phP_get_transition_frequencies(theta, gamma2, gamma3, chip, derived, phd_derived)[source]

Compute PhenomPv2 phase and amplitude transition frequencies.

Differs from the parent PhenomD method by using phP_get_fRD_fdamp() (which incorporates the in-plane spin chip) rather than the aligned-spin ringdown frequency.

Parameters:

  • theta (torch.Tensor, shape (B, 6)) – Mass-swapped reduced parameters.

  • gamma2 (torch.Tensor, shape (B, 1)) – PhenomD amplitude Lorentzian width coefficient.

  • gamma3 (torch.Tensor, shape (B, 1)) – PhenomD amplitude Lorentzian damping coefficient.

  • chip (torch.Tensor, shape (B, 1)) – In-plane spin magnitude.

  • derived (torch.Tensor, shape (B, 4)) – PhenomPv2 derived parameters.

  • phd_derived (torch.Tensor, shape (B, 4+)) – PhenomD derived parameters.

Returns:

trans_fs – Transition frequencies: f1, f2, f3, f4, fRD, fdamp in Hz.

Return type:

torch.Tensor, shape (B, 6)

phP_get_fRD_fdamp(theta, derived, phd_derived, chip)[source]

Compute ringdown and damping frequencies for PhenomPv2.

Uses the precessing final spin from FinalSpin_inplane() (which includes the in-plane chip contribution) and the radiated energy from EradRational0815() to look up fRD and fdamp from QNM tables.

Parameters:

  • theta (torch.Tensor, shape (B, 6)) – Mass-swapped reduced parameters; columns 2–3 are chi1, chi2.

  • derived (torch.Tensor, shape (B, 4)) – PhenomPv2 derived parameters.

  • phd_derived (torch.Tensor, shape (B, 4+)) – PhenomD derived parameters; column 3 is eta.

  • chip (torch.Tensor, shape (B, 1)) – In-plane spin magnitude.

Returns:

  • fRD (torch.Tensor, shape (B, 1)) – Ringdown frequency in Hz.

  • fdamp (torch.Tensor, shape (B, 1)) – Damping frequency in Hz.

FinalSpin_inplane(theta, derived, chip)[source]

Compute the final dimensionless spin including in-plane spin contribution.

Combines the aligned final spin from FinalSpin0815() with the perpendicular component S_perp = chip × (m2/M)² to produce the total final spin magnitude, preserving the sign from the aligned component.

Parameters:

  • theta (torch.Tensor, shape (B, 6)) – Mass-swapped reduced parameters; column 0 is m1 (larger mass).

  • derived (torch.Tensor, shape (B, 4)) – PhenomPv2 derived parameters; column 0 is M.

  • chip (torch.Tensor, shape (B, 1)) – In-plane spin magnitude.

Returns:

af – Final dimensionless spin.

Return type:

torch.Tensor, shape (B, 1)

FinalSpin0815(eta, chi1, chi2)[source]

Compute the aligned final spin using the Barkett et al. (0815) fit.

Delegates to FinalSpin0815_s() after forming the mass-weighted effective spin S = m1²·chi1 + m2²·chi2.

Parameters:

  • eta (torch.Tensor, shape (B, 1)) – Symmetric mass ratio.

  • chi1 (torch.Tensor, shape (B, 1)) – Dimensionless aligned spin of the larger BH.

  • chi2 (torch.Tensor, shape (B, 1)) – Dimensionless aligned spin of the smaller BH.

Returns:

af_parallel – Aligned-spin final dimensionless spin.

Return type:

torch.Tensor, shape (B, 1)

FinalSpin0815_s(eta, S)[source]

Evaluate the Barkett et al. (arXiv:0815) final-spin rational fit given S.

Parameters:

  • eta (torch.Tensor, shape (B, 1)) – Symmetric mass ratio.

  • S (torch.Tensor, shape (B, 1)) – Mass-weighted effective spin S = (m1²·chi1 + m2²·chi2) / M².

Returns:

af – Final dimensionless spin from the rational fit.

Return type:

torch.Tensor, shape (B, 1)

EradRational0815(eta, chi1, chi2)[source]

Compute radiated energy fraction using the Barkett et al. (0815) rational fit.

Delegates to EradRational0815_s() after forming the mass-weighted effective spin.

Parameters:

  • eta (torch.Tensor, shape (B, 1)) – Symmetric mass ratio.

  • chi1 (torch.Tensor, shape (B, 1)) – Dimensionless aligned spin of the larger BH.

  • chi2 (torch.Tensor, shape (B, 1)) – Dimensionless aligned spin of the smaller BH.

Returns:

Erad – Fraction of total mass radiated as gravitational waves.

Return type:

torch.Tensor, shape (B, 1)

EradRational0815_s(eta, s)[source]

Evaluate the Barkett et al. (arXiv:0815) radiated-energy rational fit.

Parameters:
Returns:

Erad – Radiated energy fraction E_rad / M_total.

Return type:

torch.Tensor, shape (B, 1)

PhenomPCoreTwistUp(f_Ms, hPhenom, eta, chi1_l, chi2_l, chip, angcoeffs, Y2m, alphaoffset, epsilonoffset)[source]

Apply the PhenomPv2 “twist-up” to convert aligned-spin PhenomD into precessing hp and hc polarisations.

Evaluates the Wigner d-matrix coefficients, computes the precessing-frame alpha and epsilon angles, and assembles the l=2 mode sum following arXiv:1408.1810 (Hannam et al.), eqs. (A1)–(A4).

Parameters:
Returns:

  • hp (torch.Tensor, shape (B, n_freq)) – Plus polarisation (complex).

  • hc (torch.Tensor, shape (B, n_freq)) – Cross polarisation (complex).

WignerdCoefficients(v, SL, eta, Sp)[source]

Compute the half-angle Wigner d-matrix coefficients cos(β/2) and sin(β/2).

Estimates the precession opening angle β from the ratio of the total in-plane spin to the total angular momentum, using the 2PN orbital angular momentum from L2PNR().

Parameters:

  • v (torch.Tensor, shape (B, n_freq)) – Orbital velocity (πMf)^(1/3).

  • SL (torch.Tensor, shape (B, 1)) – Aligned dimensionful spin S_L = chi1_l·m1² + chi2_l·m2².

  • eta (torch.Tensor, shape (B, 1)) – Symmetric mass ratio.

  • Sp (torch.Tensor, shape (B, 1)) – In-plane spin S_perp = chip · m2².

Returns:

  • cos_beta_half (torch.Tensor) – cos(β/2) for the Wigner d-matrix.

  • sin_beta_half (torch.Tensor) – sin(β/2) for the Wigner d-matrix.

PhenomPOneFrequency(f, f_Ms, fx_Ms, theta, phd_derived, coeffs, trans_fs, fcut_true)[source]

m1, m2: in solar masses phic: Orbital phase at the peak of the underlying non precessing model (rad) M: Total mass (Solar masses)

apply_time_shift_phase_correction(hptilde, hctilde, freqs_fixed, phase_fixed, f_final, offset=0)[source]

Apply time shift correction so the waveform coalesces at t=0.

Parameters:

  • hptilde – Tensor of shape (n_freq,) with plus polarization.

  • hctilde – Tensor of shape (n_freq,) with cross polarization.

  • freqs – Tensor of frequencies corresponding to hptilde/hctilde.

  • freqs_fixed – Fixed frequency grid used for spline interpolation.

  • phase_fixed – Phase values on freqs_fixed.

  • f_final – Final frequency (fRD or f_merger) to evaluate derivative.

  • offset (int) – Index offset if freqs does not start at zero.

Returns:

Tuple of corrected (hptilde, hctilde)