Force Operators

forces

Force operators for lattice field evolution using autodiff.

All potential-derived forces use JAX autodiff, guaranteeing correctness and energy conservation in symplectic integrators.

Design principles:

1. Forces are computed as \(F = -\frac{\partial H}{\partial q}\) via autodiff
2. Spatial operators (Laplacians) use finite differences (appropriate for lattices)
3. Factory functions return JIT-compiled, cached force functions
4. Support for both flat spacetime and expanding universe (FRW)

Example

from jaxlatt.core.potentials import ScalarPotential
from jaxlatt.operators.forces import make_coupled_scalar_force

potential = ScalarPotential.quartic(m=1.0, lambda_=0.1)
force_fn = make_coupled_scalar_force(potential, dx=0.5)
F = force_fn(phi, links)

make_scalar_force(potential, dx)

Create force function for free scalar field (no gauge coupling).

Computes \(F(\phi) = \nabla^2\phi - dV/d\phi\) where:

• \(\nabla^2\) is the discrete Laplacian (finite differences)
• \(dV/d\phi\) is computed via autodiff

Parameters:

Name	Type	Description	Default
potential	ScalarPotential | PotentialFunction	ScalarPotential instance or legacy potential function	required
dx	float	Lattice spacing	required

Returns:

Type	Description
Callable[[Array], Array]	JIT-compiled function: field -> force

Example

from jaxlatt.core.potentials import ScalarPotential
V = ScalarPotential.quadratic(m=1.0)
force_fn = make_scalar_force(V, dx=0.1)
F = force_fn(field)

Source code in jaxlatt/operators/forces.py

def make_scalar_force(
    potential: ScalarPotential | PotentialFunction,
    dx: float,
) -> Callable[[Array], Array]:
    """
    Create force function for free scalar field (no gauge coupling).

    Computes $F(\\phi) = \\nabla^2\\phi - dV/d\\phi$ where:

    - $\\nabla^2$ is the discrete Laplacian (finite differences)
    - $dV/d\\phi$ is computed via autodiff

    Args:
        potential: ScalarPotential instance or legacy potential function
        dx: Lattice spacing

    Returns:
        JIT-compiled function: field -> force

    Example:
        ```python
        from jaxlatt.core.potentials import ScalarPotential
        V = ScalarPotential.quadratic(m=1.0)
        force_fn = make_scalar_force(V, dx=0.1)
        F = force_fn(field)
        ```
    """
    from jaxlatt.evolution.scalar import compute_laplacian

    potential_force = _resolve_potential_force(potential)

    @jit
    def force(field: Array) -> Array:
        lap = compute_laplacian(field, dx)
        return lap + potential_force(field)

    return force

make_coupled_scalar_force(potential, dx)

Create force for scalar field coupled to gauge field (flat spacetime).

Computes \(F_\phi = D^2\phi - dV/d\phi\) where:

• \(D^2\) is the covariant Laplacian (includes gauge links)
• \(dV/d\phi\) is computed via autodiff

Parameters:

Name	Type	Description	Default
potential	ScalarPotential | PotentialFunction	ScalarPotential instance or legacy potential function	required
dx	float	Lattice spacing	required

Returns:

Type	Description
Callable[[Array, Array], Array]	JIT-compiled function: (phi, links) -> force

Example

from jaxlatt.core.potentials import ScalarPotential
V = ScalarPotential.quartic(m=1.0, lambda_=0.1)
force_fn = make_coupled_scalar_force(V, dx=0.5)
F = force_fn(phi, links)

Source code in jaxlatt/operators/forces.py

def make_coupled_scalar_force(
    potential: ScalarPotential | PotentialFunction,
    dx: float,
) -> Callable[[Array, Array], Array]:
    """
    Create force for scalar field coupled to gauge field (flat spacetime).

    Computes $F_\\phi = D^2\\phi - dV/d\\phi$ where:

    - $D^2$ is the covariant Laplacian (includes gauge links)
    - $dV/d\\phi$ is computed via autodiff

    Args:
        potential: ScalarPotential instance or legacy potential function
        dx: Lattice spacing

    Returns:
        JIT-compiled function: (phi, links) -> force

    Example:
        ```python
        from jaxlatt.core.potentials import ScalarPotential
        V = ScalarPotential.quartic(m=1.0, lambda_=0.1)
        force_fn = make_coupled_scalar_force(V, dx=0.5)
        F = force_fn(phi, links)
        ```
    """
    from jaxlatt.operators.coupled import covariant_laplacian

    potential_force = _resolve_potential_force(potential)

    @jit
    def force(phi: Array, links: Array) -> Array:
        lap = covariant_laplacian(phi, links, dx)
        return lap + potential_force(phi)

    return force

make_coupled_scalar_force_expanding(potential, dx)

Create force for scalar field in expanding FRW universe.

Computes \(F_\phi = (1/a^2) D^2\phi - a^2 \, dV/d\phi^*\) where:

• \(D^2\) is the covariant Laplacian
• \(dV/d\phi\) is computed via autodiff
• \(a\) is the scale factor

Note: The Hubble friction term \(-2H\pi\) is NOT included here; it should be applied separately in the evolution equation.

Parameters:

Name	Type	Description	Default
potential	ScalarPotential | PotentialFunction	ScalarPotential instance or legacy potential function	required
dx	float	Comoving lattice spacing	required

Returns:

Type	Description
Callable[[Array, Array, float], Array]	JIT-compiled function: (phi, links, a) -> force

Example

V = ScalarPotential.quartic(m=1.0, lambda_=0.1)
force_fn = make_coupled_scalar_force_expanding(V, dx=0.5)
F = force_fn(phi, links, a=2.0)

Source code in jaxlatt/operators/forces.py

def make_coupled_scalar_force_expanding(
    potential: ScalarPotential | PotentialFunction,
    dx: float,
) -> Callable[[Array, Array, float], Array]:
    """
    Create force for scalar field in expanding FRW universe.

    Computes $F_\\phi = (1/a^2) D^2\\phi - a^2 \\, dV/d\\phi^*$ where:

    - $D^2$ is the covariant Laplacian
    - $dV/d\\phi$ is computed via autodiff
    - $a$ is the scale factor

    Note: The Hubble friction term $-2H\\pi$ is NOT included here;
    it should be applied separately in the evolution equation.

    Args:
        potential: ScalarPotential instance or legacy potential function
        dx: Comoving lattice spacing

    Returns:
        JIT-compiled function: (phi, links, a) -> force

    Example:
        ```python
        V = ScalarPotential.quartic(m=1.0, lambda_=0.1)
        force_fn = make_coupled_scalar_force_expanding(V, dx=0.5)
        F = force_fn(phi, links, a=2.0)
        ```
    """
    from jaxlatt.operators.coupled import covariant_laplacian

    potential_force = _resolve_potential_force(potential)

    @jit
    def force(phi: Array, links: Array, a: float) -> Array:
        # Covariant Laplacian
        lap = covariant_laplacian(phi, links, dx)

        # Kinetic term: +D²φ/a²
        kinetic = lap / (a**2)

        # Potential term: -a² dV/dφ*
        # potential_force returns F = -dV/dφ
        # We want: -a² dV/dφ = a² * F
        pot = (a**2) * potential_force(phi)

        return kinetic + pot

    return force

make_rescaled_scalar_force(potential, dx)

Create force for rescaled field χ = aφ (eliminates Hubble friction).

The rescaled field evolution equation is:

\[\chi'' = \frac{1}{a^2}\nabla^2\chi - a^2 \frac{dV_\mathrm{eff}}{d\chi} - \frac{a''}{a}\chi\]

where \(V_\mathrm{eff}\) is the effective potential for \(\chi\).

Parameters:

Name	Type	Description	Default
potential	ScalarPotential | PotentialFunction	ScalarPotential for physical field φ (not χ)	required
dx	float	Comoving lattice spacing	required

Returns:

Type	Description
Callable[[Array, Array, float, float], Array]	JIT-compiled function: (chi, links, a, addot) -> force

Example

V = ScalarPotential.quartic(m=1.0, lambda_=0.1)
force_fn = make_rescaled_scalar_force(V, dx=0.5)
F = force_fn(chi, links, a=2.0, addot=0.1)

Source code in jaxlatt/operators/forces.py

def make_rescaled_scalar_force(
    potential: ScalarPotential | PotentialFunction,
    dx: float,
) -> Callable[[Array, Array, float, float], Array]:
    """
    Create force for rescaled field χ = aφ (eliminates Hubble friction).

    The rescaled field evolution equation is:

    $$\\chi'' = \\frac{1}{a^2}\\nabla^2\\chi - a^2 \\frac{dV_\\mathrm{eff}}{d\\chi} - \\frac{a''}{a}\\chi$$

    where $V_\\mathrm{eff}$ is the effective potential for $\\chi$.

    Args:
        potential: ScalarPotential for physical field φ (not χ)
        dx: Comoving lattice spacing

    Returns:
        JIT-compiled function: (chi, links, a, addot) -> force

    Example:
        ```python
        V = ScalarPotential.quartic(m=1.0, lambda_=0.1)
        force_fn = make_rescaled_scalar_force(V, dx=0.5)
        F = force_fn(chi, links, a=2.0, addot=0.1)
        ```
    """
    from jaxlatt.operators.coupled import covariant_laplacian

    # Get the potential function for computing derivatives
    if isinstance(potential, ScalarPotential):
        V_fn = lambda field: potential(field)
    else:
        V_fn = potential

    @jit
    def force(chi: Array, links: Array, a: float, addot: float) -> Array:
        # Covariant Laplacian of χ
        lap_chi = covariant_laplacian(chi, links, dx)

        # Kinetic term: (1/a²)∇²χ
        kinetic = lap_chi / (a**2)

        # Potential term via autodiff
        # For χ = aφ, we need -dV/dχ where V is expressed in terms of χ
        # V(φ) = V(χ/a), and the EOM has a⁴V(χ/a) structure
        # The force is: -a² dV/d(χ/a) * (1/a) = -a dV/dφ|_{φ=χ/a}
        def V_effective(chi_field: Array) -> Array:
            phi = chi_field / a
            return jnp.sum(V_fn(phi)).real * (a**4)

        dV_dchi = grad(V_effective, holomorphic=False)(chi)

        # Curvature term: -(a''/a)χ
        curvature = -(addot / a) * chi

        return kinetic - dV_dchi + curvature

    return force

make_gauge_force(g, dx)

Create force on electric field from pure gauge + scalar coupling.

Computes F_E for all three directions. This is geometric and doesn't particularly benefit from autodiff, but is provided here for API consistency.

Parameters:

Name	Type	Description	Default
g	float	Gauge coupling constant	required
dx	float	Lattice spacing	required

Returns:

Type	Description
Callable[[Array, Array], Array]	JIT-compiled function: (phi, links) -> force array (3, N, N, N)

Example

force_fn = make_gauge_force(g=0.5, dx=0.1)
F_E = force_fn(phi, links)

Source code in jaxlatt/operators/forces.py

def make_gauge_force(
    g: float,
    dx: float,
) -> Callable[[Array, Array], Array]:
    """
    Create force on electric field from pure gauge + scalar coupling.

    Computes F_E for all three directions. This is geometric and
    doesn't particularly benefit from autodiff, but is provided here
    for API consistency.

    Args:
        g: Gauge coupling constant
        dx: Lattice spacing

    Returns:
        JIT-compiled function: (phi, links) -> force array (3, N, N, N)

    Example:
        ```python
        force_fn = make_gauge_force(g=0.5, dx=0.1)
        F_E = force_fn(phi, links)
        ```
    """
    from jaxlatt.operators.gauge import gauge_force_all_directions

    @jit
    def force(phi: Array, links: Array) -> Array:
        return gauge_force_all_directions(phi, links, g, dx)

    return force

make_gauge_force_expanding(g, dx)

Create force on electric field in expanding universe.

Parameters:

Name	Type	Description	Default
g	float	Gauge coupling constant	required
dx	float	Comoving lattice spacing	required

Returns:

Type	Description
Callable[[Array, Array, float], Array]	JIT-compiled function: (phi, links, a) -> force array (3, N, N, N)

Source code in jaxlatt/operators/forces.py

def make_gauge_force_expanding(
    g: float,
    dx: float,
) -> Callable[[Array, Array, float], Array]:
    """
    Create force on electric field in expanding universe.

    Args:
        g: Gauge coupling constant
        dx: Comoving lattice spacing

    Returns:
        JIT-compiled function: (phi, links, a) -> force array (3, N, N, N)
    """
    from jaxlatt.operators.gauge import gauge_force_all_directions_expanding

    @jit
    def force(phi: Array, links: Array, a: float) -> Array:
        return gauge_force_all_directions_expanding(phi, links, g, dx, a)

    return force

make_forces_from_hamiltonian(potential, g, dx)

Create scalar and gauge forces from full Hamiltonian via autodiff.

This guarantees F = -∂H/∂q exactly, useful for: 1. Validating other force implementations 2. Energy conservation checks 3. Novel potentials without analytical derivatives

Note: Slightly slower than specialized implementations due to full Hamiltonian evaluation. Use make_coupled_scalar_force for production.

Parameters:

Name	Type	Description	Default
potential	ScalarPotential | PotentialFunction	ScalarPotential or legacy potential function	required
g	float	Gauge coupling constant	required
dx	float	Lattice spacing	required

Returns:

Type	Description
tuple[Callable, Callable]	Tuple of (scalar_force_fn, gauge_force_fn)

Example

V = ScalarPotential.quartic(m=1.0, lambda_=0.1)
scalar_F, gauge_F = make_forces_from_hamiltonian(V, g=0.5, dx=0.1)
F_phi = scalar_F(phi, pi, links, E)

Source code in jaxlatt/operators/forces.py

def make_forces_from_hamiltonian(
    potential: ScalarPotential | PotentialFunction,
    g: float,
    dx: float,
) -> tuple[Callable, Callable]:
    """
    Create scalar and gauge forces from full Hamiltonian via autodiff.

    This guarantees F = -∂H/∂q exactly, useful for:
    1. Validating other force implementations
    2. Energy conservation checks
    3. Novel potentials without analytical derivatives

    Note: Slightly slower than specialized implementations due to
    full Hamiltonian evaluation. Use make_coupled_scalar_force for production.

    Args:
        potential: ScalarPotential or legacy potential function
        g: Gauge coupling constant
        dx: Lattice spacing

    Returns:
        Tuple of (scalar_force_fn, gauge_force_fn)

    Example:
        ```python
        V = ScalarPotential.quartic(m=1.0, lambda_=0.1)
        scalar_F, gauge_F = make_forces_from_hamiltonian(V, g=0.5, dx=0.1)
        F_phi = scalar_F(phi, pi, links, E)
        ```
    """
    from jaxlatt.operators.coupled import covariant_gradient_squared
    from jaxlatt.operators.gauge import electric_energy_density, magnetic_energy_density

    # Get potential function
    if isinstance(potential, ScalarPotential):
        V_fn = lambda field: potential(field)
    else:
        V_fn = potential

    def total_hamiltonian(phi: Array, pi: Array, links: Array, E: Array) -> Array:
        """Compute total Hamiltonian H = T + V."""
        # Scalar kinetic energy
        T_scalar = 0.5 * jnp.sum(jnp.abs(pi) ** 2)

        # Scalar gradient energy (covariant)
        grad_sq = covariant_gradient_squared(phi, links, dx)
        T_grad = 0.5 * jnp.sum(grad_sq)

        # Scalar potential energy
        V_scalar = jnp.sum(V_fn(phi)).real

        # Electric energy
        T_E = jnp.sum(electric_energy_density(E, dx))

        # Magnetic energy
        T_B = jnp.sum(magnetic_energy_density(links, dx, g))

        return T_scalar + T_grad + V_scalar + T_E + T_B

    @jit
    def scalar_force(phi: Array, pi: Array, links: Array, E: Array) -> Array:
        """$F_\\phi = -\\partial H / \\partial \\phi^*$"""
        return -jnp.conj(
            grad(lambda p: total_hamiltonian(p, pi, links, E).real, holomorphic=False)(phi)
        )

    @jit
    def gauge_force(phi: Array, pi: Array, links: Array, E: Array) -> Array:
        """$F_U = -\\partial H / \\partial U^*$ (approximation for U(1))"""
        # For U(1) gauge theory, this is more subtle
        # We compute gradient w.r.t. link phases
        return -grad(lambda L: total_hamiltonian(phi, pi, L, E).real, holomorphic=False)(links)

    return scalar_force, gauge_force

validate_force_against_manual(autodiff_force, manual_force, rtol=1e-05, atol=1e-08)

Compare autodiff force with manual implementation.

Parameters:

Name	Type	Description	Default
autodiff_force	Array	Force computed via autodiff	required
manual_force	Array	Force computed manually	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/forces.py

def validate_force_against_manual(
    autodiff_force: Array,
    manual_force: Array,
    rtol: float = 1e-5,
    atol: float = 1e-8,
) -> dict:
    """
    Compare autodiff force with manual implementation.

    Args:
        autodiff_force: Force computed via autodiff
        manual_force: Force computed manually
        rtol: Relative tolerance
        atol: Absolute tolerance

    Returns:
        Dictionary with comparison metrics
    """
    diff = autodiff_force - manual_force
    max_abs_error = float(jnp.max(jnp.abs(diff)))
    rms_error = float(jnp.sqrt(jnp.mean(jnp.abs(diff) ** 2)))
    max_rel_error = float(jnp.max(jnp.abs(diff) / (jnp.abs(manual_force) + atol)))

    matches = jnp.allclose(autodiff_force, manual_force, rtol=rtol, atol=atol)

    return {
        "matches": bool(matches),
        "max_abs_error": max_abs_error,
        "rms_error": rms_error,
        "max_rel_error": max_rel_error,
    }