Add a module for computing null geodesics and similarity coordinates.

1 files changed, 260 insertions, 0 deletions

diff --git a/null.py b/null.py
new file mode 100644
index 0000000..05f2e9e
--- /dev/null
+++ b/null.py
@@ -0,0 +1,260 @@
+# null coordinates calculations
+
+import numpy as np
+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):
+    """
+    Right-hand-side of the null geodesic equation along an axis of symmetry.
+    Assumes zero shift.
+
+    The curve is given by X^μ(λ) = { t(λ), x(λ), 0, 0 }, the geodesic equation then is:
+        d X^μ / dλ = p^μ
+        d p^μ / dλ = -Γ^μ_αβ p^α p^β
+
+    In components:
+        dt / dλ = p^t
+        dx / dλ = p^x
+
+        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)
+
+    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
+        Γ^t_xx = -0.5 g^tt g_xx,t
+        Γ^x_tt = -0.5 g^xx g_tt,x
+        Γ^x_xt =  0.5 g^xx g_xx,t
+        Γ^x_xx =  0.5 g^xx g_xx,x
+    """
+    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]
+
+    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
+
+    mom_t = np.sqrt(-gXX / gtt) * mom_X
+
+    gtt_dt = -2.0 * alpha_dt * alpha
+    gtt_dX = -2.0 * alpha_dX * alpha
+
+    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
+
+    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)
+
+    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]
+    gXX_0   = gXX(t0, 0.0).flatten()[0]
+
+    pu_X_0 = np.sqrt(alpha_0 * alpha_0 / gXX_0) * pu_t_0
+
+    args = {
+        'fun'           : _null_geodesic_axis_rhs,
+        't_span'        : (0.0, 1e6),
+        'y0'            : (t0, 0.0, pu_t_0, pu_X_0),
+        'method'        : 'RK45',
+        'dense_output'  : True,
+        'events'        : events,
+        'args'          : (alpha, alpha_dX, alpha_dt, gXX, gXX_dX, gXX_dt),
+        'rtol'          : 1e-4,
+        'atol'          : 1e-6,
+        }
+
+    sol_fwd = None
+    if t0 < times[-1]:
+        sol = solve_ivp(**args)
+        if sol.status != 1:
+            raise ValueError('solver status from t0=%g: %d' % t0, sol.status)
+
+        sol_fwd = sol.sol
+
+    sol_back = None
+    if t0 > times[0]:
+        args['t_span'] = (0.0, -1e6)
+        sol = solve_ivp(**args)
+        if sol.status != 1:
+            raise ValueError('solver status from t0=%g: %d' % t0, sol.status)
+
+        sol_back = sol.sol
+
+    if sol_fwd is not None and sol_back is not None:
+        # combine forward and backward solutions
+        return OdeSolution(np.concatenate((sol_back.ts[::-1][:-1], sol_fwd.ts)),
+                           sol_back.interpolants[::-1] + sol_fwd.interpolants)
+    elif sol_fwd is not None:
+        return sol_fwd
+    return sol_back
+
+# terminate integration on reaching the simulation time boundaries
+def _events_bnd_null_geodesics(times, spatial_coords):
+    def event_t_bound_upper(Lambda, y, *args):
+        t = y[0]
+        return t - times[-1]
+    event_t_bound_upper.terminal  = True
+    event_t_bound_upper.direction = 1.0
+
+    def event_t_bound_lower(Lambda, y, *args):
+        t = y[0]
+        return t - times[0]
+    event_t_bound_lower.terminal  = True
+    event_t_bound_lower.direction = -1.0
+
+    def event_x_bound_upper(Lambda, y, *args):
+        x = y[1]
+        return x - spatial_coords[-1]
+    event_x_bound_upper.terminal  = True
+    event_x_bound_upper.direction = 1.0
+
+    def event_x_bound_lower(Lambda, y, sign, *args):
+        x = y[1]
+        return x - spatial_coords[0]
+    event_x_bound_lower.terminal  = True
+    event_x_bound_lower.direction = -1.0
+
+    return [event_t_bound_upper, event_t_bound_lower, event_x_bound_upper, event_x_bound_lower]
+
+def null_geodesics(times, X, gXX, alpha,
+                   pu_t_0 = None, interp_order = 6):
+    """
+    Compute null geodesics along a axis, assuming zero shift.
+
+    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
+                             are provided, must be uniform
+    :param array_like X: 1D array of spatial coordinates at which gXX and alpha
+                         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 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.
+
+    :return: A len(times)-sized list of scipy.integrate.OdeSolution, each
+             describing the null geodesic shot from the corresponding coordinate
+             time.
+    """
+
+    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)
+
+    events = _events_bnd_null_geodesics(times, X)
+
+    ret = []
+    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))
+
+    return ret
+
+def coord_similarity(times, X, T, gXX, alpha,
+                     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
+                             are provided, must be uniform
+    :param array_like X: 1D array of spatial coordinates at which gXX and alpha
+                         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 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 λ.
+    :param lambda_points: number of points in the λ direction
+
+    :return: Tuple of (T, λ, X(T, λ), t(T, λ)).
+             T and λ are 1D arrays (not necessarily uniform) specifying the
+             similarity coordinates.
+             X() and t() are 2D arrays with shapes (T.shape + λ.shape)
+             containing the values of X and t at corresponding (T, λ)
+             coordinates.
+    """
+    if g is None:
+        g = null_geodesics(times, X, gXX, alpha)
+
+    sim_coord_valid = np.isfinite(T)
+    T_valid = T[sim_coord_valid]
+    t_valid = times[sim_coord_valid]
+
+    if lambda_max is None:
+        g_valid    = np.array(g, dtype = np.object)[sim_coord_valid]
+        lambda_max = min([c.t_max for c in g_valid])
+
+    lambda_uniform = np.linspace(0., lambda_max, lambda_points)
+
+    dT_dt = diff.fd8(T_valid, 0, t_valid[1] - t_valid[0])
+
+    X_of_Tl = np.empty(T_valid.shape + lambda_uniform.shape)
+    t_of_Tl = np.empty_like(X_of_Tl)
+
+    for i in range(X_of_Tl.shape[0]):
+        gg         = g[i]
+        t_val      = t_valid[i]
+        T_val      = T_valid[i]
+
+        # the geodesics integrated above normalize the affine parameter λ so that
+        #   dt/dλ (λ=0) = 1
+        # which makes the result dependent on the time coordinate
+        # to make this coordinate-invariant we rescale the affine parameter for
+        # each geodesic such that initially
+        #  (dt/dλ) = 1 / (dT /dt)
+        # cf. Baumgarte, Gundlach, Hilditch 'Critical phenomena in the
+        # gravitational collapse of electromagnetic waves' (2019), third-to-last
+        # paragraph on page 2
+        lambda_scaled = lambda_uniform / dT_dt[i]
+
+        vals = gg(lambda_scaled)[:2]
+        t_of_Tl[i] = np.where(lambda_scaled <= gg.t_max, vals[0], np.nan)
+        X_of_Tl[i] = np.where(lambda_scaled <= gg.t_max, vals[1], np.nan)
+
+    return T_valid, lambda_uniform, X_of_Tl, t_of_Tl