FEAT: add precessing spin transformation

bilby/core/utils/constants.py

@@ -5,3 +5,4 @@
 solar_mass = 1.988409870698050731911960804878414216e30  # Kg
 radius_of_earth = 6378136.6  # m
 gravitational_constant = 6.6743e-11  # m^3 kg^-1 s^-2
+msun_time_si = gravitational_constant * solar_mass / speed_of_light**3  # s

bilby/gw/conversion.py

@@ -12,6 +12,9 @@
 from pandas import DataFrame, Series
 from scipy.stats import norm
 
+from .eos.eos import IntegrateTOV
+from .cosmology import get_cosmology, z_at_value
+from .geometry import transform_precessing_spins
 from .utils import (lalsim_SimNeutronStarEOS4ParamSDGammaCheck,
                     lalsim_SimNeutronStarEOS4ParameterSpectralDecomposition,
                     lalsim_SimNeutronStarEOS4ParamSDViableFamilyCheck,
@@ -25,14 +28,10 @@
                     lalsim_SimNeutronStarMaximumMass,
                     lalsim_SimNeutronStarRadius,
                     lalsim_SimNeutronStarLoveNumberK2)
-
 from ..compat.utils import array_module
 from ..core.likelihood import MarginalizedLikelihoodReconstructionError
-from ..core.utils import logger, solar_mass, gravitational_constant, speed_of_light, command_line_args, safe_file_dump
 from ..core.prior import DeltaFunction
-from .utils import lalsim_SimInspiralTransformPrecessingNewInitialConditions
-from .eos.eos import IntegrateTOV
-from .cosmology import get_cosmology, z_at_value
+from ..core.utils import logger, solar_mass, gravitational_constant, speed_of_light, command_line_args, safe_file_dump
 
 
 def redshift_to_luminosity_distance(redshift, cosmology=None):
@@ -152,33 +151,14 @@ def bilby_to_lalsimulation_spins(
         spin_2z = a_2 * np.cos(tilt_2)
         iota = theta_jn
     else:
-        from numbers import Number
-        args = (
-            theta_jn, phi_jl, tilt_1, tilt_2, phi_12, a_1, a_2, mass_1,
-            mass_2, reference_frequency, phase
+        func = transform_precessing_spins
+        iota, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, spin_2z = func(
+            theta_jn, phi_jl, tilt_1, tilt_2, phi_12, a_1, a_2,
+            mass_1, mass_2, reference_frequency, phase
         )
-        float_inputs = all([isinstance(arg, Number) for arg in args])
-        if float_inputs:
-            func = lalsim_SimInspiralTransformPrecessingNewInitialConditions
-        else:
-            func = transform_precessing_spins
-        iota, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, spin_2z = func(*args)
     return iota, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, spin_2z
 
 
-@np.vectorize
-def transform_precessing_spins(*args):
-    """
-    Vectorized wrapper for
-    lalsimulation.SimInspiralTransformPrecessingNewInitialConditions
-
-    For detailed documentation see
-    :code:`bilby.gw.conversion.bilby_to_lalsimulation_spins`.
-    This will be removed from the public API in a future release.
-    """
-    return lalsim_SimInspiralTransformPrecessingNewInitialConditions(*args)
-
-
 def convert_to_lal_binary_black_hole_parameters(parameters):
     """
     Convert parameters we have into parameters we need.

bilby/gw/geometry.py

@@ -1,7 +1,8 @@
 from plum import dispatch
 
 from .time import greenwich_mean_sidereal_time
-from ..compat.utils import array_module, promote_to_array
+from ..compat.utils import array_module, promote_to_array, xp_wrap
+from ..core.utils.constants import msun_time_si
 
 
 __all__ = [
@@ -375,3 +376,168 @@ def zenith_azimuth_to_theta_phi(zenith, azimuth, delta_x):
     theta = xp.arccos(omega[2])
     phi = xp.arctan2(omega[1], omega[0]) % (2 * xp.pi)
     return theta, phi
+
+
+@xp_wrap
+def transform_precessing_spins(
+    theta_jn,
+    phi_jl,
+    tilt_1,
+    tilt_2,
+    phi_12,
+    chi_1,
+    chi_2,
+    mass_1,
+    mass_2,
+    f_ref,
+    phase,
+    *,
+    xp=None,
+):
+    """
+    A direct reimplementation of
+    :code:`lalsimulation.SimInspiralTransformPrecessingNewInitialConditions`.
+
+    Parameters
+    ----------
+    theta_jn: float | xp.ndarray
+        Zenith angle between J and N (rad).
+    phi_jl: float | xp.ndarray
+        Azimuthal angle of L_N on its cone about J (rad).
+    tilt_1: float | xp.ndarray
+        Zenith angle between S1 and LNhat (rad).
+    tilt_2: float | xp.ndarray
+        Zenith angle between S2 and LNhat (rad).
+    phi_12: float | xp.ndarray
+        Difference in azimuthal angle between S1, S2 (rad).
+    chi_1: float | xp.ndarray
+        Dimensionless spin of body 1.
+    chi_2: float | xp.ndarray
+        Dimensionless spin of body 2.
+    mass_1: float | xp.ndarray
+        Mass of body 1 (solar masses).
+    mass_2: float | xp.ndarray
+        Mass of body 2 (solar masses).
+    f_ref: float | xp.ndarray
+        Reference GW frequency (Hz).
+    phase: float | xp.ndarray
+        Reference orbital phase.
+
+    Returns
+    -------
+    tuple
+        (iota, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, spin_2z)
+        - iota: Inclination angle of L_N
+        - spin_1x, spin_1y, spin_1z: Components of spin 1
+        - spin_2x, spin_2y, spin_2z: Components of spin 2
+    """
+
+    # Helper rotation functions
+    def rotate_z(angle, vec):
+        """Rotate vector about z-axis"""
+        cos_a = xp.cos(angle)
+        sin_a = xp.sin(angle)
+        x_new = cos_a * vec[0] - sin_a * vec[1]
+        y_new = sin_a * vec[0] + cos_a * vec[1]
+        return xp.asarray([x_new, y_new, vec[2]])
+
+    def rotate_y(angle, vec):
+        """Rotate vector about y-axis"""
+        cos_a = xp.cos(angle)
+        sin_a = xp.sin(angle)
+        x_new = cos_a * vec[0] + sin_a * vec[2]
+        z_new = -sin_a * vec[0] + cos_a * vec[2]
+        return xp.asarray([x_new, vec[1], z_new])
+
+    # Starting frame: LNhat is along the z-axis
+    ln_hat = xp.asarray([theta_jn * 0, theta_jn * 0, theta_jn ** 0])
+
+    # Initial spin unit vectors
+    s1_hat = xp.asarray([
+        xp.sin(tilt_1) * xp.cos(phase),
+        xp.sin(tilt_1) * xp.sin(phase),
+        xp.cos(tilt_1),
+    ])
+
+    s2_hat = xp.asarray([
+        xp.sin(tilt_2) * xp.cos(phi_12 + phase),
+        xp.sin(tilt_2) * xp.sin(phi_12 + phase),
+        xp.cos(tilt_2),
+    ])
+
+    # Compute physical parameters
+    m_total = mass_1 + mass_2
+    eta = mass_1 * mass_2 / (m_total * m_total)
+
+    # v parameter at reference point (c=G=1 units)
+    v0 = (m_total * msun_time_si * xp.pi * f_ref) ** (1 / 3)
+
+    # Compute angular momentum magnitude using PN expressions
+    # L/M = eta * v^(-1) * (1 + v^2 * L_2PN)
+    # L_2PN = 3/2 + 1/6 * eta
+    l_2pn = 1.5 + eta / 6.0
+    l_mag = eta * m_total * m_total / v0 * (1.0 + v0 * v0 * l_2pn)
+
+    # Spin vectors with proper magnitudes
+    s1 = mass_1 * mass_1 * chi_1 * s1_hat
+    s2 = mass_2 * mass_2 * chi_2 * s2_hat
+
+    # Total angular momentum J = L + S1 + S2
+    # l_vec = xp.asarray([xp.zeros_like(theta_jn), xp.zeros_like(theta_jn), l_mag])
+    l_vec = xp.asarray([l_mag * 0, l_mag * 0, l_mag])
+    j = l_vec + s1 + s2
+
+    # Normalize J to get Jhat and find its angles
+    j_norm = xp.sqrt(xp.sum(j ** 2, axis=0))
+    j_hat = j / j_norm
+
+    theta_0 = xp.arccos(j_hat[2])
+    phi_0 = xp.arctan2(j_hat[1], j_hat[0])
+
+    # Rotation 1: Rotate about z-axis by -phi_0 to put Jhat in x-z plane
+    angle = -phi_0
+    s1_hat = rotate_z(angle, s1_hat)
+    s2_hat = rotate_z(angle, s2_hat)
+
+    # Rotation 2: Rotate about y-axis by -theta_0 to put Jhat along z-axis
+    angle = -theta_0
+    ln_hat = rotate_y(angle, ln_hat)
+    s1_hat = rotate_y(angle, s1_hat)
+    s2_hat = rotate_y(angle, s2_hat)
+
+    # Rotation 3: Rotate about z-axis by (phi_jl - pi) to put L at desired azimuth
+    angle = phi_jl - xp.pi
+    ln_hat = rotate_z(angle, ln_hat)
+    s1_hat = rotate_z(angle, s1_hat)
+    s2_hat = rotate_z(angle, s2_hat)
+
+    # Compute inclination: angle between L and N
+    n = xp.asarray([theta_jn * 0, xp.sin(theta_jn), xp.cos(theta_jn)])
+    iota = xp.arccos(xp.sum(n * ln_hat, axis=0))
+
+    # Rotation 4-5: Bring L into the z-axis
+    theta_lj = xp.arccos(ln_hat[2])
+    phi_l = xp.arctan2(ln_hat[1], ln_hat[0])
+
+    angle = -phi_l
+    s1_hat = rotate_z(angle, s1_hat)
+    s2_hat = rotate_z(angle, s2_hat)
+    n = rotate_z(angle, n)
+
+    angle = -theta_lj
+    s1_hat = rotate_y(angle, s1_hat)
+    s2_hat = rotate_y(angle, s2_hat)
+    n = rotate_y(angle, n)
+
+    # Rotation 6: Bring N into y-z plane with positive y component
+    phi_n = xp.arctan2(n[1], n[0])
+
+    angle = xp.pi / 2.0 - phi_n - phase
+    s1_hat = rotate_z(angle, s1_hat)
+    s2_hat = rotate_z(angle, s2_hat)
+
+    # Return final spin components
+    spin_1 = s1_hat * chi_1
+    spin_2 = s2_hat * chi_2
+
+    return iota, *spin_1, *spin_2

bilby/gw/source.py

@@ -143,10 +143,16 @@ def gwsignal_binary_black_hole(frequency_array, mass_1, mass_2, luminosity_dista
     frequency_bounds = ((frequency_array >= minimum_frequency) *
                         (frequency_array <= maximum_frequency))
 
-    iota, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, spin_2z = bilby_to_lalsimulation_spins(
-        theta_jn=theta_jn, phi_jl=phi_jl, tilt_1=tilt_1, tilt_2=tilt_2,
-        phi_12=phi_12, a_1=a_1, a_2=a_2, mass_1=mass_1 * utils.solar_mass, mass_2=mass_2 * utils.solar_mass,
-        reference_frequency=reference_frequency, phase=phase)
+    try:
+        iota, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, spin_2z = bilby_to_lalsimulation_spins(
+            theta_jn=theta_jn, phi_jl=phi_jl, tilt_1=tilt_1, tilt_2=tilt_2,
+            phi_12=phi_12, a_1=a_1, a_2=a_2, mass_1=mass_1, mass_2=mass_2,
+            reference_frequency=reference_frequency, phase=phase)
+    except ZeroDivisionError:
+        if catch_waveform_errors:
+            return None
+        else:
+            raise
 
     eccentricity = 0.0
     longitude_ascending_nodes = 0.0
@@ -619,14 +625,21 @@ def _base_lal_cbc_fd_waveform(
                         (frequency_array <= maximum_frequency))
 
     luminosity_distance = luminosity_distance * 1e6 * utils.parsec
+
+    try:
+        iota, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, spin_2z = bilby_to_lalsimulation_spins(
+            theta_jn=theta_jn, phi_jl=phi_jl, tilt_1=tilt_1, tilt_2=tilt_2,
+            phi_12=phi_12, a_1=a_1, a_2=a_2, mass_1=mass_1, mass_2=mass_2,
+            reference_frequency=reference_frequency, phase=phase)
+    except ZeroDivisionError:
+        if catch_waveform_errors:
+            return None
+        else:
+            raise
+
     mass_1 = mass_1 * utils.solar_mass
     mass_2 = mass_2 * utils.solar_mass
 
-    iota, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, spin_2z = bilby_to_lalsimulation_spins(
-        theta_jn=theta_jn, phi_jl=phi_jl, tilt_1=tilt_1, tilt_2=tilt_2,
-        phi_12=phi_12, a_1=a_1, a_2=a_2, mass_1=mass_1, mass_2=mass_2,
-        reference_frequency=reference_frequency, phase=phase)
-
     longitude_ascending_nodes = 0.0
     mean_per_ano = 0.0
 
@@ -1113,14 +1126,21 @@ def _base_waveform_frequency_sequence(
     approximant = lalsim_GetApproximantFromString(approximant)
 
     luminosity_distance = luminosity_distance * 1e6 * utils.parsec
+
+    try:
+        iota, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, spin_2z = bilby_to_lalsimulation_spins(
+            theta_jn=theta_jn, phi_jl=phi_jl, tilt_1=tilt_1, tilt_2=tilt_2,
+            phi_12=phi_12, a_1=a_1, a_2=a_2, mass_1=mass_1, mass_2=mass_2,
+            reference_frequency=reference_frequency, phase=phase)
+    except ZeroDivisionError:
+        if catch_waveform_errors:
+            return None
+        else:
+            raise
+
     mass_1 = mass_1 * utils.solar_mass
     mass_2 = mass_2 * utils.solar_mass
 
-    iota, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, spin_2z = bilby_to_lalsimulation_spins(
-        theta_jn=theta_jn, phi_jl=phi_jl, tilt_1=tilt_1, tilt_2=tilt_2,
-        phi_12=phi_12, a_1=a_1, a_2=a_2, mass_1=mass_1, mass_2=mass_2,
-        reference_frequency=reference_frequency, phase=phase)
-
     try:
         h_plus, h_cross = lalsim_SimInspiralChooseFDWaveformSequence(
             phase, mass_1, mass_2, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y,

examples/gw_examples/injection_examples/jax_fast_tutorial.py

@@ -24,30 +24,6 @@
 jax.config.update("jax_enable_x64", True)
 
 
-def bilby_to_ripple_spins(
-    theta_jn,
-    phi_jl,
-    tilt_1,
-    tilt_2,
-    phi_12,
-    a_1,
-    a_2,
-):
-    """
-    A simplified spherical to cartesian spin conversion function.
-    This is not equivalent to the method used in `bilby.gw.conversion`
-    which comes from `lalsimulation` and is not `JAX` compatible.
-    """
-    iota = theta_jn
-    spin_1x = a_1 * jnp.sin(tilt_1) * jnp.cos(phi_jl)
-    spin_1y = a_1 * jnp.sin(tilt_1) * jnp.sin(phi_jl)
-    spin_1z = a_1 * jnp.cos(tilt_1)
-    spin_2x = a_2 * jnp.sin(tilt_2) * jnp.cos(phi_jl + phi_12)
-    spin_2y = a_2 * jnp.sin(tilt_2) * jnp.sin(phi_jl + phi_12)
-    spin_2z = a_2 * jnp.cos(tilt_2)
-    return iota, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, spin_2z
-
-
 def ripple_bbh(
     frequency,
     mass_1,
@@ -102,8 +78,19 @@ def ripple_bbh(
     dict
         Dictionary containing the plus and cross polarizations of the waveform.
     """
-    iota, *cartesian_spins = bilby_to_ripple_spins(
-        theta_jn, phi_jl, tilt_1, tilt_2, phi_12, a_1, a_2
+    reference_frequency = jnp.asarray(kwargs["reference_frequency"])
+    iota, *cartesian_spins = bilby.gw.geometry.transform_precessing_spins(
+        theta_jn,
+        phi_jl,
+        tilt_1,
+        tilt_2,
+        phi_12,
+        a_1,
+        a_2,
+        mass_1,
+        mass_2,
+        reference_frequency,
+        phase,
     )
     frequencies = jnp.maximum(frequency, kwargs["minimum_frequency"])
     theta = jnp.array(
@@ -118,7 +105,7 @@ def ripple_bbh(
         ]
     )
     wf_func = jax.jit(IMRPhenomPv2.gen_IMRPhenomPv2)
-    hp, hc = wf_func(frequencies, theta, jnp.array(20.0))
+    hp, hc = wf_func(frequencies, theta, reference_frequency)
     return dict(plus=hp, cross=hc)