# -*- coding: utf-8 -*- import itertools as it import numpy as np from . import diff from . import utils 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 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) metric_u = utils.matrix_invert(metric) dmetric = _calc_dmetric(X, Z, metric, diff_op) return _calc_christoffel(metric, metric_u, dmetric) 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 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) 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