curvature: add calculations of more variables

1 files changed, 133 insertions, 20 deletions

diff --git a/curvature.py b/curvature.py
index e9d8e17..68dd88d 100644
--- a/curvature.py
+++ b/curvature.py
@@ -1,11 +1,94 @@
 # -*- coding: utf-8 -*-
 
+import itertools as it
 import numpy as np
 
 from . import diff
 from . import utils
 
-def calc_christoffel(x, z, metric):
+def _calc_dmetric(X, Z, metric, diff_op):
+    dmetric = np.zeros((3,) + metric.shape)
+
+    res = diff_op(metric, 3, X[0, 1] - X[0, 0])
+    dmetric[0] = diff_op(metric, 3, X[0, 1] - X[0, 0])
+    dmetric[2] = diff_op(metric, 2, Z[1, 0] - Z[0, 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]
+
+    return dmetric
+
+def _calc_d2metric(X, Z, metric, dmetric, diff_op, diff2_op):
+    d2metric = np.empty((3,) + dmetric.shape)
+
+    dx = X[0, 1] - X[0, 0]
+    dz = Z[1, 0] - Z[0, 0]
+    X2 = X * X
+
+    d2metric[0, 0] = diff2_op(metric, 3, dx)
+    d2metric[2, 2] = diff2_op(metric, 2, dz)
+
+    d2metric[0, 2] = diff_op(dmetric[0], 2, dz)
+    d2metric[2, 0] = d2metric[0, 2]
+
+    d2metric[0, 1, 0, 0] = 0
+    d2metric[1, 0, 0, 0] = 0
+    d2metric[1, 2, 0, 0] = 0
+    d2metric[2, 1, 0, 0] = 0
+    d2metric[1, 1, 0, 0] = np.where(np.abs(X) > 1e-8, dmetric[0, 0, 0] / X - 2.0 * (metric[0, 0] - metric[1, 1]) / X2, d2metric[0, 0, 1, 1])
+
+    d2metric[1, 1, 1, 1] = np.where(np.abs(X) > 1e-8, dmetric[0, 1, 1] / X + 2.0 * (metric[0, 0] - metric[1, 1]) / X2, d2metric[0, 0, 0, 0])
+    d2metric[0, 1, 1, 1] = 0
+    d2metric[1, 0, 1, 1] = 0
+    d2metric[2, 1, 1, 1] = 0
+    d2metric[1, 2, 1, 1] = 0
+
+    d2metric[1, 1, 2, 2] = np.where(np.abs(X) > 1e-8, dmetric[0, 2, 2] / X, d2metric[0, 0, 2, 2])
+    d2metric[0, 1, 2, 2] = 0
+    d2metric[1, 0, 2, 2] = 0
+    d2metric[2, 1, 2, 2] = 0
+    d2metric[1, 2, 2, 2] = 0
+
+    d2metric[1, 1, 0, 1] = 0
+    d2metric[0, 1, 0, 1] = np.where(np.abs(X) > 1e-8, (dmetric[0, 0, 0] - dmetric[0, 1, 1]) / X - (metric[0, 0] - metric[1, 1]) / X2, 0.5 * (d2metric[0, 0, 0, 0] - d2metric[0, 0, 1, 1]))
+    d2metric[1, 0, 0, 1] = d2metric[0, 1, 0, 1]
+    d2metric[1, 2, 0, 1] = np.where(np.abs(X) > 1e-8, (dmetric[2, 0, 0] - dmetric[2, 1, 1]) / X, d2metric[0, 2, 0, 0] - d2metric[0, 2, 1, 1])
+    d2metric[2, 1, 0, 1] = d2metric[1, 2, 0, 1]
+
+
+    d2metric[1, 1, 0, 2] = np.where(np.abs(X) > 1e-8, dmetric[0, 0, 2] / X - metric[0, 2] / X2, 0.5 * d2metric[0, 0, 0, 2])
+    d2metric[0, 1, 0, 2] = 0
+    d2metric[1, 0, 0, 2] = 0
+    d2metric[1, 2, 0, 2] = 0
+    d2metric[2, 1, 0, 2] = 0
+
+    d2metric[1, 1, 1, 2] = 0
+    d2metric[0, 1, 1, 2] = np.where(np.abs(X) > 1e-8, dmetric[0, 0, 2] / X - metric[0, 2] / X2, 0.5 * d2metric[0, 0, 0, 2])
+    d2metric[1, 0, 1, 2] = d2metric[0, 1, 1, 2]
+    d2metric[1, 2, 1, 2] = np.where(np.abs(X) > 1e-8, dmetric[2, 0, 2] / X, d2metric[0, 2, 0, 2])
+    d2metric[2, 1, 1, 2] = d2metric[1, 2, 1, 2]
+
+    d2metric[:, :, 1, 0] = d2metric[:, :, 0, 1]
+    d2metric[:, :, 2, 0] = d2metric[:, :, 0, 2]
+    d2metric[:, :, 2, 1] = d2metric[:, :, 1, 2]
+
+    return d2metric
+
+def _calc_christoffel(metric, metric_u, dmetric):
+    Gamma = np.empty_like(dmetric)
+    for i, j, k in it.product(range(3), repeat = 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
+
+def calc_christoffel(x, z, metric, diff_op = diff.fd4):
     """
     Calculate Christoffel symbols
 
@@ -28,26 +111,56 @@ def calc_christoffel(x, z, metric):
     """
     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])
+    metric_u = utils.matrix_invert(metric)
+    dmetric  = _calc_dmetric(X, Z, metric, diff_op)
 
-    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]
+    return _calc_christoffel(metric, metric_u, dmetric)
 
-    metric_u = utils.matrix_invert(metric)
+def _calc_dchristoffel(metric, metric_u, dmetric, dmetric_u, d2metric):
+    dGamma = np.empty_like(d2metric)
+    for i, j, k, l in it.product(range(3), repeat = 4):
+        dGamma[i, j, k, l] = 0.5 * (np.einsum('k...,k...', dmetric_u[i, j], dmetric[l, k] + dmetric[k, l] - dmetric[:, k, l]) +
+                                    np.einsum('k...,k...',  metric_u[j],    d2metric[i, l, k] + d2metric[i, k, l] - d2metric[i, :, l, k]))
+    return dGamma
 
-    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])
+def calc_riemann(metric, metric_u, dmetric, d2metric):
+    dmetric_u = -np.einsum('ij...,km...,ljk...->lim...', metric_u, metric_u, dmetric)
+    Gamma     = _calc_christoffel(metric, metric_u, dmetric)
+    dGamma    = _calc_dchristoffel(metric, metric_u, dmetric, dmetric_u, d2metric)
 
-    return Gamma
+    Riemann_uddd = (np.einsum('cabd...->abcd...', dGamma) - np.einsum('dabc...->abcd...', dGamma) +
+                    np.einsum('akc...,kbd...->abcd...', Gamma, Gamma) - np.einsum('akd...,kbc...->abcd...', Gamma, Gamma))
+    Riemann = np.einsum('ik...,klmn...->ilmn...', metric, Riemann_uddd)
+
+    return Riemann
+
+def calc_ricci(metric, metric_u, dmetric, d2metric):
+    dmetric_u = -np.einsum('ij...,km...,ljk...->lim...', metric_u, metric_u, dmetric)
+    Gamma     = _calc_christoffel(metric, metric_u, dmetric)
+    dGamma    = _calc_dchristoffel(metric, metric_u, dmetric, dmetric_u, d2metric)
+
+    Ricci  = (np.einsum('mmjk...->jk...', dGamma) - np.einsum('kmjm...->jk...', dGamma) +
+              np.einsum('llm...,mjk...->jk...', Gamma, Gamma) - np.einsum('lkm...,mjl...->jk...', Gamma, Gamma))
+
+    return Ricci
+
+def calc_constraint_ham(metric, metric_u, dmetric, d2metric, curv):
+    Ricci      = calc_ricci(metric, metric_u, dmetric, d2metric)
+    R_scalar   = np.einsum('ij...,ij...', metric_u, Ricci)
+
+    curv_ud    = np.einsum('ik...,kj...->ij...', metric_u, curv)
+    curv_trace = curv_ud[0, 0] + curv_ud[1, 1] + curv_ud[2, 2]
+
+    k2 = np.einsum('ij...,ji...', curv_ud, curv_ud)
+
+    return R_scalar + curv_trace ** 2 - k2
+
+def calc_constraint_mom(metric, metric_u, dmetric, curv_m, dcurv_m):
+    Gamma      = _calc_christoffel(metric, metric_u, dmetric)
+    curv_trace = curv_m[0, 0] + curv_m[1, 1] + curv_m[2, 2]
+
+    M = np.empty((3,) + curv_trace.shape)
+    for j in range(3):
+        M[j] = np.einsum('ii...', dcurv_m[:, :, j]) + np.einsum('iik...,k...', Gamma, curv_m[:, j]) - np.einsum('ki...,ik...', Gamma[:, :, j], curv_m)
+
+    return M