Gauge Operators

gauge

Gauge field operators for U(1) lattice gauge theory.

Implements:

• Plaquettes and staples
• Magnetic field extraction
• Electric field operations
• Gauge forces and constraints

plaquette(links, i, j)

Compute spatial plaquette \(U_{ij}\) for all lattice sites.

\[U_{ij}(n) = U_i(n) \, U_j(n+\hat{i}) \, U_i^*(n+\hat{j}) \, U_j^*(n)\]

For U(1), this is a complex number with \(|U_{ij}| \leq 1\). In the continuum limit: \(U_{ij} \approx \exp(-i g \, dx^2 F_{ij})\).

Parameters:

Name	Type	Description	Default
links	Array	Link variables (3, N, N, N) complex	required
i	int	First spatial direction (0, 1, or 2) - static for JIT	required
j	int	Second spatial direction (0, 1, or 2) - static for JIT	required

Returns:

Type	Description
Array	Plaquette values at all sites (N, N, N) complex

Source code in jaxlatt/operators/gauge.py

@partial(jit, static_argnums=(1, 2))
def plaquette(links: Array, i: int, j: int) -> Array:
    """
    Compute spatial plaquette $U_{ij}$ for all lattice sites.

    $$U_{ij}(n) = U_i(n) \\, U_j(n+\\hat{i}) \\, U_i^*(n+\\hat{j}) \\, U_j^*(n)$$

    For U(1), this is a complex number with $|U_{ij}| \\leq 1$.
    In the continuum limit: $U_{ij} \\approx \\exp(-i g \\, dx^2 F_{ij})$.

    Args:
        links: Link variables (3, N, N, N) complex
        i: First spatial direction (0, 1, or 2) - static for JIT
        j: Second spatial direction (0, 1, or 2) - static for JIT

    Returns:
        Plaquette values at all sites (N, N, N) complex
    """
    Ui = links[i]
    Uj = links[j]

    # U_j(n+i): shift j-links forward in i direction
    Uj_shift_i = jnp.roll(Uj, shift=-1, axis=i)

    # U_i(n+j): shift i-links forward in j direction
    Ui_shift_j = jnp.roll(Ui, shift=-1, axis=j)

    # Plaquette: forward-forward-back-back path
    U_plaq = Ui * Uj_shift_i * jnp.conj(Ui_shift_j) * jnp.conj(Uj)

    return U_plaq

all_plaquettes(links)

Compute all three spatial plaquettes (xy, yz, zx).

Parameters:

Name	Type	Description	Default
links	Array	Link variables (3, N, N, N)	required

Returns:

Type	Description
tuple[Array, Array, Array]	Tuple of (U_01, U_12, U_20) plaquettes, each (N, N, N) complex

Source code in jaxlatt/operators/gauge.py

@jit
def all_plaquettes(links: Array) -> tuple[Array, Array, Array]:
    """
    Compute all three spatial plaquettes (xy, yz, zx).

    Args:
        links: Link variables (3, N, N, N)

    Returns:
        Tuple of (U_01, U_12, U_20) plaquettes, each (N, N, N) complex
    """
    U_xy = plaquette(links, 0, 1)
    U_yz = plaquette(links, 1, 2)
    U_zx = plaquette(links, 2, 0)

    return U_xy, U_yz, U_zx

magnetic_field(links, dx, g)

Compute magnetic field \(B_k\) from plaquettes.

\[B_k = \frac{1}{2g \, dx^2} \epsilon_{kij} \operatorname{Im}(U_{ij})\]

where \(\epsilon_{kij}\) is the Levi-Civita symbol.

Parameters:

Name	Type	Description	Default
links	Array	Link variables (3, N, N, N)	required
dx	float	Lattice spacing	required
g	float	Gauge coupling	required

Returns:

Type	Description
Array	Magnetic field components (3, N, N, N) real

Source code in jaxlatt/operators/gauge.py

@jit
def magnetic_field(links: Array, dx: float, g: float) -> Array:
    """
    Compute magnetic field $B_k$ from plaquettes.

    $$B_k = \\frac{1}{2g \\, dx^2} \\epsilon_{kij} \\operatorname{Im}(U_{ij})$$

    where $\\epsilon_{kij}$ is the Levi-Civita symbol.

    Args:
        links: Link variables (3, N, N, N)
        dx: Lattice spacing
        g: Gauge coupling

    Returns:
        Magnetic field components (3, N, N, N) real
    """
    U_xy, U_yz, U_zx = all_plaquettes(links)

    prefactor = 1.0 / (2.0 * g * dx * dx)

    # B_x from yz plaquette, B_y from zx, B_z from xy
    Bx = prefactor * jnp.imag(U_yz)
    By = prefactor * jnp.imag(U_zx)
    Bz = prefactor * jnp.imag(U_xy)

    B = jnp.stack([Bx, By, Bz], axis=0)

    # No cast: jnp.imag of a complex array already yields the matching
    # real dtype, so this follows jax_enable_x64 instead of defeating it.
    return B

staple(links, i)

Compute staple \(S_i(n)\) for link \(i\) at all sites.

The staple is the sum of 4 plaquettes touching link \(U_i(n)\), needed for computing the force \(\partial H / \partial A_i\).

Parameters:

Name	Type	Description	Default
links	Array	Link variables (3, N, N, N)	required
i	int	Direction of link - static for JIT	required

Returns:

Type	Description
Array	Staple values at all sites (N, N, N) complex

Source code in jaxlatt/operators/gauge.py

@partial(jit, static_argnums=(1,))
def staple(links: Array, i: int) -> Array:
    """
    Compute staple $S_i(n)$ for link $i$ at all sites.

    The staple is the sum of 4 plaquettes touching link $U_i(n)$,
    needed for computing the force $\\partial H / \\partial A_i$.

    Args:
        links: Link variables (3, N, N, N)
        i: Direction of link - static for JIT

    Returns:
        Staple values at all sites (N, N, N) complex
    """
    if i == 0:
        return _staple_contribution(links, 0, 1) + _staple_contribution(links, 0, 2)
    elif i == 1:
        return _staple_contribution(links, 1, 0) + _staple_contribution(links, 1, 2)
    else:  # i == 2
        return _staple_contribution(links, 2, 0) + _staple_contribution(links, 2, 1)

all_staples(links)

Compute staples for all three link directions.

Parameters:

Name	Type	Description	Default
links	Array	Link variables (3, N, N, N)	required

Returns:

Type	Description
Array	Staples for all directions (3, N, N, N) complex

Source code in jaxlatt/operators/gauge.py

@jit
def all_staples(links: Array) -> Array:
    """
    Compute staples for all three link directions.

    Args:
        links: Link variables (3, N, N, N)

    Returns:
        Staples for all directions (3, N, N, N) complex
    """
    S0 = staple(links, 0)
    S1 = staple(links, 1)
    S2 = staple(links, 2)

    return jnp.stack([S0, S1, S2], axis=0)

gauge_force(links, dx, g)

Compute force \(F_i = \partial H_\mathrm{mag} / \partial \theta_i\) for the magnetic Hamiltonian.

\(H_\mathrm{mag} = \frac{1}{2}\int B^2 \, dV\) with \(B_k = \operatorname{Im}(U_{ij})/(2g \, dx^2)\), giving

\[H_\mathrm{mag} = \Sigma_n \frac{\operatorname{Im}(P_{ij}(n))^2}{8g^2 \, dx}\]

The exact gradient is:

\[F_i(n) = \Sigma_{j \neq i} \frac{\operatorname{Im}(P_{ij}(n))\operatorname{Re}(P_{ij}(n)) - \operatorname{Im}(P_{ij}(n-\hat{e}_j))\operatorname{Re}(P_{ij}(n-\hat{e}_j))}{4g^2 \, dx}\]

Parameters:

Name	Type	Description	Default
links	Array	Link variables (3, N, N, N)	required
dx	float	Lattice spacing	required
g	float	Gauge coupling	required

Returns:

Type	Description
Array	Force components (3, N, N, N) real

Source code in jaxlatt/operators/gauge.py

@jit
def gauge_force(links: Array, dx: float, g: float) -> Array:
    """
    Compute force $F_i = \\partial H_\\mathrm{mag} / \\partial \\theta_i$ for the magnetic Hamiltonian.

    $H_\\mathrm{mag} = \\frac{1}{2}\\int B^2 \\, dV$ with $B_k = \\operatorname{Im}(U_{ij})/(2g \\, dx^2)$, giving

    $$H_\\mathrm{mag} = \\Sigma_n \\frac{\\operatorname{Im}(P_{ij}(n))^2}{8g^2 \\, dx}$$

    The exact gradient is:

    $$F_i(n) = \\Sigma_{j \\neq i} \\frac{\\operatorname{Im}(P_{ij}(n))\\operatorname{Re}(P_{ij}(n)) - \\operatorname{Im}(P_{ij}(n-\\hat{e}_j))\\operatorname{Re}(P_{ij}(n-\\hat{e}_j))}{4g^2 \\, dx}$$

    Args:
        links: Link variables (3, N, N, N)
        dx: Lattice spacing
        g: Gauge coupling

    Returns:
        Force components (3, N, N, N) real
    """
    prefactor = 1.0 / (4.0 * g**2 * dx)

    def _contrib(i: int, j: int) -> Array:
        P_fwd = plaquette(links, i, j)
        P_bwd = jnp.roll(P_fwd, shift=1, axis=j)
        return (jnp.imag(P_fwd) * jnp.real(P_fwd) - jnp.imag(P_bwd) * jnp.real(P_bwd)) * prefactor

    F0 = _contrib(0, 1) + _contrib(0, 2)
    F1 = _contrib(1, 0) + _contrib(1, 2)
    F2 = _contrib(2, 0) + _contrib(2, 1)

    return jnp.stack([F0, F1, F2], axis=0)

magnetic_energy_density(links, dx, g)

Compute total magnetic energy \(E_B = \frac{1}{2} \int B^2 \, dV\).

On the lattice:

\[E_B = \frac{1}{2} \Sigma_n \Sigma_k B_k(n)^2 \, dx^3\]

Parameters:

Name	Type	Description	Default
links	Array	Link variables	required
dx	float	Lattice spacing	required
g	float	Gauge coupling	required

Returns:

Type	Description
Array	Total magnetic energy (as JAX scalar)

Source code in jaxlatt/operators/gauge.py

@jit
def magnetic_energy_density(links: Array, dx: float, g: float) -> Array:
    """
    Compute total magnetic energy $E_B = \\frac{1}{2} \\int B^2 \\, dV$.

    On the lattice:

    $$E_B = \\frac{1}{2} \\Sigma_n \\Sigma_k B_k(n)^2 \\, dx^3$$

    Args:
        links: Link variables
        dx: Lattice spacing
        g: Gauge coupling

    Returns:
        Total magnetic energy (as JAX scalar)
    """
    B = magnetic_field(links, dx, g)
    dV = dx**3
    E_mag = 0.5 * jnp.sum(B**2) * dV
    return E_mag

electric_energy_density(E, dx)

Compute total electric energy \(E_E = \frac{1}{2} \int E^2 \, dV\).

On the lattice:

\[E_E = \frac{1}{2} \Sigma_n \Sigma_i E_i(n)^2 \, dx^3\]

Parameters:

Name	Type	Description	Default
E	Array	Electric field components (3, N, N, N)	required
dx	float	Lattice spacing	required

Returns:

Type	Description
Array	Total electric energy (as JAX scalar)

Source code in jaxlatt/operators/gauge.py

@jit
def electric_energy_density(E: Array, dx: float) -> Array:
    """
    Compute total electric energy $E_E = \\frac{1}{2} \\int E^2 \\, dV$.

    On the lattice:

    $$E_E = \\frac{1}{2} \\Sigma_n \\Sigma_i E_i(n)^2 \\, dx^3$$

    Args:
        E: Electric field components (3, N, N, N)
        dx: Lattice spacing

    Returns:
        Total electric energy (as JAX scalar)
    """
    dV = dx**3
    E_elec = 0.5 * jnp.sum(E**2) * dV
    return E_elec

divergence_3d(field, dx)

Compute divergence of a 3-component vector field.

\[\nabla \cdot F = \frac{\partial F_x}{\partial x} + \frac{\partial F_y}{\partial y} + \frac{\partial F_z}{\partial z}\]

Uses backward difference: \((F_i(n) - F_i(n-\hat{i})) / dx\).

Parameters:

Name	Type	Description	Default
field	Array	Vector field components (3, N, N, N)	required
dx	float	Lattice spacing	required

Returns:

Type	Description
Array	Divergence at each site (N, N, N)

Source code in jaxlatt/operators/gauge.py

@jit
def divergence_3d(field: Array, dx: float) -> Array:
    """
    Compute divergence of a 3-component vector field.

    $$\\nabla \\cdot F = \\frac{\\partial F_x}{\\partial x} + \\frac{\\partial F_y}{\\partial y} + \\frac{\\partial F_z}{\\partial z}$$

    Uses backward difference: $(F_i(n) - F_i(n-\\hat{i})) / dx$.

    Args:
        field: Vector field components (3, N, N, N)
        dx: Lattice spacing

    Returns:
        Divergence at each site (N, N, N)
    """
    div = (
        (field[0] - jnp.roll(field[0], 1, axis=0))
        + (field[1] - jnp.roll(field[1], 1, axis=1))
        + (field[2] - jnp.roll(field[2], 1, axis=2))
    ) / dx
    return div

gauss_constraint(links, E, dx)

Compute Gauss constraint violation \(G(n) = \nabla \cdot E\) for pure gauge (no sources).

Parameters:

Name	Type	Description	Default
links	Array	Link variables (not used, kept for interface consistency)	required
E	Array	Electric field (3, N, N, N)	required
dx	float	Lattice spacing	required

Returns:

Type	Description
Array	Gauss violation at each site (N, N, N)

Source code in jaxlatt/operators/gauge.py

@jit
def gauss_constraint(links: Array, E: Array, dx: float) -> Array:
    """
    Compute Gauss constraint violation $G(n) = \\nabla \\cdot E$ for pure gauge (no sources).

    Args:
        links: Link variables (not used, kept for interface consistency)
        E: Electric field (3, N, N, N)
        dx: Lattice spacing

    Returns:
        Gauss violation at each site (N, N, N)
    """
    return divergence_3d(E, dx)

gauss_violation_norm(links, E, dx)

Compute RMS norm of Gauss constraint violation.

Returns \(\sqrt{\langle |\nabla \cdot E|^2 \rangle}\).

Parameters:

Name	Type	Description	Default
links	Array	Link variables (not used, kept for interface consistency)	required
E	Array	Electric field (3, N, N, N)	required
dx	float	Lattice spacing	required

Returns:

Type	Description
Array	RMS Gauss violation (scalar)

Source code in jaxlatt/operators/gauge.py

@jit
def gauss_violation_norm(links: Array, E: Array, dx: float) -> Array:
    """
    Compute RMS norm of Gauss constraint violation.

    Returns $\\sqrt{\\langle |\\nabla \\cdot E|^2 \\rangle}$.

    Args:
        links: Link variables (not used, kept for interface consistency)
        E: Electric field (3, N, N, N)
        dx: Lattice spacing

    Returns:
        RMS Gauss violation (scalar)
    """
    G = gauss_constraint(links, E, dx)
    return jnp.sqrt(jnp.mean(G**2))

gauge_energy(links, E, dx, g)

Compute all energy components for pure gauge field.

Parameters:

Name	Type	Description	Default
links	Array	Gauge links (3, N, N, N)	required
E	Array	Electric field (3, N, N, N)	required
dx	float	Lattice spacing	required
g	float	Gauge coupling	required

Returns:

Type	Description
dict	Dictionary with JAX scalar 'electric', 'magnetic', and 'total' energies

Source code in jaxlatt/operators/gauge.py

@jit
def gauge_energy(links: Array, E: Array, dx: float, g: float) -> dict:
    """
    Compute all energy components for pure gauge field.

    Args:
        links: Gauge links (3, N, N, N)
        E: Electric field (3, N, N, N)
        dx: Lattice spacing
        g: Gauge coupling

    Returns:
        Dictionary with JAX scalar 'electric', 'magnetic', and 'total' energies
    """
    E_elec = 0.5 * jnp.sum(E**2)
    E_mag = magnetic_energy_density(links, dx, g)

    return {
        "electric": E_elec,
        "magnetic": E_mag,
        "total": E_elec + E_mag,
    }

evolve_links_half_step(links, E, dt)

Evolve gauge links by half time step: \(U o U \exp(i E \, d\tau/2)\).

This is the position update in the leapfrog scheme. Links are normalized to maintain unitarity.

Parameters:

Name	Type	Description	Default
links	Array	Gauge links (3, N, N, N) complex	required
E	Array	Electric field (3, N, N, N) real	required
dt	float	Time step	required

Returns:

Type	Description
Array	Updated links (3, N, N, N), normalized to unit modulus

Source code in jaxlatt/operators/gauge.py

@jit
def evolve_links_half_step(links: Array, E: Array, dt: float) -> Array:
    """
    Evolve gauge links by half time step: $U \to U \\exp(i E \\, d\\tau/2)$.

    This is the position update in the leapfrog scheme.
    Links are normalized to maintain unitarity.

    Args:
        links: Gauge links (3, N, N, N) complex
        E: Electric field (3, N, N, N) real
        dt: Time step

    Returns:
        Updated links (3, N, N, N), normalized to unit modulus
    """
    phase_kick = jnp.exp(1j * E * dt / 2)
    links_new = links * phase_kick
    return links_new / jnp.abs(links_new)

evolve_links_full_step(links, E, dt)

Evolve gauge links by full time step: \(U o U \exp(i E \, d\tau)\).

Parameters:

Name	Type	Description	Default
links	Array	Gauge links (3, N, N, N) complex	required
E	Array	Electric field (3, N, N, N) real	required
dt	float	Time step	required

Returns:

Type	Description
Array	Updated links (3, N, N, N), normalized to unit modulus

Source code in jaxlatt/operators/gauge.py

@jit
def evolve_links_full_step(links: Array, E: Array, dt: float) -> Array:
    """
    Evolve gauge links by full time step: $U \to U \\exp(i E \\, d\\tau)$.

    Args:
        links: Gauge links (3, N, N, N) complex
        E: Electric field (3, N, N, N) real
        dt: Time step

    Returns:
        Updated links (3, N, N, N), normalized to unit modulus
    """
    phase_kick = jnp.exp(1j * E * dt)
    links_new = links * phase_kick
    return links_new / jnp.abs(links_new)

gauge_force_single_direction(phi, links, direction, g, dx)

Compute force on electric field \(E_i\) including scalar current.

\[F_i = \partial H_\mathrm{mag}/\partial \theta_i + (\text{scalar current})\]

Parameters:

Name	Type	Description	Default
phi	Array	Scalar field (N, N, N)	required
links	Array	Gauge links (3, N, N, N)	required
direction	int	Direction (0=x, 1=y, 2=z) — static for JIT	required
g	float	Gauge coupling	required
dx	float	Lattice spacing	required

Returns:

Type	Description
Array	Force on \(E_i\) (real-valued, shape (N, N, N))

Source code in jaxlatt/operators/gauge.py

@partial(jit, static_argnums=(2,))
def gauge_force_single_direction(
    phi: Array,
    links: Array,
    direction: int,
    g: float,
    dx: float,
) -> Array:
    r"""
    Compute force on electric field $E_i$ including scalar current.

    $$F_i = \partial H_\mathrm{mag}/\partial \theta_i + (\text{scalar current})$$

    Args:
        phi: Scalar field (N, N, N)
        links: Gauge links (3, N, N, N)
        direction: Direction (0=x, 1=y, 2=z) — static for JIT
        g: Gauge coupling
        dx: Lattice spacing

    Returns:
        Force on $E_i$ (real-valued, shape (N, N, N))
    """
    prefactor = 1.0 / (4.0 * g**2 * dx)
    pure_gauge_force = jnp.zeros(links.shape[1:], dtype=jnp.float32)
    for j in [k for k in range(3) if k != direction]:
        P_fwd = plaquette(links, direction, j)
        P_bwd = jnp.roll(P_fwd, shift=1, axis=j)
        pure_gauge_force = (
            pure_gauge_force
            + (jnp.imag(P_fwd) * jnp.real(P_fwd) - jnp.imag(P_bwd) * jnp.real(P_bwd)) * prefactor
        )

    U_i = links[direction]
    phi_forward = jnp.roll(phi, -1, axis=direction)
    current = jnp.imag(jnp.conj(phi) * U_i * phi_forward) / dx**2

    return pure_gauge_force + current

gauge_force_all_directions(phi, links, g, dx)

Compute coupled gauge-force for all three directions (flat spacetime).

Parameters:

Name	Type	Description	Default
phi	Array	Scalar field array.	required
links	Array	Gauge links array.	required
g	float	Gauge coupling.	required
dx	float	Lattice spacing.	required

Returns:

Type	Description
Array	Force array of shape (3, N, N, N).

Source code in jaxlatt/operators/gauge.py

@jit
def gauge_force_all_directions(
    phi: Array,
    links: Array,
    g: float,
    dx: float,
) -> Array:
    """Compute coupled gauge-force for all three directions (flat spacetime).

    Args:
        phi: Scalar field array.
        links: Gauge links array.
        g: Gauge coupling.
        dx: Lattice spacing.

    Returns:
        Force array of shape ``(3, N, N, N)``.
    """
    F0 = gauge_force_single_direction(phi, links, 0, g, dx)
    F1 = gauge_force_single_direction(phi, links, 1, g, dx)
    F2 = gauge_force_single_direction(phi, links, 2, g, dx)
    return jnp.stack([F0, F1, F2], axis=0)

gauge_force_expanding(phi, links, direction, g, dx, a)

Compute force on electric field in one direction for an FRW background.

Parameters:

Name	Type	Description	Default
phi	Array	Scalar field	required
links	Array	Gauge links	required
direction	int	Direction (0, 1, 2) — static for JIT	required
g	float	Gauge coupling	required
dx	float	Comoving lattice spacing	required
a	float	Scale factor	required

Returns:

Type	Description
Array	Force on \(E_i\) (shape (N, N, N))

Source code in jaxlatt/operators/gauge.py

@partial(jit, static_argnums=(2,))
def gauge_force_expanding(
    phi: Array,
    links: Array,
    direction: int,
    g: float,
    dx: float,
    a: float,
) -> Array:
    """Compute force on electric field in one direction for an FRW background.

    Args:
        phi: Scalar field
        links: Gauge links
        direction: Direction (0, 1, 2) — static for JIT
        g: Gauge coupling
        dx: Comoving lattice spacing
        a: Scale factor

    Returns:
        Force on $E_i$ (shape (N, N, N))
    """
    gauge_staple = staple(links, direction)
    U_i = links[direction]

    pure_gauge_force = g * dx * jnp.imag(U_i * jnp.conj(gauge_staple))

    phi_forward = jnp.roll(phi, -1, axis=direction)
    current = jnp.imag(jnp.conj(phi) * U_i * phi_forward) / dx**2

    return pure_gauge_force + current

gauge_force_all_directions_expanding(phi, links, g, dx, a)

Compute coupled gauge-force for all three directions in an FRW background.

Parameters:

Name	Type	Description	Default
phi	Array	Scalar field array.	required
links	Array	Gauge links array.	required
g	float	Gauge coupling.	required
dx	float	Comoving lattice spacing.	required
a	float	Scale factor.	required

Returns:

Type	Description
Array	Force array of shape (3, N, N, N).

Source code in jaxlatt/operators/gauge.py

@jit
def gauge_force_all_directions_expanding(
    phi: Array,
    links: Array,
    g: float,
    dx: float,
    a: float,
) -> Array:
    """Compute coupled gauge-force for all three directions in an FRW background.

    Args:
        phi: Scalar field array.
        links: Gauge links array.
        g: Gauge coupling.
        dx: Comoving lattice spacing.
        a: Scale factor.

    Returns:
        Force array of shape ``(3, N, N, N)``.
    """
    F0 = gauge_force_expanding(phi, links, 0, g, dx, a)
    F1 = gauge_force_expanding(phi, links, 1, g, dx, a)
    F2 = gauge_force_expanding(phi, links, 2, g, dx, a)
    return jnp.stack([F0, F1, F2], axis=0)

Coupled Scalar-Gauge Operators

coupled

Coupled scalar-gauge operators for Abelian Higgs model.

Implements gauge-covariant operations for charged scalar fields:

• Covariant derivatives and Laplacian
• Covariant gradient energy
• Scalar forces with gauge coupling
• Charge density and current
• Gauss law with scalar sources

covariant_derivative(phi, links, direction, dx)

Compute covariant derivative \(D_i \phi = (U_i(n) \phi(n+\hat{i}) - \phi(n)) / dx\).

For a charged field, the covariant derivative includes the gauge link to maintain gauge invariance.

Parameters:

Name	Type	Description	Default
phi	Array	Complex scalar field (N, N, N)	required
links	Array	Gauge links (3, N, N, N)	required
direction	int	Direction (0=x, 1=y, 2=z) - static for JIT	required
dx	float	Lattice spacing	required

Returns:

Type	Description
Array	\(D_i \phi\): Covariant derivative (N, N, N)

Source code in jaxlatt/operators/coupled.py

@partial(jit, static_argnums=(2,))
def covariant_derivative(phi: Array, links: Array, direction: int, dx: float) -> Array:
    """
    Compute covariant derivative $D_i \\phi = (U_i(n) \\phi(n+\\hat{i}) - \\phi(n)) / dx$.

    For a charged field, the covariant derivative includes
    the gauge link to maintain gauge invariance.

    Args:
        phi: Complex scalar field (N, N, N)
        links: Gauge links (3, N, N, N)
        direction: Direction (0=x, 1=y, 2=z) - static for JIT
        dx: Lattice spacing

    Returns:
        $D_i \\phi$: Covariant derivative (N, N, N)
    """
    phi_shifted = jnp.roll(phi, -1, axis=direction)
    U_i = links[direction]

    return (U_i * phi_shifted - phi) / dx

covariant_gradient_squared(phi, links, dx)

Compute \(\Sigma_i |D_i \phi|^2\) at each site.

This is the kinetic energy density for the scalar field in the presence of gauge fields.

Parameters:

Name	Type	Description	Default
phi	Array	Complex scalar field (N, N, N)	required
links	Array	Gauge links (3, N, N, N)	required
dx	float	Lattice spacing	required

Returns:

Type	Description
Array	\(|\nabla_\mathrm{cov} \phi|^2\): Covariant gradient squared (N, N, N)

Source code in jaxlatt/operators/coupled.py

@jit
def covariant_gradient_squared(phi: Array, links: Array, dx: float) -> Array:
    """
    Compute $\\Sigma_i |D_i \\phi|^2$ at each site.

    This is the kinetic energy density for the scalar field
    in the presence of gauge fields.

    Args:
        phi: Complex scalar field (N, N, N)
        links: Gauge links (3, N, N, N)
        dx: Lattice spacing

    Returns:
        $|\\nabla_\\mathrm{cov} \\phi|^2$: Covariant gradient squared (N, N, N)
    """
    D0 = covariant_derivative(phi, links, 0, dx)
    D1 = covariant_derivative(phi, links, 1, dx)
    D2 = covariant_derivative(phi, links, 2, dx)

    return jnp.abs(D0) ** 2 + jnp.abs(D1) ** 2 + jnp.abs(D2) ** 2

covariant_laplacian(phi, links, dx)

Compute covariant Laplacian \(\Delta_\mathrm{cov} \phi\) for a gauge-coupled scalar field.

\[\Delta_\mathrm{cov} \phi = \Sigma_i \frac{U_i(n)\phi(n+\hat{i}) + U_i^\dagger(n-\hat{i})\phi(n-\hat{i}) - 2\phi(n)}{dx^2}\]

This is the gauge-covariant kinetic term in the scalar equation of motion.

Parameters:

Name	Type	Description	Default
phi	Array	Complex scalar field (N, N, N)	required
links	Array	Gauge links \(U_i\) (3, N, N, N)	required
dx	float	Lattice spacing	required

Returns:

Type	Description
Array	Covariant Laplacian (N, N, N)

Source code in jaxlatt/operators/coupled.py

@jit
def covariant_laplacian(phi: Array, links: Array, dx: float) -> Array:
    """
    Compute covariant Laplacian $\\Delta_\\mathrm{cov} \\phi$ for a gauge-coupled scalar field.

    $$\\Delta_\\mathrm{cov} \\phi = \\Sigma_i \\frac{U_i(n)\\phi(n+\\hat{i}) + U_i^\\dagger(n-\\hat{i})\\phi(n-\\hat{i}) - 2\\phi(n)}{dx^2}$$

    This is the gauge-covariant kinetic term in the scalar equation of motion.

    Args:
        phi: Complex scalar field (N, N, N)
        links: Gauge links $U_i$ (3, N, N, N)
        dx: Lattice spacing

    Returns:
        Covariant Laplacian (N, N, N)
    """
    laplacian = jnp.zeros_like(phi)

    # Direction 0 (x)
    laplacian = laplacian + (
        links[0] * jnp.roll(phi, -1, axis=0)
        + jnp.roll(jnp.conj(links[0]), 1, axis=0) * jnp.roll(phi, 1, axis=0)
        - 2 * phi
    )

    # Direction 1 (y)
    laplacian = laplacian + (
        links[1] * jnp.roll(phi, -1, axis=1)
        + jnp.roll(jnp.conj(links[1]), 1, axis=1) * jnp.roll(phi, 1, axis=1)
        - 2 * phi
    )

    # Direction 2 (z)
    laplacian = laplacian + (
        links[2] * jnp.roll(phi, -1, axis=2)
        + jnp.roll(jnp.conj(links[2]), 1, axis=2) * jnp.roll(phi, 1, axis=2)
        - 2 * phi
    )

    return laplacian / (dx**2)

scalar_force(phi, links, m, lambda_, dx)

Compute force on scalar field: \(F_\phi = -\delta H / \delta \phi^*\).

\[F_\phi = \Delta_\mathrm{cov} \phi - m^2 \phi - \lambda |\phi|^2 \phi\]

Includes:

• Covariant Laplacian (gauge-covariant kinetic term)
• Mass term
• Self-interaction term

Parameters:

Name	Type	Description	Default
phi	Array	Complex scalar field (N, N, N)	required
links	Array	Gauge links (3, N, N, N)	required
m	float	Mass parameter	required
lambda_	float	Self-coupling	required
dx	float	Lattice spacing	required

Returns:

Type	Description
Array	\(F_\phi\): Force on phi (N, N, N)

Source code in jaxlatt/operators/coupled.py

@jit
def scalar_force(
    phi: Array,
    links: Array,
    m: float,
    lambda_: float,
    dx: float,
) -> Array:
    """
    Compute force on scalar field: $F_\\phi = -\\delta H / \\delta \\phi^*$.

    $$F_\\phi = \\Delta_\\mathrm{cov} \\phi - m^2 \\phi - \\lambda |\\phi|^2 \\phi$$

    Includes:

    - Covariant Laplacian (gauge-covariant kinetic term)
    - Mass term
    - Self-interaction term

    Args:
        phi: Complex scalar field (N, N, N)
        links: Gauge links (3, N, N, N)
        m: Mass parameter
        lambda_: Self-coupling
        dx: Lattice spacing

    Returns:
        $F_\\phi$: Force on phi (N, N, N)
    """
    laplacian = covariant_laplacian(phi, links, dx)
    potential_force = scalar_potential_force(phi, m, lambda_)
    return laplacian + potential_force

scalar_charge_density(phi, pi, g)

Compute scalar field charge density \(\rho = g \operatorname{Im}(\phi^* \pi)\).

This appears as a source in the Gauss law: \(\nabla \cdot E = -g \rho\).

Parameters:

Name	Type	Description	Default
phi	Array	Complex scalar field (N, N, N)	required
pi	Array	Conjugate momentum (N, N, N)	required
g	float	Gauge coupling	required

Returns:

Type	Description
Array	\(\rho\): Charge density (N, N, N)

Source code in jaxlatt/operators/coupled.py

@jit
def scalar_charge_density(phi: Array, pi: Array, g: float) -> Array:
    """
    Compute scalar field charge density $\\rho = g \\operatorname{Im}(\\phi^* \\pi)$.

    This appears as a source in the Gauss law: $\\nabla \\cdot E = -g \\rho$.

    Args:
        phi: Complex scalar field (N, N, N)
        pi: Conjugate momentum (N, N, N)
        g: Gauge coupling

    Returns:
        $\\rho$: Charge density (N, N, N)
    """
    return g * jnp.imag(jnp.conj(phi) * pi)

scalar_current_density(phi, links, direction, g, dx)

Compute scalar field current density in direction \(i\).

\[j_i = \frac{g}{dx} \operatorname{Im}[\phi^*(n) \, U_i(n) \, \phi(n+\hat{i})]\]

This is the physical U(1) conserved current (matches \(j_i = g \operatorname{Im}(\phi^* D_i \phi)\) in the continuum limit).

Parameters:

Name	Type	Description	Default
phi	Array	Complex scalar field (N, N, N)	required
links	Array	Gauge links (3, N, N, N)	required
direction	int	Direction (0=x, 1=y, 2=z) - static for JIT	required
g	float	Gauge coupling	required
dx	float	Lattice spacing	required

Returns:

Type	Description
Array	\(j_i\): Current density in direction \(i\) (N, N, N)

Source code in jaxlatt/operators/coupled.py

@partial(jit, static_argnums=(2,))
def scalar_current_density(phi: Array, links: Array, direction: int, g: float, dx: float) -> Array:
    """
    Compute scalar field current density in direction $i$.

    $$j_i = \\frac{g}{dx} \\operatorname{Im}[\\phi^*(n) \\, U_i(n) \\, \\phi(n+\\hat{i})]$$

    This is the physical U(1) conserved current (matches $j_i = g \\operatorname{Im}(\\phi^* D_i \\phi)$
    in the continuum limit).

    Args:
        phi: Complex scalar field (N, N, N)
        links: Gauge links (3, N, N, N)
        direction: Direction (0=x, 1=y, 2=z) - static for JIT
        g: Gauge coupling
        dx: Lattice spacing

    Returns:
        $j_i$: Current density in direction $i$ (N, N, N)
    """
    phi_shifted = jnp.roll(phi, -1, axis=direction)
    U_i = links[direction]

    return (g / dx) * jnp.imag(jnp.conj(phi) * U_i * phi_shifted)

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

Compute Gauss constraint with scalar field source.

The lattice Hamiltonian uses \(E\) conjugate to the link angle \(\theta_i\) (not the physical gauge field \(A_i = \theta_i / (g \, dx)\)). In this convention the conserved Gauss law is:

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

This is exact: the scalar-current term in the gauge force is \(\operatorname{Im}(\phi^* U_i \phi_{n+\hat{i}}) / dx^2\), whose lattice divergence is \(\operatorname{Im}(\phi^* \Delta_\mathrm{cov} \phi) / dx\), matching \(d/d\tau[(1/dx) \operatorname{Im}(\phi^* \pi)] = (1/dx) \operatorname{Im}(\phi^* \Delta_\mathrm{cov} \phi)\).

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:

Name	Type	Description
G	Array	Gauss constraint violation (N, N, N)

Source code in jaxlatt/operators/coupled.py

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

    The lattice Hamiltonian uses $E$ conjugate to the link angle $\\theta_i$ (not the
    physical gauge field $A_i = \\theta_i / (g \\, dx)$).  In this convention the
    conserved Gauss law is:

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

    This is exact: the scalar-current term in the gauge force is
    $\\operatorname{Im}(\\phi^* U_i \\phi_{n+\\hat{i}}) / dx^2$, whose lattice divergence is
    $\\operatorname{Im}(\\phi^* \\Delta_\\mathrm{cov} \\phi) / dx$,
    matching $d/d\\tau[(1/dx) \\operatorname{Im}(\\phi^* \\pi)] = (1/dx) \\operatorname{Im}(\\phi^* \\Delta_\\mathrm{cov} \\phi)$.

    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:
        G: Gauss constraint violation (N, N, N)
    """
    div_E = divergence_3d(E, dx)
    # Coefficient is 1/dx, not g — see docstring for derivation
    rho_theta = jnp.imag(jnp.conj(phi) * pi) / dx

    return div_E + rho_theta

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

Compute RMS norm of Gauss constraint violation.

Returns \(\sqrt{\langle |\nabla \cdot E + \operatorname{Im}(\phi^* \pi)/dx|^2 \rangle}\).

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	\(\|G\|_2\): RMS Gauss violation (scalar)

Source code in jaxlatt/operators/coupled.py

@jit
def gauss_violation_norm_with_source(
    links: Array,
    E: Array,
    phi: Array,
    pi: Array,
    dx: float,
) -> Array:
    """
    Compute RMS norm of Gauss constraint violation.

    Returns $\\sqrt{\\langle |\\nabla \\cdot E + \\operatorname{Im}(\\phi^* \\pi)/dx|^2 \\rangle}$.

    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:
        $\\|G\\|_2$: RMS Gauss violation (scalar)
    """
    G = gauss_constraint_with_source(links, E, phi, pi, dx)
    return jnp.sqrt(jnp.mean(G**2))

coupled_energy(phi, pi, links, E, dx, m, lambda_, g)

Compute total energy of coupled scalar-gauge system.

\[H = \int d^3x \left[ \frac{1}{2}|\pi|^2 + \frac{1}{2}|D_i \phi|^2 + V(\phi) + \frac{1}{2}E_i^2 + \frac{1}{2}B_i^2 \right]\]

where the terms are scalar kinetic, scalar gradient, scalar potential, electric energy, and magnetic energy respectively.

Parameters:

Name	Type	Description	Default
phi	Array	Scalar field (N, N, N)	required
pi	Array	Conjugate momentum (N, N, N)	required
links	Array	Gauge links (3, N, N, N)	required
E	Array	Electric field (3, N, N, N)	required
dx	float	Lattice spacing	required
m	float	Scalar mass	required
lambda_	float	Self-coupling	required
g	float	Gauge coupling	required

Returns:

Type	Description
dict	dict with JAX scalar energy components and total

Source code in jaxlatt/operators/coupled.py

@jit
def coupled_energy(
    phi: Array,
    pi: Array,
    links: Array,
    E: Array,
    dx: float,
    m: float,
    lambda_: float,
    g: float,
) -> dict:
    """
    Compute total energy of coupled scalar-gauge system.

    $$H = \\int d^3x \\left[ \\frac{1}{2}|\\pi|^2 + \\frac{1}{2}|D_i \\phi|^2 + V(\\phi) + \\frac{1}{2}E_i^2 + \\frac{1}{2}B_i^2 \\right]$$

    where the terms are scalar kinetic, scalar gradient, scalar potential,
    electric energy, and magnetic energy respectively.

    Args:
        phi: Scalar field (N, N, N)
        pi: Conjugate momentum (N, N, N)
        links: Gauge links (3, N, N, N)
        E: Electric field (3, N, N, N)
        dx: Lattice spacing
        m: Scalar mass
        lambda_: Self-coupling
        g: Gauge coupling

    Returns:
        dict with JAX scalar energy components and total
    """
    # Scalar kinetic energy: Σ (1/2)|π|² — canonical
    T_scalar = scalar_kinetic_energy_density(pi)
    E_scalar_kinetic = jnp.sum(T_scalar)

    # Scalar gradient energy: Σ (1/2)|D_iφ|² — canonical
    grad_sq = covariant_gradient_squared(phi, links, dx)
    E_scalar_gradient = 0.5 * jnp.sum(grad_sq)

    # Scalar potential energy: Σ V(φ) — canonical
    V = scalar_potential_energy(phi, m, lambda_)
    E_scalar_potential = jnp.sum(V)

    # Electric energy: Σ (1/2)E² — canonical
    E_electric = 0.5 * jnp.sum(E**2)

    # Magnetic energy: Σ Im(P)²/(8g²dx) — has its own dx scaling
    E_magnetic = magnetic_energy_density(links, dx, g)

    E_total = E_scalar_kinetic + E_scalar_gradient + E_scalar_potential + E_electric + E_magnetic

    return {
        "scalar_kinetic": E_scalar_kinetic,
        "scalar_gradient": E_scalar_gradient,
        "scalar_potential": E_scalar_potential,
        "electric": E_electric,
        "magnetic": E_magnetic,
        "total": E_total,
    }