ml4gw 0.6.0__tar.gz → 0.6.1__tar.gz

{ml4gw-0.6.0 → ml4gw-0.6.1}/PKG-INFO

@@ -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 → ml4gw-0.6.1}/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-0.6.0 → ml4gw-0.6.1}/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-0.6.0 → ml4gw-0.6.1}/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 → ml4gw-0.6.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "ml4gw"
-version = "0.6.0"
+version = "0.6.1"
 description = "Tools for training torch models on gravitational wave data"
 readme = "README.md"
 authors = [