teukolsky_data.py: add a method to evaluate cartesian components

1 files changed, 38 insertions, 0 deletions

diff --git a/teukolsky_data.py b/teukolsky_data.py
index a819e5c..00ec5c4 100644
--- a/teukolsky_data.py
+++ b/teukolsky_data.py
@@ -131,6 +131,44 @@ class TeukolskyData(object):
     def eval_lapse(self, r, theta, diff_order = None):
         return self._eval_var(self._libtd.td_eval_lapse, r, theta, diff_order)
 
+    def eval_data_cart(self, x, z):
+        X, Z  = np.meshgrid(x, z)
+        X2    = X * X
+        Z2    = Z * Z
+        R     = np.sqrt(X2 + Z2)
+        R2    = R * R
+        R3    = R2 * R
+        Theta = np.where(R > 1e-15, np.arccos(Z / R), 0.0)
+
+        psi = self.eval_psi(R, Theta)
+        krr = self.eval_krr(R, Theta)
+        kpp = self.eval_kpp(R, Theta)
+        krt = self.eval_krt(R, Theta)
+
+        krr_x = self.eval_krr(np.array([0.0]), np.array([np.pi / 2]))[0]
+        krr_z = self.eval_krr(np.array([0.0]), np.array([0.0]))[0]
+
+        psi4 = (psi ** 4)
+        ktt = -(krr + kpp)
+
+        metric = np.zeros((3, 3) + X.shape)
+        curv   = np.zeros_like(metric)
+
+        metric[0, 0] = psi4
+        metric[1, 1] = psi4
+        metric[2, 2] = psi4
+
+        curv[1, 1] = psi4 * kpp
+
+        curv[0, 0] = psi4 * np.where(R > 1e-15,
+                X2 / R2 * krr + Z2 / R2 * ktt + 2.0 * Z * np.abs(X) / R3 * krt, krr_x)
+        curv[2, 2] = psi4 * np.where(R > 1e-15,
+                Z2 / R2 * krr + X2 / R2 * ktt - 2.0 * Z * np.abs(X) / R3 * krt, krr_z)
+        curv[0, 2] = psi4 * np.where(R > 1e-15,
+                X * Z / R2 * krr - X * Z / R2 * ktt + (-X * np.abs(X) + np.sign(X) * Z2) / R3 * krt, 0.0)
+        curv[2, 0] = curv[0, 2]
+
+        return metric, curv
 
     @property
     def amplitude(self):