Constraints

constraints

Constraint enforcement for gauge theories.

This module provides comprehensive tools for:

• Computing Gauss law constraint violations
• Projecting fields onto constraint manifold
• Gauge transformations
• Baumgarte-Shapiro constraint stabilization
• Diagnostic utilities

All Jacobians computed automatically via JAX autodiff.

Note on Gauss Constraint Implementations

The canonical implementations are:

• Pure gauge (no sources): jaxlatt.operators.gauge.gauss_constraint Signature: (links, E, dx) -> G = div(E)

• With scalar sources: jaxlatt.operators.coupled.gauss_constraint_with_source Signature: (links, E, phi, pi, dx) -> G = div(E) + Im(phi* pi)/dx

This module re-exports these for convenience and provides additional constraint enforcement tools (projection, gauge transformations, etc.).

gauss_constraint_with_source(links, E, phi, pi, dx)

Compute Gauss constraint with scalar field source.

This is the canonical implementation for coupled scalar-gauge systems. Gauss law:

\[G = \nabla \cdot E + \frac{\mathrm{Im}(\phi^* \pi)}{dx} = 0\]

Parameters:

Name	Type	Description	Default
links	Array	Gauge links (not used, kept for interface consistency)	required
E	Array	Electric field (3, N, N, N)	required
phi	Array	Complex scalar field (N, N, N)	required
pi	Array	Conjugate momentum (N, N, N)	required
dx	float	Lattice spacing	required

Returns:

Type	Description
Array	Gauss constraint violation \(G\), shape (N, N, N)

See Also

jaxlatt.operators.coupled.gauss_constraint_with_source

Source code in jaxlatt/core/constraints.py

def gauss_constraint_with_source(
    links: Array,
    E: Array,
    phi: Array,
    pi: Array,
    dx: float,
) -> Array:
    """
    Compute Gauss constraint with scalar field source.

    This is the canonical implementation for coupled scalar-gauge systems.
    Gauss law:

    $$G = \\nabla \\cdot E + \\frac{\\mathrm{Im}(\\phi^* \\pi)}{dx} = 0$$

    Args:
        links: Gauge links (not used, kept for interface consistency)
        E: Electric field `(3, N, N, N)`
        phi: Complex scalar field `(N, N, N)`
        pi: Conjugate momentum `(N, N, N)`
        dx: Lattice spacing

    Returns:
        Gauss constraint violation $G$, shape `(N, N, N)`

    See Also:
        `jaxlatt.operators.coupled.gauss_constraint_with_source`
    """
    return _gauss_constraint_with_source(links, E, phi, pi, dx)

project_gauss_constraint(E, phi, pi, dx, tolerance=1e-10, max_iterations=100)

Project electric field to satisfy Gauss constraint \(\nabla \cdot E + \mathrm{Im}(\phi^*\pi)/dx = 0\).

Uses iterative Jacobi method to solve Poisson equation:

\[\nabla^2 \phi_{\mathrm{gauge}} = \nabla \cdot E + \frac{\mathrm{Im}(\phi^*\pi)}{dx}\]

\[E_{\mathrm{corrected}} = E - \nabla \phi_{\mathrm{gauge}}\]

This ensures Gauss law is satisfied to numerical precision without changing the transverse part of \(E\).

Parameters:

Name	Type	Description	Default
E	Array	Electric field to project, shape (3, N, N, N).	required
phi	Array	Scalar field \(\phi\), shape (N, N, N).	required
pi	Array	Conjugate momentum \(\pi\), shape (N, N, N).	required
dx	float	Lattice spacing.	required
tolerance	float	Convergence tolerance for the iterative solver.	1e-10
max_iterations	int	Maximum number of Jacobi iterations.	100

Returns:

Type	Description
Array	Electric field satisfying Gauss constraint, shape (3, N, N, N).

Source code in jaxlatt/core/constraints.py

def project_gauss_constraint(
    E: Array,
    phi: Array,
    pi: Array,
    dx: float,
    tolerance: float = 1e-10,
    max_iterations: int = 100,
) -> Array:
    """
    Project electric field to satisfy Gauss constraint $\\nabla \\cdot E + \\mathrm{Im}(\\phi^*\\pi)/dx = 0$.

    Uses iterative Jacobi method to solve Poisson equation:

    $$\\nabla^2 \\phi_{\\mathrm{gauge}} = \\nabla \\cdot E + \\frac{\\mathrm{Im}(\\phi^*\\pi)}{dx}$$

    $$E_{\\mathrm{corrected}} = E - \\nabla \\phi_{\\mathrm{gauge}}$$

    This ensures Gauss law is satisfied to numerical precision without
    changing the transverse part of $E$.

    Args:
        E: Electric field to project, shape `(3, N, N, N)`.
        phi: Scalar field $\\phi$, shape `(N, N, N)`.
        pi: Conjugate momentum $\\pi$, shape `(N, N, N)`.
        dx: Lattice spacing.
        tolerance: Convergence tolerance for the iterative solver.
        max_iterations: Maximum number of Jacobi iterations.

    Returns:
        Electric field satisfying Gauss constraint, shape `(3, N, N, N)`.
    """
    # Source term: div(E) + Im(phi*pi)/dx (this is what we want to set to zero)
    div_E = divergence_3d(E, dx)
    source = div_E + jnp.imag(jnp.conj(phi) * pi) / dx

    # Solve nabla^2 phi_gauge = source using Jacobi iteration
    phi_gauge = jnp.zeros_like(phi, dtype=jnp.float32)

    # Jacobi iteration: solve nabla^2 phi = source
    # In discrete form: (phi_{i+1} + phi_{i-1} + ... - 6*phi_i) / dx^2 = source
    # Rearrange: phi_i = (phi_{i+1} + phi_{i-1} + ... - source*dx^2) / 6
    for _ in range(max_iterations):
        # Neighbor sum
        neighbor_sum = jnp.zeros_like(phi_gauge)
        for i in range(3):
            neighbor_sum = (
                neighbor_sum + jnp.roll(phi_gauge, -1, axis=i) + jnp.roll(phi_gauge, 1, axis=i)
            )

        # Jacobi update
        phi_gauge_new = (neighbor_sum - source * dx**2) / 6.0

        # Check convergence
        diff = phi_gauge_new - phi_gauge
        error = jnp.sqrt(jnp.mean(diff**2))

        phi_gauge = phi_gauge_new

        if error < tolerance:
            break

    # Compute gradient of phi_gauge
    E_correction = jnp.zeros_like(E)
    for i in range(3):
        phi_forward = jnp.roll(phi_gauge, -1, axis=i)
        grad_i = (phi_forward - phi_gauge) / dx
        E_correction = E_correction.at[i].set(grad_i)

    # Project E
    E_projected = E - E_correction

    return E_projected

project_gauss_coupled_lattice(lattice, tolerance=1e-10, max_iterations=100)

Project electric field in coupled lattice to satisfy Gauss constraint.

Creates a new lattice with \(E\) corrected to enforce \(\nabla \cdot E + g\rho = 0\).

Parameters:

Name	Type	Description	Default
lattice	CoupledLattice	Lattice to project.	required
tolerance	float	Convergence tolerance.	1e-10
max_iterations	int	Maximum iterations for solver.	100

Returns:

Type	Description
CoupledLattice	CoupledLattice with projected electric field.

Source code in jaxlatt/core/constraints.py

def project_gauss_coupled_lattice(
    lattice: CoupledLattice,
    tolerance: float = 1e-10,
    max_iterations: int = 100,
) -> CoupledLattice:
    """
    Project electric field in coupled lattice to satisfy Gauss constraint.

    Creates a new lattice with $E$ corrected to enforce $\\nabla \\cdot E + g\\rho = 0$.

    Args:
        lattice: Lattice to project.
        tolerance: Convergence tolerance.
        max_iterations: Maximum iterations for solver.

    Returns:
        `CoupledLattice` with projected electric field.
    """
    E_projected = project_gauss_constraint(
        lattice.E,
        lattice.phi,
        lattice.pi,
        lattice.dx,
        tolerance,
        max_iterations,
    )

    return CoupledLattice(
        phi=lattice.phi,
        pi=lattice.pi,
        links=lattice.links,
        E=E_projected,
        m=lattice.m,
        lambda_=lattice.lambda_,
        g=lattice.g,
        dx=lattice.dx,
        size=lattice.size,
        length=lattice.length,
    )

stabilized_constraint_evolution(G, dG_dt, alpha=0.1, beta=0.01)

Baumgarte-Shapiro constraint damping.

Instead of enforcing \(dG/dt = 0\), enforce:

\[\frac{d^2 G}{dt^2} + \alpha \frac{dG}{dt} + \beta G = 0\]

This exponentially damps constraint violations.

Parameters:

Name	Type	Description	Default
G	Array	Current constraint violation	required
dG_dt	Array	Time derivative of constraint	required
alpha	float	Linear damping coefficient \(\alpha\)	0.1
beta	float	Constraint restoring coefficient \(\beta\)	0.01

Returns:

Type	Description
Array	Modified \(d^2 G/dt^2\) with damping

Source code in jaxlatt/core/constraints.py

@jit
def stabilized_constraint_evolution(
    G: Array,
    dG_dt: Array,
    alpha: float = 0.1,
    beta: float = 0.01,
) -> Array:
    """
    Baumgarte-Shapiro constraint damping.

    Instead of enforcing $dG/dt = 0$, enforce:

    $$\\frac{d^2 G}{dt^2} + \\alpha \\frac{dG}{dt} + \\beta G = 0$$

    This exponentially damps constraint violations.

    Args:
        G: Current constraint violation
        dG_dt: Time derivative of constraint
        alpha: Linear damping coefficient $\\alpha$
        beta: Constraint restoring coefficient $\\beta$

    Returns:
        Modified $d^2 G/dt^2$ with damping
    """
    return -alpha * dG_dt - beta * G

gauge_transform_scalar(phi, alpha, g)

Apply gauge transformation to scalar field.

Under a U(1) gauge transformation with phase \(\alpha(x)\):

\[\phi(x) \to e^{ig\alpha(x)} \phi(x)\]

Parameters:

Name	Type	Description	Default
phi	Array	Scalar field \(\phi\), shape (N, N, N).	required
alpha	Array	Gauge transformation parameter \(\alpha\), shape (N, N, N).	required
g	float	Gauge coupling.	required

Returns:

Type	Description
Array	Transformed scalar field.

Source code in jaxlatt/core/constraints.py

@jit
def gauge_transform_scalar(phi: Array, alpha: Array, g: float) -> Array:
    """
    Apply gauge transformation to scalar field.

    Under a U(1) gauge transformation with phase $\\alpha(x)$:

    $$\\phi(x) \\to e^{ig\\alpha(x)} \\phi(x)$$

    Args:
        phi: Scalar field $\\phi$, shape `(N, N, N)`.
        alpha: Gauge transformation parameter $\\alpha$, shape `(N, N, N)`.
        g: Gauge coupling.

    Returns:
        Transformed scalar field.
    """
    return jnp.exp(1j * g * alpha) * phi

gauge_transform_links(links, alpha, g)

Apply gauge transformation to gauge links.

Under a U(1) gauge transformation:

\[U_i(n) \to e^{ig\alpha(n)} U_i(n) e^{-ig\alpha(n+\hat{i})}\]

This keeps the gauge-covariant derivative properly transformed.

Parameters:

Name	Type	Description	Default
links	Array	Gauge links, shape (3, N, N, N).	required
alpha	Array	Gauge transformation parameter \(\alpha\), shape (N, N, N).	required
g	float	Gauge coupling.	required

Returns:

Type	Description
Array	Transformed gauge links.

Source code in jaxlatt/core/constraints.py

@jit
def gauge_transform_links(links: Array, alpha: Array, g: float) -> Array:
    """
    Apply gauge transformation to gauge links.

    Under a U(1) gauge transformation:

    $$U_i(n) \\to e^{ig\\alpha(n)} U_i(n) e^{-ig\\alpha(n+\\hat{i})}$$

    This keeps the gauge-covariant derivative properly transformed.

    Args:
        links: Gauge links, shape `(3, N, N, N)`.
        alpha: Gauge transformation parameter $\\alpha$, shape `(N, N, N)`.
        g: Gauge coupling.

    Returns:
        Transformed gauge links.
    """
    links_new = jnp.zeros_like(links)

    for i in range(3):
        alpha_n = alpha
        alpha_n_plus_i = jnp.roll(alpha, -1, axis=i)
        phase = jnp.exp(1j * g * (alpha_n - alpha_n_plus_i))
        links_new = links_new.at[i].set(phase * links[i])

    return links_new

gauge_transform_coupled_lattice(lattice, alpha)

Apply gauge transformation to entire coupled lattice.

Transforms all fields simultaneously:

• \(\phi \to e^{ig\alpha} \phi\)
• \(\pi \to e^{ig\alpha} \pi\)
• \(U_i(n) \to e^{ig\alpha(n)} U_i(n) e^{-ig\alpha(n+\hat{i})}\)
• \(E\) unchanged (gauge invariant)

Parameters:

Name	Type	Description	Default
lattice	CoupledLattice	Lattice to transform.	required
alpha	Array	Gauge transformation parameter \(\alpha\), shape (N, N, N).	required

Returns:

Type	Description
CoupledLattice	Transformed CoupledLattice.

Source code in jaxlatt/core/constraints.py

def gauge_transform_coupled_lattice(lattice: CoupledLattice, alpha: Array) -> CoupledLattice:
    """
    Apply gauge transformation to entire coupled lattice.

    Transforms all fields simultaneously:

    - $\\phi \\to e^{ig\\alpha} \\phi$
    - $\\pi \\to e^{ig\\alpha} \\pi$
    - $U_i(n) \\to e^{ig\\alpha(n)} U_i(n) e^{-ig\\alpha(n+\\hat{i})}$
    - $E$ unchanged (gauge invariant)

    Args:
        lattice: Lattice to transform.
        alpha: Gauge transformation parameter $\\alpha$, shape `(N, N, N)`.

    Returns:
        Transformed `CoupledLattice`.
    """
    phi_new = gauge_transform_scalar(lattice.phi, alpha, lattice.g)
    pi_new = gauge_transform_scalar(lattice.pi, alpha, lattice.g)
    links_new = gauge_transform_links(lattice.links, alpha, lattice.g)

    return CoupledLattice(
        phi=phi_new,
        pi=pi_new,
        links=links_new,
        E=lattice.E,  # E is gauge invariant
        m=lattice.m,
        lambda_=lattice.lambda_,
        g=lattice.g,
        dx=lattice.dx,
        size=lattice.size,
        length=lattice.length,
    )

check_gauge_invariance(observable_before, observable_after, tolerance=1e-10)

Check if an observable is gauge invariant.

Parameters:

Name	Type	Description	Default
observable_before	float	Observable value before gauge transformation.	required
observable_after	float	Observable value after gauge transformation.	required
tolerance	float	Allowed numerical tolerance.	1e-10

Returns:

Type	Description
Array	Boolean JAX scalar indicating whether the observable is gauge invariant.

Source code in jaxlatt/core/constraints.py

@jit
def check_gauge_invariance(
    observable_before: float, observable_after: float, tolerance: float = 1e-10
) -> Array:
    """
    Check if an observable is gauge invariant.

    Args:
        observable_before: Observable value before gauge transformation.
        observable_after: Observable value after gauge transformation.
        tolerance: Allowed numerical tolerance.

    Returns:
        Boolean JAX scalar indicating whether the observable is gauge invariant.
    """
    return jnp.abs(observable_before - observable_after) < tolerance

compute_constraint_violation_norm(phi, pi, E, links, dx)

Compute various norms of constraint violation.

Parameters:

Name	Type	Description	Default
phi	Array	Scalar field	required
pi	Array	Conjugate momentum	required
E	Array	Electric field	required
links	Array	Gauge links	required
dx	float	Lattice spacing	required

Returns:

Type	Description
dict[str, float]	Dict with L1, L2, Linf norms of G

Source code in jaxlatt/core/constraints.py

def compute_constraint_violation_norm(
    phi: Array,
    pi: Array,
    E: Array,
    links: Array,
    dx: float,
) -> dict[str, float]:
    """
    Compute various norms of constraint violation.

    Args:
        phi: Scalar field
        pi: Conjugate momentum
        E: Electric field
        links: Gauge links
        dx: Lattice spacing

    Returns:
        Dict with L1, L2, Linf norms of G
    """
    G = gauss_constraint_with_source(links, E, phi, pi, dx)

    return {
        "L1_norm": float(jnp.mean(jnp.abs(G))),
        "L2_norm": float(jnp.sqrt(jnp.mean(G**2))),
        "Linf_norm": float(jnp.max(jnp.abs(G))),
        "total_violation": float(jnp.sum(jnp.abs(G)) * dx**3),
    }

check_constraint_preservation(phi_old, pi_old, E_old, links_old, phi_new, pi_new, E_new, links_new, dx)

Check if time evolution preserves Gauss constraint.

Parameters:

Name	Type	Description	Default
phi_old	Array	Scalar field at previous time.	required
pi_old	Array	Conjugate momentum at previous time.	required
E_old	Array	Electric field at previous time.	required
links_old	Array	Gauge links at previous time.	required
phi_new	Array	Scalar field at current time.	required
pi_new	Array	Conjugate momentum at current time.	required
E_new	Array	Electric field at current time.	required
links_new	Array	Gauge links at current time.	required
dx	float	Lattice spacing.	required

Returns:

Type	Description
dict[str, float]	Dictionary containing Gauss-constraint violation metrics before and after
dict[str, float]	the update, including absolute drift and relative drift rate.

Source code in jaxlatt/core/constraints.py

def check_constraint_preservation(
    phi_old: Array,
    pi_old: Array,
    E_old: Array,
    links_old: Array,
    phi_new: Array,
    pi_new: Array,
    E_new: Array,
    links_new: Array,
    dx: float,
) -> dict[str, float]:
    """
    Check if time evolution preserves Gauss constraint.

    Args:
        phi_old: Scalar field at previous time.
        pi_old: Conjugate momentum at previous time.
        E_old: Electric field at previous time.
        links_old: Gauge links at previous time.
        phi_new: Scalar field at current time.
        pi_new: Conjugate momentum at current time.
        E_new: Electric field at current time.
        links_new: Gauge links at current time.
        dx: Lattice spacing.

    Returns:
        Dictionary containing Gauss-constraint violation metrics before and after
        the update, including absolute drift and relative drift rate.
    """
    G_old = gauss_constraint_with_source(links_old, E_old, phi_old, pi_old, dx)
    G_new = gauss_constraint_with_source(links_new, E_new, phi_new, pi_new, dx)

    violation_old = float(jnp.sqrt(jnp.mean(G_old**2)))
    violation_new = float(jnp.sqrt(jnp.mean(G_new**2)))
    drift = violation_new - violation_old

    return {
        "violation_before": violation_old,
        "violation_after": violation_new,
        "drift": drift,
        "drift_rate": drift / (violation_old + 1e-10),
    }

charge_current_continuity(lattice)

Compute charge and current densities for continuity equation check.

The continuity equation for conserved charge:

\[\frac{d\rho}{dt} + \nabla \cdot j = 0\]

where:

• \(\rho = g \, \mathrm{Im}(\phi^* \pi)\) (charge density)
• \(j_i = \frac{g}{2 \, dx} \mathrm{Im}[\phi^*(n) U_i(n) \phi(n+\hat{i})]\) (current density)

Parameters:

Name	Type	Description	Default
lattice	CoupledLattice	Coupled lattice state.	required

Returns:

Type	Description
tuple[ndarray, ndarray]	Tuple (rho, div_j) with charge density and current divergence arrays.

Source code in jaxlatt/core/constraints.py

def charge_current_continuity(
    lattice: CoupledLattice,
) -> tuple[jnp.ndarray, jnp.ndarray]:
    """
    Compute charge and current densities for continuity equation check.

    The continuity equation for conserved charge:

    $$\\frac{d\\rho}{dt} + \\nabla \\cdot j = 0$$

    where:

    - $\\rho = g \\, \\mathrm{Im}(\\phi^* \\pi)$ (charge density)
    - $j_i = \\frac{g}{2 \\, dx} \\mathrm{Im}[\\phi^*(n) U_i(n) \\phi(n+\\hat{i})]$ (current density)

    Args:
        lattice: Coupled lattice state.

    Returns:
        Tuple `(rho, div_j)` with charge density and current divergence arrays.
    """
    from jaxlatt.operators.coupled import scalar_current_density

    # Charge density
    rho = scalar_charge_density(lattice.phi, lattice.pi, lattice.g)

    # Current divergence
    div_j = jnp.zeros_like(lattice.phi, dtype=jnp.float32)

    for i in range(3):
        j_i = scalar_current_density(lattice.phi, lattice.links, i, lattice.g, lattice.dx)
        j_i_backward = jnp.roll(j_i, 1, axis=i)
        div_j = div_j + (j_i - j_i_backward) / lattice.dx

    return rho, div_j