diff --git a/bilby/core/utils/constants.py b/bilby/core/utils/constants.py index 8dbec27da..e734923a5 100644 --- a/bilby/core/utils/constants.py +++ b/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 diff --git a/bilby/gw/conversion.py b/bilby/gw/conversion.py index 96cd02dd1..ae908da1c 100644 --- a/bilby/gw/conversion.py +++ b/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. diff --git a/bilby/gw/geometry.py b/bilby/gw/geometry.py index fa4be3d86..461ae7ae4 100644 --- a/bilby/gw/geometry.py +++ b/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,175 @@ 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.stack([x_new, y_new, vec[2]], axis=0) + + 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.stack([x_new, vec[1], z_new], axis=0) + + # Starting frame: LNhat is along the z-axis + ln_hat = xp.stack([ + xp.zeros_like(theta_jn), + xp.zeros_like(theta_jn), + xp.ones_like(theta_jn) + ], axis=0) + + # Initial spin unit vectors + s1_hat = xp.stack([ + xp.sin(tilt_1) * xp.cos(phase), + xp.sin(tilt_1) * xp.sin(phase), + xp.cos(tilt_1) + ], axis=0) + + s2_hat = xp.stack([ + xp.sin(tilt_2) * xp.cos(phi_12 + phase), + xp.sin(tilt_2) * xp.sin(phi_12 + phase), + xp.cos(tilt_2) + ], axis=0) + + # 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.stack([xp.zeros_like(theta_jn), xp.zeros_like(theta_jn), l_mag], axis=0) + j = l_vec + s1 + s2 + + # Normalize J to get Jhat and find its angles + j_norm = xp.sqrt(xp.sum(j * j, 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.stack([ + xp.zeros_like(theta_jn), + xp.sin(theta_jn), + xp.cos(theta_jn) + ], axis=0) + 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 diff --git a/bilby/gw/source.py b/bilby/gw/source.py index 11411b468..07d3f8327 100644 --- a/bilby/gw/source.py +++ b/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, diff --git a/examples/gw_examples/injection_examples/jax_fast_tutorial.py b/examples/gw_examples/injection_examples/jax_fast_tutorial.py index 22a64c935..7a51a6162 100644 --- a/examples/gw_examples/injection_examples/jax_fast_tutorial.py +++ b/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) diff --git a/test/gw/geometry_test.py b/test/gw/geometry_test.py new file mode 100644 index 000000000..505965d48 --- /dev/null +++ b/test/gw/geometry_test.py @@ -0,0 +1,112 @@ +import array_api_compat as aac +import numpy as np +import pytest +from bilby.core.prior import Uniform +from bilby.gw.prior import BBHPriorDict +from bilby.gw.geometry import transform_precessing_spins + + +@pytest.mark.array_backend +@pytest.mark.usefixtures("xp_class") +class TestTransformPrecessingSpins: + def test_transform_precessing_spins(self): + """ + Verify that our port of this function matches the lalsimulation version. + """ + import lal + from lalsimulation import SimInspiralTransformPrecessingNewInitialConditions + + priors = BBHPriorDict() + priors["mass_1"] = Uniform(1, 1000) + priors["mass_2"] = Uniform(1, 1000) + priors["reference_frequency"] = Uniform(10, 100) + + # some default priors are problematic for some array backends + for key in ["luminosity_distance", "chirp_mass", "mass_ratio"]: + del priors[key] + + for _ in range(100): + point = priors.sample(random_state=self.rng) + bilby_transformed = transform_precessing_spins( + point["theta_jn"], + point["phi_jl"], + point["tilt_1"], + point["tilt_2"], + point["phi_12"], + point["a_1"], + point["a_2"], + point["mass_1"], + point["mass_2"], + point["reference_frequency"], + point["phase"], + ) + assert aac.get_namespace(*bilby_transformed) == self.xp + bilby_transformed = np.asarray(bilby_transformed) + lalsim_transformed = np.asarray(SimInspiralTransformPrecessingNewInitialConditions( + float(point["theta_jn"]), + float(point["phi_jl"]), + float(point["tilt_1"]), + float(point["tilt_2"]), + float(point["phi_12"]), + float(point["a_1"]), + float(point["a_2"]), + float(point["mass_1"] * lal.MSUN_SI), + float(point["mass_2"] * lal.MSUN_SI), + float(point["reference_frequency"]), + float(point["phase"]), + )) + # the different array backends have some precision loss + np.testing.assert_allclose(bilby_transformed, lalsim_transformed, atol=1e-5) + + @pytest.mark.array_backend + def test_transform_precessing_spins_vectorized(self): + """ + Run the tests with vectorization, note that this returns a tuple of arrays. + """ + import lal + from lalsimulation import SimInspiralTransformPrecessingNewInitialConditions + + priors = BBHPriorDict() + priors["mass_1"] = Uniform(1, 1000) + priors["mass_2"] = Uniform(1, 1000) + priors["reference_frequency"] = Uniform(10, 100) + + # some default priors are problematic for some array backends + for key in ["luminosity_distance", "chirp_mass", "mass_ratio"]: + del priors[key] + + points = priors.sample(100, random_state=self.rng) + bilby_transformed = transform_precessing_spins( + points["theta_jn"], + points["phi_jl"], + points["tilt_1"], + points["tilt_2"], + points["phi_12"], + points["a_1"], + points["a_2"], + points["mass_1"], + points["mass_2"], + points["reference_frequency"], + points["phase"], + ) + assert aac.get_namespace(*bilby_transformed) == self.xp + bilby_transformed = np.asarray(bilby_transformed) + lalsim_transformed = list() + for ii in range(len(points["theta_jn"])): + point = {key: points[key][ii] for key in points.keys()} + lalsim_transformed.append(np.asarray(SimInspiralTransformPrecessingNewInitialConditions( + float(point["theta_jn"]), + float(point["phi_jl"]), + float(point["tilt_1"]), + float(point["tilt_2"]), + float(point["phi_12"]), + float(point["a_1"]), + float(point["a_2"]), + float(point["mass_1"] * lal.MSUN_SI), + float(point["mass_2"] * lal.MSUN_SI), + float(point["reference_frequency"]), + float(point["phase"]), + ))) + lalsim_transformed = np.asarray(lalsim_transformed).T + # the different array backends have some precision loss + np.testing.assert_allclose(bilby_transformed, lalsim_transformed, atol=1e-5)