Skip to content

Commit

Permalink
Improved Commenting and corrected constraint function
Browse files Browse the repository at this point in the history
  • Loading branch information
@ThomasHelfer
ThomasHelfer committed Dec 27, 2023
1 parent 0f2bab8 commit 971b3b4
Show file tree
Hide file tree
Showing 3 changed files with 116 additions and 66 deletions.
21 changes: 9 additions & 12 deletions GeneralRelativity/CCZ4Geometry.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,23 +32,18 @@ def compute_ricci_Z(
boxtildechi = 0

covdtilde2chi = torch.zeros_like(vars["h"])
# FOR2(k, l)
# {
# covdtilde2chi[k][l] = d2.chi[k][l];
# FOR1(m) { covdtilde2chi[k][l] -= chris.ULL[m][k][l] * d1.chi[m]; }
# }
for k, l in FOR2():
# covdtilde2chi[k][l] = d2.chi[k][l];
covdtilde2chi[..., k, l] = d2["chi"][..., k, l]
# FOR1(m) { covdtilde2chi[k][l] -= chris.ULL[m][k][l] * d1.chi[m]; }
for m in FOR1():
covdtilde2chi[..., k, l] -= chris["ULL"][..., m, k, l] * d1["chi"][..., m]
# FOR2(k, l) { boxtildechi += h_UU[k][l] * covdtilde2chi[k][l]; }
# FOR2(k, l) { boxtildechi += h_UU[k][l] * covdtilde2chi[k][l]; }
for k, l in FOR2():
boxtildechi += h_UU[..., k, l] * covdtilde2chi[..., k, l]
# data_t dchi_dot_dchi = 0;
# {
# FOR2(m, n) { dchi_dot_dchi += h_UU[m][n] * d1.chi[m] * d1.chi[n]; }
# }
dchi_dot_dchi = torch.zeros_like(vars["chi"][...])
# FOR2(m, n) { dchi_dot_dchi += h_UU[m][n] * d1.chi[m] * d1.chi[n]; }
for m, n in FOR2():
dchi_dot_dchi += h_UU[..., m, n] * d1["chi"][..., m] * d1["chi"][..., n]

Expand All @@ -62,12 +57,13 @@ def compute_ricci_Z(
for k in FOR1():
# ricci_tilde += 0.5 * (vars.h[k][i] * d1.Gamma[k][j] +
# vars.h[k][j] * d1.Gamma[k][i]);
# ricci_tilde += 0.5 * (vars.Gamma[k] - 2 * Z_over_chi[k]) *
# (chris.LLL[i][j][k] + chris.LLL[j][i][k]);

ricci_tilde += 0.5 * (
vars["h"][..., k, i] * d1["Gamma"][..., k, j]
+ vars["h"][..., k, j] * d1["Gamma"][..., k, i]
)
# ricci_tilde += 0.5 * (vars.Gamma[k] - 2 * Z_over_chi[k]) *
# (chris.LLL[i][j][k] + chris.LLL[j][i][k]);
ricci_tilde += (
0.5
* (vars["Gamma"][..., k] - 2 * Z_over_chi[..., k])
Expand Down Expand Up @@ -139,6 +135,7 @@ def compute_ricci(
d2: Dict[str, torch.Tensor],
h_UU: torch.Tensor,
chris: Dict[str, torch.Tensor],
GR_SPACEDIM: int = 4,
) -> Dict[str, torch.Tensor]:
"""
Compute the Ricci tensor using the provided variables, derivatives, and Christoffel symbols.
Expand All @@ -154,4 +151,4 @@ def compute_ricci(
dict: Dictionary containing the components of the Ricci tensor.
"""
Z0 = torch.zeros_like(torch.zeros_like(d1["chi"]))
return compute_ricci_Z(vars, d1, d2, h_UU, chris, Z0)
return compute_ricci_Z(vars, d1, d2, h_UU, chris, Z0, GR_SPACEDIM)
138 changes: 84 additions & 54 deletions GeneralRelativity/Constraints.py
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import torch
from typing import Dict
from GeneralRelativity.DimensionDefinitions import FOR1, FOR2
from GeneralRelativity.TensorAlgebra import compute_trace
from typing import Dict, List
from GeneralRelativity.DimensionDefinitions import FOR1, FOR2, FOR3
from GeneralRelativity.TensorAlgebra import compute_trace, raise_all
from GeneralRelativity.CCZ4Geometry import compute_ricci


Expand All @@ -11,78 +11,108 @@ def constraint_equations(
d2: Dict[str, torch.Tensor],
h_UU: torch.Tensor,
chris: Dict[str, torch.Tensor],
m_cosmological_constant: float = 0.0,
GR_SPACEDIM: int = 4,
cosmological_constant: float = 0,
) -> Dict[str, torch.Tensor]:
"""
Compute the Hamiltonian and momentum constraints.
Calculates the constraint equations in the context of general relativity.
Args:
vars (dict): Dictionary of tensor variables.
d1 (dict): Dictionary of first derivatives of tensor variables.
d2 (dict): Dictionary of second derivatives of tensor variables.
h_UU (torch.Tensor): Inverse metric tensor.
chris (dict): Dictionary containing Christoffel symbols.
m_cosmological_constant (float): Cosmological constant (default is 0).
GR_SPACEDIM (int): Spatial dimension (default is 4).
vars (Dict[str, torch.Tensor]): Dictionary containing variables.
Expected to have keys like 'A', 'K', 'chi', etc.
d1 (Dict[str, torch.Tensor]): Dictionary containing first derivatives.
d2 (Dict[str, torch.Tensor]): Dictionary containing second derivatives.
h_UU (torch.Tensor): The inverse metric tensor.
chris (Dict[str, torch.Tensor]): Dictionary containing Christoffel symbols,
expected to have keys like 'ULL'.
GR_SPACEDIM (int, optional): The spatial dimension of the general relativity
calculations. Defaults to 4.
cosmological_constant (float, optional): The cosmological constant. Defaults to 0.
Returns:
dict: Dictionary containing the Hamiltonian and momentum constraints.
Dict[str, torch.Tensor]: A dictionary containing the computed constraint equations.
Keys include 'Ham', 'Ham_abs_terms', 'Mom', 'Mom_abs_terms'.
"""

out = {}
ricci = compute_ricci(vars, d1, d2, h_UU, chris)
out = {
"Ham": torch.zeros_like(vars["chi"]),
"Ham_abs_terms": torch.zeros_like(vars["chi"]),
"Mom": torch.zeros_like(vars["shift"]),
"Mom_abs_terms": torch.zeros_like(vars["shift"]),
}

# auto ricci = CCZ4Geometry::compute_ricci(vars, d1, d2, h_UU, chris);
# auto A_UU = TensorAlgebra::raise_all(vars.A, h_UU);
# data_t tr_A2 = TensorAlgebra::compute_trace(vars.A, A_UU);
ricci = compute_ricci(vars, d1, d2, h_UU, chris, GR_SPACEDIM)
A_UU = raise_all(vars["A"], h_UU)
tr_A2 = compute_trace(vars["A"], A_UU)

# out.Ham = ricci.scalar +
# (GR_SPACEDIM - 1.) * vars.K * vars.K / GR_SPACEDIM - tr_A2;
# out.Ham -= 2 * m_cosmological_constant;
out["Ham"] = (
ricci["scalar"]
+ (GR_SPACEDIM - 1.0) * vars["K"] * vars["K"] / GR_SPACEDIM
- tr_A2
ricci["scalar"] + (GR_SPACEDIM - 1) * vars["K"] ** 2 / GR_SPACEDIM - tr_A2
)
out["Ham"] -= 2 * m_cosmological_constant
out["Ham"] -= 2 * cosmological_constant

# out.Ham_abs_terms =
# abs(ricci.scalar) + abs(tr_A2) +
# abs((GR_SPACEDIM - 1.) * vars.K * vars.K / GR_SPACEDIM);
# out.Ham_abs_terms += 2.0 * abs(m_cosmological_constant);
out["Ham_abs_terms"] = (
torch.abs(ricci["scalar"])
+ torch.abs(tr_A2)
+ torch.abs((GR_SPACEDIM - 1.0) * vars["K"] * vars["K"] / GR_SPACEDIM)
+ torch.abs((GR_SPACEDIM - 1) * vars["K"] ** 2 / GR_SPACEDIM)
)
out["Ham_abs_terms"] += 2.0 * torch.abs(m_cosmological_constant)

if "Mom" in vars or "Mom_abs_terms" in vars:
covd_A = torch.zeros_like(vars["A"])
for i, j, k in FOR3():
covd_A[..., i, j, k] = d1["A"][..., j, k, i]
for l in FOR1():
covd_A[..., i, j, k] -= (
chris["ULL"][..., l, i, j] * vars["A"][..., l, k]
+ chris["ULL"][..., l, i, k] * vars["A"][..., l, j]
)

out["Mom"] = torch.zeros((GR_SPACEDIM,))
out["Mom_abs_terms"] = torch.zeros((GR_SPACEDIM,))
for i in FOR1():
out["Mom"][..., i] = -(GR_SPACEDIM - 1.0) * d1["K"][..., i] / GR_SPACEDIM
out["Mom_abs_terms"][..., i] = torch.abs(out["Mom"][..., i])
out["Ham_abs_terms"] += 2.0 * abs(cosmological_constant)

covd_A_term = torch.zeros((GR_SPACEDIM,))
d1_chi_term = torch.zeros((GR_SPACEDIM,))
chi_regularised = torch.clamp(vars["chi"], min=1e-6)
for i, j, k in FOR3():
covd_A_term[..., i] += h_UU[..., j, k] * covd_A[..., k, j, i]
d1_chi_term[..., i] -= (
GR_SPACEDIM
* h_UU[..., j, k]
* vars["A"][..., i, j]
* d1["chi"][..., k]
/ (2 * chi_regularised)
# Tensor<2, data_t> covd_A[CH_SPACEDIM];
# FOR3(i, j, k)
# {
covd_A = torch.zeros_like(d1["A"])
for i, j, k in FOR3():
# covd_A[i][j][k] = d1.A[j][k][i];
# FOR1(l)
# {
covd_A[..., i, j, k] = d1["A"][..., j, k, i]
for l in FOR1():
# covd_A[i][j][k] += -chris.ULL[l][i][j] * vars.A[l][k] -
# chris.ULL[l][i][k] * vars.A[l][j];
covd_A[..., i, j, k] -= (
chris["ULL"][..., l, i, j] * vars["A"][..., l, k]
- chris["ULL"][..., l, i, k] * vars["A"][..., l, j]
)

for i in FOR1():
out["Mom"][..., i] += covd_A_term[..., i] + d1_chi_term[..., i]
out["Mom_abs_terms"][..., i] += torch.abs(covd_A_term[..., i]) + torch.abs(
d1_chi_term[..., i]
)
for i in FOR1():
# out.Mom[i] = -(GR_SPACEDIM - 1.) * d1.K[i] / GR_SPACEDIM;
# out.Mom_abs_terms[i] = abs(out.Mom[i]);
out["Mom"][..., i] = -(GR_SPACEDIM - 1.0) * d1["K"][..., i] / GR_SPACEDIM
out["Mom_abs_terms"][..., i] = torch.abs(out["Mom"][..., i])
# Tensor<1, data_t> covd_A_term = 0.0;
# Tensor<1, data_t> d1_chi_term = 0.0;
# const data_t chi_regularised = simd_max(1e-6, vars.chi);
covd_A_term = torch.zeros_like(d1["chi"])
d1_chi_term = torch.zeros_like(d1["chi"])
chi_regularised = torch.maximum(torch.tensor(1e-6), vars["chi"])
for i, j, k in FOR3():
# covd_A_term[i] += h_UU[j][k] * covd_A[k][j][i];
# d1_chi_term[i] += -GR_SPACEDIM * h_UU[j][k] * vars.A[i][j] *
# d1.chi[k] / (2 * chi_regularised);
covd_A_term[..., i] += h_UU[..., j, k] * covd_A[..., k, j, i]
d1_chi_term[..., i] += (
-GR_SPACEDIM
* h_UU[..., j, k]
* vars["A"][..., i, j]
* d1["chi"][..., k]
/ (2 * chi_regularised)
)

for i in FOR1():
# out.Mom[i] += covd_A_term[i] + d1_chi_term[i];
# out.Mom_abs_terms[i] += abs(covd_A_term[i]) + abs(d1_chi_term[i]);
out["Mom"][..., i] += covd_A_term[..., i] + d1_chi_term[..., i]
out["Mom_abs_terms"][..., :, i] += torch.abs(covd_A_term[..., i]) + torch.abs(
d1_chi_term[..., i]
)

return out
23 changes: 23 additions & 0 deletions GeneralRelativity/TensorAlgebra.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -140,3 +140,26 @@ def raise_all_metric(
inverse_metric[..., i, k] * inverse_metric[..., j, l] * tensor_LL[..., k, l]
)
return tensor_UU


def raise_all(tensor: torch.Tensor, inverse_metric: torch.Tensor) -> torch.Tensor:
"""
Raises the indices of a (0,1,2)-Tensor using the inverse metric tensor.
Args:
tensor_LL (torch.Tensor): The 2-Tensor with lower indices.
inverse_metric (torch.Tensor): The inverse metric tensor (2-Tensor).
Returns:
torch.Tensor: The resulting tensor with indices raised (2-Tensor).
"""
if len(tensor.shape) == 4:
return tensor
elif len(tensor.shape) == 5:
return raise_all_vector(tensor, inverse_metric)
elif len(tensor.shape) == 6:
return raise_all_metric(tensor, inverse_metric)
else:
raise ValueError(
"Raise_all can only deal with scalar,vector or tensor dimensions."
)

0 comments on commit 971b3b4

Please sign in to comment.