null: fix the null-vector constraint with non-zero shift

1 files changed, 2 insertions, 1 deletions

diff --git a/null.py b/null.py
index 13ed6d2..407987f 100644
--- a/null.py
+++ b/null.py
@@ -62,7 +62,8 @@ def _null_geodesic_axis_rhs(Lambda, f,
     gu = np.linalg.inv(g)
 
     # enforce the momentum to be a null vector
-    mom_t = np.sqrt(-gXX / gtt) * mom_X
+    termX = gtX * mom_X
+    mom_t = (-termX + np.sqrt(termX * termX - gtt * gXX * mom_X * mom_X)) / gtt
 
     G = np.zeros((4, 4, 4))
     G[0, 0, 0] =  0.5 * gtt_dt