doublenull: add a function for inverse double-null coordinates

I.e. computing T/X(U, V).

3 files changed, 103 insertions, 0 deletions

diff --git a/doublenull.py b/doublenull.py
index e3fe0d1..3794b9b 100644
--- a/doublenull.py
+++ b/doublenull.py
@@ -185,3 +185,78 @@ def null_coordinates(times, spatial_coords, u_rays, v_rays,
         v_of_tx[i] = interp1d(Xv, v_rays, bounds_error = False)(spatial_coords)
 
     return (u_of_tx, v_of_tx)
+
+def null_coordinates_inv(t, X, uv, gXX, gXt, gtt, *,
+                         uv_times = None,
+                         interp_order = 6):
+    """
+    Compute values of (t, X) on a uniform grid in double-null
+    coordinates (U, V).
+
+    :param array_like t: 1D array of coordinate times at which the spacetime
+                         curvature is provided
+    :param array_like X: 1D array of spatial coordinates
+    :param array_like uv: 1D array of values of U/V values on the symmetry axis.
+                          Every value coresponds to the appropriate element in
+                          uv_times, i.e. uv[i] = U(t = uv_times[i], X = 0).
+    :param array_like gXX: 2D array containing the values of the XX component
+                           of the spacetime metric, where X is the spatial
+                           coordinate along which the rays are traced.
+                           gXX[i, j] is the value at spacetime point
+                           (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 array_like uv_times: 1D array of coordinate times corresponding to
+                                uv; may be None, in which case t is used instead
+
+    :return: tuple containing two 2D arrays with, respectively, values of T and
+             X as functions of U and V. Tuv[i, j] is the value of T at
+             U=uv[i], V=uv[j].
+    """
+    t_span = [t[0], t[-1]]
+    dt     = t[1] - t[0]
+
+    if np.any(np.abs((t[1:] - t[:-1]) - dt) > 1e-6 * dt):
+        raise ValueError("Non-uniform t-grid")
+
+    X_span = [X[0], X[-1]]
+    dX     = X[1] - X[0]
+    if np.any(np.abs((X[1:] - X[:-1]) - dX) > 1e-6 * dX):
+        raise ValueError("Non-uniform X-grid")
+
+    uv_times = t if uv_times is None else uv_times
+    if uv.shape != uv_times.shape:
+        raise ValueError("Shape of uv must match uv_times: %s != %s" % (uv.shape, uv_times.shape))
+
+    pos, neg = _null_curves(t_span, dt, X_span, dX, uv_times,
+                            gXX, gXt, gtt,
+                            Curves.TEMPORAL, interp_order)
+
+    # interpolator for V(t, X)
+    # FIXME: integrates the curves again at different times
+    # (uniform in U/V vs uniform in X/t)
+    if uv_times is t:
+        uv_t = uv
+    else:
+        uv_t = interp1d(uv_times, uv, bounds_error = False)(t)
+    _, vtx_vals = null_coordinates(t, X, uv_t, uv_t, gXX, gXt, gtt,
+                                   kind = Curves.TEMPORAL)
+    Vtx = interp.Interp2D_C([t[0], X[0]], [dt, dX], vtx_vals, interp_order)
+
+    Xuv = np.empty(uv.shape * 2)
+    Tuv = np.empty_like(Xuv)
+    for i in range(uv.shape[0]):
+        # X(t) at constant U=uv[i]
+        xt_u = pos[i]
+
+        # uniform t grid along the curve
+        t_uniform = np.linspace(xt_u.t_min, xt_u.t_max, 2 * uv.shape[0])
+
+        # values of V(t, X(t)) at this t-grid along the curve
+        v_vals = Vtx(t_uniform, xt_u(t_uniform))
+
+        # finally invert V(t) into t(V) on a uniform V-grid
+        Tuv[i] = interp1d(v_vals, t_uniform, bounds_error = False)(uv)
+        Xuv[i] = xt_u(Tuv[i])
+
+    return Tuv, Xuv
diff --git a/test/test_doublenull.npz b/test/test_doublenull.npz
index db5fc85..bf7b384 100644
--- a/test/test_doublenull.npz
+++ b/test/test_doublenull.npz
diff --git a/test/test_doublenull.py b/test/test_doublenull.py
index 204870d..41fbc91 100644
--- a/test/test_doublenull.py
+++ b/test/test_doublenull.py
@@ -4,6 +4,9 @@ from unittest import TestCase
 import numpy as np
 from   numpy.testing import assert_allclose
 
+from   scipy.integrate   import cumulative_trapezoid
+from   scipy.interpolate import interp1d
+
 from math_utils      import array_utils
 from nr_analysis_axi import doublenull, invars
 
@@ -43,3 +46,28 @@ class TestDoubleNull(TestCase):
         assert_allclose(coords[1], self.data['vxt'], 1e-12)
 
         #self.update_refs(uxt = coords[0], vxt = coords[1])
+
+    def test_null_coordinates_inv(self):
+        t = self.data['t']
+        x = self.data['x']
+        idx = x.shape[0] // 2
+
+        gxx   = self.data['gxx']
+        alpha = self.data['alpha']
+
+        # proper time τ(t)
+        tau_t = cumulative_trapezoid(alpha[:, idx], t, initial = 0.0)
+
+        # uniform grid in τ
+        tau_uniform = np.linspace(tau_t[0], tau_t[-1], t.shape[0])
+
+        # t on the uniform τ grid
+        t_tau = interp1d(tau_t, t)(tau_uniform)
+
+        Tuv, Xuv = doublenull.null_coordinates_inv(t, x, tau_uniform,
+                                                   gxx, None, -(alpha ** 2),
+                                                   uv_times = t_tau)
+        assert_allclose(Tuv, self.data['Tuv'], 1e-12)
+        assert_allclose(Xuv, self.data['Xuv'], 1e-12)
+
+        #self.update_refs(Tuv = Tuv, Xuv = Xuv)