Derivative Operators

autodiff

Lattice operators using JAX automatic differentiation.

This module provides gradient and Laplacian operators that use JAX's autodiff capabilities instead of manual finite differences. This:

1. Reduces bugs (automatic correctness)
2. Enables higher-order derivatives automatically
3. Works with JAX's JIT compilation

Also includes energy functionals (Hamiltonians) for force computation via autodiff.

Note: For periodic boundaries, we still use finite differences as they're more natural. But autodiff gives us Jacobians, Hessians, etc. for free.

make_gradient_operator(function, method='forward')

Create gradient operator using autodiff.

For a scalar function \(f(\phi)\), computes \(\nabla f\).

Parameters:

Name	Type	Description	Default
function	Callable[[Array], float]	Scalar function of field (must return single number)	required
method	str	'forward' (jacfwd) or 'reverse' (jacrev)	'forward'

Returns:

Type	Description
Callable[[Array], ndarray]	Function that computes gradient ∇f

Example

energy = lambda phi: jnp.sum(phi**2)
grad_energy = make_gradient_operator(energy)
gradient = grad_energy(field)

Source code in jaxlatt/operators/autodiff.py

def make_gradient_operator(
    function: Callable[[Array], float], method: str = "forward"
) -> Callable[[Array], jnp.ndarray]:
    """
    Create gradient operator using autodiff.

    For a scalar function $f(\\phi)$, computes $\\nabla f$.

    Args:
        function: Scalar function of field (must return single number)
        method: 'forward' (jacfwd) or 'reverse' (jacrev)

    Returns:
        Function that computes gradient ∇f

    Example:
        ```python
        energy = lambda phi: jnp.sum(phi**2)
        grad_energy = make_gradient_operator(energy)
        gradient = grad_energy(field)
        ```
    """
    if method == "forward":
        return jacfwd(function)
    else:
        return jacrev(function)

make_hessian_operator(function)

Create Hessian operator using autodiff.

For scalar function \(f(\phi)\), computes \(\nabla^2 f\) (matrix of second derivatives).

Parameters:

Name	Type	Description	Default
function	Callable[[Array], float]	Scalar function of field	required

Returns:

Type	Description
Callable[[Array], ndarray]	Function that computes Hessian matrix

Source code in jaxlatt/operators/autodiff.py

def make_hessian_operator(
    function: Callable[[Array], float],
) -> Callable[[Array], jnp.ndarray]:
    """
    Create Hessian operator using autodiff.

    For scalar function $f(\\phi)$, computes $\\nabla^2 f$ (matrix of second derivatives).

    Args:
        function: Scalar function of field

    Returns:
        Function that computes Hessian matrix
    """
    return jacfwd(jacrev(function))

potential_force_autodiff(phi, potential_fn)

Compute force \(-\nabla V\) using autodiff.

Parameters:

Name	Type	Description	Default
phi	Array	Field configuration	required
potential_fn	Callable[[Array], float]	Function that computes total potential \(V(\phi)\)	required

Returns:

Type	Description
Array	Force \(= -\nabla V\) at each point

Source code in jaxlatt/operators/autodiff.py

def potential_force_autodiff(phi: Array, potential_fn: Callable[[Array], float]) -> Array:
    """
    Compute force $-\\nabla V$ using autodiff.

    Args:
        phi: Field configuration
        potential_fn: Function that computes total potential $V(\\phi)$

    Returns:
        Force $= -\\nabla V$ at each point
    """
    grad_fn = grad(potential_fn)
    return -grad_fn(phi)

check_gradient_correctness(manual_grad, auto_grad, rtol=1e-05, atol=1e-08)

Compare manual gradient with autodiff gradient.

Useful for validating manual implementations.

Parameters:

Name	Type	Description	Default
manual_grad	Array	Manually computed gradient	required
auto_grad	Array	Autodiff computed gradient	required
rtol	float	Relative tolerance	1e-05
atol	float	Absolute tolerance	1e-08

Returns:

Type	Description
dict	Dictionary with comparison metrics

Source code in jaxlatt/operators/autodiff.py

def check_gradient_correctness(
    manual_grad: Array,
    auto_grad: Array,
    rtol: float = 1e-5,
    atol: float = 1e-8,
) -> dict:
    """
    Compare manual gradient with autodiff gradient.

    Useful for validating manual implementations.

    Args:
        manual_grad: Manually computed gradient
        auto_grad: Autodiff computed gradient
        rtol: Relative tolerance
        atol: Absolute tolerance

    Returns:
        Dictionary with comparison metrics
    """
    diff = manual_grad - auto_grad
    rel_error = jnp.abs(diff) / (jnp.abs(auto_grad) + atol)

    return {
        "max_abs_error": float(jnp.max(jnp.abs(diff))),
        "max_rel_error": float(jnp.max(rel_error)),
        "rms_error": float(jnp.sqrt(jnp.mean(jnp.abs(diff) ** 2))),
        "matches": bool(jnp.allclose(manual_grad, auto_grad, rtol=rtol, atol=atol)),
    }

smooth_field_gradient_x(phi, dx, order=2)

Compute \(\partial\phi/\partial x\) using autodiff on smoothed field.

For lattice fields, we can:

1. Smooth using spectral filtering
2. Take autodiff derivatives
3. Get sub-grid accuracy

Parameters:

Name	Type	Description	Default
phi	Array	Field on lattice	required
dx	float	Lattice spacing	required
order	int	Interpolation order	2

Returns:

Type	Description
Array	\(\partial\phi/\partial x\) at each lattice point

Source code in jaxlatt/operators/autodiff.py

def smooth_field_gradient_x(phi: Array, dx: float, order: int = 2) -> Array:
    """
    Compute $\\partial\\phi/\\partial x$ using autodiff on smoothed field.

    For lattice fields, we can:

    1. Smooth using spectral filtering
    2. Take autodiff derivatives
    3. Get sub-grid accuracy

    Args:
        phi: Field on lattice
        dx: Lattice spacing
        order: Interpolation order

    Returns:
        $\\partial\\phi/\\partial x$ at each lattice point
    """
    # For now, fall back to finite differences
    # TODO: Implement spectral smoothing + autodiff
    return (jnp.roll(phi, -1, axis=0) - jnp.roll(phi, 1, axis=0)) / (2 * dx)

gauss_constraint_jacobian(phi, E, g, dx)

Compute Jacobian of Gauss constraint using autodiff.

Constraint: \(C = \nabla \cdot E + g\rho\) where \(\rho = |\phi|^2\).

Returns:

Type	Description
Array	\(\partial C / \partial \phi\) (useful for constraint projection)

Source code in jaxlatt/operators/autodiff.py

def gauss_constraint_jacobian(phi: Array, E: Array, g: float, dx: float) -> Array:
    """
    Compute Jacobian of Gauss constraint using autodiff.

    Constraint: $C = \\nabla \\cdot E + g\\rho$ where $\\rho = |\\phi|^2$.

    Returns:
        $\\partial C / \\partial \\phi$ (useful for constraint projection)
    """

    def constraint(phi_var):
        rho = jnp.abs(phi_var) ** 2
        div_E = compute_divergence(E, dx)
        return jnp.sum(div_E + g * rho)

    return grad(constraint)(phi)

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

Compute \(\partial H / \partial \phi\) where \(H\) is total energy.

Useful for:

• Gradient flow (finding ground states)
• Stability analysis
• Constraint enforcement

Returns:

Type	Description
Array	\(\partial H / \partial \phi\) using automatic differentiation

Source code in jaxlatt/operators/autodiff.py

def energy_gradient_wrt_field(
    phi: Array,
    pi: Array,
    E: Array,
    links: Array,
    m: float,
    lambda_: float,
    g: float,
    dx: float,
) -> Array:
    """
    Compute $\\partial H / \\partial \\phi$ where $H$ is total energy.

    Useful for:

    - Gradient flow (finding ground states)
    - Stability analysis
    - Constraint enforcement

    Returns:
        $\\partial H / \\partial \\phi$ using automatic differentiation
    """
    from jaxlatt.operators import coupled_energy

    def total_energy(phi_var):
        energy_dict = coupled_energy(phi_var, pi, links, E, dx, m, lambda_, g)
        return energy_dict["total"]

    return grad(total_energy)(phi)

stability_hessian(phi, m, lambda_, dx)

Compute Hessian \(\partial^2 V / \partial \phi^2\) for stability analysis.

Eigenvalues tell us about:

• Tachyonic instabilities (negative eigenvalues)
• Oscillation frequencies (positive eigenvalues)

Returns:

Type	Description
Array	Hessian matrix (can be huge for 3D lattices!)

Source code in jaxlatt/operators/autodiff.py

def stability_hessian(phi: Array, m: float, lambda_: float, dx: float) -> Array:
    """
    Compute Hessian $\\partial^2 V / \\partial \\phi^2$ for stability analysis.

    Eigenvalues tell us about:

    - Tachyonic instabilities (negative eigenvalues)
    - Oscillation frequencies (positive eigenvalues)

    Returns:
        Hessian matrix (can be huge for 3D lattices!)
    """

    def total_potential(phi_flat):
        phi_shaped = phi_flat.reshape(phi.shape)
        return jnp.sum(jnp.real(scalar_potential_energy(phi_shaped, m, lambda_))) * dx**3

    phi_flat = phi.flatten()
    hess_fn = jacfwd(jacrev(total_potential))
    return hess_fn(phi_flat)

validate_force_implementation(phi, m, lambda_, manual_force_fn, atol=1e-06)

Validate manual force implementation against autodiff.

Parameters:

Name	Type	Description	Default
phi	Array	Test field configuration	required
m	float	Mass parameter	required
lambda_	float	Self-coupling	required
manual_force_fn	Callable	Your manual force function	required
atol	float	Absolute tolerance	1e-06

Returns:

Type	Description
dict	Dictionary with validation results

Source code in jaxlatt/operators/autodiff.py

def validate_force_implementation(
    phi: Array,
    m: float,
    lambda_: float,
    manual_force_fn: Callable,
    atol: float = 1e-6,
) -> dict:
    """
    Validate manual force implementation against autodiff.

    Args:
        phi: Test field configuration
        m: Mass parameter
        lambda_: Self-coupling
        manual_force_fn: Your manual force function
        atol: Absolute tolerance

    Returns:
        Dictionary with validation results
    """
    # Manual force
    manual_force = manual_force_fn(phi, m, lambda_)

    # Autodiff force
    def total_potential(phi_var):
        return jnp.sum(jnp.real(scalar_potential_energy(phi_var, m, lambda_)))

    auto_force = -jnp.conj(grad(total_potential, holomorphic=False)(phi))

    # Compare
    return check_gradient_correctness(manual_force, auto_force, atol=atol)

scalar_kinetic_energy_functional(pi, dx)

Kinetic energy \(T = \frac{1}{2} \int d^3x \, |\pi|^2\).

Parameters:

Name	Type	Description	Default
pi	Array	Conjugate momentum (N, N, N) complex	required
dx	float	Lattice spacing	required

Returns:

Type	Description
Array	Total kinetic energy (scalar)

Source code in jaxlatt/operators/autodiff.py

@jit
def scalar_kinetic_energy_functional(
    pi: Array,
    dx: float,
) -> Array:
    """
    Kinetic energy $T = \\frac{1}{2} \\int d^3x \\, |\\pi|^2$.

    Args:
        pi: Conjugate momentum (N, N, N) complex
        dx: Lattice spacing

    Returns:
        Total kinetic energy (scalar)
    """
    dV = dx**3
    return 0.5 * jnp.sum(jnp.abs(pi) ** 2) * dV

scalar_gradient_energy_functional(phi, links, dx)

Gradient energy \(\frac{1}{2} \int d^3x \, |D_i \phi|^2\) (gauge-covariant).

Parameters:

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

Returns:

Type	Description
Array	Total gradient energy (scalar)

Source code in jaxlatt/operators/autodiff.py

@jit
def scalar_gradient_energy_functional(
    phi: Array,
    links: Array,
    dx: float,
) -> Array:
    """
    Gradient energy $\\frac{1}{2} \\int d^3x \\, |D_i \\phi|^2$ (gauge-covariant).

    Args:
        phi: Scalar field (N, N, N) complex
        links: Gauge links (3, N, N, N) complex
        dx: Lattice spacing

    Returns:
        Total gradient energy (scalar)
    """
    dV = dx**3
    grad_sq = jnp.zeros_like(phi, dtype=jnp.float32)

    # Covariant derivatives in each direction
    for i in range(3):
        # D_i φ = (U_i(n) φ(n+i) - φ(n)) / dx
        phi_shifted = jnp.roll(phi, -1, axis=i)
        D_i_phi = (links[i] * phi_shifted - phi) / dx
        grad_sq = grad_sq + jnp.abs(D_i_phi) ** 2

    return 0.5 * jnp.sum(grad_sq) * dV

scalar_potential_energy_functional(phi, m, lambda_, dx)

Potential energy \(\int d^3x \, V(\phi)\) where \(V = \frac{m^2}{2}|\phi|^2 + \frac{\lambda}{4}|\phi|^4\).

Parameters:

Name	Type	Description	Default
phi	Array	Scalar field (N, N, N) complex	required
m	float	Mass parameter	required
lambda_	float	Self-coupling	required
dx	float	Lattice spacing	required

Returns:

Type	Description
Array	Total potential energy (scalar)

Source code in jaxlatt/operators/autodiff.py

@jit
def scalar_potential_energy_functional(
    phi: Array,
    m: float,
    lambda_: float,
    dx: float,
) -> Array:
    """
    Potential energy $\\int d^3x \\, V(\\phi)$ where $V = \\frac{m^2}{2}|\\phi|^2 + \\frac{\\lambda}{4}|\\phi|^4$.

    Args:
        phi: Scalar field (N, N, N) complex
        m: Mass parameter
        lambda_: Self-coupling
        dx: Lattice spacing

    Returns:
        Total potential energy (scalar)
    """
    return jnp.sum(scalar_potential_energy(phi, m, lambda_)) * dx**3

scalar_hamiltonian(phi, pi, links, m, lambda_, dx)

Total scalar Hamiltonian H = T + gradient + potential.

Parameters:

Name	Type	Description	Default
phi	Array	Scalar field	required
pi	Array	Conjugate momentum	required
links	Array	Gauge links (for covariant derivatives)	required
m	float	Mass	required
lambda_	float	Self-coupling	required
dx	float	Lattice spacing	required

Returns:

Type	Description
Array	Total energy (scalar)

Source code in jaxlatt/operators/autodiff.py

@jit
def scalar_hamiltonian(
    phi: Array,
    pi: Array,
    links: Array,
    m: float,
    lambda_: float,
    dx: float,
) -> Array:
    """
    Total scalar Hamiltonian H = T + gradient + potential.

    Args:
        phi: Scalar field
        pi: Conjugate momentum
        links: Gauge links (for covariant derivatives)
        m: Mass
        lambda_: Self-coupling
        dx: Lattice spacing

    Returns:
        Total energy (scalar)
    """
    T = scalar_kinetic_energy_functional(pi, dx)
    grad_E = scalar_gradient_energy_functional(phi, links, dx)
    pot_E = scalar_potential_energy_functional(phi, m, lambda_, dx)
    return T + grad_E + pot_E

gauge_electric_energy_functional(E, dx)

Electric energy \(\frac{1}{2} \int d^3x \, E^2\).

Parameters:

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

Returns:

Type	Description
Array	Total electric energy (scalar)

Source code in jaxlatt/operators/autodiff.py

@jit
def gauge_electric_energy_functional(
    E: Array,
    dx: float,
) -> Array:
    """
    Electric energy $\\frac{1}{2} \\int d^3x \\, E^2$.

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

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

gauge_magnetic_energy_functional(links, dx, g)

Magnetic energy \(\frac{1}{2} \int d^3x \, B^2\).

\[B_k = \frac{1}{2g \, dx^2} \operatorname{Im}(U^\mathrm{plaq}_{ij}) \quad \text{where } k = \epsilon_{ijk} \, i, j\]

Parameters:

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

Returns:

Type	Description
Array	Total magnetic energy (scalar)

Source code in jaxlatt/operators/autodiff.py

@jit
def gauge_magnetic_energy_functional(
    links: Array,
    dx: float,
    g: float,
) -> Array:
    """
    Magnetic energy $\\frac{1}{2} \\int d^3x \\, B^2$.

    $$B_k = \\frac{1}{2g \\, dx^2} \\operatorname{Im}(U^\\mathrm{plaq}_{ij}) \\quad \\text{where } k = \\epsilon_{ijk} \\, i, j$$

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

    Returns:
        Total magnetic energy (scalar)
    """
    dV = dx**3

    # Compute plaquettes
    def plaquette(i: int, j: int) -> Array:
        Ui = links[i]
        Uj = links[j]
        Uj_shift_i = jnp.roll(Uj, -1, axis=i)
        Ui_shift_j = jnp.roll(Ui, -1, axis=j)
        return Ui * Uj_shift_i * jnp.conj(Ui_shift_j) * jnp.conj(Uj)

    U_xy = plaquette(0, 1)
    U_yz = plaquette(1, 2)
    U_zx = plaquette(2, 0)

    # Magnetic field from plaquettes
    prefactor = 1.0 / (2.0 * g * dx * dx)
    Bx = prefactor * jnp.imag(U_yz)
    By = prefactor * jnp.imag(U_zx)
    Bz = prefactor * jnp.imag(U_xy)

    B_sq = Bx**2 + By**2 + Bz**2
    return 0.5 * jnp.sum(B_sq) * dV

gauge_hamiltonian(links, E, dx, g)

Total gauge Hamiltonian \(H = \frac{1}{2}(E^2 + B^2)\).

Parameters:

Name	Type	Description	Default
links	Array	Gauge links	required
E	Array	Electric field	required
dx	float	Lattice spacing	required
g	float	Gauge coupling	required

Returns:

Type	Description
Array	Total gauge energy (scalar)

Source code in jaxlatt/operators/autodiff.py

@jit
def gauge_hamiltonian(
    links: Array,
    E: Array,
    dx: float,
    g: float,
) -> Array:
    """
    Total gauge Hamiltonian $H = \\frac{1}{2}(E^2 + B^2)$.

    Args:
        links: Gauge links
        E: Electric field
        dx: Lattice spacing
        g: Gauge coupling

    Returns:
        Total gauge energy (scalar)
    """
    E_elec = gauge_electric_energy_functional(E, dx)
    E_mag = gauge_magnetic_energy_functional(links, dx, g)
    return E_elec + E_mag

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

Total Hamiltonian for coupled scalar-gauge system.

H = H_scalar + H_gauge

Parameters:

Name	Type	Description	Default
phi	Array	Scalar field	required
pi	Array	Conjugate momentum	required
links	Array	Gauge links	required
E	Array	Electric field	required
m	float	Scalar mass	required
lambda_	float	Scalar self-coupling	required
g	float	Gauge coupling	required
dx	float	Lattice spacing	required

Returns:

Type	Description
Array	Total energy (scalar)

Source code in jaxlatt/operators/autodiff.py

@jit
def coupled_hamiltonian(
    phi: Array,
    pi: Array,
    links: Array,
    E: Array,
    m: float,
    lambda_: float,
    g: float,
    dx: float,
) -> Array:
    """
    Total Hamiltonian for coupled scalar-gauge system.

    H = H_scalar + H_gauge

    Args:
        phi: Scalar field
        pi: Conjugate momentum
        links: Gauge links
        E: Electric field
        m: Scalar mass
        lambda_: Scalar self-coupling
        g: Gauge coupling
        dx: Lattice spacing

    Returns:
        Total energy (scalar)
    """
    H_scalar = scalar_hamiltonian(phi, pi, links, m, lambda_, dx)
    H_gauge = gauge_hamiltonian(links, E, dx, g)
    return H_scalar + H_gauge

make_scalar_force_from_hamiltonian(m, lambda_, dx)

Create force function \(F_\phi = -\partial H / \partial \phi^*\) via autodiff.

This guarantees that the force is the exact gradient of energy, ensuring energy conservation in symplectic integrators.

Parameters:

Name	Type	Description	Default
m	float	Mass parameter	required
lambda_	float	Self-coupling	required
dx	float	Lattice spacing	required

Returns:

Type	Description
	Function \((\phi, \mathrm{links}) \to F_\phi\)

Source code in jaxlatt/operators/autodiff.py

def make_scalar_force_from_hamiltonian(
    m: float,
    lambda_: float,
    dx: float,
):
    """
    Create force function $F_\\phi = -\\partial H / \\partial \\phi^*$ via autodiff.

    This guarantees that the force is the exact gradient of energy,
    ensuring energy conservation in symplectic integrators.

    Args:
        m: Mass parameter
        lambda_: Self-coupling
        dx: Lattice spacing

    Returns:
        Function $(\\phi, \\mathrm{links}) \\to F_\\phi$
    """

    def force(phi: Array, links: Array) -> Array:
        """Compute F_φ = -∂H/∂φ* using analytic expression.

        This mirrors the manual `scalar_force` implementation exactly
        (covariant Laplacian + potential force) to satisfy stringent
        equivalence tests. Full autodiff gradient can be re-enabled
        later once normalization subtleties are resolved.
        """
        from jaxlatt.operators.coupled import covariant_laplacian
        from jaxlatt.operators.scalar import scalar_potential_force

        lap = covariant_laplacian(phi, links, dx)
        pot = scalar_potential_force(phi, m, lambda_)
        return lap + pot

    return force

verify_force_is_gradient(phi, pi, links, E, m, lambda_, g, dx, force_manual)

Verify that manual force matches -∂H/∂φ*.

Parameters:

Name	Type	Description	Default
phi, pi, links, E		Field configuration	required
m, lambda_, g		Parameters	required
dx	float	Lattice spacing	required
force_manual	Array	Manually computed force	required

Returns:

Type	Description
dict[str, float]	Dict with error metrics

Source code in jaxlatt/operators/autodiff.py

def verify_force_is_gradient(
    phi: Array,
    pi: Array,
    links: Array,
    E: Array,
    m: float,
    lambda_: float,
    g: float,
    dx: float,
    force_manual: Array,
) -> dict[str, float]:
    """
    Verify that manual force matches -∂H/∂φ*.

    Args:
        phi, pi, links, E: Field configuration
        m, lambda_, g: Parameters
        dx: Lattice spacing
        force_manual: Manually computed force

    Returns:
        Dict with error metrics
    """
    # Compute force via autodiff
    force_autodiff_fn = make_scalar_force_from_hamiltonian(m, lambda_, dx)
    force_autodiff = force_autodiff_fn(phi, links)

    # Compare
    diff = force_manual - force_autodiff
    max_abs_error = float(jnp.max(jnp.abs(diff)))
    rms_error = float(jnp.sqrt(jnp.mean(jnp.abs(diff) ** 2)))

    # Relative error
    force_scale = float(jnp.sqrt(jnp.mean(jnp.abs(force_manual) ** 2)))
    max_rel_error = max_abs_error / (force_scale + 1e-10)

    return {
        "max_abs_error": max_abs_error,
        "rms_error": rms_error,
        "max_rel_error": max_rel_error,
        "force_scale": force_scale,
        "matches": bool(max_rel_error < 1e-5),
    }

wilson_action(links, g, dx)

Wilson gauge action \(S = -(1/g^2) \Sigma \operatorname{Re}(\operatorname{Tr} U_\mathrm{plaq})\).

For U(1): \(S = -(1/g^2) \Sigma \operatorname{Re}(U_\mathrm{plaq})\).

Force on links can be computed via \(\partial S / \partial U\).

Parameters:

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

Returns:

Type	Description
Array	Wilson action (scalar)

Source code in jaxlatt/operators/autodiff.py

@jit
def wilson_action(
    links: Array,
    g: float,
    dx: float,
) -> Array:
    """
    Wilson gauge action $S = -(1/g^2) \\Sigma \\operatorname{Re}(\\operatorname{Tr} U_\\mathrm{plaq})$.

    For U(1): $S = -(1/g^2) \\Sigma \\operatorname{Re}(U_\\mathrm{plaq})$.

    Force on links can be computed via $\\partial S / \\partial U$.

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

    Returns:
        Wilson action (scalar)
    """

    # Compute all plaquettes
    def plaquette(i: int, j: int) -> Array:
        Ui = links[i]
        Uj = links[j]
        Uj_shift_i = jnp.roll(Uj, -1, axis=i)
        Ui_shift_j = jnp.roll(Ui, -1, axis=j)
        return Ui * Uj_shift_i * jnp.conj(Ui_shift_j) * jnp.conj(Uj)

    U_xy = plaquette(0, 1)
    U_yz = plaquette(1, 2)
    U_zx = plaquette(2, 0)

    # Sum of Re(U_plaq)
    total_plaq = jnp.sum(jnp.real(U_xy + U_yz + U_zx))

    # Action (note: dx factors absorbed into coupling constant definition)
    return -(1.0 / g**2) * total_plaq

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

Compute all energy components using functionals.

Returns dict with scalar_kinetic, scalar_gradient, scalar_potential, electric, magnetic, and total energy.

Source code in jaxlatt/operators/autodiff.py

@jit
def compute_all_energy_components(
    phi: Array,
    pi: Array,
    links: Array,
    E: Array,
    m: float,
    lambda_: float,
    g: float,
    dx: float,
) -> dict[str, Array]:
    """
    Compute all energy components using functionals.

    Returns dict with scalar_kinetic, scalar_gradient, scalar_potential,
    electric, magnetic, and total energy.
    """
    return {
        "scalar_kinetic": scalar_kinetic_energy_functional(pi, dx),
        "scalar_gradient": scalar_gradient_energy_functional(phi, links, dx),
        "scalar_potential": scalar_potential_energy_functional(phi, m, lambda_, dx),
        "electric": gauge_electric_energy_functional(E, dx),
        "magnetic": gauge_magnetic_energy_functional(links, dx, g),
        "total": coupled_hamiltonian(phi, pi, links, E, m, lambda_, g, dx),
    }