Merge pull request #19 from ThibeauWouters/clean

TaylorF2 and IMRPhenomD_NRTidalv2

README.md

@@ -6,15 +6,15 @@ A small `jax` package for differentiable and fast gravitational wave data analys
 
 ### Installation
 
-Both waveforms have been tested extensively and match lalsuite implementations to machine precision across all the parameter space. 
+Both waveforms have been tested extensively and match `lalsuite` implementations to machine precision across all the parameter space. 
 
 Ripple can be installed using 
 
 ```
 pip3 install ripplegw
 ```
 
-Note that by default we do not include enable float64 in jax since we want allow users to use float32 to improve performance.
+Note that by default we do not include enable float64 in `jax`` since we want allow users to use float32 to improve performance.
 If you require float64, please include the following code at the start of the script:
 
 ```
@@ -29,10 +29,12 @@ See https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html fo
 - IMRPhenomXAS (aligned spin)
 - IMRPhenomD (aligned spin)
 - IMRPhenomPv2 (Still finalizing sampling checks)
+- TaylorF2 with tidal effects
+- IMRPhenomD_NRTidalv2, verified for the low spin regime (chi1, chi2 < 0.05), further testing is required for higher spins
 
 ### Generating a waveform and its derivative
 
-Generating a waveform is increadibly easy. Below is an example of calling the PhenomXAS waveform model
+Generating a waveform is incredibly easy. Below is an example of calling the PhenomXAS waveform model
 to get the h_+ and h_x polarizations of the waveform model
 
 We start with some basic imports:

src/ripple/__init__.py

@@ -51,6 +51,233 @@ def ms_to_Mc_eta(m):
     return (m1 * m2) ** (3 / 5) / (m1 + m2) ** (1 / 5), m1 * m2 / (m1 + m2) ** 2
 
 
+# TODO in code below, reduce copy-paste
+def lambdas_to_lambda_tildes_from_q(params: Array):
+    """
+    Convert from individual tidal parameters to domainant tidal term. (Code taken from Bilby)
+
+    See, e.g., Wade et al., https://arxiv.org/pdf/1402.5156.pdf.
+
+    Parameters
+    ==========
+    lambda_1: float
+        Tidal parameter of more massive neutron star.
+    lambda_2: float
+        Tidal parameter of less massive neutron star.
+    mass_1: float
+        Mass of more massive neutron star.
+    mass_2: float
+        Mass of less massive neutron star.
+
+    Returns
+    =======
+    lambda_tilde: float
+        Dominant tidal term.
+    """
+    lambda_1, lambda_2, q = params
+    eta = q / (1 + q) ** 2
+    lambda_plus = lambda_1 + lambda_2
+    lambda_minus = lambda_1 - lambda_2
+    lambda_tilde = (
+        8
+        / 13
+        * (
+            (1 + 7 * eta - 31 * eta**2) * lambda_plus
+            + (1 - 4 * eta) ** 0.5 * (1 + 9 * eta - 11 * eta**2) * lambda_minus
+        )
+    )
+
+    delta_lambda_tilde = (
+        1
+        / 2
+        * (
+            (1 - 4 * eta) ** 0.5
+            * (1 - 13272 / 1319 * eta + 8944 / 1319 * eta**2)
+            * lambda_plus
+            + (1 - 15910 / 1319 * eta + 32850 / 1319 * eta**2 + 3380 / 1319 * eta**3)
+            * lambda_minus
+        )
+    )
+
+    return lambda_tilde, delta_lambda_tilde
+
+
+def lambdas_to_lambda_tildes(params: Array):
+    """
+    Convert from individual tidal parameters to domainant tidal term. (Code taken from Bilby)
+
+    See, e.g., Wade et al., https://arxiv.org/pdf/1402.5156.pdf.
+
+    Parameters
+    ==========
+    lambda_1: float
+        Tidal parameter of more massive neutron star.
+    lambda_2: float
+        Tidal parameter of less massive neutron star.
+    mass_1: float
+        Mass of more massive neutron star.
+    mass_2: float
+        Mass of less massive neutron star.
+
+    Returns
+    =======
+    lambda_tilde: float
+        Dominant tidal term.
+    """
+    lambda_1, lambda_2, mass_1, mass_2 = params
+    _, eta = ms_to_Mc_eta(jnp.array([mass_1, mass_2]))
+    lambda_plus = lambda_1 + lambda_2
+    lambda_minus = lambda_1 - lambda_2
+    lambda_tilde = (
+        8
+        / 13
+        * (
+            (1 + 7 * eta - 31 * eta**2) * lambda_plus
+            + (1 - 4 * eta) ** 0.5 * (1 + 9 * eta - 11 * eta**2) * lambda_minus
+        )
+    )
+
+    delta_lambda_tilde = (
+        1
+        / 2
+        * (
+            (1 - 4 * eta) ** 0.5
+            * (1 - 13272 / 1319 * eta + 8944 / 1319 * eta**2)
+            * lambda_plus
+            + (1 - 15910 / 1319 * eta + 32850 / 1319 * eta**2 + 3380 / 1319 * eta**3)
+            * lambda_minus
+        )
+    )
+
+    return lambda_tilde, delta_lambda_tilde
+
+
+def lambda_tildes_to_lambdas(params: Array):
+    """
+    Convert from dominant tidal terms to individual tidal parameters. Code taken from bilby.
+
+    See, e.g., Wade et al., https://arxiv.org/pdf/1402.5156.pdf.
+
+    Parameters
+    ==========
+    lambda_tilde: float
+        Dominant tidal term.
+    delta_lambda_tilde: float
+        Secondary tidal term.
+    mass_1: float
+        Mass of more massive neutron star.
+    mass_2: float
+        Mass of less massive neutron star.
+
+    Returns
+    =======
+    lambda_1: float
+        Tidal parameter of more massive neutron star.
+    lambda_2: float
+        Tidal parameter of less massive neutron star.
+
+    """
+
+    lambda_tilde, delta_lambda_tilde, mass_1, mass_2 = params
+
+    _, eta = ms_to_Mc_eta(jnp.array([mass_1, mass_2]))
+    coefficient_1 = 1 + 7 * eta - 31 * eta**2
+    coefficient_2 = (1 - 4 * eta) ** 0.5 * (1 + 9 * eta - 11 * eta**2)
+    coefficient_3 = (1 - 4 * eta) ** 0.5 * (
+        1 - 13272 / 1319 * eta + 8944 / 1319 * eta**2
+    )
+    coefficient_4 = (
+        1 - 15910 / 1319 * eta + 32850 / 1319 * eta**2 + 3380 / 1319 * eta**3
+    )
+    lambda_1 = (
+        13 * lambda_tilde / 8 * (coefficient_3 - coefficient_4)
+        - 2 * delta_lambda_tilde * (coefficient_1 - coefficient_2)
+    ) / (
+        (coefficient_1 + coefficient_2) * (coefficient_3 - coefficient_4)
+        - (coefficient_1 - coefficient_2) * (coefficient_3 + coefficient_4)
+    )
+    lambda_2 = (
+        13 * lambda_tilde / 8 * (coefficient_3 + coefficient_4)
+        - 2 * delta_lambda_tilde * (coefficient_1 + coefficient_2)
+    ) / (
+        (coefficient_1 - coefficient_2) * (coefficient_3 + coefficient_4)
+        - (coefficient_1 + coefficient_2) * (coefficient_3 - coefficient_4)
+    )
+
+    return lambda_1, lambda_2
+
+
+def lambda_tildes_to_lambdas_from_q(params: Array):
+    """
+    Convert from dominant tidal terms to individual tidal parameters. Code taken from bilby.
+
+    See, e.g., Wade et al., https://arxiv.org/pdf/1402.5156.pdf.
+
+    Parameters
+    ==========
+    lambda_tilde: float
+        Dominant tidal term.
+    delta_lambda_tilde: float
+        Secondary tidal term.
+    mass_1: float
+        Mass of more massive neutron star.
+    mass_2: float
+        Mass of less massive neutron star.
+
+    Returns
+    =======
+    lambda_1: float
+        Tidal parameter of more massive neutron star.
+    lambda_2: float
+        Tidal parameter of less massive neutron star.
+
+    """
+
+    lambda_tilde, delta_lambda_tilde, q = params
+
+    eta = q / (1 + q) ** 2
+
+    coefficient_1 = 1 + 7 * eta - 31 * eta**2
+    coefficient_2 = (1 - 4 * eta) ** 0.5 * (1 + 9 * eta - 11 * eta**2)
+    coefficient_3 = (1 - 4 * eta) ** 0.5 * (
+        1 - 13272 / 1319 * eta + 8944 / 1319 * eta**2
+    )
+    coefficient_4 = (
+        1 - 15910 / 1319 * eta + 32850 / 1319 * eta**2 + 3380 / 1319 * eta**3
+    )
+    lambda_1 = (
+        13 * lambda_tilde / 8 * (coefficient_3 - coefficient_4)
+        - 2 * delta_lambda_tilde * (coefficient_1 - coefficient_2)
+    ) / (
+        (coefficient_1 + coefficient_2) * (coefficient_3 - coefficient_4)
+        - (coefficient_1 - coefficient_2) * (coefficient_3 + coefficient_4)
+    )
+    lambda_2 = (
+        13 * lambda_tilde / 8 * (coefficient_3 + coefficient_4)
+        - 2 * delta_lambda_tilde * (coefficient_1 + coefficient_2)
+    ) / (
+        (coefficient_1 - coefficient_2) * (coefficient_3 + coefficient_4)
+        - (coefficient_1 + coefficient_2) * (coefficient_3 - coefficient_4)
+    )
+
+    return lambda_1, lambda_2
+
+
+def get_chi_eff(params: Array) -> float:
+    """Compute effective spin.
+
+    Args:
+        params (Array): Parameters: mass1, mass2, spin1, spin2
+
+    Returns:
+        float: Effective spin.
+    """
+    m1, m2, chi1, chi2 = params
+
+    chi_eff = (m1 * chi1 + m2 * chi2) / (m1 + m2)
+    return chi_eff
+
+
 def get_f_isco(m):
     r"""
     Computes the ISCO frequency for a black hole.

src/ripple/constants.py

@@ -7,6 +7,9 @@
 MSUN = 1.988409902147041637325262574352366540e30  # kg
 """Solar mass"""
 
+MRSUN = 1.476625038050124729627979840144936351e3
+"""Geometrized nominal solar mass, m"""
+
 G = 6.67430e-11  # m^3 / kg / s^2
 """Newton's gravitational constant"""
 
@@ -16,7 +19,9 @@
 """Pi"""
 PI = 3.141592653589793238462643383279502884
 
-gt = G * MSUN / (C ** 3.0)
+TWO_PI = 6.283185307179586476925286766559005768
+
+gt = G * MSUN / (C**3.0)
 """
 G MSUN / C^3 in seconds
 """
@@ -25,3 +30,8 @@
 """
 Meters per Mpc.
 """
+
+clightGpc = C / 3.0856778570831e22
+"""
+Speed of light in vacuum (:math:`c`), in gigaparsecs per second
+"""

src/ripple/noise.py

@@ -2,7 +2,6 @@
 Detector noise power spectral densities.
 """
 
-
 import importlib.resources as pkg_resources
 from typing import Callable, Tuple
 

src/ripple/waveforms/IMRPhenomD.py

@@ -143,11 +143,12 @@ def get_inspiral_phase(fM_s: Array, theta: Array, coeffs: Array) -> Array:
         + phi2 * ((PI * fM_s) ** -1.0)
         + phi3 * ((PI * fM_s) ** -(2.0 / 3.0))
         + phi4 * ((PI * fM_s) ** -(1.0 / 3.0))
-        + phi5_log * jnp.log(v) + phi5
+        + phi5_log * jnp.log(v)
+        + phi5
         + phi6_log * jnp.log(v) * ((PI * fM_s) ** (1.0 / 3.0))
         + phi6 * ((PI * fM_s) ** (1.0 / 3.0))
         + phi7 * ((PI * fM_s) ** (2.0 / 3.0))
-    ) * (3.0 / (128.0 * eta)) - PI/4.0
+    ) * (3.0 / (128.0 * eta)) - PI / 4.0
     phi_Ins = (
         phi_TF2
         + (
@@ -169,9 +170,7 @@ def get_IIa_raw_phase(fM_s: Array, theta: Array, coeffs: Array) -> Array:
     eta = m1_s * m2_s / (M_s**2.0)
 
     phi_IIa_raw = (
-        coeffs[11] * fM_s
-        + coeffs[12] * jnp.log(fM_s)
-        - coeffs[13] * (fM_s**-3.0) / 3.0
+        coeffs[11] * fM_s + coeffs[12] * jnp.log(fM_s) - coeffs[13] * (fM_s**-3.0) / 3.0
     ) / eta
 
     return phi_IIa_raw