Allow time domain approximants

Closes #308 (closed)

This changes the lalsimulation function being called and enforces frequency bounds on the waveforms returned.

I noticed there's a lot of duplicate code in the source functions. I'll make a separate issue to rationalise this.

@michael.puerrer can you check this?

UPDATE: I got carried away and made all the lal source functions call a single base function. This is likely a more maintainable solution.

bilby/gw/source.py

@@ -4,8 +4,9 @@ import numpy as np
 
 from ..core import utils
 from ..core.utils import logger
-from .utils import (lalsim_SimInspiralTransformPrecessingNewInitialConditions,
-                    lalsim_GetApproximantFromString,
+from .conversion import transform_precessing_spins
+from .utils import (lalsim_GetApproximantFromString,
+                    lalsim_SimInspiralFD,
                     lalsim_SimInspiralChooseFDWaveform,
                     lalsim_SimInspiralWaveformParamsInsertTidalLambda1,
                     lalsim_SimInspiralWaveformParamsInsertTidalLambda2,
@@ -59,61 +60,73 @@ def lal_binary_black_hole(
     -------
     dict: A dictionary with the plus and cross polarisation strain modes
     """
-
-    waveform_kwargs = dict(waveform_approximant='IMRPhenomPv2',
-                           reference_frequency=50.0,
-                           minimum_frequency=20.0)
+    waveform_kwargs = dict(
+        waveform_approximant='IMRPhenomPv2', reference_frequency=50.0,
+        minimum_frequency=20.0, maximum_frequency=frequency_array[-1])
     waveform_kwargs.update(kwargs)
-    waveform_approximant = waveform_kwargs['waveform_approximant']
-    reference_frequency = waveform_kwargs['reference_frequency']
-    minimum_frequency = waveform_kwargs['minimum_frequency']
-
-    if mass_2 > mass_1:
-        return None
-
-    luminosity_distance = luminosity_distance * 1e6 * utils.parsec
-    mass_1 = mass_1 * utils.solar_mass
-    mass_2 = mass_2 * utils.solar_mass
-
-    if tilt_1 == 0 and tilt_2 == 0:
-        spin_1x = 0
-        spin_1y = 0
-        spin_1z = a_1
-        spin_2x = 0
-        spin_2y = 0
-        spin_2z = a_2
-        iota = theta_jn
-    else:
-        iota, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, spin_2z = (
-            lalsim_SimInspiralTransformPrecessingNewInitialConditions(
-                theta_jn, phi_jl, tilt_1, tilt_2, phi_12, a_1, a_2, mass_1,
-                mass_2, reference_frequency, phase))
-
-    longitude_ascending_nodes = 0.0
-    eccentricity = 0.0
-    mean_per_ano = 0.0
-
-    waveform_dictionary = None
+    return _base_lal_cbc_fd_waveform(
+        frequency_array=frequency_array, mass_1=mass_1, mass_2=mass_2,
+        luminosity_distance=luminosity_distance, theta_jn=theta_jn, phase=phase,
+        a_1=a_1, a_2=a_2, tilt_1=tilt_1, tilt_2=tilt_2, phi_12=phi_12,
+        phi_jl=phi_jl, **waveform_kwargs)
 
-    approximant = lalsim_GetApproximantFromString(waveform_approximant)
 
-    maximum_frequency = frequency_array[-1]
-    delta_frequency = frequency_array[1] - frequency_array[0]
-
-    hplus, hcross = lalsim_SimInspiralChooseFDWaveform(
-        mass_1, mass_2, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y,
-        spin_2z, luminosity_distance, iota, phase,
-        longitude_ascending_nodes, eccentricity, mean_per_ano, delta_frequency,
-        minimum_frequency, maximum_frequency, reference_frequency,
-        waveform_dictionary, approximant)
+def lal_binary_neutron_star(
+        frequency_array, mass_1, mass_2, luminosity_distance, chi_1, chi_2,
+        theta_jn, phase, lambda_1, lambda_2, **kwargs):
+    """ A Binary Neutron Star waveform model using lalsimulation
 
-    h_plus = hplus.data.data
-    h_cross = hcross.data.data
+    Parameters
+    ----------
+    frequency_array: array_like
+        The frequencies at which we want to calculate the strain
+    mass_1: float
+        The mass of the heavier object in solar masses
+    mass_2: float
+        The mass of the lighter object in solar masses
+    luminosity_distance: float
+        The luminosity distance in megaparsec
+    chi_1: float
+        Dimensionless aligned spin
+    chi_2: float
+        Dimensionless aligned spin
+    theta_jn: float
+        Orbital inclination
+    phase: float
+        The phase at coalescence
+    ra: float
+        The right ascension of the binary
+    dec: float
+        The declination of the object
+    geocent_time: float
+        The time at coalescence
+    psi: float
+        Orbital polarisation
+    lambda_1: float
+        Dimensionless tidal deformability of mass_1
+    lambda_2: float
+        Dimensionless tidal deformability of mass_2
 
-    h_plus = h_plus[:len(frequency_array)]
-    h_cross = h_cross[:len(frequency_array)]
+    kwargs: dict
+        Optional keyword arguments
 
-    return {'plus': h_plus, 'cross': h_cross}
+    Returns
+    -------
+    dict: A dictionary with the plus and cross polarisation strain modes
+    """
+    waveform_kwargs = dict(
+        waveform_approximant='TaylorF2', reference_frequency=50.0,
+        minimum_frequency=20.0, maximum_frequency=frequency_array[-1])
+    a_1 = abs(chi_1)
+    a_2 = abs(chi_2)
+    tilt_1 = np.arccos(chi_1)
+    tilt_2 = np.arccos(chi_2)
+    waveform_kwargs.update(kwargs)
+    return _base_lal_cbc_fd_waveform(
+        frequency_array=frequency_array, mass_1=mass_1, mass_2=mass_2,
+        luminosity_distance=luminosity_distance, theta_jn=theta_jn, phase=phase,
+        a_1=a_1, a_2=a_2, tilt_1=tilt_1, tilt_2=tilt_2, lambda_1=lambda_1,
+        lambda_2=lambda_2, **waveform_kwargs)
 
 
 def lal_eccentric_binary_black_hole_no_spins(
@@ -144,13 +157,69 @@ def lal_eccentric_binary_black_hole_no_spins(
     -------
     dict: A dictionary with the plus and cross polarisation strain modes
     """
-
-    waveform_kwargs = dict(waveform_approximant='EccentricFD', reference_frequency=10.0,
-                           minimum_frequency=10.0)
+    waveform_kwargs = dict(
+        waveform_approximant='EccentricFD', reference_frequency=10.0,
+        minimum_frequency=10.0, maximum_frequency=frequency_array[-1])
     waveform_kwargs.update(kwargs)
+    return _base_lal_cbc_fd_waveform(
+        frequency_array=frequency_array, mass_1=mass_1, mass_2=mass_2,
+        luminosity_distance=luminosity_distance, theta_jn=theta_jn, phase=phase,
+        eccentricity=eccentricity, **waveform_kwargs)
+
+
+def _base_lal_cbc_fd_waveform(
+        frequency_array, mass_1, mass_2, luminosity_distance, theta_jn, phase,
+        a_1=0.0, a_2=0.0, tilt_1=0.0, tilt_2=0.0, phi_12=0.0, phi_jl=0.0,
+        lambda_1=0.0, lambda_2=0.0, eccentricity=0.0, **waveform_kwargs):
+    """ Generate a cbc waveform model using lalsimulation
+
+    Parameters
+    ----------
+    frequency_array: array_like
+        The frequencies at which we want to calculate the strain
+    mass_1: float
+        The mass of the heavier object in solar masses
+    mass_2: float
+        The mass of the lighter object in solar masses
+    luminosity_distance: float
+        The luminosity distance in megaparsec
+    a_1: float
+        Dimensionless primary spin magnitude
+    tilt_1: float
+        Primary tilt angle
+    phi_12: float
+        Azimuthal angle between the component spins
+    a_2: float
+        Dimensionless secondary spin magnitude
+    tilt_2: float
+        Secondary tilt angle
+    phi_jl: float
+        Azimuthal angle between the total and orbital angular momenta
+    theta_jn: float
+        Orbital inclination
+    phase: float
+        The phase at coalescence
+    eccentricity: float
+        Binary eccentricity
+    lambda_1: float
+        Tidal deformability of the more massive object
+    lambda_2: float
+        Tidal deformability of the less massive object
+    kwargs: dict
+        Optional keyword arguments
+
+    Returns
+    -------
+    dict: A dictionary with the plus and cross polarisation strain modes
+    """
     waveform_approximant = waveform_kwargs['waveform_approximant']
     reference_frequency = waveform_kwargs['reference_frequency']
     minimum_frequency = waveform_kwargs['minimum_frequency']
+    maximum_frequency = waveform_kwargs['maximum_frequency']
+    delta_frequency = frequency_array[1] - frequency_array[0]
+
+    frequency_bounds = ((frequency_array >= minimum_frequency) *
+                        (frequency_array <= maximum_frequency))
 
     if mass_2 > mass_1:
         return None
@@ -159,33 +228,42 @@ def lal_eccentric_binary_black_hole_no_spins(
     mass_1 = mass_1 * utils.solar_mass
     mass_2 = mass_2 * utils.solar_mass
 
-    spin_1x = 0.0
-    spin_1y = 0.0
-    spin_1z = 0.0
-    spin_2x = 0.0
-    spin_2y = 0.0
-    spin_2z = 0.0
-    iota = theta_jn
+    if tilt_1 in [0.0, np.pi] and tilt_2 in [0, np.pi]:
+        spin_1x = 0.0
+        spin_1y = 0.0
+        spin_1z = a_1 * np.cos(tilt_1)
+        spin_2x = 0.0
+        spin_2y = 0.0
+        spin_2z = a_2 * np.cos(tilt_2)
+        iota = theta_jn
+    else:
+        iota, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, spin_2z = (
+            transform_precessing_spins(
+                theta_jn, phi_jl, tilt_1, tilt_2, phi_12, a_1, a_2, mass_1,
+                mass_2, reference_frequency, phase))
 
     longitude_ascending_nodes = 0.0
     mean_per_ano = 0.0
 
-    waveform_dictionary = None
+    waveform_dictionary = lal.CreateDict()
+    lalsim_SimInspiralWaveformParamsInsertTidalLambda1(waveform_dictionary, lambda_1)
+    lalsim_SimInspiralWaveformParamsInsertTidalLambda2(waveform_dictionary, lambda_2)
 
     approximant = lalsim_GetApproximantFromString(waveform_approximant)
 
-    maximum_frequency = frequency_array[-1]
-    delta_frequency = frequency_array[1] - frequency_array[0]
-
-    hplus, hcross = lalsim_SimInspiralChooseFDWaveform(
+    if lalsim.SimInspiralImplementedFDApproximants(approximant):
+        wf_func = lalsim_SimInspiralChooseFDWaveform
+    else:
+        wf_func = lalsim_SimInspiralFD
+    hplus, hcross = wf_func(
         mass_1, mass_2, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y,
         spin_2z, luminosity_distance, iota, phase,
         longitude_ascending_nodes, eccentricity, mean_per_ano, delta_frequency,
         minimum_frequency, maximum_frequency, reference_frequency,
         waveform_dictionary, approximant)
 
-    h_plus = hplus.data.data
-    h_cross = hcross.data.data
+    h_plus = hplus.data.data[:len(frequency_array)] * frequency_bounds
+    h_cross = hcross.data.data[:len(frequency_array)] * frequency_bounds
 
     return {'plus': h_plus, 'cross': h_cross}
 
@@ -249,101 +327,6 @@ def supernova_pca_model(
     return {'plus': h_plus, 'cross': h_cross}
 
 
-def lal_binary_neutron_star(
-        frequency_array, mass_1, mass_2, luminosity_distance, chi_1, chi_2,
-        theta_jn, phase, lambda_1, lambda_2, **kwargs):
-    """ A Binary Neutron Star waveform model using lalsimulation
-
-    Parameters
-    ----------
-    frequency_array: array_like
-        The frequencies at which we want to calculate the strain
-    mass_1: float
-        The mass of the heavier object in solar masses
-    mass_2: float
-        The mass of the lighter object in solar masses
-    luminosity_distance: float
-        The luminosity distance in megaparsec
-    chi_1: float
-        Dimensionless aligned spin
-    chi_2: float
-        Dimensionless aligned spin
-    theta_jn: float
-        Orbital inclination
-    phase: float
-        The phase at coalescence
-    ra: float
-        The right ascension of the binary
-    dec: float
-        The declination of the object
-    geocent_time: float
-        The time at coalescence
-    psi: float
-        Orbital polarisation
-    lambda_1: float
-        Dimensionless tidal deformability of mass_1
-    lambda_2: float
-        Dimensionless tidal deformability of mass_2
-
-    kwargs: dict
-        Optional keyword arguments
-
-    Returns
-    -------
-    dict: A dictionary with the plus and cross polarisation strain modes
-    """
-
-    waveform_kwargs = dict(waveform_approximant='TaylorF2', reference_frequency=50.0,
-                           minimum_frequency=20.0)
-    waveform_kwargs.update(kwargs)
-    waveform_approximant = waveform_kwargs['waveform_approximant']
-    reference_frequency = waveform_kwargs['reference_frequency']
-    minimum_frequency = waveform_kwargs['minimum_frequency']
-
-    if mass_2 > mass_1:
-        return None
-
-    luminosity_distance = luminosity_distance * 1e6 * utils.parsec
-    mass_1 = mass_1 * utils.solar_mass
-    mass_2 = mass_2 * utils.solar_mass
-
-    spin_1x = 0
-    spin_1y = 0
-    spin_1z = chi_1
-    spin_2x = 0
-    spin_2y = 0
-    spin_2z = chi_2
-    iota = theta_jn
-
-    longitude_ascending_nodes = 0.0
-    eccentricity = 0.0
-    mean_per_ano = 0.0
-
-    waveform_dictionary = lal.CreateDict()
-    lalsim_SimInspiralWaveformParamsInsertTidalLambda1(waveform_dictionary, lambda_1)
-    lalsim_SimInspiralWaveformParamsInsertTidalLambda2(waveform_dictionary, lambda_2)
-
-    approximant = lalsim_GetApproximantFromString(waveform_approximant)
-
-    maximum_frequency = frequency_array[-1]
-    delta_frequency = frequency_array[1] - frequency_array[0]
-
-    hplus, hcross = lalsim_SimInspiralChooseFDWaveform(
-        mass_1, mass_2, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y,
-        spin_2z, luminosity_distance, iota, phase,
-        longitude_ascending_nodes, eccentricity, mean_per_ano, delta_frequency,
-        minimum_frequency, maximum_frequency, reference_frequency,
-        waveform_dictionary, approximant)
-
-    h_plus = hplus.data.data
-    h_cross = hcross.data.data
-
-    h_plus = h_plus[:len(frequency_array)]
-    h_cross = h_cross[:len(frequency_array)]
-
-    return {'plus': h_plus, 'cross': h_cross}
-
-
 def roq(frequency_array, mass_1, mass_2, luminosity_distance, a_1, tilt_1,
         phi_12, a_2, tilt_2, phi_jl, theta_jn, phase, **waveform_arguments):
     """
@@ -419,7 +402,7 @@ def roq(frequency_array, mass_1, mass_2, luminosity_distance, a_1, tilt_1,
         iota = theta_jn
     else:
         iota, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, spin_2z = \
-            lalsim_SimInspiralTransformPrecessingNewInitialConditions(
+            transform_precessing_spins(
                 theta_jn, phi_jl, tilt_1, tilt_2, phi_12, a_1, a_2, mass_1,
                 mass_2, reference_frequency, phase)