null: add support for non-zero shift

1 files changed, 92 insertions, 52 deletions

diff --git a/null.py b/null.py
index 77ba068..eb49c4d 100644
--- a/null.py
+++ b/null.py
@@ -6,11 +6,13 @@ from scipy.integrate import solve_ivp, OdeSolution
 from . import diff
 from . import interp
 
-def _null_geodesic_axis_rhs(Lambda, f, alpha_i, alpha_dX_i, alpha_dt_i,
-                            gXX_i, gXX_dX_i, gXX_dt_i):
+def _null_geodesic_axis_rhs(Lambda, f,
+                            gtt_i, gtt_dX_i, gtt_dt_i,
+                            gXX_i, gXX_dX_i, gXX_dt_i,
+                            gtX_i, gtX_dX_i, gtX_dt_i):
     """
-    Right-hand-side of the null geodesic equation along an axis of symmetry.
-    Assumes zero shift.
+    Right-hand-side of the null geodesic equation along an axis of symmetry or
+    reflection-symmetric equatorial plane.
 
     The curve is given by X^μ(λ) = { t(λ), x(λ), 0, 0 }, the geodesic equation then is:
         d X^μ / dλ = p^μ
@@ -23,6 +25,10 @@ def _null_geodesic_axis_rhs(Lambda, f, alpha_i, alpha_dX_i, alpha_dt_i,
         dp^t / dλ = -(Γ^t_tt (p^t)^2 + 2 Γ^t_xt p^x p^t + Γ^t_xx (p^x)^2)
         dp^x / dλ = -(Γ^x_tt (p^t)^2 + 2 Γ^x_xt p^x p^t + Γ^x_xx (p^x)^2)
 
+    As p^μ should be null, it also holds
+        0 = g_μν p^μ p^ν = g_tt (p^t)^2 + 2 g_tx p^t p^x + g_xx (p^x)^2
+     -> p^t = (-g_tx p^x ± √((g_tx p^x)^2 - g_tt g_xx (p^x)^2)) / g_tt
+
     With zero shift, it can be easily computed that
         Γ^t_tt =  0.5 g^tt g_tt,t
         Γ^t_xt =  0.5 g^tt g_tt,x
@@ -33,46 +39,63 @@ def _null_geodesic_axis_rhs(Lambda, f, alpha_i, alpha_dX_i, alpha_dt_i,
     """
     t, X, mom_t, mom_X = f
 
-    alpha    = alpha_i   (t, X).flatten()[0]
-    alpha_dt = alpha_dt_i(t, X).flatten()[0]
-    alpha_dX = alpha_dX_i(t, X).flatten()[0]
+    gtt    = gtt_i   (t, X).flatten()[0]
+    gtt_dt = gtt_dt_i(t, X).flatten()[0]
+    gtt_dX = gtt_dX_i(t, X).flatten()[0]
 
     gXX    = gXX_i   (t, X).flatten()[0]
     gXX_dt = gXX_dt_i(t, X).flatten()[0]
     gXX_dX = gXX_dX_i(t, X).flatten()[0]
 
-    gtt = -alpha * alpha
-    gutt = 1.0 / gtt
-    guXX = 1.0 / gXX
-
+    gtX    = gtX_i   (t, X).flatten()[0]
+    gtX_dt = gtX_dt_i(t, X).flatten()[0]
+    gtX_dX = gtX_dX_i(t, X).flatten()[0]
+
+    # set the unavailable (but unneeded) diagonal elements
+    # to NaN, to ensure they are actually unneeded
+    # off-diagonal elements other than gtX should be zero
+    # due to symmetry
+    g = np.array([[gtt, gtX,    0.0, 0.0],
+                  [gtX, gXX,    0.0, 0.0],
+                  [0.0, 0.0, np.inf, 0.0],
+                  [0.0, 0.0, 0.0, np.inf]])
+    gu = np.linalg.inv(g)
+
+    # enforce the momentum to be a null vector
     mom_t = np.sqrt(-gXX / gtt) * mom_X
 
-    gtt_dt = -2.0 * alpha_dt * alpha
-    gtt_dX = -2.0 * alpha_dX * alpha
+    G = np.zeros((4, 4, 4))
+    G[0, 0, 0] =  0.5 * gtt_dt
+    G[0, 1, 1] = -0.5 * gXX_dt + gtX_dX
+    G[0, 0, 1] =  0.5 * gtt_dX
+    G[0, 1, 0] =  G[0, 0, 1]
 
-    Gttt =  0.5 * gutt * gtt_dt
-    GtXt =  0.5 * gutt * gtt_dX
-    GtXX = -0.5 * gutt * gXX_dt
-    GXtt = -0.5 * guXX * gtt_dX
-    GXXt =  0.5 * guXX * gXX_dt
-    GXXX =  0.5 * guXX * gXX_dX
+    G[1, 1, 1] =  0.5 * gXX_dX
+    G[1, 0, 0] = -0.5 * gtt_dX + gtX_dt
+    G[1, 0, 1] =  0.5 * gXX_dt
+    G[1, 1, 0] =  G[1, 0, 1]
 
-    rhs_mom_t = -(Gttt * mom_t * mom_t + 2.0 * GtXt * mom_X * mom_t + GtXX * mom_X * mom_X)
-    rhs_mom_X = -(GXtt * mom_t * mom_t + 2.0 * GXXt * mom_X * mom_t + GXXX * mom_X * mom_X)
+    Gu = np.einsum('ij,jkl->ikl', gu, G)
+
+    rhs_mom_t = -(Gu[0, 0, 0] * mom_t * mom_t + 2.0 * Gu[0, 1, 0] * mom_X * mom_t + Gu[0, 1, 1] * mom_X * mom_X)
+    rhs_mom_X = -(Gu[1, 0, 0] * mom_t * mom_t + 2.0 * Gu[1, 1, 0] * mom_X * mom_t + Gu[1, 1, 1] * mom_X * mom_X)
 
     return np.array([mom_t, mom_X, rhs_mom_t, rhs_mom_X])
 
 def _null_geodesic_kernel(times, t0, events, pu_t_0,
-                          alpha, alpha_dX, alpha_dt,
-                          gXX,     gXX_dX,   gXX_dt):
-    alpha_0 = alpha(t0, 0.0).flatten()[0]
+                          gtt,     gtt_dX,   gtt_dt,
+                          gXX,     gXX_dX,   gXX_dt,
+                          gtX,     gtX_dX,   gtX_dt):
+    gtt_0   = gtt(t0, 0.0).flatten()[0]
     gXX_0   = gXX(t0, 0.0).flatten()[0]
 
-    pu_X_0 = np.sqrt(alpha_0 * alpha_0 / gXX_0) * pu_t_0
+    pu_X_0 = np.sqrt(-gtt_0 / gXX_0) * pu_t_0
 
     args = {
         # workaround for older versions of scipy not supporting args for fun
-        'fun'           : lambda t,x:_null_geodesic_axis_rhs(t,x,alpha, alpha_dX, alpha_dt, gXX, gXX_dX, gXX_dt),
+        'fun'           : lambda t, x: \
+            _null_geodesic_axis_rhs(t, x, gtt, gtt_dX, gtt_dt, gXX, gXX_dX, gXX_dt,
+                                    gtX, gtX_dX, gtX_dt),
         't_span'        : (0.0, 1e6),
         'y0'            : (t0, 0.0, pu_t_0, pu_X_0),
         'method'        : 'RK45',
@@ -135,7 +158,8 @@ def _events_bnd_null_geodesics(times, spatial_coords):
 
     return [event_t_bound_upper, event_t_bound_lower, event_x_bound_upper, event_x_bound_lower]
 
-def null_geodesics(times, X, gXX, alpha,
+def null_geodesics(times, X, gXX, gtt,
+                   gtX = None,
                    pu_t_0 = None, interp_order = 6):
     """
     Compute null geodesics along a axis, assuming zero shift.
@@ -143,22 +167,25 @@ def null_geodesics(times, X, gXX, alpha,
     For every coordinate time times[i], integrate a null geodesic from
     (t=times[i], x=0) forward and backward in time.
 
-    :param array_like times: 1D array of coordinate times at which gXX and alpha
+    :param array_like times: 1D array of coordinate times at which gXX and gtt
                              are provided, must be uniform
-    :param array_like X: 1D array of spatial coordinates at which gXX and alpha
+    :param array_like X: 1D array of spatial coordinates at which gXX and gtt
                          are provided, must be uniform
     :param array_like gXX: 2D array containing the values of the XX component of
                            the spacetime metric. gXX[i, j] is the value at
                            spacetime point (t=times[i], x=X[j])
-    :param array_like alpha: 2D array containing the values of the lapse,
-                             analogous to gXX.
+    :param array_like gtt: 2D array containing the values of the tt component of
+                           the spacetime metric, analogous to gXX.
+    :param array_like gtX: 2D array containing the values of the tX component of
+                           the spacetime metric, analogous to gXX. Assumed to be
+                           zero if not provided.
     :param array_like pu_t_0: 1D array of the same shape as times containing
                               initial time component of the 4-momentum p^t
                               for the null geodesic shot from each time
                               coordinate.
                               Defaults to 1.0 for each geodesic if not supplied.
     :param int interp_order: order of the Lagrange interpolation used for gXX
-                             and alpha.
+                             and gtt.
 
     :return: A len(times)-sized list of scipy.integrate.OdeSolution, each
              describing the null geodesic shot from the corresponding coordinate
@@ -168,17 +195,25 @@ def null_geodesics(times, X, gXX, alpha,
     dt = times[1] - times[0]
     dX = X[1]     - X[0]
 
-    gXX_dt   = diff.fd8(gXX,   0, dt)
-    gXX_dX   = diff.fd8(gXX,   1, dX)
-    alpha_dt = diff.fd8(alpha, 0, dt)
-    alpha_dX = diff.fd8(alpha, 1, dX)
-
-    gXX_interp      = interp.Interp2D_C([times[0], X[0]], [dt, dX], gXX,      interp_order)
-    gXX_dX_interp   = interp.Interp2D_C([times[0], X[0]], [dt, dX], gXX_dX,   interp_order)
-    gXX_dt_interp   = interp.Interp2D_C([times[0], X[0]], [dt, dX], gXX_dt,   interp_order)
-    alpha_interp    = interp.Interp2D_C([times[0], X[0]], [dt, dX], alpha,    interp_order)
-    alpha_dX_interp = interp.Interp2D_C([times[0], X[0]], [dt, dX], alpha_dX, interp_order)
-    alpha_dt_interp = interp.Interp2D_C([times[0], X[0]], [dt, dX], alpha_dt, interp_order)
+    if gtX is None:
+        gtX = np.zeros_like(gXX)
+
+    gXX_dt   = diff.fd8(gXX, 0, dt)
+    gXX_dX   = diff.fd8(gXX, 1, dX)
+    gtt_dt   = diff.fd8(gtt, 0, dt)
+    gtt_dX   = diff.fd8(gtt, 1, dX)
+    gtX_dt   = diff.fd8(gtX, 0, dt)
+    gtX_dX   = diff.fd8(gtX, 1, dX)
+
+    gXX_interp      = interp.Interp2D_C([times[0], X[0]], [dt, dX], gXX,    interp_order)
+    gXX_dX_interp   = interp.Interp2D_C([times[0], X[0]], [dt, dX], gXX_dX, interp_order)
+    gXX_dt_interp   = interp.Interp2D_C([times[0], X[0]], [dt, dX], gXX_dt, interp_order)
+    gtt_interp      = interp.Interp2D_C([times[0], X[0]], [dt, dX], gtt,    interp_order)
+    gtt_dX_interp   = interp.Interp2D_C([times[0], X[0]], [dt, dX], gtt_dX, interp_order)
+    gtt_dt_interp   = interp.Interp2D_C([times[0], X[0]], [dt, dX], gtt_dt, interp_order)
+    gtX_interp      = interp.Interp2D_C([times[0], X[0]], [dt, dX], gtX,    interp_order)
+    gtX_dX_interp   = interp.Interp2D_C([times[0], X[0]], [dt, dX], gtX_dX, interp_order)
+    gtX_dt_interp   = interp.Interp2D_C([times[0], X[0]], [dt, dX], gtX_dt, interp_order)
 
     events = _events_bnd_null_geodesics(times, X)
 
@@ -186,27 +221,32 @@ def null_geodesics(times, X, gXX, alpha,
     for i, t0 in enumerate(times):
         pu_t_0_val = 1.0 if pu_t_0 is None else pu_t_0[i]
         ret.append(_null_geodesic_kernel(times, t0, events, pu_t_0_val,
-                                         alpha_interp, alpha_dX_interp, alpha_dt_interp,
-                                           gXX_interp,   gXX_dX_interp,   gXX_dt_interp))
+                                           gtt_interp, gtt_dX_interp, gtt_dt_interp,
+                                           gXX_interp, gXX_dX_interp, gXX_dt_interp,
+                                           gtX_interp, gtX_dX_interp, gtX_dt_interp))
 
     return ret
 
-def coord_similarity(times, X, T, gXX, alpha,
-                     g = None, lambda_max = None, lambda_points = 1 << 8):
+def coord_similarity(times, X, T, gXX, gtt,
+                     gtX = None, g = None,
+                     lambda_max = None, lambda_points = 1 << 8):
     """
     Compute null similarity coordinates, assuming zero shift.
 
-    :param array_like times: 1D array of coordinate times at which gXX and alpha
+    :param array_like times: 1D array of coordinate times at which gXX and gtt
                              are provided, must be uniform
-    :param array_like X: 1D array of spatial coordinates at which gXX and alpha
+    :param array_like X: 1D array of spatial coordinates at which gXX and gtt
                          are provided, must be uniform
     :param array_like T: 1D array of slow time T corresponding to coordinate
                          times.
     :param array_like gXX: 2D array containing the values of the XX component of
                            the spacetime metric. gXX[i, j] is the value at
                            spacetime point (t=times[i], x=X[j])
-    :param array_like alpha: 2D array containing the values of the lapse,
-                             analogous to gXX.
+    :param array_like gtt: 2D array containing the values of the tt component of
+                           the spacetime metric analogous to gXX.
+    :param array_like gtX: 2D array containing the values of the tX component of
+                           the spacetime metric analogous to gXX. Assumed to be
+                           zero if not provided.
     :param g: Optional list of pre-computed geodesics, as returned by null_geodesics().
               Computed by this function if not supplied.
     :param lambda_max: maximum value of the affine parameter λ.
@@ -220,7 +260,7 @@ def coord_similarity(times, X, T, gXX, alpha,
              coordinates.
     """
     if g is None:
-        g = null_geodesics(times, X, gXX, alpha)
+        g = null_geodesics(times, X, gXX, gtt, gtX)
 
     sim_coord_valid = np.isfinite(T)
     T_valid = T[sim_coord_valid]