added constraint code

GeneralRelativity/CCZ4Geometry.py

@@ -11,6 +11,7 @@ def compute_ricci_Z(
     h_UU: torch.Tensor,
     chris: Dict[str, torch.Tensor],
     Z_over_chi: torch.Tensor,
+    GR_SPACEDIM: int = 4,
 ) -> Dict[str, torch.Tensor]:
     """
     Compute the Ricci tensor Z using the provided variables, derivatives, and Christoffel symbols.
@@ -28,7 +29,6 @@ def compute_ricci_Z(
     """
     out = {"LL": torch.zeros_like(vars["h"]), "scalar": 0}
 
-    GR_SPACEDIM = 4
     boxtildechi = 0
 
     covdtilde2chi = torch.zeros_like(vars["h"])

GeneralRelativity/Constraints.py

@@ -0,0 +1,88 @@
+import torch
+from typing import Dict
+from GeneralRelativity.DimensionDefinitions import FOR1, FOR2
+from GeneralRelativity.TensorAlgebra import compute_trace
+from GeneralRelativity.CCZ4Geometry import compute_ricci
+
+
+def constraint_equations(
+    vars: Dict[str, torch.Tensor],
+    d1: Dict[str, torch.Tensor],
+    d2: Dict[str, torch.Tensor],
+    h_UU: torch.Tensor,
+    chris: Dict[str, torch.Tensor],
+    m_cosmological_constant: float = 0.0,
+    GR_SPACEDIM: int = 4,
+) -> Dict[str, torch.Tensor]:
+    """
+    Compute the Hamiltonian and momentum constraints.
+
+    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).
+
+    Returns:
+        dict: Dictionary containing the Hamiltonian and momentum constraints.
+    """
+
+    out = {}
+    ricci = compute_ricci(vars, d1, d2, h_UU, chris)
+
+    A_UU = raise_all(vars["A"], h_UU)
+    tr_A2 = compute_trace(vars["A"], A_UU)
+
+    out["Ham"] = (
+        ricci["scalar"]
+        + (GR_SPACEDIM - 1.0) * vars["K"] * vars["K"] / GR_SPACEDIM
+        - tr_A2
+    )
+    out["Ham"] -= 2 * 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)
+    )
+    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])
+
+        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)
+            )
+
+        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]
+            )
+
+    return out

GeneralRelativity/TensorAlgebra.py

@@ -3,6 +3,10 @@
 import glob
 import torch
 from GeneralRelativity.DimensionDefinitions import FOR1, FOR2, FOR3, FOR4
+from typing import Dict
+from GeneralRelativity.DimensionDefinitions import FOR1, FOR2
+from GeneralRelativity.TensorAlgebra import compute_trace
+import torch
 
 
 def compute_christoffel(d1_metric: torch.tensor, h_UU: torch.tensor) -> torch.tensor:
@@ -99,7 +103,9 @@ def compute_trace(tensor_LL, inverse_metric):
     return trace
 
 
-def raise_all(tensor_L: torch.Tensor, inverse_metric: torch.Tensor) -> torch.Tensor:
+def raise_all_vector(
+    tensor_L: torch.Tensor, inverse_metric: torch.Tensor
+) -> torch.Tensor:
     """
     Raises the index of a covector (tensor with a lower index) using the inverse metric.
 
@@ -114,3 +120,24 @@ def raise_all(tensor_L: torch.Tensor, inverse_metric: torch.Tensor) -> torch.Ten
     for i, j in FOR2():
         tensor_U[..., i] += inverse_metric[..., i, j] * tensor_L[..., j]
     return tensor_U
+
+
+def raise_all_metric(
+    tensor_LL: torch.Tensor, inverse_metric: torch.Tensor
+) -> torch.Tensor:
+    """
+    Raises the indices of a 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).
+    """
+    tensor_UU = torch.zeros_like(tensor_LL)
+    for i, j, k, l in FOR4():
+        tensor_UU[..., i, j] += (
+            inverse_metric[..., i, k] * inverse_metric[..., j, l] * tensor_LL[..., k, l]
+        )
+    return tensor_UU