horizon: split off the Christoffel symbol calculation

2 files changed, 57 insertions, 19 deletions

diff --git a/curvature.py b/curvature.py
new file mode 100644
index 0000000..e9d8e17
--- /dev/null
+++ b/curvature.py
@@ -0,0 +1,53 @@
+# -*- coding: utf-8 -*-
+
+import numpy as np
+
+from . import diff
+from . import utils
+
+def calc_christoffel(x, z, metric):
+    """
+    Calculate Christoffel symbols
+
+     i    1 il /                       \
+    Γ   = -γ   | ∂ γ   + ∂ γ   - ∂ γ   |
+     jk   2    \  j kl    k jl    l jk /
+
+    using finite differences.
+
+    :param array_like x: 1D array of x coordinates.
+    :param array_like z: 1D array of z-coordinates.
+    :param array_like metric: 4D array of spatial metric values at the grid
+                              formed by x and z. metric[i, j, k, l] is the ijth
+                              component of the metric at the point (X=x[l],
+                              Z=z[k]).
+    :rtype: array_like, shape (3, 3, 3, z.shape[0], x.shape[0])
+    :return: Christoffel symbols, first axis is the upper index, following two
+             axes are the two lower indices, final two axes correspond to the z
+             and x grid points respectively.
+    """
+    X, Z  = np.meshgrid(x, z)
+
+    dmetric = np.zeros((3,) + metric.shape)
+    dmetric[0] = diff.fd4(metric, 3, x[1] - x[0])
+    dmetric[2] = diff.fd4(metric, 2, z[1] - z[0])
+
+    dmetric[1, 0, 0] = 0.0
+    dmetric[1, 1, 1] = 0.0
+    dmetric[1, 2, 2] = 0.0
+    dmetric[1, 0, 1] = np.where(np.abs(X) > 1e-8, (metric[0, 0] - metric[1, 1]) / X, dmetric[0, 0, 0] - dmetric[0, 1, 1])
+    dmetric[1, 1, 0] = dmetric[1, 0, 1]
+    dmetric[1, 0, 2] = 0.0
+    dmetric[1, 2, 0] = 0.0
+    dmetric[1, 1, 2] = np.where(np.abs(X) > 1e-8, metric[0, 2] / X, dmetric[0, 0, 2])
+    dmetric[1, 2, 1] = dmetric[1, 1, 2]
+
+    metric_u = utils.matrix_invert(metric)
+
+    Gamma = np.empty_like(dmetric)
+    for i in range(3):
+        for j in range(3):
+            for k in range(3):
+                Gamma[i, j, k] = 0.5 * np.einsum('k...,k...', metric_u[i], dmetric[j, k] + dmetric[k, j] - dmetric[:, k, j])
+
+    return Gamma
diff --git a/horizon.py b/horizon.py
index 3f74242..f1b5123 100644
--- a/horizon.py
+++ b/horizon.py
@@ -1,3 +1,6 @@
+# -*- coding: utf-8 -*-
+
+from . import curvature
 from . import diff
 from . import utils
 
@@ -34,25 +37,7 @@ def calc_expansion(x, z, metric, curv, surf, direction = 1):
 
     metric_u = utils.matrix_invert(metric)
 
-    dmetric = np.zeros((3,) + metric.shape)
-    dmetric[0] = diff.fd4(metric, 3, dX[0])
-    dmetric[2] = diff.fd4(metric, 2, dX[2])
-
-    dmetric[1, 0, 0] = 0.0
-    dmetric[1, 1, 1] = 0.0
-    dmetric[1, 2, 2] = 0.0
-    dmetric[1, 0, 1] = np.where(np.abs(X) > 1e-8, (metric[0, 0] - metric[1, 1]) / X, dmetric[0, 0, 0] - dmetric[0, 1, 1])
-    dmetric[1, 1, 0] = dmetric[1, 0, 1]
-    dmetric[1, 0, 2] = 0.0
-    dmetric[1, 2, 0] = 0.0
-    dmetric[1, 1, 2] = np.where(np.abs(X) > 1e-8, metric[0, 2] / X, dmetric[0, 0, 2])
-    dmetric[1, 2, 1] = dmetric[1, 1, 2]
-
-    Gamma = np.empty_like(dmetric)
-    for i in range(3):
-        for j in range(3):
-            for k in range(3):
-                Gamma[i, j, k] = 0.5 * np.einsum('k...,k...', metric_u[i], dmetric[j, k] + dmetric[k, j] - dmetric[:, k, j])
+    Gamma = curvature.calc_christoffel(x, z, metric)
 
     trK = np.einsum('ij...,ij...', metric_u, curv)