Add RHSes for the apparent horizon equation in axial symmetry.

1 files changed, 128 insertions, 0 deletions

diff --git a/horizon_axi.py b/horizon_axi.py
new file mode 100644
index 0000000..8807f3a
--- /dev/null
+++ b/horizon_axi.py
@@ -0,0 +1,128 @@
+import numpy as np
+import scipy.interpolate
+
+def _invert_matrix(mat):
+    inv = np.empty_like(mat)
+    for i in xrange(mat.shape[-1]):
+        try:
+            inv[:, :, i] = np.linalg.inv(mat[:, :, i])
+        except:
+            #print i, mat[:, :, i]
+            inv[:, :, i] = np.array([[1.0, 0.0, 0.0],
+                                     [0.0, 1.0, 0.0],
+                                     [0.0, 0.0, 1.0]])
+    return inv
+
+def _tensor_transform_cart_spherical_d2(theta, r, tens):
+    z = np.zeros_like(theta)
+    o = z + 1.0
+    jac = np.array([[np.sin(theta), r * np.cos(theta),   z ],
+                    [z,             z,                   r * np.sin(theta)],
+                    [np.cos(theta), -r * np.sin(theta),  z ]])
+    ret = np.zeros_like(tens)
+    for i in xrange(3):
+        for j in xrange(3):
+            for k in xrange(3):
+                for l in xrange(3):
+                    ret[i, j] += jac[k, i] * jac[l, j] * tens[k, l]
+
+    return ret;
+
+def _calc_metric(X, Z, metric, curv, r, theta):
+    R     = np.sqrt(X ** 2 + Z ** 2)
+    Theta = np.where(np.abs(X) > 1e-12, np.arccos(Z / R), 0)
+
+    x = r * np.sin(theta)
+    z = r * np.cos(theta)
+
+    metric_src_spherical = _tensor_transform_cart_spherical_d2(Theta, R, metric)
+    curv_src_spherical   = _tensor_transform_cart_spherical_d2(Theta, R, curv)
+
+    metric_val = np.empty((3, 3) + theta.shape)
+    curv_val   = np.empty_like(metric_val)
+    dg_cart    = np.empty((2,) + metric_val.shape)
+    for i in xrange(3):
+        for j in xrange(3):
+            metric_interp = scipy.interpolate.RectBivariateSpline(Z[:, 0], X[0], metric_src_spherical[i, j])
+            metric_val[i, j] = metric_interp.ev(z, x)
+            dg_cart[0, i, j] = metric_interp.ev(z, x, 0, 1)
+            dg_cart[1, i, j] = metric_interp.ev(z, x, 1, 0)
+
+            curv_interp    = scipy.interpolate.RectBivariateSpline(Z[:, 0], X[0], curv_src_spherical[i, j])
+            curv_val[i, j] = curv_interp.ev(z, x)
+
+    dg = np.empty_like(dg_cart)
+    dg[0] = dg_cart[0] * np.sin(theta) + dg_cart[1] * np.cos(theta)
+    dg[1] = z * dg_cart[0] - x * dg_cart[1]
+
+    metric_u_val = _invert_matrix(metric_val)
+
+    Gamma_r_val     = np.zeros_like(metric_val)
+    Gamma_theta_val = np.zeros_like(metric_val)
+    for i in xrange(3):
+        for j in xrange(3):
+            val_r     = np.zeros_like(theta)
+            val_theta = np.zeros_like(theta)
+            if i != 2:
+                val_r     += dg[i, j, 0]
+                val_theta += dg[i, j, 1]
+            if j != 2:
+                val_r     += dg[j, i, 0]
+                val_theta += dg[j, i, 1]
+            val_r     -= dg[0, i, j]
+            val_theta -= dg[1, i, j]
+            Gamma_r_val[i, j]     = 0.5 * (metric_u_val[0, 0] * val_r     + metric_u_val[0, 1] * val_theta)
+            Gamma_theta_val[i, j] = 0.5 * (metric_u_val[1, 1] * val_theta + metric_u_val[0, 1] * val_r)
+
+    return metric_val, metric_u_val, curv_val, Gamma_r_val, Gamma_theta_val
+
+
+def _horiz_axi_F(theta, u, args):
+    r, dr = u
+
+    X, Z, metric, curv = args
+    metric_val, metric_u_val, curv_val, Gamma_r_val, Gamma_theta_val = _calc_metric(X, Z, metric, curv, r, theta)
+
+    u2 = metric_u_val[1, 1] * dr * dr - 2 * metric_u_val[0, 1] * dr + metric_u_val[0, 0]
+
+    df_ij = np.empty_like(metric_u_val)
+    for i in xrange(3):
+        for j in xrange(3):
+            df_ij[i, j] = (metric_u_val[0, i] - metric_u_val[1, i] * dr) * (metric_u_val[0, j] - metric_u_val[1, j] * dr)
+
+    term1 = u2 * metric_u_val - df_ij
+    term2 = Gamma_r_val - dr * Gamma_theta_val + np.sqrt(u2) * curv_val
+
+    return -np.einsum('ij...,ij...', term1, term2) / (metric_u_val[0, 0] * metric_u_val[1, 1] - metric_u_val[0, 1] * metric_u_val[0, 1])
+
+def _horiz_axi_dF_r(theta, u, args):
+    eps = 1e-4
+    r, R = u
+
+    rp = r + eps
+    rm = r - eps
+    Fp = _horiz_axi_F(theta, (rp, R), args)
+    Fm = _horiz_axi_F(theta, (rm, R), args)
+    return (Fp - Fm) / (2 * eps)
+
+def _horiz_axi_dF_R(theta, u, args):
+    eps = 1e-4
+    r, R = u
+
+    Rp = R + eps
+    Rm = R - eps
+    Fp = _horiz_axi_F(theta, (r, Rp), args)
+    Fm = _horiz_axi_F(theta, (r, Rm), args)
+    return (Fp - Fm) / (2 * eps)
+
+"""
+The right-hand side function (and its derivatives wrt u and u') defining the apparent horizon equation in axial
+symmetry. Intended to be passed to the nonlin_solve_1d solver.
+
+The extra args is a tuple of:
+    - a meshgrid of the cartesian r coordinates
+    - a meshgrid of the cartesian z coordinates
+    - the spatial metric in cartesian coordinates at the provided coordinates -- a numpy aray of (3, 3) + z.shape
+    - the extrinsic curvature in cartesian coordinates at the provided coordinates -- a numpy aray of (3, 3) + z.shape
+"""
+Fs = (_horiz_axi_F, _horiz_axi_dF_r, _horiz_axi_dF_R)