From bb7f5696ef576bf635a25548d354426fc12c3eb4 Mon Sep 17 00:00:00 2001
From: "Samson Leong (wkstn)" <55839002+SSL32081@users.noreply.github.com>
Date: Mon, 14 Apr 2025 22:13:23 +0800
Subject: [PATCH 1/3] Add aligned-spin conditions to PhenomPv2 angle
 calculations

---
 src/ripplegw/waveforms/IMRPhenomPv2_utils.py | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/src/ripplegw/waveforms/IMRPhenomPv2_utils.py b/src/ripplegw/waveforms/IMRPhenomPv2_utils.py
index 42aac9ae..fc680b93 100644
--- a/src/ripplegw/waveforms/IMRPhenomPv2_utils.py
+++ b/src/ripplegw/waveforms/IMRPhenomPv2_utils.py
@@ -16,6 +16,8 @@
 from ..typing import Array
 from .IMRPhenomD_QNMdata import QNMData_a, QNMData_fRD, QNMData_fdamp
 
+MAX_TOL_ATAN = 1.0e-15
+
 
 # helper functions for LALtoPhenomP:
 def ROTATEZ(angle, x, y, z):
@@ -92,7 +94,12 @@ def convert_spins(
 
     thetaJ_sf = jnp.arccos(J0z_sf / J0)
 
-    phiJ_sf = jnp.arctan2(J0y_sf, J0x_sf)
+    no_inplane_J0 = jnp.abs(J0x_sf) < MAX_TOL_ATAN and jnp.abs(J0y_sf) < MAX_TOL_ATAN
+    phiJ_sf = jax.where(
+        no_inplane_J0,
+        jnp.pi / 2.0 - phiRef,  # This is the aligned-spin case
+        jnp.arctan2(J0y_sf, J0x_sf),
+    )
 
     phi_aligned = -phiJ_sf
 
@@ -117,7 +124,12 @@ def convert_spins(
     tmp_x, tmp_y, tmp_z = ROTATEY(-thetaJ_sf, tmp_x, tmp_y, tmp_z)
     tmp_x, tmp_y, tmp_z = ROTATEZ(kappa, tmp_x, tmp_y, tmp_z)
 
-    alpha0 = jnp.arctan2(tmp_y, tmp_x)
+    no_inplane_LN = jnp.abs(tmp_x) < MAX_TOL_ATAN and jnp.abs(tmp_y) < MAX_TOL_ATAN
+    alpha0 = jnp.where(
+        no_inplane_LN,
+        jnp.pi,  # This is the aligned-spin case
+        jnp.arctan2(tmp_y, tmp_x),
+    )
 
     # Finally we determine thetaJ, by rotating N
     tmp_x, tmp_y, tmp_z = Nx_sf, Ny_sf, Nz_sf

From d3ed2c0171f68fa5be596e2cbc5fc4cd5e90f023 Mon Sep 17 00:00:00 2001
From: "Samson Leong (wkstn)" <55839002+SSL32081@users.noreply.github.com>
Date: Tue, 15 Apr 2025 15:10:07 +0800
Subject: [PATCH 2/3] Fix jax.where to jnp.where

---
 src/ripplegw/waveforms/IMRPhenomPv2_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/ripplegw/waveforms/IMRPhenomPv2_utils.py b/src/ripplegw/waveforms/IMRPhenomPv2_utils.py
index fc680b93..095a2701 100644
--- a/src/ripplegw/waveforms/IMRPhenomPv2_utils.py
+++ b/src/ripplegw/waveforms/IMRPhenomPv2_utils.py
@@ -95,7 +95,7 @@ def convert_spins(
     thetaJ_sf = jnp.arccos(J0z_sf / J0)
 
     no_inplane_J0 = jnp.abs(J0x_sf) < MAX_TOL_ATAN and jnp.abs(J0y_sf) < MAX_TOL_ATAN
-    phiJ_sf = jax.where(
+    phiJ_sf = jnp.where(
         no_inplane_J0,
         jnp.pi / 2.0 - phiRef,  # This is the aligned-spin case
         jnp.arctan2(J0y_sf, J0x_sf),

From 9c2f2213a8bf8a09bbe0be8c90807e7143e2cba5 Mon Sep 17 00:00:00 2001
From: "Samson Leong (wkstn)" <55839002+SSL32081@users.noreply.github.com>
Date: Wed, 16 Apr 2025 03:05:09 +0800
Subject: [PATCH 3/3] Switch to use `&` instead of `and`

---
 src/ripplegw/waveforms/IMRPhenomPv2_utils.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/ripplegw/waveforms/IMRPhenomPv2_utils.py b/src/ripplegw/waveforms/IMRPhenomPv2_utils.py
index 095a2701..5aa51b28 100644
--- a/src/ripplegw/waveforms/IMRPhenomPv2_utils.py
+++ b/src/ripplegw/waveforms/IMRPhenomPv2_utils.py
@@ -94,7 +94,7 @@ def convert_spins(
 
     thetaJ_sf = jnp.arccos(J0z_sf / J0)
 
-    no_inplane_J0 = jnp.abs(J0x_sf) < MAX_TOL_ATAN and jnp.abs(J0y_sf) < MAX_TOL_ATAN
+    no_inplane_J0 = (jnp.abs(J0x_sf) < MAX_TOL_ATAN) & (jnp.abs(J0y_sf) < MAX_TOL_ATAN)
     phiJ_sf = jnp.where(
         no_inplane_J0,
         jnp.pi / 2.0 - phiRef,  # This is the aligned-spin case
@@ -124,7 +124,7 @@ def convert_spins(
     tmp_x, tmp_y, tmp_z = ROTATEY(-thetaJ_sf, tmp_x, tmp_y, tmp_z)
     tmp_x, tmp_y, tmp_z = ROTATEZ(kappa, tmp_x, tmp_y, tmp_z)
 
-    no_inplane_LN = jnp.abs(tmp_x) < MAX_TOL_ATAN and jnp.abs(tmp_y) < MAX_TOL_ATAN
+    no_inplane_LN = (jnp.abs(tmp_x) < MAX_TOL_ATAN) & (jnp.abs(tmp_y) < MAX_TOL_ATAN)
     alpha0 = jnp.where(
         no_inplane_LN,
         jnp.pi,  # This is the aligned-spin case