From 70e2bbfe72f0bff021295dc18586ec970f6d08a5 Mon Sep 17 00:00:00 2001 From: Colm Talbot Date: Fri, 26 Jun 2026 12:02:06 -0400 Subject: [PATCH 01/11] FEAT: add precessing spin transformation --- bilby/gw/conversion.py | 27 +-- bilby/gw/geometry.py | 171 ++++++++++++++++++ .../injection_examples/jax_fast_tutorial.py | 28 +-- test/gw/geometry_test.py | 107 +++++++++++ 4 files changed, 285 insertions(+), 48 deletions(-) create mode 100644 test/gw/geometry_test.py diff --git a/bilby/gw/conversion.py b/bilby/gw/conversion.py index 96cd02dd1..6a442bec1 100644 --- a/bilby/gw/conversion.py +++ b/bilby/gw/conversion.py @@ -33,6 +33,7 @@ from .utils import lalsim_SimInspiralTransformPrecessingNewInitialConditions from .eos.eos import IntegrateTOV from .cosmology import get_cosmology, z_at_value +from .geometry import transform_precessing_spins def redshift_to_luminosity_distance(redshift, cosmology=None): @@ -152,32 +153,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): """ diff --git a/bilby/gw/geometry.py b/bilby/gw/geometry.py index fa4be3d86..da788d428 100644 --- a/bilby/gw/geometry.py +++ b/bilby/gw/geometry.py @@ -375,3 +375,174 @@ 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 + + +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, +): + """ + 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 + """ + + xp = array_module(theta_jn) + + # 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/examples/gw_examples/injection_examples/jax_fast_tutorial.py b/examples/gw_examples/injection_examples/jax_fast_tutorial.py index 22a64c935..8e0e53530 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,8 @@ 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 + 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, f_ref, phase ) frequencies = jnp.maximum(frequency, kwargs["minimum_frequency"]) theta = jnp.array( diff --git a/test/gw/geometry_test.py b/test/gw/geometry_test.py new file mode 100644 index 000000000..3d87c628e --- /dev/null +++ b/test/gw/geometry_test.py @@ -0,0 +1,107 @@ +import numpy as np +import pytest + + +@pytest.mark.array_backend +def test_transform_precessing_spins(xp): + """ + Verify that our port of this function matches the lalsimulation version. + """ + import lal + from bilby.core.prior import Uniform + from bilby.gw.prior import BBHPriorDict + from bilby.gw.geometry import transform_precessing_spins + 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(xp=xp) + bilby_transformed = np.asarray(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"], + )) + 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"]), + )) + np.testing.assert_allclose(bilby_transformed, lalsim_transformed, rtol=1e-10) + + +@pytest.mark.array_backend +def test_transform_precessing_spins_vectorized(xp): + """ + Run the tests with vectorization, note that this returns a tuple of arrays. + """ + import lal + from bilby.core.prior import Uniform + from bilby.gw.prior import BBHPriorDict + from bilby.gw.geometry import transform_precessing_spins + 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, xp=xp) + bilby_transformed = np.asarray(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"], + )) + 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 + np.testing.assert_allclose(bilby_transformed, lalsim_transformed, rtol=1e-10) From baca325a1c6a8a05219a7d592f8719c7a3aa19ee Mon Sep 17 00:00:00 2001 From: Colm Talbot Date: Wed, 1 Jul 2026 09:18:07 -0500 Subject: [PATCH 02/11] FMT: fix precommits --- bilby/core/utils/constants.py | 1 + bilby/gw/conversion.py | 11 ++++------- bilby/gw/geometry.py | 3 ++- .../injection_examples/jax_fast_tutorial.py | 15 +++++++++++++-- 4 files changed, 20 insertions(+), 10 deletions(-) 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 6a442bec1..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,15 +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 .geometry import transform_precessing_spins +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): @@ -161,7 +159,6 @@ def bilby_to_lalsimulation_spins( return iota, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, spin_2z - 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 da788d428..75afdbdca 100644 --- a/bilby/gw/geometry.py +++ b/bilby/gw/geometry.py @@ -2,6 +2,7 @@ from .time import greenwich_mean_sidereal_time from ..compat.utils import array_module, promote_to_array +from ...core.utils.constants import msun_time_si __all__ = [ @@ -472,7 +473,7 @@ def rotate_y(angle, vec): 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) + 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) diff --git a/examples/gw_examples/injection_examples/jax_fast_tutorial.py b/examples/gw_examples/injection_examples/jax_fast_tutorial.py index 8e0e53530..7a51a6162 100644 --- a/examples/gw_examples/injection_examples/jax_fast_tutorial.py +++ b/examples/gw_examples/injection_examples/jax_fast_tutorial.py @@ -78,8 +78,19 @@ def ripple_bbh( dict Dictionary containing the plus and cross polarizations of the waveform. """ + 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, f_ref, phase + 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( @@ -94,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) From ae9ae26528c30196ebc845b6f9b930c4612a522a Mon Sep 17 00:00:00 2001 From: Colm Talbot Date: Mon, 6 Jul 2026 13:07:28 -0400 Subject: [PATCH 03/11] TYPO: fix relative import --- bilby/gw/geometry.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bilby/gw/geometry.py b/bilby/gw/geometry.py index 75afdbdca..544d4acaa 100644 --- a/bilby/gw/geometry.py +++ b/bilby/gw/geometry.py @@ -2,7 +2,7 @@ from .time import greenwich_mean_sidereal_time from ..compat.utils import array_module, promote_to_array -from ...core.utils.constants import msun_time_si +from ..core.utils.constants import msun_time_si __all__ = [ From c8e856958754f7eabc70d432033b188f1a007323 Mon Sep 17 00:00:00 2001 From: Colm Talbot Date: Mon, 6 Jul 2026 13:16:11 -0400 Subject: [PATCH 04/11] TEST: Refactor transform_precessing_spins tests --- test/gw/geometry_test.py | 206 ++++++++++++++++++++------------------- 1 file changed, 105 insertions(+), 101 deletions(-) diff --git a/test/gw/geometry_test.py b/test/gw/geometry_test.py index 3d87c628e..02fe1897a 100644 --- a/test/gw/geometry_test.py +++ b/test/gw/geometry_test.py @@ -1,107 +1,111 @@ +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 -def test_transform_precessing_spins(xp): - """ - Verify that our port of this function matches the lalsimulation version. - """ - import lal - from bilby.core.prior import Uniform - from bilby.gw.prior import BBHPriorDict - from bilby.gw.geometry import transform_precessing_spins - 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(xp=xp) - bilby_transformed = np.asarray(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"], - )) - 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"]), - )) +@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(rng=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"]), + )) + np.testing.assert_allclose(bilby_transformed, lalsim_transformed, rtol=1e-10) + + + @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, rng=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 np.testing.assert_allclose(bilby_transformed, lalsim_transformed, rtol=1e-10) - - -@pytest.mark.array_backend -def test_transform_precessing_spins_vectorized(xp): - """ - Run the tests with vectorization, note that this returns a tuple of arrays. - """ - import lal - from bilby.core.prior import Uniform - from bilby.gw.prior import BBHPriorDict - from bilby.gw.geometry import transform_precessing_spins - 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, xp=xp) - bilby_transformed = np.asarray(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"], - )) - 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 - np.testing.assert_allclose(bilby_transformed, lalsim_transformed, rtol=1e-10) From 6b7f8380deb1702ee2bed435d85a88e42f148116 Mon Sep 17 00:00:00 2001 From: Colm Talbot Date: Mon, 6 Jul 2026 13:23:26 -0400 Subject: [PATCH 05/11] MAINT: add xp_wrap to transform_precessing_spins --- bilby/gw/geometry.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/bilby/gw/geometry.py b/bilby/gw/geometry.py index 544d4acaa..461ae7ae4 100644 --- a/bilby/gw/geometry.py +++ b/bilby/gw/geometry.py @@ -1,7 +1,7 @@ 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 @@ -378,6 +378,7 @@ def zenith_azimuth_to_theta_phi(zenith, azimuth, delta_x): return theta, phi +@xp_wrap def transform_precessing_spins( theta_jn, phi_jl, @@ -390,6 +391,8 @@ def transform_precessing_spins( mass_2, f_ref, phase, + *, + xp=None, ): """ A direct reimplementation of @@ -429,8 +432,6 @@ def transform_precessing_spins( - spin_2x, spin_2y, spin_2z: Components of spin 2 """ - xp = array_module(theta_jn) - # Helper rotation functions def rotate_z(angle, vec): """Rotate vector about z-axis""" From 56c50c4e072aed2f8cba17f3e19a227842d910c1 Mon Sep 17 00:00:00 2001 From: Colm Talbot Date: Mon, 6 Jul 2026 14:09:16 -0400 Subject: [PATCH 06/11] TEST: fix precessing spin tests --- test/gw/geometry_test.py | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/test/gw/geometry_test.py b/test/gw/geometry_test.py index 02fe1897a..dfd5ae2cc 100644 --- a/test/gw/geometry_test.py +++ b/test/gw/geometry_test.py @@ -15,18 +15,18 @@ def test_transform_precessing_spins(self): """ 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(rng=self.rng) + point = priors.sample(random_state=self.rng) bilby_transformed = transform_precessing_spins( point["theta_jn"], point["phi_jl"], @@ -56,8 +56,7 @@ def test_transform_precessing_spins(self): float(point["phase"]), )) np.testing.assert_allclose(bilby_transformed, lalsim_transformed, rtol=1e-10) - - + @pytest.mark.array_backend def test_transform_precessing_spins_vectorized(self): """ @@ -65,17 +64,17 @@ def test_transform_precessing_spins_vectorized(self): """ 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, rng=self.rng) + + points = priors.sample(100, random_state=self.rng) bilby_transformed = transform_precessing_spins( points["theta_jn"], points["phi_jl"], From 14354f9bc60f376503dc77a0f1486b7ea17705c9 Mon Sep 17 00:00:00 2001 From: Colm Talbot Date: Mon, 6 Jul 2026 14:25:57 -0400 Subject: [PATCH 07/11] TEST: unpack transformed values for backend test --- test/gw/geometry_test.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/gw/geometry_test.py b/test/gw/geometry_test.py index dfd5ae2cc..2ae45caf8 100644 --- a/test/gw/geometry_test.py +++ b/test/gw/geometry_test.py @@ -40,7 +40,7 @@ def test_transform_precessing_spins(self): point["reference_frequency"], point["phase"], ) - assert aac.get_namespace(bilby_transformed) == self.xp + assert aac.get_namespace(*bilby_transformed) == self.xp bilby_transformed = np.asarray(bilby_transformed) lalsim_transformed = np.asarray(SimInspiralTransformPrecessingNewInitialConditions( float(point["theta_jn"]), @@ -88,7 +88,7 @@ def test_transform_precessing_spins_vectorized(self): points["reference_frequency"], points["phase"], ) - assert aac.get_namespace(bilby_transformed) == self.xp + 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"])): From 71a6af0d6d34b884d3072afb9f38454eb4182c30 Mon Sep 17 00:00:00 2001 From: Colm Talbot Date: Mon, 6 Jul 2026 14:49:03 -0400 Subject: [PATCH 08/11] TEST: Adjust tolerance in geometry tests Updated tolerance for assert_allclose due to precision loss in different array backends. --- test/gw/geometry_test.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/test/gw/geometry_test.py b/test/gw/geometry_test.py index 2ae45caf8..4061a78db 100644 --- a/test/gw/geometry_test.py +++ b/test/gw/geometry_test.py @@ -55,7 +55,8 @@ def test_transform_precessing_spins(self): float(point["reference_frequency"]), float(point["phase"]), )) - np.testing.assert_allclose(bilby_transformed, lalsim_transformed, rtol=1e-10) + # the different array backends have some precision loss + np.testing.assert_allclose(bilby_transformed, lalsim_transformed, rtol=1e-5) @pytest.mark.array_backend def test_transform_precessing_spins_vectorized(self): @@ -107,4 +108,5 @@ def test_transform_precessing_spins_vectorized(self): float(point["phase"]), ))) lalsim_transformed = np.asarray(lalsim_transformed).T - np.testing.assert_allclose(bilby_transformed, lalsim_transformed, rtol=1e-10) + # the different array backends have some precision loss + np.testing.assert_allclose(bilby_transformed, lalsim_transformed, rtol=1e-5) From 78082525169bd6490ce03e70ba2ec8434d154377 Mon Sep 17 00:00:00 2001 From: Colm Talbot Date: Mon, 6 Jul 2026 15:07:31 -0400 Subject: [PATCH 09/11] TEST: switch test to absolute tolerance --- test/gw/geometry_test.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/gw/geometry_test.py b/test/gw/geometry_test.py index 4061a78db..505965d48 100644 --- a/test/gw/geometry_test.py +++ b/test/gw/geometry_test.py @@ -56,7 +56,7 @@ def test_transform_precessing_spins(self): float(point["phase"]), )) # the different array backends have some precision loss - np.testing.assert_allclose(bilby_transformed, lalsim_transformed, rtol=1e-5) + np.testing.assert_allclose(bilby_transformed, lalsim_transformed, atol=1e-5) @pytest.mark.array_backend def test_transform_precessing_spins_vectorized(self): @@ -109,4 +109,4 @@ def test_transform_precessing_spins_vectorized(self): ))) lalsim_transformed = np.asarray(lalsim_transformed).T # the different array backends have some precision loss - np.testing.assert_allclose(bilby_transformed, lalsim_transformed, rtol=1e-5) + np.testing.assert_allclose(bilby_transformed, lalsim_transformed, atol=1e-5) From 93a2206966447770400d14913e91ecb92c499f66 Mon Sep 17 00:00:00 2001 From: Colm Talbot Date: Mon, 6 Jul 2026 15:54:05 -0400 Subject: [PATCH 10/11] Fix mass units in conversion function --- bilby/gw/source.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/bilby/gw/source.py b/bilby/gw/source.py index 11411b468..9e86ec227 100644 --- a/bilby/gw/source.py +++ b/bilby/gw/source.py @@ -145,7 +145,7 @@ def gwsignal_binary_black_hole(frequency_array, mass_1, mass_2, luminosity_dista 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, + 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) eccentricity = 0.0 @@ -619,14 +619,15 @@ def _base_lal_cbc_fd_waveform( (frequency_array <= maximum_frequency)) luminosity_distance = luminosity_distance * 1e6 * utils.parsec - 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) + mass_1 = mass_1 * utils.solar_mass + mass_2 = mass_2 * utils.solar_mass + longitude_ascending_nodes = 0.0 mean_per_ano = 0.0 @@ -1113,14 +1114,15 @@ def _base_waveform_frequency_sequence( approximant = lalsim_GetApproximantFromString(approximant) luminosity_distance = luminosity_distance * 1e6 * utils.parsec - 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) + mass_1 = mass_1 * utils.solar_mass + mass_2 = mass_2 * utils.solar_mass + try: h_plus, h_cross = lalsim_SimInspiralChooseFDWaveformSequence( phase, mass_1, mass_2, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, From 1c0807aff2f163082acc5c230e5ac4cdfbdd6ca1 Mon Sep 17 00:00:00 2001 From: Colm Talbot Date: Fri, 10 Jul 2026 10:59:21 -0400 Subject: [PATCH 11/11] Implement error handling for spin calculations Added error handling for ZeroDivisionError when calculating spins. --- bilby/gw/source.py | 42 ++++++++++++++++++++++++++++++------------ 1 file changed, 30 insertions(+), 12 deletions(-) diff --git a/bilby/gw/source.py b/bilby/gw/source.py index 9e86ec227..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, mass_2=mass_2, - 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 @@ -620,10 +626,16 @@ def _base_lal_cbc_fd_waveform( luminosity_distance = luminosity_distance * 1e6 * utils.parsec - 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: + 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 @@ -1115,10 +1127,16 @@ def _base_waveform_frequency_sequence( luminosity_distance = luminosity_distance * 1e6 * utils.parsec - 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: + 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