ml4gw 0.6.1__py3-none-any.whl → 0.6.2__py3-none-any.whl

ml4gw/transforms/spline_interpolation.py

@@ -263,7 +263,7 @@ class SplineInterpolate(torch.nn.Module):
             return torch.linalg.solve(self.BxT_Bx, Z_Bx.mT).mT
 
         # Adding batch/channel dimension handling
-        # ByT @ Z @ Bx
+        # ByT @ Z @ BxW
         ByT_Z_Bx = torch.einsum("ij,bcik,kl->bcjl", self.By, Z, self.Bx)
         # (ByT @ By)^-1 @ (ByT @ Z @ Bx) = By^-1 @ Z @ Bx
         E = torch.linalg.solve(self.ByT_By, ByT_Z_Bx)

ml4gw/waveforms/cbc/phenom_d.py

@@ -157,20 +157,30 @@ class IMRPhenomD(TaylorF2):
         chi22,
         xi,
         distance,
+        fRD=None,  # used for passing ringdown frequency from phenom_p
+        fDM=None,  # used for passing damping frequency from phenom_p
     ):
         ins_amp, ins_Damp = self.phenom_d_inspiral_amp(
             Mf, eta, eta2, Seta, xi, chi1, chi2, chi12, chi22
         )
         int_amp, int_Damp = self.phenom_d_int_amp(
-            Mf, eta, eta2, Seta, chi1, chi2, chi12, chi22, xi
+            Mf, eta, eta2, Seta, chi1, chi2, chi12, chi22, xi, fRD, fDM
         )
         mrd_amp, mrd_Damp = self.phenom_d_mrd_amp(
-            Mf, eta, eta2, chi1, chi2, xi
+            Mf, eta, eta2, chi1, chi2, xi, fRD, fDM
         )
 
         gamma2 = self.gamma2_fun(eta, eta2, xi)
         gamma3 = self.gamma3_fun(eta, eta2, xi)
-        fRD, fDM = self.fring_fdamp(eta, eta2, chi1, chi2)
+
+        # merger ringdown
+        if (fRD is None) != (fDM is None):
+            raise ValueError(
+                "Both fRD and fDM must either be provided or both be None"
+            )
+        if (fRD is None) and (fDM is None):
+            fRD, fDM = self.fring_fdamp(eta, eta2, chi1, chi2)
+
         Mf_peak = self.fmaxCalc(fRD, fDM, gamma2, gamma3)
         # Geometric peak and joining frequencies
         Mf_peak = (torch.ones_like(Mf).mT * Mf_peak).mT
@@ -201,10 +211,27 @@ class IMRPhenomD(TaylorF2):
         return amp, Damp
 
     def phenom_d_int_amp(
-        self, Mf, eta, eta2, Seta, chi1, chi2, chi12, chi22, xi
+        self,
+        Mf,
+        eta,
+        eta2,
+        Seta,
+        chi1,
+        chi2,
+        chi12,
+        chi22,
+        xi,
+        fRD=None,  # used for passing ringdown frequency from phenom_p
+        fDM=None,  # used for passing damping frequency from phenom_p
     ):
         # merger ringdown
-        fRD, fDM = self.fring_fdamp(eta, eta2, chi1, chi2)
+        if (fRD is None) != (fDM is None):
+            raise ValueError(
+                "Both fRD and fDM must either be provided or both be None"
+            )
+        if (fRD is None) and (fDM is None):
+            fRD, fDM = self.fring_fdamp(eta, eta2, chi1, chi2)
+
         # Geometric frequency definition from PhenomD header file
         AMP_fJoin_INS = 0.014
 
@@ -220,7 +247,9 @@ class IMRPhenomD(TaylorF2):
         v1, d1 = self.phenom_d_inspiral_amp(
             Mf1, eta, eta2, Seta, xi, chi1, chi2, chi12, chi22
         )
-        v3, d2 = self.phenom_d_mrd_amp(Mf3, eta, eta2, chi1, chi2, xi)
+        v3, d2 = self.phenom_d_mrd_amp(
+            Mf3, eta, eta2, chi1, chi2, xi, fRD, fDM
+        )
         v2 = (
             torch.ones_like(Mf).mT * self.AmpIntColFitCoeff(eta, eta2, xi)
         ).mT
@@ -239,9 +268,18 @@ class IMRPhenomD(TaylorF2):
         )
         return amp, Damp
 
-    def phenom_d_mrd_amp(self, Mf, eta, eta2, chi1, chi2, xi):
+    # fRD and fDM are to be passed for generating phenom_p waveforms
+    # and remain None for phenom_d
+    def phenom_d_mrd_amp(
+        self, Mf, eta, eta2, chi1, chi2, xi, fRD=None, fDM=None
+    ):
         # merger ringdown
-        fRD, fDM = self.fring_fdamp(eta, eta2, chi1, chi2)
+        if (fRD is None) != (fDM is None):
+            raise ValueError(
+                "Both fRD and fDM must either be provided or both be None"
+            )
+        if (fRD is None) and (fDM is None):
+            fRD, fDM = self.fring_fdamp(eta, eta2, chi1, chi2)
 
         gamma1 = self.gamma1_fun(eta, eta2, xi)
         gamma2 = self.gamma2_fun(eta, eta2, xi)
@@ -384,17 +422,26 @@ class IMRPhenomD(TaylorF2):
 
         return amp, Damp
 
-    def phenom_d_phase(self, Mf, mass_1, mass_2, eta, eta2, chi1, chi2, xi):
+    # fRD and fDM are to be passed for generating phenom_p waveforms
+    # and remain None for phenom_d
+    def phenom_d_phase(
+        self, Mf, mass_1, mass_2, eta, eta2, chi1, chi2, xi, fRD=None, fDM=None
+    ):
         ins_phase, ins_Dphase = self.phenom_d_inspiral_phase(
             Mf, mass_1, mass_2, eta, eta2, xi, chi1, chi2
         )
         int_phase, int_Dphase = self.phenom_d_int_phase(Mf, eta, eta2, xi)
         mrd_phase, mrd_Dphase = self.phenom_d_mrd_phase(
-            Mf, eta, eta2, chi1, chi2, xi
+            Mf, eta, eta2, chi1, chi2, xi, fRD, fDM
         )
 
         # merger ringdown
-        fRD, fDM = self.fring_fdamp(eta, eta2, chi1, chi2)
+        if (fRD is None) != (fDM is None):
+            raise ValueError(
+                "Both fRD and fDM must either be provided or both be None"
+            )
+        if (fRD is None) and (fDM is None):
+            fRD, fDM = self.fring_fdamp(eta, eta2, chi1, chi2)
         # definitions in Eq. (35) of arXiv:1508.07253
         # PHI_fJoin_INS in header LALSimIMRPhenomD.h
         # C1 continuity at intermediate region i.e. f_1
@@ -415,7 +462,7 @@ class IMRPhenomD(TaylorF2):
             fRDJoin, eta, eta2, xi
         )
         mrd_phase_rd, mrd_Dphase_rd = self.phenom_d_mrd_phase(
-            fRDJoin, eta, eta2, chi1, chi2, xi
+            fRDJoin, eta, eta2, chi1, chi2, xi, fRD, fDM
         )
         PhiIntTempVal = (int_phase_rd.mT / eta).mT + C1Int + C2Int * fRDJoin
         # C2MRD = int_Dphase_rd - mrd_Dphase_rd
@@ -454,7 +501,11 @@ class IMRPhenomD(TaylorF2):
 
         return phasing, Dphasing
 
-    def phenom_d_mrd_phase(self, Mf, eta, eta2, chi1, chi2, xi):
+    # fRD and fDM are to be passed for generating phenom_p waveforms
+    # and remain None for phenom_d
+    def phenom_d_mrd_phase(
+        self, Mf, eta, eta2, chi1, chi2, xi, fRD=None, fDM=None
+    ):
         alpha1 = self.alpha1Fit(eta, eta2, xi)
         alpha2 = self.alpha2Fit(eta, eta2, xi)
         alpha3 = self.alpha3Fit(eta, eta2, xi)
@@ -462,9 +513,14 @@ class IMRPhenomD(TaylorF2):
         alpha5 = self.alpha5Fit(eta, eta2, xi)
 
         # merger ringdown
-        fRD, fDM = self.fring_fdamp(eta, eta2, chi1, chi2)
-        f_minus_alpha5_fRD = (Mf.t() - alpha5 * fRD).t()
+        if (fRD is None) != (fDM is None):
+            raise ValueError(
+                "Both fRD and fDM must either be provided or both be None"
+            )
+        if (fRD is None) and (fDM is None):
+            fRD, fDM = self.fring_fdamp(eta, eta2, chi1, chi2)
 
+        f_minus_alpha5_fRD = (Mf.t() - alpha5 * fRD).t()
         # Leading 1/eta is not multiplied at this stage
         mrd_phasing = (Mf.t() * alpha1).t()
         mrd_phasing -= (1 / Mf.t() * alpha2).t()

ml4gw/waveforms/cbc/phenom_p.py

@@ -1,13 +1,21 @@
+"""
+    Based on the JAX implementation of IMRPhenomPv2 from
+    https://github.com/tedwards2412/ripple/blob/main/src/ripplegw/waveforms/IMRPhenomPv2.py
+"""
+
 from typing import Dict, Optional, Tuple
 
 import torch
 from jaxtyping import Float
 from torch import Tensor
 
-from ml4gw.constants import MPC_SEC, MTSUN_SI, PI
-from ml4gw.types import BatchTensor, FrequencySeries1d
-from ml4gw.waveforms.conversion import rotate_y, rotate_z
-
+from ...constants import MPC_SEC, MTSUN_SI, PI
+from ...types import BatchTensor, FrequencySeries1d
+from ..conversion import (
+    chirp_mass_and_mass_ratio_to_components,
+    rotate_y,
+    rotate_z,
+)
 from .phenom_d import IMRPhenomD
 
 
@@ -75,15 +83,15 @@ class IMRPhenomPv2(IMRPhenomD):
         if tc is None:
             tc = torch.zeros_like(chirp_mass)
 
-        m2 = chirp_mass * (1.0 + mass_ratio) ** 0.2 / mass_ratio**0.6
-        m1 = m2 * mass_ratio
+        m1, m2 = chirp_mass_and_mass_ratio_to_components(
+            chirp_mass, mass_ratio
+        )
 
         # # flip m1 m2. For some reason LAL uses this convention for PhenomPv2
         m1, m2 = m2, m1
         s1x, s2x = s2x, s1x
         s1y, s2y = s2y, s1y
         s1z, s2z = s2z, s1z
-
         (
             chi1_l,
             chi2_l,
@@ -320,42 +328,61 @@ class IMRPhenomPv2(IMRPhenomD):
         phic: Orbital phase at peak of the underlying non precessing model
         M: Total mass (Solar masses)
         """
-
         M_s = M * MTSUN_SI
         Mf = torch.outer(M_s, fs)
-        fRD, _ = self.phP_get_fRD_fdamp(m1, m2, chi1, chi2, chip)
+        fRD, fDM = self.phP_get_fRD_fdamp(m1, m2, chi1, chi2, chip)
+        # pass M_s * ringdown and M_s * damping frequency to PhenomD functions
+        MfRD, MfDM = M_s * fRD, M_s * fDM
 
-        phase, _ = self.phenom_d_phase(Mf, m1, m2, eta, eta2, chi1, chi2, xi)
+        phase, _ = self.phenom_d_phase(
+            Mf, m1, m2, eta, eta2, chi1, chi2, xi, MfRD, MfDM
+        )
         phase = (phase.mT - (phic + PI / 4.0)).mT
+        # why are they subtracting 2*phic?
+        # https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/lib/LALSimIMRPhenomP.c#L1316
+
         Amp = self.phenom_d_amp(
-            Mf, m1, m2, eta, eta2, Seta, chi1, chi2, chi12, chi22, xi, distance
+            Mf,
+            m1,
+            m2,
+            eta,
+            eta2,
+            Seta,
+            chi1,
+            chi2,
+            chi12,
+            chi22,
+            xi,
+            distance,
+            MfRD,
+            MfDM,
         )[0]
         Amp0 = self.get_Amp0(Mf, eta)
         dist_s = distance * MPC_SEC
         Amp = ((Amp0 * Amp).mT * (M_s**2.0) / dist_s).mT
-        # phase -= 2. * phic; # line 1316 ???
+
         hPhenom = Amp * (torch.exp(-1j * phase))
 
-        fRDs = torch.outer(
-            fRD, torch.linspace(0.5, 1.5, 101, device=fRD.device)
-        )
-        delta_fRds = torch.median(torch.diff(fRDs, axis=1), axis=1)[0]
+        # calculating derivative of phase with frequency following
+        # https://git.ligo.org/lscsoft/lalsuite/-/blame/master/lalsimulation/lib/LALSimIMRPhenomP.c?page=2#L1057 # noqa: E501
+        n_fixed = 1000
+        x = torch.linspace(0.8, 1.2, n_fixed, device=fRD.device)
+        fRDs = torch.outer(fRD, x)
+        delta_fRds = (1.2 * fRD - 0.8 * fRD) / (n_fixed - 1)
         MfRDs = torch.zeros_like(fRDs)
         for i in range(fRD.shape[0]):
             MfRDs[i, :] = torch.outer(M_s, fRDs[i, :])[i, :]
         RD_phase = self.phenom_d_phase(
-            MfRDs, m1, m2, eta, eta2, chi1, chi2, xi
+            MfRDs, m1, m2, eta, eta2, chi1, chi2, xi, MfRD, MfDM
         )[0]
         diff = torch.diff(RD_phase, axis=1)
         diffRDphase = (diff[:, 1:] + diff[:, :-1]) / (
             2 * delta_fRds.unsqueeze(1)
         )
-        diffRDphase = -diffRDphase[:, 50]
-        # MfRD = torch.outer(M_s, fRD)
-        # Dphase = torch.diag(
-        #     -self.phenom_d_phase(
-        #         MfRD, m1, m2, eta, eta2, chi1, chi2, xi)[1] * M_s
-        # ).view(-1, 1)
+        # reshape x to have same shape as diffRDphase
+        x = x[1:-1].unsqueeze(0).expand(diffRDphase.shape)
+        # interpolate at x = 1, as thats the same as f = fRD
+        diffRDphase = -self.interpolate(torch.tensor([1]), x, diffRDphase)
         return hPhenom, diffRDphase
 
     # Utility functions
@@ -752,9 +779,9 @@ class IMRPhenomPv2(IMRPhenomD):
         eta2 = eta * eta
         # m1 > m2, the convention used in phenomD
         # (not the convention of internal phenomP)
-        mass_ratio = m1 / m2
+        q_factor = m1 / M
         af_parallel = self.FinalSpin0815(eta, eta2, chi1_l, chi2_l)
-        Sperp = chip * mass_ratio * mass_ratio
+        Sperp = chip * q_factor * q_factor
         af = torch.copysign(
             torch.ones_like(af_parallel), af_parallel
         ) * torch.sqrt(Sperp * Sperp + af_parallel * af_parallel)

{ml4gw-0.6.1.dist-info → ml4gw-0.6.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.3
 Name: ml4gw
-Version: 0.6.1
+Version: 0.6.2
 Summary: Tools for training torch models on gravitational wave data
 Author: Alec Gunny
 Author-email: alec.gunny@ligo.org

{ml4gw-0.6.1.dist-info → ml4gw-0.6.2.dist-info}/RECORD

@@ -28,7 +28,7 @@ ml4gw/transforms/scaler.py,sha256=souOt-hOO4M6dqPNXOspfmeU2V9622yGoIMNvju5JZI,25
 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/spline_interpolation.py,sha256=GkyAVLrtZODIIDLkBdAngO9jqEHRzvEFTdxjNM7U1Bc,13526
+ml4gw/transforms/spline_interpolation.py,sha256=oXih-gLMbIbI36DPKLTk39WcjiNUJN_rcQia_k3OrMY,13527
 ml4gw/transforms/transform.py,sha256=BuzTbPFxp18OEGP9Tu9jBGtvqy3len1cqvqg5X37DiY,2512
 ml4gw/transforms/waveforms.py,sha256=LkYCvxPqYhHa2yYZTvPE6j0E4HFy16b5ndCRQb7WfcA,3196
 ml4gw/transforms/whitening.py,sha256=Aw_ogq93CYCATiHWBqSZ-qsUtaHAMA3k009ZRtQTtHA,9596
@@ -40,12 +40,12 @@ ml4gw/waveforms/adhoc/__init__.py,sha256=XVwP4t8TMUj87WY3yMGRTkXsv7_lVr1w8p8iKBW
 ml4gw/waveforms/adhoc/ringdown.py,sha256=m8IBQTxKBBGFqBtWGEO4KG3DEYR8TTnNyGVdVLaMKa8,3316
 ml4gw/waveforms/adhoc/sine_gaussian.py,sha256=-MtrI7ydwBTk4K0O4tdkC8-w5OifQszdnWN9__I4XzY,3569
 ml4gw/waveforms/cbc/__init__.py,sha256=hGbPsFNAIveYJnff8qKY8RWeBPFtZoYcnGHxraPWtWI,99
-ml4gw/waveforms/cbc/phenom_d.py,sha256=3JhN77Xl5foqjmdDJSw4GrTYcaOr1RxXdL1XjUAb8IY,47053
+ml4gw/waveforms/cbc/phenom_d.py,sha256=0pcVAt7b1cjTbphdClPCjenv2sC8bp6oXGGlEUyW-mY,48973
 ml4gw/waveforms/cbc/phenom_d_data.py,sha256=WA1FBxUp9fo1IQaV_OLJ_5g5gI166mY1FtG9n25he9U,53447
-ml4gw/waveforms/cbc/phenom_p.py,sha256=Y8L2r3UPkJeQqJNwknWBmcG_nO2Z_aXJ_DfWc_lzJhg,26720
+ml4gw/waveforms/cbc/phenom_p.py,sha256=nCVu6Tb-xeplch3MUrso_jiThyiNeuIFPnizBan5RfM,27563
 ml4gw/waveforms/cbc/taylorf2.py,sha256=_s-faE8yWMULMxGd4VvzPI54R3G-O2TF2G4-T2m2rDM,10510
 ml4gw/waveforms/conversion.py,sha256=zPkaGkMVqsdrF0fS3ZscyP-2jX8YK40d4smUoJb4gj4,6903
 ml4gw/waveforms/generator.py,sha256=dO6RQ96EC87p2q0tEkxA62XkkJc1xARFO1SKcGvyDhM,1272
-ml4gw-0.6.1.dist-info/METADATA,sha256=APoH6HIVkaBFuDLHQQ58NVmu9VJ5zBuo26cLwjHBrE4,5785
-ml4gw-0.6.1.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-ml4gw-0.6.1.dist-info/RECORD,,
+ml4gw-0.6.2.dist-info/METADATA,sha256=BhP0NjJl1RhuBauJBL3Vu9Rw2iZTIWkJyyLzLsg0D_o,5785
+ml4gw-0.6.2.dist-info/WHEEL,sha256=RaoafKOydTQ7I_I3JTrPCg6kUmTgtm4BornzOqyEfJ8,88
+ml4gw-0.6.2.dist-info/RECORD,,

{ml4gw-0.6.1.dist-info → ml4gw-0.6.2.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: poetry-core 1.9.1
+Generator: poetry-core 2.0.0
 Root-Is-Purelib: true
 Tag: py3-none-any