Adds TimeDomainCBCWaveformGenerator (#188)

* improve generator to use better waveform conditioning

* finalize waveform generator

* add conversion function to parameter sampler

* proper indexing when slicing

* add proper device handling

* add kwargs to phenomD and taylorf2 to soak up unused parameters

* update testing

* match conditioning with gwsignal

* add fft into generator

* add tests for waveform generator

* fix pre-commit

* remove coefficients tests since there is not lalsuite equivalent

* update number of samples in conftest

* adjust tests for off by one error

* add iirfilter for highpassing

* remove notebook

* use relative imoprts

* fix type hint

* bump test tolerances

* clean up tests and account for discrepancy in argma

* reduce number of samples to 1000

* increas tolerance

* remove unused high_pass_time_series

ml4gw/transforms/__init__.py

@@ -1,3 +1,4 @@
+from .iirfilter import IIRFilter
 from .pearson import ShiftedPearsonCorrelation
 from .qtransform import QScan, SingleQTransform
 from .scaler import ChannelWiseScaler

ml4gw/waveforms/cbc/coefficients.py

@@ -0,0 +1,35 @@
+import torch
+
+from ml4gw.constants import C, G
+from ml4gw.types import BatchTensor
+
+
+def taylor_t2_timing_0pn_coeff(total_mass: BatchTensor, eta: BatchTensor):
+    """
+    https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/lib/LALSimInspiralPNCoefficients.c#L1528
+    """
+
+    output = total_mass * G / C**3
+    return -5.0 * output / (256.0 * eta)
+
+
+def taylor_t2_timing_2pn_coeff(eta: BatchTensor):
+    """
+    https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/lib/LALSimInspiralPNCoefficients.c#L1545
+    """
+    return 7.43 / 2.52 + 11.0 / 3.0 * eta
+
+
+def taylor_t2_timing_4pn_coeff(eta: BatchTensor):
+    """
+    https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/lib/LALSimInspiralPNCoefficients.c#L1560
+    """
+    return 30.58673 / 5.08032 + 54.29 / 5.04 * eta + 61.7 / 7.2 * eta**2
+
+
+def taylor_t3_frequency_0pn_coeff(total_mass: BatchTensor):
+    """
+    https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/lib/LALSimInspiralPNCoefficients.c#L1723
+    """
+    output = total_mass * G / C**3.0
+    return 1.0 / (8.0 * torch.pi * output)

ml4gw/waveforms/cbc/phenom_d.py

@@ -25,6 +25,7 @@ def forward(
         phic: BatchTensor,
         inclination: BatchTensor,
         f_ref: float,
+        **kwargs
     ):
         """
         IMRPhenomD waveform

ml4gw/waveforms/cbc/phenom_p.py

@@ -39,6 +39,7 @@ def forward(
         inclination: BatchTensor,
         f_ref: float,
         tc: Optional[BatchTensor] = None,
+        **kwargs,
     ):
         """
         IMRPhenomPv2 waveform

ml4gw/waveforms/cbc/taylorf2.py

@@ -22,6 +22,7 @@ def forward(
         phic: BatchTensor,
         inclination: BatchTensor,
         f_ref: float,
+        **kwargs
     ):
         """
         TaylorF2 up to 3.5 PN in phase. Newtonian SPA amplitude.

ml4gw/waveforms/cbc/utils.py

@@ -0,0 +1,111 @@
+"""
+Utilities for conditioning waveforms
+See https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/lib/LALSimInspiral.c # noqa
+"""
+import torch
+
+from ml4gw.constants import MRSUN, MSUN, MTSUN_SI, C, G
+from ml4gw.types import BatchTensor
+from ml4gw.waveforms.cbc import coefficients
+
+
+def chirp_time_bound(
+    fstart: BatchTensor,
+    mass_1: BatchTensor,
+    mass_2: BatchTensor,
+    s1: BatchTensor,
+    s2: BatchTensor,
+) -> BatchTensor:
+    """
+    https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/lib/LALSimInspiral.c#L4969
+    """
+
+    total_mass = mass_1 + mass_2
+    reduced_mass = mass_1 * mass_2 / total_mass
+    eta = reduced_mass / total_mass
+    chi = torch.max(s1.abs(), s2.abs()).abs()
+
+    c0 = torch.abs(coefficients.taylor_t2_timing_0pn_coeff(total_mass, eta))
+
+    c2 = coefficients.taylor_t2_timing_2pn_coeff(eta)
+    c3 = (226.0 / 15.0) * chi
+    c4 = coefficients.taylor_t2_timing_4pn_coeff(eta)
+
+    v = (torch.pi * total_mass * fstart * G) ** (1.0 / 3.0)
+    v /= C
+
+    return (c0 * (v**-8) * (1 + (c2 + (c3 + c4 * v) * v) * v * v)).float()
+
+
+def chirp_start_frequency_bound(
+    tchirp: BatchTensor,
+    mass_1: BatchTensor,
+    mass_2: BatchTensor,
+):
+    """
+    https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/lib/LALSimInspiral.c#L5104
+    """
+    total_mass = mass_1 + mass_2
+    mu = mass_1 * mass_2 / total_mass
+
+    eta = mu / total_mass
+    c0 = coefficients.taylor_t3_frequency_0pn_coeff(total_mass)
+    return (
+        c0
+        * pow(5.0 * total_mass * (MTSUN_SI / MSUN) / (eta * tchirp), 3.0 / 8.0)
+    ).float()
+
+
+def final_black_hole_spin_bound(
+    s1z: BatchTensor, s2z: BatchTensor
+) -> BatchTensor:
+    """
+    https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/lib/LALSimInspiral.c#L5081
+    """
+    maximum_black_hole_spin = 0.998
+    s = 0.686 + 0.15 * (s1z + s2z)
+    s = torch.maximum(s, torch.abs(s1z)).maximum(torch.abs(s2z))
+    s = torch.clamp(s, max=maximum_black_hole_spin)
+    return s
+
+
+def merge_time_bound(mass_1: BatchTensor, mass_2: BatchTensor) -> BatchTensor:
+    """
+    https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/lib/LALSimInspiral.c#L5007
+    """
+
+    n_orbits = 1
+    total_mass = mass_1 + mass_2
+    r = 9.0 * total_mass * MRSUN / MSUN
+    v = C / 3.0
+    return (n_orbits * (2.0 * torch.pi * r / v)).float()
+
+
+def ringdown_time_bound(
+    total_mass: BatchTensor, s: BatchTensor
+) -> BatchTensor:
+    """
+    https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/lib/LALSimInspiral.c#L5032
+    """
+    n_efolds = 11
+
+    f1 = 1.5251
+    f2 = -1.1568
+    f3 = 0.1292
+    q1 = 0.7000
+    q2 = 1.4187
+    q3 = -0.4990
+
+    omega = (f1 + f2 * (1.0 - s) ** f3) / (total_mass * MTSUN_SI / MSUN)
+    Q = q1 + q2 * (1.0 - s) ** q3
+    tau = 2.0 * Q / omega
+    return (n_efolds * tau).float()
+
+
+def frequency_isco(mass_1: BatchTensor, mass_2: BatchTensor):
+    return (
+        1.0
+        / (
+            (9.0**1.5) * torch.pi * (mass_1 + mass_2) * MTSUN_SI / MSUN
+        ).float()
+    )