doublenull: add a new null curve tracking mode

Trace from the spatial origin and each time point backwards and forwards
in time.

1 files changed, 105 insertions, 47 deletions

diff --git a/doublenull.py b/doublenull.py
index 6d11bd4..241a3a5 100644
--- a/doublenull.py
+++ b/doublenull.py
@@ -1,13 +1,100 @@
+from enum import Enum, auto
+
 import numpy as np
 from scipy.interpolate import RectBivariateSpline, interp1d
 from scipy.integrate import solve_ivp
 
-def calc_null_curves(times, spatial_coords, gXX, gXt, gtt, reverse = False):
+def _photon_dXdt(t, coord, sign, gXX, gtt, gXt):
+    """
+    Null curve equation RHS (not a geodesic - not affinely parametrized).
+
+    gXX dX^2 + 2 gXt dX dt - gtt dt^2 = 0
+     => dx / dt = (-gXt ± √(gXt^2 - gXX gtt)) / gXX
+    """
+    gXt_val = gXt(t, coord)
+    gXX_val = gXX(t, coord)
+    gtt_val = gtt(t, coord)
+    return ((-gXt_val + sign * np.sqrt((gXt_val ** 2) - gtt_val * gXX_val)) / gXX_val).flatten()[0]
+
+# terminate integration on reaching the outer boundaries
+def _events_bnd_null_curves(spatial_coords):
+    def event_x_bound_upper(t, x, sign, *args):
+        return x[0] - spatial_coords[-1]
+    event_x_bound_upper.terminal  = True
+    event_x_bound_upper.direction = 1.0
+
+    def event_x_bound_lower(t, x, sign, *args):
+        return x[0] - spatial_coords[0]
+    event_x_bound_lower.terminal  = True
+    event_x_bound_lower.direction = -1.0
+
+    return [event_x_bound_upper, event_x_bound_lower]
+
+def _calc_null_kernel_time(times, origin, sign, gXX, gXt, gtt, events):
+    idx = np.where(times == origin)[0][0]
+
+    t_fwd  = times[idx:]
+    t_back = times[:idx + 1][::-1]
+
+    ray_fwd = [0.0]
+    if len(t_fwd) > 1:
+        sol = solve_ivp(_photon_dXdt, (t_fwd[0], t_fwd[-1]), (0,),
+                        method = 'RK45', t_eval = t_fwd,
+                        args = (sign, gXX, gtt, gXt),
+                        dense_output = True, events = events,
+                        rtol = 1e-6, atol = 1e-8)
+        ray_fwd = sol.y[0]
+        ray_fwd = np.concatenate((ray_fwd, [ray_fwd[-1]] * (len(t_fwd) - len(ray_fwd))))
+
+    ray_back = []
+    if len(t_back) > 1:
+        sol = solve_ivp(_photon_dXdt, (t_back[0], t_back[-1]), (0,),
+                            method = 'RK45', t_eval = t_back,
+                            args = (sign, gXX, gtt, gXt),
+                            dense_output = True, events = events,
+                            rtol = 1e-6, atol = 1e-8)
+        ray_back = sol.y[0]
+        ray_back = np.concatenate((ray_back, [ray_back[-1]] * (len(t_back) - len(ray_back))))
+
+    return np.concatenate((ray_back[::-1][:-1], ray_fwd))
+
+def _calc_null_kernel_space(times, origin, sign, gXX, gXt, gtt, events):
+    sol = solve_ivp(_photon_dXdt, (times[0], times[-1]), (origin,),
+                    method = 'RK45', t_eval = times, args = (sign, gXX, gtt, gXt),
+                    dense_output = True, events = events, rtol = 1e-6, atol = 1e-8)
+
+    t, x = sol.t, sol.y[0]
+
+    if len(t) < len(times):
+        x_ext = np.empty_like(times)
+        x_ext[:x.shape[0]] = x
+        x_ext[x.shape[0]:] = x[-1] + sign * (times[x.shape[0]:] - t[-1])
+
+        x = x_ext
+
+    return x
+
+class Curves(Enum):
+    SPATIAL_FORWARD = auto()
+    "photons are traced forward in time from (t=times[0], X=spatial_coords[i])"
+
+    SPATIAL_BACK    = auto()
+    "photons are traced backward in time from (t=times[-1], X=spatial_coords[i])"
+
+    TEMPORAL        = auto()
+    """
+    photons are traced both forward and backward in time from
+     (t=times[i], X=0); the forward and backward segment are combined
+     to form the resulting curve
+    """
+
+def calc_null_curves(times, spatial_coords, gXX, gXt, gtt,
+                     kind = Curves.SPATIAL_FORWARD):
     """
     Compute null curves along a given axis.
 
     Shoot a null ray from each point in spatial_coords, in the positive and
-    negative direction and compute its trajectory.
+    negative spatial direction and compute its trajectory.
 
     :param array_like times: 1D array of coordinate times at which the spacetime
                              curvature is provided
@@ -19,8 +106,7 @@ def calc_null_curves(times, spatial_coords, gXX, gXt, gtt, reverse = False):
                            (t=times[i], X=spatial_coords[j]).
     :param array_like gXt: same as gXX, but for the Xt component of the metric
     :param array_like gtt: same as gXX, but for the tt component of the metric
-    :param bool reverse:   when true, the null curves are traced from the last
-                           time coordinate backwards in time
+    :param Curves kind:    specifies where to integrate the curves from
     :return: Tuple of (ray_times, rays_pos, rays_neg). rays_*[i, j] contains the
              X-coordinate of the ray shot from (t=ray_times[0],
              X=spatial_coords[i]) at time t=ray_times[j].
@@ -30,56 +116,28 @@ def calc_null_curves(times, spatial_coords, gXX, gXt, gtt, reverse = False):
     gXt_interp = RectBivariateSpline(times, spatial_coords, gXt)
     gtt_interp = RectBivariateSpline(times, spatial_coords, gtt)
 
-    if reverse:
-        ray_times = times[::-1]
+    if kind == Curves.TEMPORAL:
+        kernel  = _calc_null_kernel_time
+        origins = times
     else:
-        ray_times = times
-
-    # terminate integration on reaching the outer boundaries
-    def event_x_bound_upper(t, x, sign):
-        return x[0] - spatial_coords[-1]
-    event_x_bound_upper.terminal  = True
-    event_x_bound_upper.direction = 1.0
-
-    def event_x_bound_lower(t, x, sign):
-        return x[0] - spatial_coords[0]
-    event_x_bound_lower.terminal  = True
-    event_x_bound_lower.direction = -1.0
+        kernel  = _calc_null_kernel_space
+        origins = spatial_coords
 
-    events = [event_x_bound_upper, event_x_bound_lower]
+    ray_times = times[::-1] if kind == Curves.SPATIAL_BACK else times
+    events    = _events_bnd_null_curves(spatial_coords)
 
-    # null geodesic equation RHS:
-    #    gXX dX^2 + 2 gXt dX dt - gtt dt^2 = 0
-    # => dx / dt = (-gXt ± √(gXt^2 - gXX gtt)) / gXX
-    def dXdt(t, coord, sign, gXX = gXX_interp, gtt = gtt_interp, gXt = gXt_interp):
-        gXt_val = gXt(t, coord)
-        gXX_val = gXX(t, coord)
-        gtt_val = gtt(t, coord)
-        return ((-gXt_val + sign * np.sqrt((gXt_val ** 2) - gtt_val * gXX_val)) / gXX_val).flatten()[0]
-
-    rays_pos = np.empty((spatial_coords.shape[0], times.shape[0]))
+    rays_pos = np.empty((origins.shape[0], times.shape[0]))
     rays_neg = np.empty_like(rays_pos)
-    for j, X0 in enumerate(spatial_coords):
-        for tgt, sign in ((rays_pos, 1.0), (rays_neg, -1.0)):
-            ret = solve_ivp(dXdt, (ray_times[0], ray_times[-1]), (X0,),
-                            method = 'RK45', t_eval = ray_times, args = (sign,),
-                            dense_output = True, events = events, rtol = 1e-6, atol = 1e-8)
-
-            t, x = ret.t, ret.y[0]
-
-            if len(t) < len(ray_times):
-                x_ext = np.empty_like(ray_times)
-                x_ext[:x.shape[0]] = x
-                x_ext[x.shape[0]:] = x[-1] + sign * (ray_times[x.shape[0]:] - t[-1])
-
-                x = x_ext
-
-            tgt[j] = x
+    for tgt, sign in ((rays_pos, 1.0), (rays_neg, -1.0)):
+        for i, origin in enumerate(origins):
+            tgt[i] = kernel(ray_times, origin, sign,
+                            gXX_interp, gXt_interp, gtt_interp,
+                            events)
 
     return (ray_times, rays_pos, rays_neg)
 
 def calc_null_coordinates(times, spatial_coords, u_rays, v_rays,
-                          gXX, gXt, gtt):
+                          gXX, gXt, gtt, kind = Curves.SPATIAL_FORWARD):
     """
     Compute double-null coordinates (u, v) as functions of
     position and time.
@@ -102,7 +160,7 @@ def calc_null_coordinates(times, spatial_coords, u_rays, v_rays,
              v as functions of t and X. u/v[i, j] is the value of u/v at
              t=times[i], X=spatial_coords[j].
     """
-    _, X_of_ut, X_of_vt = calc_null_curves(times, spatial_coords, gXX, gXt, gtt)
+    _, X_of_ut, X_of_vt = calc_null_curves(times, spatial_coords, gXX, gXt, gtt, kind = kind)
 
     u_of_tx = np.empty((times.shape[0], spatial_coords.shape[0]))
     v_of_tx = np.empty_like(u_of_tx)