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

ml4gw/waveforms/cbc/phenom_d.py

@@ -590,18 +590,19 @@ class IMRPhenomD(TaylorF2):
         return fRD, fDM
 
     def fmaxCalc(self, fRD, fDM, gamma2, gamma3):
-        res = torch.zeros_like(gamma2)
-        res = torch.abs(fRD + (-fDM * gamma3) / gamma2) * (gamma2 > 1).to(
-            torch.int
-        ) + torch.abs(
-            fRD
-            + (fDM * (-1 + torch.sqrt(1 - gamma2 * gamma2)) * gamma3) / gamma2
-        ) * (
-            gamma2 <= 1
-        ).to(
-            torch.int
-        )
-        return res
+        mask = gamma2 <= 1
+        # calculate result for gamma2 <= 1 case
+        sqrt_term = torch.sqrt(1 - gamma2.pow(2))
+        result_case1 = fRD + (fDM * (-1 + sqrt_term) * gamma3) / gamma2
+
+        # calculate result for gamma2 > 1 case
+        # i.e. don't add sqrt term
+        result_case2 = fRD + (-fDM * gamma3) / gamma2
+
+        # combine results using mask
+        result = torch.where(mask, result_case1, result_case2)
+
+        return torch.abs(result)
 
     def _linear_interp_finspin(self, finspin):
         # chi is a batch of final spins i.e. torch.Size([n])

ml4gw/waveforms/cbc/taylorf2.py

@@ -4,6 +4,7 @@ from jaxtyping import Float
 from ml4gw.constants import MPC_SEC, MTSUN_SI, PI
 from ml4gw.constants import EulerGamma as GAMMA
 from ml4gw.types import BatchTensor, FrequencySeries1d
+from ml4gw.waveforms.conversion import chirp_mass_and_mass_ratio_to_components
 
 
 class TaylorF2(torch.nn.Module):
@@ -60,8 +61,11 @@ class TaylorF2(torch.nn.Module):
             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
+
+        mass1, mass2 = chirp_mass_and_mass_ratio_to_components(
+            chirp_mass, mass_ratio
+        )
+
         cfac = torch.cos(inclination)
         pfac = 0.5 * (1.0 + cfac * cfac)
 

ml4gw/waveforms/conversion.py

@@ -32,6 +32,23 @@ def XLALSimInspiralL_2PN(eta: BatchTensor):
     return 1.5 + eta / 6.0
 
 
+def chirp_mass_and_mass_ratio_to_components(
+    chirp_mass: BatchTensor, mass_ratio: BatchTensor
+):
+    """
+    Compute component masses from chirp mass and mass ratio.
+    Args:
+        chirp_mass: Tensor of chirp mass values
+        mass_ratio:
+            Tensor of mass ratio values, `m2 / m1`,
+            where m1 >= m2, so that mass_ratio <= 1
+    """
+    total_mass = chirp_mass * (1 + mass_ratio) ** 1.2 / mass_ratio**0.6
+    mass_1 = total_mass / (1 + mass_ratio)
+    mass_2 = mass_1 * mass_ratio
+    return mass_1, mass_2
+
+
 def bilby_spins_to_lalsim(
     theta_jn: BatchTensor,
     phi_jl: BatchTensor,

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

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

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

@@ -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=vA60SjOvWSIcsU83-KEw2hnU3ATo4eW8A2mMmuMXo7Y,46941
+ml4gw/waveforms/cbc/phenom_d.py,sha256=3JhN77Xl5foqjmdDJSw4GrTYcaOr1RxXdL1XjUAb8IY,47053
 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/taylorf2.py,sha256=ySYLGTT_c3k4NzPDsQ9v822kzvU6TwYpELJEWlCDGQE,10428
-ml4gw/waveforms/conversion.py,sha256=F5fsNeqf6KHY66opDIj8fN9bwUcwrt9f7PCaxLAi9Jk,6367
+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.0.dist-info/METADATA,sha256=6bwcfu6ojmrxgtMnFVViy9FanSmMXjhnN33yAzViFzo,5785
-ml4gw-0.6.0.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-ml4gw-0.6.0.dist-info/RECORD,,
+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,,