horizon: make the finite difference operator specifiable

1 files changed, 5 insertions, 5 deletions

diff --git a/horizon.py b/horizon.py
index f1b5123..739714f 100644
--- a/horizon.py
+++ b/horizon.py
@@ -6,7 +6,7 @@ from . import utils
 
 import numpy as np
 
-def calc_expansion(x, z, metric, curv, surf, direction = 1):
+def calc_expansion(x, z, metric, curv, surf, direction = 1, diff_op = diff.fd4):
     """
     Calculate expansion of null geodesics on a sequence of surfaces [1]. The
     surfaces are specified as level surfaces of F(r, θ) = r - h(θ).
@@ -37,7 +37,7 @@ def calc_expansion(x, z, metric, curv, surf, direction = 1):
 
     metric_u = utils.matrix_invert(metric)
 
-    Gamma = curvature.calc_christoffel(x, z, metric)
+    Gamma = curvature.calc_christoffel(x, z, metric, diff_op)
 
     trK = np.einsum('ij...,ij...', metric_u, curv)
 
@@ -46,9 +46,9 @@ def calc_expansion(x, z, metric, curv, surf, direction = 1):
         F[i] -= surf.eval(Theta[i])
 
     dF = np.empty((3,) + F.shape)
-    dF[0] = diff.fd4(F, 1, dX[0])
+    dF[0] = diff_op(F, 1, dX[0])
     dF[1] = 0.0
-    dF[2] = diff.fd4(F, 0, dX[2])
+    dF[2] = diff_op(F, 0, dX[2])
 
     s_l  = direction * dF[:]
     s_u  = np.einsum('ij...,j...->i...', metric_u, s_l)
@@ -60,7 +60,7 @@ def calc_expansion(x, z, metric, curv, surf, direction = 1):
             if i == 1 or j == 1:
                 continue
             diff_dir = 1 if (i == 0) else 0
-            ds_u[i, j] = diff.fd4(s_u[j], diff_dir, dX[i])
+            ds_u[i, j] = diff_op(s_u[j], diff_dir, dX[i])
     ds_u[1, 1] = np.where(np.abs(X) > 1e-8, s_u[0] / X, ds_u[0, 0])
 
     Div_s_u = np.einsum('ii...', ds_u) + np.einsum('iki...,k...', Gamma, s_u)