horizon: properly calculate rhs at θ=0

1 files changed, 119 insertions, 58 deletions

diff --git a/horizon.py b/horizon.py
index c137906..ea491ea 100644
--- a/horizon.py
+++ b/horizon.py
@@ -82,6 +82,9 @@ def _tensor_transform_cart_spherical_d2(theta, r, x, z, tens):
     return np.einsum('ai...,bj...,ab...->ij...', jac, jac, tens)
 
 def _tensor_transform_cart_spherical_u2(theta, r, x, z, tens):
+    eps = 1e-10 # small number to replace NaNs with just huge numbers
+    x = np.where(np.abs(x) > eps, x, eps)
+
     zero = np.zeros_like(theta)
     st = np.sin(theta)
     ct = np.cos(theta)
@@ -120,6 +123,9 @@ def _dtensor_transform_cart_spherical_d2(theta, r, x, z, tens, dtens):
            np.einsum('bj...,ck...,ai...,abc...->ijk...', jac, jac, jac, dtens)
 
 def _christoffel_transform_cart_spherical(theta, r, x, z, Gamma_cart):
+    eps = 1e-10 # small number to replace NaNs with just huge numbers
+    x = np.where(np.abs(x) > eps, x, eps)
+
     zero = np.zeros_like(theta)
     st = np.sin(theta)
     ct = np.cos(theta)
@@ -152,6 +158,9 @@ def _christoffel_transform_cart_spherical(theta, r, x, z, Gamma_cart):
            np.einsum('ak...,kbc...->abc...', jac_inv, djac)
 
 def _dchristoffel_transform_cart_spherical(theta, r, x, z, Gamma_cart, dGamma_cart):
+    eps = 1e-10 # small number to replace NaNs with just huge numbers
+    x = np.where(np.abs(x) > eps, x, eps)
+
     zero = np.zeros_like(theta)
     st = np.sin(theta)
     ct = np.cos(theta)
@@ -318,6 +327,28 @@ class AHCalc(object):
     _dGamma_cart    = None
     _dcurv_cart     = None
 
+    class _RHSFunc:
+        _hc = None
+
+        def __init__(self, hc):
+            self._hc = hc
+
+        def __call__(self, val = 0.0):
+            rhs      = lambda x, y: self._hc._ah_rhs(x, y, val = val)
+            rhs_var0 = lambda x, y: self._hc._var_diff(x, y, rhs, 0)
+            rhs_var1 = lambda x, y: self._hc._var_diff(x, y, rhs, 1)
+
+            return (rhs, rhs_var0, rhs_var1)
+
+        def __getitem__(self, key):
+            return self()[key]
+
+        def __len__(self):
+            return len(self())
+
+        def __iter__(self):
+            return iter(self())
+
     def __init__(self, X, Z, metric_cart, curv_cart):
         self._diff_op      = diff.fd8
         self._diff2_op     = diff.fd28
@@ -328,7 +359,8 @@ class AHCalc(object):
         self._metric_cart  = metric_cart
         self._curv_cart    = curv_cart
 
-        self.Fs = [self._F, self._dF_r_fd, self._dF_R_fd]
+        self.Fs = [self._ah_rhs, self._dF_r_fd, self._dF_R_fd]
+        self.rhs           = self._RHSFunc(self)
 
     @property
     def X(self):
@@ -396,19 +428,22 @@ class AHCalc(object):
             metric_cart_dst[idx] = self._interp_var(theta, r, x, z, self.metric_cart[idx])
         return _tensor_transform_cart_spherical_d2(theta, r, x, z, metric_cart_dst)
 
-    def _spherical_metric_u(self, theta, r, x, z):
-        metric_u_cart_dst = np.empty((3, 3) + theta.shape)
-        for idx in it.product(range(3), repeat = 2):
-            metric_u_cart_dst[idx] = self._interp_var(theta, r, x, z, self.metric_u_cart[idx])
-        return _tensor_transform_cart_spherical_u2(theta, r, x, z, metric_u_cart_dst)
-
-    def _spherical_dmetric_u(self, theta, r, x, z, metric, metric_u, Gamma):
+    def _spherical_dmetric(self, theta, r, x, z, metric, Gamma):
         # extract the derivatives of the metric from the Christoffel symbols
         Gamma_l = np.einsum('ij...,ikl...->jkl...', metric, Gamma)
         dmetric = np.empty_like(Gamma_l)
         for i, j, k in it.product(range(3), repeat = 3):
             dmetric[i, j, k] = Gamma_l[j, i, k] + Gamma_l[k, i, j]
 
+        return dmetric
+
+    def _spherical_metric_u(self, theta, r, x, z):
+        metric_u_cart_dst = np.empty((3, 3) + theta.shape)
+        for idx in it.product(range(3), repeat = 2):
+            metric_u_cart_dst[idx] = self._interp_var(theta, r, x, z, self.metric_u_cart[idx])
+        return _tensor_transform_cart_spherical_u2(theta, r, x, z, metric_u_cart_dst)
+
+    def _spherical_dmetric_u(self, theta, r, x, z, dmetric, metric_u):
         return -np.einsum('ik...,jl...,mkl...->mij...', metric_u, metric_u, dmetric)
 
     def _spherical_curv(self, theta, r, x, z):
@@ -442,62 +477,87 @@ class AHCalc(object):
 
         return _dchristoffel_transform_cart_spherical(theta, r, x, z, Gamma_cart_dst, dGamma_cart_dst)
 
-    def _F(self, theta, H):
+    def _spherical_trK(self, theta, r, x, z):
+        metric_u = np.empty((3, 3) + theta.shape)
+        curv     = np.empty((3, 3) + theta.shape)
+        for idx in it.product(range(3), repeat = 2):
+            metric_u[idx] = self._interp_var(theta, r, x, z, self.metric_u_cart[idx])
+            curv[idx]     = self._interp_var(theta, r, x, z, self.curv_cart[idx])
+
+        return np.einsum('ij...,ij...', metric_u, curv)
+
+    def _var_diff(self, x, y, func, idx, eps = 1e-8):
+        y_var      = list(y)
+        y_var[idx] = y[idx] + eps
+        f_plus      = func(x, y_var)
+
+        y_var[idx]  = y[idx] - eps
+        f_minus      = func(x, y_var)
+
+        return (f_plus - f_minus) / (2 * eps)
+
+    def _ah_rhs(self, theta, H, val = 0.0):
         r, dr = H
 
+        r2 = r * r
+        st = np.sin(theta)
+        ct = np.cos(theta)
+
         x = r * np.sin(theta)
         z = r * np.cos(theta)
 
+        metric   = self._spherical_metric   (theta, r, x, z)
         metric_u = self._spherical_metric_u(theta, r, x, z)
         curv     = self._spherical_curv    (theta, r, x, z)
         Gamma    = self._spherical_Gamma   (theta, r, x, z)
+        trK      = self._spherical_trK     (theta, r, x, z)
+
+        dmetric            = self._spherical_dmetric(theta, r, x, z, metric, Gamma)
+        dmetric_u          = self._spherical_dmetric_u(theta, r, x, z, dmetric, metric_u)
+        metric_u_cart_yy   = self._interp_var(theta, r, x, z, self.metric_u_cart[1, 1])
+        dmetric_cart_yy_dz = self._interp_var(theta, r, x, z, self.dmetric_cart[2, 1, 1])
 
         u2 = metric_u[1, 1] * dr * dr - 2 * metric_u[0, 1] * dr + metric_u[0, 0]
         u  = np.sqrt(u2)
-
-        df_ij = np.empty_like(metric_u)
-        for i, j in it.product(range(3), repeat = 2):
-            df_ij[i, j] = (metric_u[0, i] - metric_u[1, i] * dr) * (metric_u[0, j] - metric_u[1, j] * dr)
-
-        term1 = u2 * metric_u - df_ij
-        term2 = Gamma[0] - dr * Gamma[1] + u * curv
-
-        ret = -np.einsum('ij...,ij...', term1, term2) / (metric_u[0, 0] * metric_u[1, 1] - metric_u[0, 1] * metric_u[0, 1])
-
-        # handle the symmetry axis
-        g = theta / np.pi
-        for idx in np.where(np.abs(g.round() - g) < 1e-12)[0]:
-            stencil    = 8
-
-            coords     = np.empty(stencil * 2)
-            vals       = np.empty(stencil * 2)
-            fact       = np.empty(stencil * 2)
-            if idx + stencil < ret.shape[0]:
-                for i in range(stencil):
-                    coords[2 * i]     = theta[idx + i + 1]
-                    coords[2 * i + 1] = 2 * theta[idx] - coords[2 * i]
-
-                    vals[2 * i]       = ret[idx + i + 1]
-                    vals[2 * i + 1]   = ret[idx + i + 1]
-            elif idx - stencil >= 0:
-                for i in range(stencil):
-                    coords[2 * i]     = theta[idx - i - 1]
-                    coords[2 * i + 1] = 2 * theta[idx] - coords[2 * i]
-
-                    vals[2 * i]       = ret[idx - i - 1]
-                    vals[2 * i + 1]   = ret[idx - i - 1]
-            else:
-                continue
-
-            for i in range(len(fact)):
-                f = 1.0
-                for j in range(len(fact)):
-                    if i == j:
-                        continue
-                    f *= (theta[idx] - coords[j]) / (coords[i] - coords[j])
-                fact[i] = f
-
-            ret[idx] = np.dot(fact, vals)
+        u3 = u * u2
+
+        du_F    = np.zeros((3,) + u.shape)
+        du_F[0] = metric_u[0, 0] - metric_u[0, 1] * dr
+        du_F[1] = metric_u[0, 1] - metric_u[1, 1] * dr
+
+        # F_term = (∂^i F)(∂^j F)(Γ_{ij}^k ∂_k F + u K_{ij})
+        F_term = np.zeros_like(u)
+        for i, j in it.product(range(2), repeat = 2):
+            F_term += du_F[i] * du_F[j] * (Gamma[0, i, j] - dr * Gamma[1, i, j] + u * curv[i, j])
+
+        # X_term = γ^{ij} Γ_{ij}^k ∂_k F
+        # this is the only term that is singular at θ = nπ
+        X_term = np.zeros_like(u)
+        for i, j in it.product(range(2), repeat = 2):
+            X_term += metric_u[i, j] * (Gamma[0, i, j] - dr * Gamma[1, i, j])
+
+        # γ^{φφ} Γ_{φφ}^r ∂_r F
+        X_term_phi_reg_r  = metric_u[2, 2] * Gamma[0, 2, 2]
+        X_term_phi_sing_r = -0.5 * (
+                metric_u[0, 0] / r * (2. + r *
+                    metric_u_cart_yy *  dmetric_cart_yy_dz) +
+                dmetric_u[1, 0, 1] * 2
+                )
+        # γ^{φφ} Γ_{φφ}^θ
+        # at the axis this term goes to γ^{θθ} * d2r, i.e. the rhs we are computing;
+        # when multiplied by prefactors they all cancel out, so we get:
+        #  rhs = <regular part> - rhs
+        # in other words, at the axis we need to set X_term_phi_theta to 0 and
+        # instead divide rhs by 2
+        X_term_phi_reg_theta  = -metric_u[2, 2] * Gamma[1, 2, 2] * dr
+
+        X_term += np.where(theta != 0.0, X_term_phi_reg_r + X_term_phi_reg_theta,
+                           X_term_phi_sing_r + 0.0)
+
+        denom = metric_u[0, 0] * metric_u[1, 1] - metric_u[0, 1] * metric_u[0, 1]
+        ret   = -(u2 * X_term + u3 * trK - F_term + val * u3) / denom
+
+        ret  *= np.where(theta == 0.0, 0.5, 1.0)
 
         return ret
 
@@ -507,8 +567,8 @@ class AHCalc(object):
 
         rp = r + eps
         rm = r - eps
-        Fp = self._F(theta, (rp, R))
-        Fm = self._F(theta, (rm, R))
+        Fp = self._ah_rhs(theta, (rp, R))
+        Fm = self._ah_rhs(theta, (rm, R))
         return (Fp - Fm) / (2 * eps)
 
     def _dF_r_exact(self, theta, H):
@@ -522,7 +582,8 @@ class AHCalc(object):
         curv      = self._spherical_curv     (theta, r, x, z)
         dcurv     = self._spherical_dcurv    (theta, r, x, z)
         Gamma     = self._spherical_Gamma    (theta, r, x, z)
-        dmetric_u = self._spherical_dmetric_u(theta, r, x, z, metric, metric_u, Gamma)
+        dmetric   = self._spherical_dmetric  (theta, r, x, z, metric, Gamma)
+        dmetric_u = self._spherical_dmetric_u(theta, r, x, z, dmetric, metric_u)
         dGamma    = self._spherical_dGamma   (theta, r, x, z)
 
         u2     = metric_u[1, 1] * dr * dr - 2 * metric_u[0, 1] * dr + metric_u[0, 0]
@@ -554,8 +615,8 @@ class AHCalc(object):
 
         Rp = R + eps
         Rm = R - eps
-        Fp = self._F(theta, (r, Rp))
-        Fm = self._F(theta, (r, Rm))
+        Fp = self._ah_rhs(theta, (r, Rp))
+        Fm = self._ah_rhs(theta, (r, Rm))
         return (Fp - Fm) / (2 * eps)
 
     def _dF_R_exact(self, theta, H):