ml4gw 0.4.2__py3-none-any.whl → 0.5.1__py3-none-any.whl

ml4gw/waveforms/taylorf2.py

@@ -1,306 +1,331 @@
 import torch
-from torchtyping import TensorType
+from jaxtyping import Float
 
-GAMMA = 0.577215664901532860606512090082402431
-"""Euler-Mascheroni constant. Same as lal.GAMMA"""
+from ml4gw.constants import MPC_SEC, MTSUN_SI, PI
+from ml4gw.constants import EulerGamma as GAMMA
+from ml4gw.types import BatchTensor, FrequencySeries1d
 
-MSUN_SI = 1.988409870698050731911960804878414216e30
-"""Solar mass in kg. Same as lal.MSUN_SI"""
 
-MTSUN_SI = 4.925490947641266978197229498498379006e-6
-"""1 solar mass in seconds. Same value as lal.MTSUN_SI"""
+class TaylorF2(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
 
-PI = 3.141592653589793238462643383279502884
-"""Archimedes constant. Same as lal.PI"""
-
-MPC_SEC = 1.02927125e14
-"""
-1 Mpc in seconds.
-"""
+    def forward(
+        self,
+        f: FrequencySeries1d,
+        chirp_mass: BatchTensor,
+        mass_ratio: BatchTensor,
+        chi1: BatchTensor,
+        chi2: BatchTensor,
+        distance: BatchTensor,
+        phic: BatchTensor,
+        inclination: BatchTensor,
+        f_ref: float,
+    ):
+        """
+        TaylorF2 up to 3.5 PN in phase. Newtonian SPA amplitude.
 
+        Args:
+            f:
+                Frequency series in Hz.
+            chirp_mass:
+                Chirp mass in solar masses
+            mass_ratio:
+                Mass ratio m1/m2
+            chi1:
+                Spin of m1
+            chi2:
+                Spin of m2
+            distance:
+                Luminosity distance
+            phic:
+                Phase at coalescence
+            inclination:
+                Inclination angle
+            f_ref:
+                Reference frequency
 
-def taylorf2_phase(
-    Mf: TensorType,
-    mass1: TensorType,
-    mass2: TensorType,
-    chi1: TensorType,
-    chi2: TensorType,
-) -> TensorType:
-    """
-    Calculate the inspiral phase for the TaylorF2.
-    """
-    M = mass1 + mass2
-    eta = mass1 * mass2 / M / M
-    m1byM = mass1 / M
-    m2byM = mass2 / M
-    chi1sq = chi1 * chi1
-    chi2sq = chi2 * chi2
+        Returns:
+            hc, hp: Tuple[torch.Tensor, torch.Tensor]
+                Cross and plus polarizations
+        """
 
-    v0 = torch.ones_like(Mf)
-    v1 = (PI * Mf) ** (1.0 / 3.0)
-    v2 = v1 * v1
-    v3 = v2 * v1
-    v4 = v3 * v1
-    v5 = v4 * v1
-    v6 = v5 * v1
-    v7 = v6 * v1
-    logv = torch.log(v1)
-    v5_logv = v5 * logv
-    v6_logv = v6 * logv
+        # shape assumed (n_batch, params)
+        if (
+            chirp_mass.shape[0] != mass_ratio.shape[0]
+            or chirp_mass.shape[0] != chi1.shape[0]
+            or chi1.shape[0] != chi2.shape[0]
+            or chi2.shape[0] != distance.shape[0]
+            or distance.shape[0] != phic.shape[0]
+            or phic.shape[0] != inclination.shape[0]
+        ):
+            raise RuntimeError("Tensors should have same batch size")
+        mass2 = chirp_mass * (1.0 + mass_ratio) ** 0.2 / mass_ratio**0.6
+        mass1 = mass_ratio * mass2
+        cfac = torch.cos(inclination)
+        pfac = 0.5 * (1.0 + cfac * cfac)
 
-    # Phase coeffeciencts from https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/lib/LALSimInspiralPNCoefficients.c  # noqa E501
-    pfaN = 3.0 / (128.0 * eta)
-    pfa_v0 = 1.0
-    pfa_v1 = 0.0
-    pfa_v2 = 5.0 * (74.3 / 8.4 + 11.0 * eta) / 9.0
-    pfa_v3 = -16.0 * PI
-    # SO contributions at 1.5 PN
-    pfa_v3 += (
-        m1byM * (25.0 + 38.0 / 3.0 * m1byM) * chi1
-        + m2byM * (25.0 + 38.0 / 3.0 * m2byM) * chi2
-    )
-    pfa_v4 = (
-        5.0
-        * (3058.673 / 7.056 + 5429.0 / 7.0 * eta + 617.0 * eta * eta)
-        / 72.0
-    )
-    # SO, SS, S1,2-squared contributions
-    pfa_v4 += (
-        247.0 / 4.8 * eta * chi1 * chi2
-        + -721.0 / 4.8 * eta * chi1 * chi2
-        + (-720.0 / 9.6 * m1byM * m1byM + 1.0 / 9.6 * m1byM * m1byM) * chi1sq
-        + (-720.0 / 9.6 * m2byM * m2byM + 1.0 / 9.6 * m2byM * m2byM) * chi2sq
-        + (240.0 / 9.6 * m1byM * m1byM + -7.0 / 9.6 * m1byM * m1byM) * chi1sq
-        + (240.0 / 9.6 * m2byM * m2byM + -7.0 / 9.6 * m2byM * m2byM) * chi2sq
-    )
-    pfa_v5logv = 5.0 / 3.0 * (772.9 / 8.4 - 13.0 * eta) * PI
-    pfa_v5 = 5.0 / 9.0 * (772.9 / 8.4 - 13.0 * eta) * PI
-    # SO coefficient for 2.5 PN
-    pfa_v5logv += 3.0 * (
-        -m1byM
-        * (
-            1391.5 / 8.4
-            - 10.0 / 3.0 * m1byM * (1.0 - m1byM)
-            + m1byM * (1276.0 / 8.1 + 170.0 / 9.0 * m1byM * (1.0 - m1byM))
-        )
-        * chi1
-        - m2byM
-        * (
-            1391.5 / 8.4
-            - 10.0 / 3.0 * m2byM * (1.0 - m2byM)
-            + m2byM * (1276.0 / 8.1 + 170.0 / 9.0 * m2byM * (1.0 - m2byM))
-        )
-        * chi2
-    )
-    pfa_v5 += (
-        -m1byM
-        * (
-            1391.5 / 8.4
-            - 10.0 / 3.0 * m1byM * (1.0 - m1byM)
-            + m1byM * (1276.0 / 8.1 + 170.0 / 9.0 * m1byM * (1.0 - m1byM))
-        )
-        * chi1
-        + -m2byM
-        * (
-            1391.5 / 8.4
-            - 10.0 / 3.0 * m2byM * (1.0 - m2byM)
-            + m2byM * (1276.0 / 8.1 + 170.0 / 9.0 * m2byM * (1.0 - m2byM))
-        )
-        * chi2
-    )
-    pfa_v6logv = -684.8 / 2.1
-    pfa_v6 = (
-        11583.231236531 / 4.694215680
-        - 640.0 / 3.0 * PI * PI
-        - 684.8 / 2.1 * GAMMA
-        + eta * (-15737.765635 / 3.048192 + 225.5 / 1.2 * PI * PI)
-        + eta * eta * 76.055 / 1.728
-        - eta * eta * eta * 127.825 / 1.296
-        + pfa_v6logv * torch.log(torch.tensor(4.0))
-    )
-    # SO + S1-S2 + S-squared contribution at 3 PN
-    pfa_v6 += (
-        PI * m1byM * (1490.0 / 3.0 + m1byM * 260.0) * chi1
-        + PI * m2byM * (1490.0 / 3.0 + m2byM * 260.0) * chi2
-        + (326.75 / 1.12 + 557.5 / 1.8 * eta) * eta * chi1 * chi2
-        + (
-            (4703.5 / 8.4 + 2935.0 / 6.0 * m1byM - 120.0 * m1byM * m1byM)
-            * m1byM
-            * m1byM
-            + (
-                -4108.25 / 6.72
-                - 108.5 / 1.2 * m1byM
-                + 125.5 / 3.6 * m1byM * m1byM
-            )
-            * m1byM
-            * m1byM
-        )
-        * chi1sq
-        + (
-            (4703.5 / 8.4 + 2935.0 / 6.0 * m2byM - 120.0 * m2byM * m2byM)
-            * m2byM
-            * m2byM
-            + (
-                -4108.25 / 6.72
-                - 108.5 / 1.2 * m2byM
-                + 125.5 / 3.6 * m2byM * m2byM
-            )
-            * m2byM
-            * m2byM
+        htilde = self.taylorf2_htilde(
+            f, mass1, mass2, chi1, chi2, distance, phic, f_ref
         )
-        * chi2sq
-    )
-    pfa_v7 = PI * (
-        770.96675 / 2.54016 + 378.515 / 1.512 * eta - 740.45 / 7.56 * eta * eta
-    )
-    # SO contribution at 3.5 PN
-    pfa_v7 += (
-        m1byM
-        * (
-            -17097.8035 / 4.8384
-            + eta * 28764.25 / 6.72
-            + eta * eta * 47.35 / 1.44
-            + m1byM
-            * (
-                -7189.233785 / 1.524096
-                + eta * 458.555 / 3.024
-                - eta * eta * 534.5 / 7.2
-            )
-        )
-    ) * chi1 + (
-        m2byM
-        * (
-            -17097.8035 / 4.8384
-            + eta * 28764.25 / 6.72
-            + eta * eta * 47.35 / 1.44
-            + m2byM
-            * (
-                -7189.233785 / 1.524096
-                + eta * 458.555 / 3.024
-                - eta * eta * 534.5 / 7.2
-            )
-        )
-    ) * chi2
-    # construct power series
-    phasing = (v7.T * pfa_v7).T
-    phasing += (v6.T * pfa_v6 + v6_logv.T * pfa_v6logv).T
-    phasing += (v5.T * pfa_v5 + v5_logv.T * pfa_v5logv).T
-    phasing += (v4.T * pfa_v4).T
-    phasing += (v3.T * pfa_v3).T
-    phasing += (v2.T * pfa_v2).T
-    phasing += (v1.T * pfa_v1).T
-    phasing += (v0.T * pfa_v0).T
-    # Divide by 0PN v-dependence
-    phasing /= v5
-    # Multiply by 0PN coefficient
-    phasing = (phasing.T * pfaN).T
-
-    # Derivative of phase w.r.t Mf
-    # dPhi/dMf = dPhi/dv dv/dMf
-    Dphasing = (2.0 * v7.T * pfa_v7).T
-    Dphasing += (v6.T * (pfa_v6 + pfa_v6logv)).T
-    Dphasing += (v6_logv.T * pfa_v6logv).T
-    Dphasing += (v5.T * pfa_v5logv).T
-    Dphasing += (-1.0 * v4.T * pfa_v4).T
-    Dphasing += (-2.0 * v3.T * pfa_v3).T
-    Dphasing += (-3.0 * v2.T * pfa_v2).T
-    Dphasing += (-4.0 * v1.T * pfa_v1).T
-    Dphasing += -5.0 * v0
-    Dphasing /= 3.0 * v1 * v7
-    Dphasing *= PI
-    Dphasing = (Dphasing.T * pfaN).T
 
-    return phasing, Dphasing
+        hp = (htilde.mT * pfac).mT
+        hc = -1j * (htilde.mT * cfac).mT
 
+        return hc, hp
 
-def taylorf2_amplitude(
-    Mf: TensorType, mass1, mass2, eta, distance
-) -> TensorType:
-    mass1_s = mass1 * MTSUN_SI
-    mass2_s = mass2 * MTSUN_SI
-    v = (PI * Mf) ** (1.0 / 3.0)
-    v10 = v**10
+    def taylorf2_htilde(
+        self,
+        f: FrequencySeries1d,
+        mass1: BatchTensor,
+        mass2: BatchTensor,
+        chi1: BatchTensor,
+        chi2: BatchTensor,
+        distance: BatchTensor,
+        phic: BatchTensor,
+        f_ref: float,
+    ) -> Float[FrequencySeries1d, " batch"]:
+        mass1_s = mass1 * MTSUN_SI
+        mass2_s = mass2 * MTSUN_SI
+        M_s = mass1_s + mass2_s
+        eta = mass1_s * mass2_s / M_s / M_s
 
-    # Flux and energy coefficient at newtonian
-    FTaN = 32.0 * eta * eta / 5.0
-    dETaN = 2 * (-eta / 2.0)
+        Mf = torch.outer(M_s, f)
+        Mf_ref = torch.outer(M_s, f_ref * torch.ones_like(f))
 
-    amp0 = -4.0 * mass1_s * mass2_s * (PI / 12.0) ** 0.5
+        Psi, _ = self.taylorf2_phase(Mf, mass1, mass2, chi1, chi2)
+        Psi_ref, _ = self.taylorf2_phase(Mf_ref, mass1, mass2, chi1, chi2)
 
-    amp0 /= distance * MPC_SEC
-    flux = (v10.T * FTaN).T
-    dEnergy = (v.T * dETaN).T
-    amp = torch.sqrt(-dEnergy / flux) * v
-    amp = (amp.T * amp0).T
+        Psi = (Psi.mT - 2 * phic).mT
+        Psi -= Psi_ref
 
-    return amp
+        amp0 = self.taylorf2_amplitude(Mf, mass1, mass2, eta, distance)
+        h0 = amp0 * torch.exp(-1j * (Psi - PI / 4))
+        return h0
 
+    def taylorf2_amplitude(
+        self,
+        Mf: BatchTensor,
+        mass1: BatchTensor,
+        mass2: BatchTensor,
+        eta: BatchTensor,
+        distance: BatchTensor,
+    ) -> Float[FrequencySeries1d, " batch"]:
+        mass1_s = mass1 * MTSUN_SI
+        mass2_s = mass2 * MTSUN_SI
+        v = (PI * Mf) ** (1.0 / 3.0)
+        v10 = v**10
 
-def taylorf2_htilde(
-    f: TensorType,
-    mass1: TensorType,
-    mass2: TensorType,
-    chi1: TensorType,
-    chi2: TensorType,
-    distance: TensorType,
-    phic: TensorType,
-    f_ref: float,
-):
-    mass1_s = mass1 * MTSUN_SI
-    mass2_s = mass2 * MTSUN_SI
-    M_s = mass1_s + mass2_s
-    eta = mass1_s * mass2_s / M_s / M_s
+        # Flux and energy coefficient at newtonian
+        FTaN = 32.0 * eta * eta / 5.0
+        dETaN = 2 * (-eta / 2.0)
 
-    Mf = torch.outer(M_s, f)
-    Mf_ref = torch.outer(M_s, f_ref * torch.ones_like(f))
+        amp0 = -4.0 * mass1_s * mass2_s * (PI / 12.0) ** 0.5
 
-    Psi, _ = taylorf2_phase(Mf, mass1, mass2, chi1, chi2)
-    Psi_ref, _ = taylorf2_phase(Mf_ref, mass1, mass2, chi1, chi2)
+        amp0 /= distance * MPC_SEC
+        flux = (v10.mT * FTaN).mT
+        dEnergy = (v.mT * dETaN).mT
+        amp = torch.sqrt(-dEnergy / flux) * v
+        amp = (amp.mT * amp0).mT
 
-    Psi = (Psi.T - 2 * phic).T
-    Psi -= Psi_ref
+        return amp
 
-    amp0 = taylorf2_amplitude(Mf, mass1, mass2, eta, distance)
-    h0 = amp0 * torch.exp(-1j * (Psi - PI / 4))
-    return h0
+    def taylorf2_phase(
+        self,
+        Mf: BatchTensor,
+        mass1: BatchTensor,
+        mass2: BatchTensor,
+        chi1: BatchTensor,
+        chi2: BatchTensor,
+    ) -> Float[FrequencySeries1d, " batch"]:
+        """
+        Calculate the inspiral phase for the TaylorF2.
+        """
+        M = mass1 + mass2
+        eta = mass1 * mass2 / M / M
+        m1byM = mass1 / M
+        m2byM = mass2 / M
+        chi1sq = chi1 * chi1
+        chi2sq = chi2 * chi2
 
+        v0 = torch.ones_like(Mf)
+        v1 = (PI * Mf) ** (1.0 / 3.0)
+        v2 = v1 * v1
+        v3 = v2 * v1
+        v4 = v3 * v1
+        v5 = v4 * v1
+        v6 = v5 * v1
+        v7 = v6 * v1
+        logv = torch.log(v1)
+        v5_logv = v5 * logv
+        v6_logv = v6 * logv
 
-def TaylorF2(
-    f: TensorType,
-    mass1: TensorType,
-    mass2: TensorType,
-    chi1: TensorType,
-    chi2: TensorType,
-    distance: TensorType,
-    phic: TensorType,
-    inclination: TensorType,
-    f_ref: float,
-):
-    """
-    TaylorF2 up to 3.5 PN in phase. Newtonian SPA amplitude.
-
-    Returns:
-    --------
-      hp, hc
-    """
-    # shape assumed (n_batch, params)
-    if (
-        mass1.shape[0] != mass2.shape[0]
-        or mass2.shape[0] != chi1.shape[0]
-        or chi1.shape[0] != chi2.shape[0]
-        or chi2.shape[0] != distance.shape[0]
-        or distance.shape[0] != phic.shape[0]
-        or phic.shape[0] != inclination.shape[0]
-    ):
-        raise RuntimeError("Tensors should have same batch size")
-    cfac = torch.cos(inclination)
-    pfac = 0.5 * (1.0 + cfac * cfac)
-
-    htilde = taylorf2_htilde(
-        f, mass1, mass2, chi1, chi2, distance, phic, f_ref
-    )
+        # Phase coeffeciencts from https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/lib/LALSimInspiralPNCoefficients.c  # noqa E501
+        pfaN = 3.0 / (128.0 * eta)
+        pfa_v0 = 1.0
+        pfa_v1 = 0.0
+        pfa_v2 = 5.0 * (74.3 / 8.4 + 11.0 * eta) / 9.0
+        pfa_v3 = -16.0 * PI
+        # SO contributions at 1.5 PN
+        pfa_v3 += (
+            m1byM * (25.0 + 38.0 / 3.0 * m1byM) * chi1
+            + m2byM * (25.0 + 38.0 / 3.0 * m2byM) * chi2
+        )
+        pfa_v4 = (
+            5.0
+            * (3058.673 / 7.056 + 5429.0 / 7.0 * eta + 617.0 * eta * eta)
+            / 72.0
+        )
+        # SO, SS, S1,2-squared contributions
+        pfa_v4 += (
+            247.0 / 4.8 * eta * chi1 * chi2
+            + -721.0 / 4.8 * eta * chi1 * chi2
+            + (-720.0 / 9.6 * m1byM * m1byM + 1.0 / 9.6 * m1byM * m1byM)
+            * chi1sq
+            + (-720.0 / 9.6 * m2byM * m2byM + 1.0 / 9.6 * m2byM * m2byM)
+            * chi2sq
+            + (240.0 / 9.6 * m1byM * m1byM + -7.0 / 9.6 * m1byM * m1byM)
+            * chi1sq
+            + (240.0 / 9.6 * m2byM * m2byM + -7.0 / 9.6 * m2byM * m2byM)
+            * chi2sq
+        )
+        pfa_v5logv = 5.0 / 3.0 * (772.9 / 8.4 - 13.0 * eta) * PI
+        pfa_v5 = 5.0 / 9.0 * (772.9 / 8.4 - 13.0 * eta) * PI
+        # SO coefficient for 2.5 PN
+        pfa_v5logv += 3.0 * (
+            -m1byM
+            * (
+                1391.5 / 8.4
+                - 10.0 / 3.0 * m1byM * (1.0 - m1byM)
+                + m1byM * (1276.0 / 8.1 + 170.0 / 9.0 * m1byM * (1.0 - m1byM))
+            )
+            * chi1
+            - m2byM
+            * (
+                1391.5 / 8.4
+                - 10.0 / 3.0 * m2byM * (1.0 - m2byM)
+                + m2byM * (1276.0 / 8.1 + 170.0 / 9.0 * m2byM * (1.0 - m2byM))
+            )
+            * chi2
+        )
+        pfa_v5 += (
+            -m1byM
+            * (
+                1391.5 / 8.4
+                - 10.0 / 3.0 * m1byM * (1.0 - m1byM)
+                + m1byM * (1276.0 / 8.1 + 170.0 / 9.0 * m1byM * (1.0 - m1byM))
+            )
+            * chi1
+            + -m2byM
+            * (
+                1391.5 / 8.4
+                - 10.0 / 3.0 * m2byM * (1.0 - m2byM)
+                + m2byM * (1276.0 / 8.1 + 170.0 / 9.0 * m2byM * (1.0 - m2byM))
+            )
+            * chi2
+        )
+        pfa_v6logv = -684.8 / 2.1
+        pfa_v6 = (
+            11583.231236531 / 4.694215680
+            - 640.0 / 3.0 * PI * PI
+            - 684.8 / 2.1 * GAMMA
+            + eta * (-15737.765635 / 3.048192 + 225.5 / 1.2 * PI * PI)
+            + eta * eta * 76.055 / 1.728
+            - eta * eta * eta * 127.825 / 1.296
+            + pfa_v6logv * torch.log(torch.tensor(4.0))
+        )
+        # SO + S1-S2 + S-squared contribution at 3 PN
+        pfa_v6 += (
+            PI * m1byM * (1490.0 / 3.0 + m1byM * 260.0) * chi1
+            + PI * m2byM * (1490.0 / 3.0 + m2byM * 260.0) * chi2
+            + (326.75 / 1.12 + 557.5 / 1.8 * eta) * eta * chi1 * chi2
+            + (
+                (4703.5 / 8.4 + 2935.0 / 6.0 * m1byM - 120.0 * m1byM * m1byM)
+                * m1byM
+                * m1byM
+                + (
+                    -4108.25 / 6.72
+                    - 108.5 / 1.2 * m1byM
+                    + 125.5 / 3.6 * m1byM * m1byM
+                )
+                * m1byM
+                * m1byM
+            )
+            * chi1sq
+            + (
+                (4703.5 / 8.4 + 2935.0 / 6.0 * m2byM - 120.0 * m2byM * m2byM)
+                * m2byM
+                * m2byM
+                + (
+                    -4108.25 / 6.72
+                    - 108.5 / 1.2 * m2byM
+                    + 125.5 / 3.6 * m2byM * m2byM
+                )
+                * m2byM
+                * m2byM
+            )
+            * chi2sq
+        )
+        pfa_v7 = PI * (
+            770.96675 / 2.54016
+            + 378.515 / 1.512 * eta
+            - 740.45 / 7.56 * eta * eta
+        )
+        # SO contribution at 3.5 PN
+        pfa_v7 += (
+            m1byM
+            * (
+                -17097.8035 / 4.8384
+                + eta * 28764.25 / 6.72
+                + eta * eta * 47.35 / 1.44
+                + m1byM
+                * (
+                    -7189.233785 / 1.524096
+                    + eta * 458.555 / 3.024
+                    - eta * eta * 534.5 / 7.2
+                )
+            )
+        ) * chi1 + (
+            m2byM
+            * (
+                -17097.8035 / 4.8384
+                + eta * 28764.25 / 6.72
+                + eta * eta * 47.35 / 1.44
+                + m2byM
+                * (
+                    -7189.233785 / 1.524096
+                    + eta * 458.555 / 3.024
+                    - eta * eta * 534.5 / 7.2
+                )
+            )
+        ) * chi2
+        # construct power series
+        phasing = (v7.mT * pfa_v7).mT
+        phasing += (v6.mT * pfa_v6 + v6_logv.mT * pfa_v6logv).mT
+        phasing += (v5.mT * pfa_v5 + v5_logv.mT * pfa_v5logv).mT
+        phasing += (v4.mT * pfa_v4).mT
+        phasing += (v3.mT * pfa_v3).mT
+        phasing += (v2.mT * pfa_v2).mT
+        phasing += (v1.mT * pfa_v1).mT
+        phasing += (v0.mT * pfa_v0).mT
+        # Divide by 0PN v-dependence
+        phasing /= v5
+        # Multiply by 0PN coefficient
+        phasing = (phasing.mT * pfaN).mT
 
-    hp = (htilde.T * pfac).T
-    hc = -1j * (htilde.T * cfac).T
+        # Derivative of phase w.r.t Mf
+        # dPhi/dMf = dPhi/dv dv/dMf
+        Dphasing = (2.0 * v7.mT * pfa_v7).mT
+        Dphasing += (v6.mT * (pfa_v6 + pfa_v6logv)).mT
+        Dphasing += (v6_logv.mT * pfa_v6logv).mT
+        Dphasing += (v5.mT * pfa_v5logv).mT
+        Dphasing += (-1.0 * v4.mT * pfa_v4).mT
+        Dphasing += (-2.0 * v3.mT * pfa_v3).mT
+        Dphasing += (-3.0 * v2.mT * pfa_v2).mT
+        Dphasing += (-4.0 * v1.mT * pfa_v1).mT
+        Dphasing += -5.0 * v0
+        Dphasing /= 3.0 * v1 * v7
+        Dphasing *= PI
+        Dphasing = (Dphasing.mT * pfaN).mT
 
-    return hp, hc
+        return phasing, Dphasing

{ml4gw-0.4.2.dist-info → ml4gw-0.5.1.dist-info}/METADATA

@@ -1,18 +1,20 @@
 Metadata-Version: 2.1
 Name: ml4gw
-Version: 0.4.2
+Version: 0.5.1
 Summary: Tools for training torch models on gravitational wave data
 Author: Alec Gunny
 Author-email: alec.gunny@ligo.org
-Requires-Python: >=3.8,<3.12
+Requires-Python: >=3.8,<3.13
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: jaxtyping (>=0.2,<0.3)
+Requires-Dist: numpy (<2.0.0)
 Requires-Dist: torch (>=2.0,<3.0)
 Requires-Dist: torchaudio (>=2.0,<3.0)
-Requires-Dist: torchtyping (>=0.1,<0.2)
 Description-Content-Type: text/markdown
 
 # ML4GW

ml4gw-0.5.1.dist-info/RECORD

@@ -0,0 +1,47 @@
+ml4gw/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ml4gw/augmentations.py,sha256=pZH9tjEpXV0AIqvHHDkpUE-BorG02beOz2pmSipw2EY,1232
+ml4gw/constants.py,sha256=W9beA9RDRdIug1I2H7VLPEPv_DFsQWWoYRmzxv7FWgM,891
+ml4gw/dataloading/__init__.py,sha256=EHBBqU7y2-Np5iQ_xyufxamUEM1pPEquqFo7oaJnaJE,149
+ml4gw/dataloading/chunked_dataset.py,sha256=FpDc4gFxt-PMyXs5qSWLuTGXMTuS1B-hH8gUOCOGxZk,5260
+ml4gw/dataloading/hdf5_dataset.py,sha256=UB1Eog8l7m4M78Owst7oYQZICb0DRJer9WVLVn4hl_I,6645
+ml4gw/dataloading/in_memory_dataset.py,sha256=kleMA9ABUKA6J0tCdz78tbX9lM6uxVSLhqgHbSa1iWY,9550
+ml4gw/distributions.py,sha256=tUuaOiX5enjKLYWD7uiN8rdRVQcrIKps64xBkTl8fMs,4991
+ml4gw/gw.py,sha256=To_hQz9tUp02ADllGLxFCPsNcfbb-kbvfgGpooxcOII,17693
+ml4gw/nn/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ml4gw/nn/autoencoder/__init__.py,sha256=ZaT1XhJTHpMuPQqu5E__Jezeh9uwtjcXlT7IZ18byq4,161
+ml4gw/nn/autoencoder/base.py,sha256=4d5Ej30IUzZh3XbldzWlCpp3p0_91YUvKeRID8ZEZGA,3225
+ml4gw/nn/autoencoder/convolutional.py,sha256=2BXDuPWYC-151RO_FL0ogdrqSVTfo4YNrY80lPwrmFA,5419
+ml4gw/nn/autoencoder/skip_connection.py,sha256=fpXxxIIl0CXY4mAfUZQuvI542pEBSwpg90TNG2rbZY8,1411
+ml4gw/nn/autoencoder/utils.py,sha256=m_ivYGNwdrhA7cFxJVD4gqM8AHiWIGmlQI3pFNRklXQ,355
+ml4gw/nn/norm.py,sha256=JIOMXQbUtoWlrhncGsqW6f1-DiGDx9zQH2O3CvQml3U,3594
+ml4gw/nn/resnet/__init__.py,sha256=vBI0IftVP_EYAeDlqomtkGqUYE-RE_S4WNioUhniw9s,64
+ml4gw/nn/resnet/resnet_1d.py,sha256=IQ-EIIzAXd-NWuLwt7JTXLWg5bO3FGJpuFAZwZ78jaI,13218
+ml4gw/nn/resnet/resnet_2d.py,sha256=aK4I0FOZk62JxnYFz0t1O0s5s7J7yRNYSM1flRypvVc,13301
+ml4gw/nn/streaming/__init__.py,sha256=zgjGR2L8t0txXLnil9ceZT0tM8Y2FC8yPxqIKYH0o1A,80
+ml4gw/nn/streaming/online_average.py,sha256=aI8hkT7I3thXkda9tsXxYrzump9swelSXPdSTwPlJWY,4719
+ml4gw/nn/streaming/snapshotter.py,sha256=B9qtbHxnPszAHQ5WQppWJLRuMnnYIxGk7MRUlgja7Is,4476
+ml4gw/spectral.py,sha256=Mt3-yz4a83z0X7M1sVp00_vB947w-9OjU0iNdEkbQcU,19145
+ml4gw/transforms/__init__.py,sha256=24pdP_hIg1wfrtZxxRBPhcEXsCbvVKtNKp7JL8SEogE,362
+ml4gw/transforms/pearson.py,sha256=Ep3mMsY15AF55taRaWNjpHRTvtr1StShUDfqk0dN-qo,3235
+ml4gw/transforms/qtransform.py,sha256=umBSpykfmPftjfyMqbniiP2mTh62q4hoYPA55qneJ4o,17702
+ml4gw/transforms/scaler.py,sha256=fLZo-m6_yFY3UDoLEaS_YgCnYggxlcKstXcM7749TiU,2433
+ml4gw/transforms/snr_rescaler.py,sha256=3XXCTaXc2dzzpXRZx7iqRwImvYtRSJLM5fHdBGfpoUs,2351
+ml4gw/transforms/spectral.py,sha256=gTHUeC0gGYbzgBZHb_FxC_4zdhl5H-XCiLg1hrvKB70,4393
+ml4gw/transforms/spectrogram.py,sha256=HS3Rf5iB7JjhlSESRDdFGUwCtIBdvUaJUDulkB4Lmos,6162
+ml4gw/transforms/transform.py,sha256=BuzTbPFxp18OEGP9Tu9jBGtvqy3len1cqvqg5X37DiY,2512
+ml4gw/transforms/waveforms.py,sha256=LkYCvxPqYhHa2yYZTvPE6j0E4HFy16b5ndCRQb7WfcA,3196
+ml4gw/transforms/whitening.py,sha256=Aw_ogq93CYCATiHWBqSZ-qsUtaHAMA3k009ZRtQTtHA,9596
+ml4gw/types.py,sha256=CcctqDcNajR7khGT6BD-WYsfRKpiP0udoSAB0k1qcFw,863
+ml4gw/utils/interferometer.py,sha256=lRS0N3SwUTknhYXX57VACJ99jK1P9M19oUWN_i_nQN0,1814
+ml4gw/utils/slicing.py,sha256=ilRz_5sJzwmd5VyBlrj81tvyC3uCnXYjd0TO2fzFMr8,13563
+ml4gw/waveforms/__init__.py,sha256=dnxfRGX_B3zQPB3_3srLyjZXRxTn4miZqYIRe7PYyrU,170
+ml4gw/waveforms/generator.py,sha256=dO6RQ96EC87p2q0tEkxA62XkkJc1xARFO1SKcGvyDhM,1272
+ml4gw/waveforms/phenom_d.py,sha256=vA60SjOvWSIcsU83-KEw2hnU3ATo4eW8A2mMmuMXo7Y,46941
+ml4gw/waveforms/phenom_d_data.py,sha256=WA1FBxUp9fo1IQaV_OLJ_5g5gI166mY1FtG9n25he9U,53447
+ml4gw/waveforms/phenom_p.py,sha256=VybpPlc2_yMGywnPz5B79QAygAj-WAeHZTPiZHets28,26951
+ml4gw/waveforms/ringdown.py,sha256=m8IBQTxKBBGFqBtWGEO4KG3DEYR8TTnNyGVdVLaMKa8,3316
+ml4gw/waveforms/sine_gaussian.py,sha256=-MtrI7ydwBTk4K0O4tdkC8-w5OifQszdnWN9__I4XzY,3569
+ml4gw/waveforms/taylorf2.py,sha256=ySYLGTT_c3k4NzPDsQ9v822kzvU6TwYpELJEWlCDGQE,10428
+ml4gw-0.5.1.dist-info/METADATA,sha256=P2uoQtMX_K5SSwAzTY5tyNvWYszxaDADTS54iDOQYKw,5785
+ml4gw-0.5.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+ml4gw-0.5.1.dist-info/RECORD,,