path: root/null.py

# 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,
                            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 or
    reflection-symmetric equatorial plane.

    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)

    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^x = (-g_tx p^t ± √((g_tx p^t)^2 - g_tt g_xx (p^t)^2)) / g_xx
    This constraint needs to be enforced at runtime to ensure the curve remains
    null.

    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

    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]

    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
    # 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
    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]

    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]

    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,
                          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]
    gtX_0   = gtX(t0, 0.0).flatten()[0]

    term_t = gtX_0 * pu_t_0
    pu_X_0 = (-term_t + np.sqrt(term_t * term_t - gtt_0 * gXX_0 * pu_t_0 * pu_t_0)) / gXX_0

    args = {
        # workaround for older versions of scipy not supporting args for fun
        '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',
        'dense_output'  : True,
        'events'        : events,
        '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, *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, gtt,
                   gtX = None,
                   pu_t_0 = None, interp_order = 6, integrate_times=None):
    """
    Compute null geodesics along a axis, assuming zero shift.

    For every coordinate time integrate_times[i], integrate a null geodesic from
    (t=integrate_times[i], x=0) forward and backward in time.

    :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 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 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 gtt.
    :param array_like integrate_times: 1D array of coordinate times from which
        each geodesic will be integrated; or to put it differntly, of coordinate
        times at which each geodesic hits X=0.

    :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]

    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)

    ret = []
    if integrate_times is None:
        integrate_times = times
    for i, t0 in enumerate(integrate_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,
                                           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, 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 gtt
                             are provided, must be uniform
    :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 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 λ.
    :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, gtt, gtX)

    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