1 files changed, 3 insertions, 2 deletions
diff --git a/null.py b/null.py
index c0f08ee..22f1860 100644
--- a/null.py
+++ b/null.py
@@ -62,8 +62,9 @@ def _null_geodesic_axis_rhs(Lambda, f,
     gu = np.linalg.inv(g)
 
     # enforce the momentum to be a null vector
-    termX = gtX * mom_X
-    mom_t = (-termX + np.sqrt(termX * termX - gtt * gXX * mom_X * mom_X)) / gtt
+    # otherwise the curve stops being null rapidly:
+    term_sqrt  = np.sqrt(gtX * gtX - gtt * gXX)
+    mom_X      = (- gtX + term_sqrt ) * mom_t / gXX
 
     G = np.zeros((4, 4, 4))
     G[0, 0, 0] =  0.5 * gtt_dt