Evolution Module

evolution

Evolution module for lattice field theory.

Provides symplectic leapfrog integrators for:

• Pure gauge fields (flat + expanding space)
• Coupled scalar-gauge fields (flat + expanding space)
• Simple scalar fields (flat space only)

Also provides RK4 (4th-order Runge-Kutta) integrators as a high-accuracy alternative.

All leapfrog integrators preserve energy, symplectic structure, and constraints. RK4 offers higher accuracy but does not preserve symplectic structure.

make_radiation_rho(rho_initial, a_initial)

Create a radiation-era density function with \(\rho(a) \propto a^{-4}\).

Parameters:

Name	Type	Description	Default
rho_initial	float	Reference density at a_initial.	required
a_initial	float	Reference scale factor.	required

Returns:

Type	Description
	Callable rho_func(a) that evaluates radiation density at scale factor a.

Source code in jaxlatt/core/cosmology/frw.py

def make_radiation_rho(rho_initial: float, a_initial: float):
    r"""Create a radiation-era density function with $\rho(a) \propto a^{-4}$.

    Args:
        rho_initial: Reference density at `a_initial`.
        a_initial: Reference scale factor.

    Returns:
        Callable `rho_func(a)` that evaluates radiation density at scale factor `a`.
    """

    def rho_func(a: float) -> float:
        return rho_initial * (a_initial / a) ** 4

    return rho_func

compute_energy_densities_expanding(lattice, a)

Compute COMOVING energy densities in expanding universe.

Phase A: Returns only comoving densities (energy per comoving volume). Physical densities with correct a-scaling deferred to Phase B.

Parameters:

Name	Type	Description	Default
lattice	CoupledLattice	Field state	required
a	float	Scale factor (unused in Phase A, kept for API compatibility)	required

Returns:

Type	Description
dict	dict with: "total": total comoving energy density "kinetic": scalar kinetic energy density (comoving) "gradient": scalar gradient energy density (comoving) "potential": scalar potential energy density (comoving) "electric": electric energy density (comoving) "magnetic": magnetic energy density (comoving)

Source code in jaxlatt/evolution/leapfrog.py

def compute_energy_densities_expanding(
    lattice: "CoupledLattice",
    a: float,
) -> dict:
    """
    Compute COMOVING energy densities in expanding universe.

    Phase A: Returns only comoving densities (energy per comoving volume).
    Physical densities with correct a-scaling deferred to Phase B.

    Args:
        lattice: Field state
        a: Scale factor (unused in Phase A, kept for API compatibility)

    Returns:
        dict with:
            "total": total comoving energy density
            "kinetic": scalar kinetic energy density (comoving)
            "gradient": scalar gradient energy density (comoving)
            "potential": scalar potential energy density (comoving)
            "electric": electric energy density (comoving)
            "magnetic": magnetic energy density (comoving)
    """
    from jaxlatt.operators.coupled import coupled_energy

    energy_dict = coupled_energy(
        lattice.phi,
        lattice.pi,
        lattice.links,
        lattice.E,
        lattice.dx,
        lattice.m,
        lattice.lambda_,
        lattice.g,
    )

    volume = (lattice.dx * lattice.size[0]) ** 3

    # Return COMOVING energy densities only (Phase A)
    return {
        "total": energy_dict["total"] / volume,
        "kinetic": energy_dict["scalar_kinetic"] / volume,
        "gradient": energy_dict["scalar_gradient"] / volume,
        "potential": energy_dict["scalar_potential"] / volume,
        "electric": energy_dict["electric"] / volume,
        "magnetic": energy_dict["magnetic"] / volume,
    }

coupled_evolve(lattice, dt, steps, save_every=1, verbose=False)

Evolve coupled system over multiple time steps.

Eager ergonomics layer over :func:_coupled_evolve_core, which is the entry point to use inside a JAX transformation.

Parameters:

Name	Type	Description	Default
lattice	CoupledLattice	Initial state	required
dt	float	Time step	required
steps	int	Number of steps	required
save_every	int	Save snapshot every N steps	1
verbose	bool	Show progress bar	False

Returns:

Name	Type	Description
times	Array	Saved snapshot times
snapshots	list[CoupledLattice]	List of saved lattice states

Source code in jaxlatt/evolution/leapfrog.py

def coupled_evolve(
    lattice: "CoupledLattice",
    dt: float,
    steps: int,
    save_every: int = 1,
    verbose: bool = False,
) -> tuple[Array, list["CoupledLattice"]]:
    """
    Evolve coupled system over multiple time steps.

    Eager ergonomics layer over :func:`_coupled_evolve_core`, which is the entry
    point to use inside a JAX transformation.

    Args:
        lattice: Initial state
        dt: Time step
        steps: Number of steps
        save_every: Save snapshot every N steps
        verbose: Show progress bar

    Returns:
        times: Saved snapshot times
        snapshots: List of saved lattice states
    """
    num_segments = steps // save_every
    times = snapshot_times(num_segments, save_every, dt)

    if verbose:
        snapshots = [lattice]
        carry = (lattice.phi, lattice.pi, lattice.links, lattice.E)
        with _evolution_progress("Evolving coupled lattice...", num_segments) as (progress, task):
            for _ in range(num_segments):
                carry = lax.fori_loop(
                    0,
                    save_every,
                    lambda _, st: coupled_leapfrog_step_arrays(
                        *st, lattice.m, lattice.lambda_, lattice.g, lattice.dx, dt
                    ),
                    carry,
                )
                snapshots.append(lattice.update(*carry))
                progress.update(task, advance=1)
        return times, snapshots

    _, s_phi, s_pi, s_links, s_E = _coupled_evolve_core(lattice, dt, steps, save_every)
    snapshots = [
        lattice.update(s_phi[i], s_pi[i], s_links[i], s_E[i]) for i in range(num_segments + 1)
    ]
    return times, snapshots

coupled_evolve_expanding(lattice, universe, rho_func, pressure_func, dt, steps, save_every=1, verbose=False)

Evolve coupled fields in expanding universe.

Eager ergonomics layer over :func:_coupled_evolve_expanding_core: it adds the host-side overflow diagnostics and unstacks the trajectory into Python lists. Being eager is deliberate — to trace, jit/grad/vmap the core instead. To keep peak memory bounded on long runs, use :func:coupled_evolve_expanding_chunked.

rho_func and pressure_func must be JAX-traceable (no Python control flow on their argument).

Parameters:

Name	Type	Description	Default
lattice	CoupledLattice	Initial field configuration	required
universe	FRWUniverse	Initial cosmological state	required
rho_func	Callable[[float], float]	Function rho(a) — must be JAX-traceable	required
pressure_func	Callable[[float], float]	Function p(a) — must be JAX-traceable	required
dt	float	Time step (conformal time)	required
steps	int	Number of steps	required
save_every	int	Save snapshot every N steps	1
verbose	bool	If True, use Python loop instead of lax.scan (for debugging)	False

Returns:

Name	Type	Description
times	Array	Conformal times at each snapshot
lattices	list[CoupledLattice]	List of saved lattice states
universes	list[FRWUniverse]	List of saved universe states

r

Source code in jaxlatt/evolution/leapfrog.py

def coupled_evolve_expanding(
    lattice: "CoupledLattice",
    universe: "FRWUniverse",
    rho_func: Callable[[float], float],
    pressure_func: Callable[[float], float],
    dt: float,
    steps: int,
    save_every: int = 1,
    verbose: bool = False,
) -> tuple[Array, list["CoupledLattice"], list["FRWUniverse"]]:
    """
    Evolve coupled fields in expanding universe.

    Eager ergonomics layer over :func:`_coupled_evolve_expanding_core`: it adds
    the host-side overflow diagnostics and unstacks the trajectory into Python
    lists. Being eager is deliberate — to trace, ``jit``/``grad``/``vmap`` the
    core instead. To keep peak memory bounded on long runs, use
    :func:`coupled_evolve_expanding_chunked`.

    rho_func and pressure_func must be JAX-traceable (no Python control flow
    on their argument).

    Args:
        lattice: Initial field configuration
        universe: Initial cosmological state
        rho_func: Function rho(a) — must be JAX-traceable
        pressure_func: Function p(a) — must be JAX-traceable
        dt: Time step (conformal time)
        steps: Number of steps
        save_every: Save snapshot every N steps
        verbose: If True, use Python loop instead of lax.scan (for debugging)

    Returns:
        times: Conformal times at each snapshot
        lattices: List of saved lattice states
        universes: List of saved universe states
    r"""
    num_segments = steps // save_every

    if verbose:
        lattices = [lattice]
        universes = [universe]
        taus = [universe.tau]
        carry = (lattice, universe)
        with _evolution_progress("Evolving expanding universe...", num_segments) as (
            progress,
            task,
        ):
            for _ in range(num_segments):
                carry = lax.fori_loop(
                    0,
                    save_every,
                    lambda _, st: coupled_leapfrog_step_expanding(
                        st[0], st[1], rho_func, pressure_func, dt
                    ),
                    carry,
                )
                lat_snap, univ_snap = carry
                lattices.append(lat_snap)
                universes.append(univ_snap)
                taus.append(univ_snap.tau)
                where = format_tau(univ_snap.tau)
                warn_if_overflow("Scalar field", jnp.max(jnp.abs(lat_snap.phi)), context=where)
                warn_if_overflow("Momentum", jnp.max(jnp.abs(lat_snap.pi)), context=where)
                warn_if_overflow("Electric field", jnp.max(jnp.abs(lat_snap.E)), context=where)
                progress.update(task, advance=1)
        return jnp.stack(taus), lattices, universes

    times, stacked_lats, stacked_univs = _coupled_evolve_expanding_core(
        lattice, universe, rho_func, pressure_func, dt, steps, save_every
    )

    # Single post-scan overflow check — avoids per-step D2H syncs.
    # No-ops under jit/grad/vmap; see jaxlatt._diagnostics.
    warn_if_overflow("Scalar field", jnp.max(jnp.abs(stacked_lats.phi)))
    warn_if_overflow("Momentum", jnp.max(jnp.abs(stacked_lats.pi)))
    warn_if_overflow("Electric field", jnp.max(jnp.abs(stacked_lats.E)))

    n = num_segments + 1
    return times, unstack(stacked_lats, n), unstack(stacked_univs, n)

coupled_evolve_expanding_chunked(lattice, universe, rho_func, pressure_func, dt, steps, save_every=1, chunk_segments=10)

Stream an expanding-universe run chunk by chunk.

Yields (times, stacked_lattices, stacked_universes) for at most chunk_segments snapshots at a time, so peak device memory is set by chunk_segments rather than by steps. The eager evolvers stack the whole trajectory, which at N=256 with 100 snapshots is tens of GB.

Only the first chunk includes the initial state, so concatenating the chunks reproduces :func:coupled_evolve_expanding exactly.

Parameters:

Name	Type	Description	Default
chunk_segments	int	Snapshots materialised per chunk.	10

Yields:

Type	Description
Array	(times, stacked_lattices, stacked_universes) per chunk. Write each
CoupledLattice	out (see :func:jaxlatt.io.append_snapshots) before taking the next.

Example

for times, lats, univs in coupled_evolve_expanding_chunked(
    lat, univ, rho, p, dt=1e-3, steps=100_000, save_every=1000
):
    append_snapshots("run.h5", times, lats)

Source code in jaxlatt/evolution/leapfrog.py

def coupled_evolve_expanding_chunked(
    lattice: "CoupledLattice",
    universe: "FRWUniverse",
    rho_func: Callable[[float], float],
    pressure_func: Callable[[float], float],
    dt: float,
    steps: int,
    save_every: int = 1,
    chunk_segments: int = 10,
) -> Iterator[tuple[Array, "CoupledLattice", "FRWUniverse"]]:
    """Stream an expanding-universe run chunk by chunk.

    Yields ``(times, stacked_lattices, stacked_universes)`` for at most
    ``chunk_segments`` snapshots at a time, so peak device memory is set by
    ``chunk_segments`` rather than by ``steps``. The eager evolvers stack the
    whole trajectory, which at N=256 with 100 snapshots is tens of GB.

    Only the first chunk includes the initial state, so concatenating the chunks
    reproduces :func:`coupled_evolve_expanding` exactly.

    Args:
        chunk_segments: Snapshots materialised per chunk.

    Yields:
        ``(times, stacked_lattices, stacked_universes)`` per chunk. Write each
        out (see :func:`jaxlatt.io.append_snapshots`) before taking the next.

    Example:
        ```python
        for times, lats, univs in coupled_evolve_expanding_chunked(
            lat, univ, rho, p, dt=1e-3, steps=100_000, save_every=1000
        ):
            append_snapshots("run.h5", times, lats)
        ```
    """
    if chunk_segments < 1:
        raise ValueError("chunk_segments must be >= 1")

    total_segments = steps // save_every
    state = (lattice, universe)
    done = 0
    first = True

    while done < total_segments:
        n = min(chunk_segments, total_segments - done)

        def step(st):
            return coupled_leapfrog_step_expanding(st[0], st[1], rho_func, pressure_func, dt)

        state, (stacked_lats, stacked_univs) = scan_segments(
            step, state, n, save_every, include_initial=first
        )
        warn_if_overflow("Scalar field", jnp.max(jnp.abs(stacked_lats.phi)))
        warn_if_overflow("Momentum", jnp.max(jnp.abs(stacked_lats.pi)))
        warn_if_overflow("Electric field", jnp.max(jnp.abs(stacked_lats.E)))
        yield stacked_univs.tau, stacked_lats, stacked_univs

        done += n
        first = False

    if total_segments == 0:
        stacked = jax.tree_util.tree_map(lambda x: jnp.asarray(x)[None], (lattice, universe))
        yield stacked[1].tau, stacked[0], stacked[1]

coupled_leapfrog_step(lattice, dt)

Single leapfrog step for coupled scalar-gauge evolution.

Equations of motion:

• \(d\phi/dt = \pi\)
• \(d\pi/dt = F_\phi\)
• \(dA_i/dt = E_i\)
• \(dE_i/dt = F_A\) (with scalar current)

Parameters:

Name	Type	Description	Default
lattice	CoupledLattice	Current state	required
dt	float	Time step	required

Returns:

Type	Description
CoupledLattice	CoupledLattice at t+dt

Source code in jaxlatt/evolution/leapfrog.py

def coupled_leapfrog_step(lattice: "CoupledLattice", dt: float) -> "CoupledLattice":
    r"""
    Single leapfrog step for coupled scalar-gauge evolution.

    Equations of motion:

    - $d\phi/dt = \pi$
    - $d\pi/dt = F_\phi$
    - $dA_i/dt = E_i$
    - $dE_i/dt = F_A$ (with scalar current)

    Args:
        lattice: Current state
        dt: Time step

    Returns:
        CoupledLattice at t+dt
    """
    phi_new, pi_new, links_new, E_new = coupled_leapfrog_step_arrays(
        lattice.phi,
        lattice.pi,
        lattice.links,
        lattice.E,
        lattice.m,
        lattice.lambda_,
        lattice.g,
        lattice.dx,
        dt,
    )

    return lattice.update(phi_new, pi_new, links_new, E_new)

coupled_leapfrog_step_expanding(lattice, universe, rho_func, pressure_func, dt)

Leapfrog step with cosmological expansion and Hubble friction.

Modified equations in conformal time:

• \(\phi' = \pi\)
• \(\pi' = F_\phi/a^2 - 2H\pi\) (Hubble friction on \(\pi\), treated via PC)
• \(A'_i = E_i\)
• \(E'_i = -F_E/a^2 - 2H E_i\) (Hubble friction on \(E\), treated via PC)
• \(a' = a^2 H\) (scale factor evolution)

The Hubble friction terms use a predictor-corrector (trapezoidal) scheme: \(H_n \cdot \pi_n\) at step \(n\) and \(H_{n+1} \cdot \pi_\mathrm{pred}\) at the predicted step, giving 2nd-order accuracy in \(dt \cdot H\). Forces are evaluated at the half-step position as usual.

Parameters:

Name	Type	Description	Default
lattice	CoupledLattice	Current field state (comoving coordinates)	required
universe	FRWUniverse	Current cosmological state	required
rho_func	Callable[[float], float]	Function \(\rho(a)\) giving physical energy density	required
pressure_func	Callable[[float], float]	Function \(p(a)\) giving physical pressure	required
dt	float	Time step (conformal time)	required

Returns:

Name	Type	Description
lattice_new	CoupledLattice	Updated field state
universe_new	FRWUniverse	Updated cosmological state

Source code in jaxlatt/evolution/leapfrog.py

@eqx.filter_jit
def coupled_leapfrog_step_expanding(
    lattice: "CoupledLattice",
    universe: "FRWUniverse",
    rho_func: Callable[[float], float],
    pressure_func: Callable[[float], float],
    dt: float,
) -> tuple["CoupledLattice", "FRWUniverse"]:
    r"""
    Leapfrog step with cosmological expansion and Hubble friction.

    Modified equations in conformal time:

    - $\phi' = \pi$
    - $\pi' = F_\phi/a^2 - 2H\pi$ (Hubble friction on $\pi$, treated via PC)
    - $A'_i = E_i$
    - $E'_i = -F_E/a^2 - 2H E_i$ (Hubble friction on $E$, treated via PC)
    - $a' = a^2 H$ (scale factor evolution)

    The Hubble friction terms use a predictor-corrector (trapezoidal) scheme:
    $H_n \cdot \pi_n$ at step $n$ and $H_{n+1} \cdot \pi_\mathrm{pred}$ at the
    predicted step, giving 2nd-order accuracy in $dt \cdot H$. Forces are
    evaluated at the half-step position as usual.

    Args:
        lattice: Current field state (comoving coordinates)
        universe: Current cosmological state
        rho_func: Function $\rho(a)$ giving physical energy density
        pressure_func: Function $p(a)$ giving physical pressure
        dt: Time step (conformal time)

    Returns:
        lattice_new: Updated field state
        universe_new: Updated cosmological state
    """
    # Advance Friedmann first to obtain H_{n+1} for the corrector
    universe_new = friedmann_step_predictor_corrector(universe, rho_func, pressure_func, dt)

    phi_new, pi_new, links_new, E_new = _coupled_leapfrog_expanding_arrays(
        lattice.phi,
        lattice.pi,
        lattice.links,
        lattice.E,
        universe.a,
        universe.H,
        universe_new.H,
        lattice.m,
        lattice.lambda_,
        lattice.g,
        lattice.dx,
        dt,
    )

    lattice_new = lattice.update(phi_new, pi_new, links_new, E_new)

    return lattice_new, universe_new

coupled_leapfrog_step_rescaled(lattice, universe, rho_func, pressure_func, dt)

Leapfrog step using rescaled field \(\chi = a\phi\) eliminating explicit friction.

Algorithm:

• \(\chi = a\phi\)
• \(\pi_\chi = \chi'\)
• \(F_\chi = (1/a^2)\nabla^2\chi - a^2 m^2 \chi - \lambda|\chi|^2\chi + (a''/a)\chi\)
• \(\chi_{n+1/2} = \chi_n + (dt/2)\,\pi_{\chi,n}\)
• \(\pi_{\chi,n+1} = \pi_{\chi,n} + dt\, F_\chi(\chi_{n+1/2})\)
• \(\chi_{n+1} = \chi_{n+1/2} + (dt/2)\,\pi_{\chi,n+1}\)
• Back-transform: \(\phi = \chi/a\), \(\pi = (\pi_\chi - a'\phi)/a\)

Gauge fields still evolved with friction (future: rescale).

Parameters:

Name	Type	Description	Default
lattice	CoupledLattice	Current field state	required
universe	FRWUniverse	Current cosmological state	required
rho_func	Callable[[float], float]	Function \(\rho(a)\) for scale factor evolution	required
pressure_func	Callable[[float], float]	Function \(p(a)\) for scale factor evolution	required
dt	float	Time step	required

Returns:

Name	Type	Description
lattice_new	CoupledLattice	Updated field state
universe_new	FRWUniverse	Updated cosmological state

Source code in jaxlatt/evolution/leapfrog.py

@eqx.filter_jit
def coupled_leapfrog_step_rescaled(
    lattice: "CoupledLattice",
    universe: "FRWUniverse",
    rho_func: Callable[[float], float],
    pressure_func: Callable[[float], float],
    dt: float,
) -> tuple["CoupledLattice", "FRWUniverse"]:
    r"""
    Leapfrog step using rescaled field $\chi = a\phi$ eliminating explicit friction.

    Algorithm:

    - $\chi = a\phi$
    - $\pi_\chi = \chi'$
    - $F_\chi = (1/a^2)\nabla^2\chi - a^2 m^2 \chi - \lambda|\chi|^2\chi + (a''/a)\chi$
    - $\chi_{n+1/2} = \chi_n + (dt/2)\,\pi_{\chi,n}$
    - $\pi_{\chi,n+1} = \pi_{\chi,n} + dt\, F_\chi(\chi_{n+1/2})$
    - $\chi_{n+1} = \chi_{n+1/2} + (dt/2)\,\pi_{\chi,n+1}$
    - Back-transform: $\phi = \chi/a$, $\pi = (\pi_\chi - a'\phi)/a$

    Gauge fields still evolved with friction (future: rescale).

    Args:
        lattice: Current field state
        universe: Current cosmological state
        rho_func: Function $\rho(a)$ for scale factor evolution
        pressure_func: Function $p(a)$ for scale factor evolution
        dt: Time step

    Returns:
        lattice_new: Updated field state
        universe_new: Updated cosmological state
    """
    a = universe.a
    adot = universe.adot

    # Compute a'' for curvature term
    rho_current = rho_func(a)
    p_current = pressure_func(a)
    addot = friedmann_acceleration(rho_current, p_current, universe.M_pl, a)

    # Rescaled variables
    chi = a * lattice.phi
    pi_chi = a * lattice.pi + adot * lattice.phi  # χ' = a φ' + a' φ

    # Half-step for χ
    chi_half = chi + 0.5 * dt * pi_chi

    # Update gauge links half-step
    links_half = evolve_links_half_step(lattice.links, lattice.E, dt)

    # Force at half-step
    F_chi = scalar_force_rescaled(
        chi_half, links_half, lattice.m, lattice.lambda_, lattice.dx, a, addot
    )

    # Full update for π_χ
    pi_chi_new = pi_chi + dt * F_chi
    chi_new = chi_half + 0.5 * dt * pi_chi_new

    # Evolve gauge momenta with original friction form (vectorized, no Python loop)
    H = universe.H
    F_gauge = gauge_force_all_directions_expanding(
        lattice.phi, lattice.links, lattice.g, lattice.dx, a
    )
    E_new = lattice.E + dt * ((F_gauge / (a**2)) - 2.0 * H * lattice.E)

    # Second half-step links
    links_new = evolve_links_half_step(links_half, E_new, dt)

    # Update scale factor (predictor-corrector) after field step
    universe_new = friedmann_step_predictor_corrector(universe, rho_func, pressure_func, dt)
    a_new = universe_new.a
    adot_new = universe_new.adot

    # Back-transform to φ, π at new time
    phi_new = chi_new / a_new
    pi_new = (pi_chi_new - adot_new * phi_new) / a_new

    lattice_new = lattice.update(phi_new, pi_new, links_new, E_new)

    return lattice_new, universe_new

gauge_evolve(lattice, dt, steps, save_every=1, verbose=False)

Evolve gauge field using leapfrog integrator.

Eager ergonomics layer over :func:_gauge_evolve_core, which is the entry point to use inside a JAX transformation.

Parameters:

Name	Type	Description	Default
lattice	GaugeLattice	Initial gauge lattice state	required
dt	float	Time step	required
steps	int	Number of steps	required
save_every	int	Save snapshot every N steps	1
verbose	bool	Show progress bar	False

Returns:

Name	Type	Description
times	Array	Array of snapshot times
snapshots	list[GaugeLattice]	List of GaugeLattice states

Source code in jaxlatt/evolution/leapfrog.py

def gauge_evolve(
    lattice: "GaugeLattice",
    dt: float,
    steps: int,
    save_every: int = 1,
    verbose: bool = False,
) -> tuple[Array, list["GaugeLattice"]]:
    """
    Evolve gauge field using leapfrog integrator.

    Eager ergonomics layer over :func:`_gauge_evolve_core`, which is the entry
    point to use inside a JAX transformation.

    Args:
        lattice: Initial gauge lattice state
        dt: Time step
        steps: Number of steps
        save_every: Save snapshot every N steps
        verbose: Show progress bar

    Returns:
        times: Array of snapshot times
        snapshots: List of GaugeLattice states
    """
    num_segments = steps // save_every
    times = snapshot_times(num_segments, save_every, dt)

    if verbose:
        snapshots = [lattice]
        carry = (lattice.links, lattice.E)
        with _evolution_progress("Evolving gauge lattice...", num_segments) as (progress, task):
            for _ in range(num_segments):
                carry = lax.fori_loop(
                    0,
                    save_every,
                    lambda _, st: _gauge_leapfrog_step_arrays(
                        st[0], st[1], lattice.dx, lattice.g, dt
                    ),
                    carry,
                )
                snapshots.append(lattice.update(links=carry[0], E=carry[1]))
                progress.update(task, advance=1)
        return times, snapshots

    _, stacked_links, stacked_E = _gauge_evolve_core(lattice, dt, steps, save_every)
    snapshots = [
        lattice.update(links=stacked_links[i], E=stacked_E[i]) for i in range(num_segments + 1)
    ]
    return times, snapshots

gauge_leapfrog_step(lattice, dt)

Single leapfrog step for pure gauge U(1) evolution.

Algorithm

1. E(t+dt/2) = E(t) - (dt/2) * F[U(t)]
2. U(t+dt) = U(t) * exp(i * dt * E(t+dt/2))
3. E(t+dt) = E(t+dt/2) - (dt/2) * F[U(t+dt)]

Preserves unitarity, Gauss constraint, and energy.

Parameters:

Name	Type	Description	Default
lattice	GaugeLattice	Current gauge lattice state	required
dt	float	Time step	required

Returns:

Type	Description
GaugeLattice	Updated GaugeLattice at t+dt

Source code in jaxlatt/evolution/leapfrog.py

def gauge_leapfrog_step(lattice: "GaugeLattice", dt: float) -> "GaugeLattice":
    """
    Single leapfrog step for pure gauge U(1) evolution.

    Algorithm:
        1. E(t+dt/2) = E(t) - (dt/2) * F[U(t)]
        2. U(t+dt) = U(t) * exp(i * dt * E(t+dt/2))
        3. E(t+dt) = E(t+dt/2) - (dt/2) * F[U(t+dt)]

    Preserves unitarity, Gauss constraint, and energy.

    Args:
        lattice: Current gauge lattice state
        dt: Time step

    Returns:
        Updated GaugeLattice at t+dt
    """
    links_new, E_new = _gauge_leapfrog_step_arrays(
        lattice.links, lattice.E, lattice.dx, lattice.g, dt
    )
    return lattice.update(links=links_new, E=E_new)

scalar_force_expanding(phi, links, m, lambda_, dx, a)

Compute force on scalar field in expanding universe.

Force for \(\pi\) update (excluding friction \(-2H\pi\) applied separately):

\[F_\phi = (\nabla^2 \phi)/a^2 - a^2 \, dV/d\phi^*\]

Uses autodiff via scalar_potential_force for correct Wirtinger derivatives on complex fields.

Parameters:

Name	Type	Description	Default
phi	Array	Scalar field (comoving)	required
links	Array	Gauge links	required
m	float	Mass parameter	required
lambda_	float	Self-coupling	required
dx	float	Lattice spacing (comoving)	required
a	float	Scale factor	required

Returns:

Type	Description
Array	\(F_\phi\): Force on \(\phi\) (excluding friction term)

Source code in jaxlatt/evolution/leapfrog.py

def scalar_force_expanding(
    phi: Array,
    links: Array,
    m: float,
    lambda_: float,
    dx: float,
    a: float,
) -> Array:
    r"""
    Compute force on scalar field in expanding universe.

    Force for $\pi$ update (excluding friction $-2H\pi$ applied separately):

    $$F_\phi = (\nabla^2 \phi)/a^2 - a^2 \, dV/d\phi^*$$

    Uses autodiff via `scalar_potential_force` for correct Wirtinger derivatives
    on complex fields.

    Args:
        phi: Scalar field (comoving)
        links: Gauge links
        m: Mass parameter
        lambda_: Self-coupling
        dx: Lattice spacing (comoving)
        a: Scale factor

    Returns:
        $F_\phi$: Force on $\phi$ (excluding friction term)
    """
    from jaxlatt.operators.scalar import scalar_potential_force

    # Covariant Laplacian
    laplacian = covariant_laplacian(phi, links, dx)

    # Kinetic term: +∇²φ/a²
    kinetic_force = laplacian / (a**2)

    # Potential force via autodiff: scalar_potential_force returns -dV/dφ*
    # We want: -a² dV/dφ* = a² * scalar_potential_force
    pot_force = scalar_potential_force(phi, m, lambda_)
    potential_force = (a**2) * pot_force

    return kinetic_force + potential_force

scalar_force_rescaled(chi, links, m, lambda_, dx, a, addot)

Force for rescaled field \(\chi = a\phi\) eliminating explicit Hubble friction.

Derived equation for \(\chi\):

\[\chi'' = \frac{1}{a^2} \nabla^2 \chi - a^2 m^2 \chi - \lambda |\chi|^2 \chi + \frac{a''}{a} \chi\]

Parameters:

Name	Type	Description	Default
chi	Array	Rescaled field \(\chi = a\phi\)	required
links	Array	Gauge links	required
m, lambda_		Scalar parameters	required
dx	float	Lattice spacing (comoving)	required
a	float	Scale factor	required
addot	float	Conformal second derivative \(a''\) (from Friedmann)	required

Returns:

Type	Description
Array	\(F_\chi\) array same shape as \(\chi\)

Source code in jaxlatt/evolution/leapfrog.py

def scalar_force_rescaled(
    chi: Array,
    links: Array,
    m: float,
    lambda_: float,
    dx: float,
    a: float,
    addot: float,
) -> Array:
    r"""
    Force for rescaled field $\chi = a\phi$ eliminating explicit Hubble friction.

    Derived equation for $\chi$:

    $$\chi'' = \frac{1}{a^2} \nabla^2 \chi - a^2 m^2 \chi - \lambda |\chi|^2 \chi + \frac{a''}{a} \chi$$

    Args:
        chi: Rescaled field $\chi = a\phi$
        links: Gauge links
        m, lambda_: Scalar parameters
        dx: Lattice spacing (comoving)
        a: Scale factor
        addot: Conformal second derivative $a''$ (from Friedmann)

    Returns:
        $F_\chi$ array same shape as $\chi$
    """
    # ∇²χ (covariant Laplacian on χ)
    laplacian_chi = covariant_laplacian(chi, links, dx)

    # (1/a²) ∇²χ term
    kinetic_term = laplacian_chi / (a**2)

    # Potential term: −a² m² χ − λ |χ|² χ  (restoring forces)
    chi_sq = jnp.abs(chi) ** 2
    potential_term = -(a**2) * (m**2) * chi - lambda_ * chi_sq * chi

    # Curvature term: +(a''/a) χ
    curvature_term = jnp.where(a > 0, (addot / a) * chi, jnp.zeros_like(chi))

    return kinetic_term + potential_term + curvature_term

coupled_rk4_step(lattice, dt)

Single RK4 step for coupled scalar-gauge field evolution (flat space).

This is a 4th-order accurate integrator that requires 4 force evaluations per step, compared to 2 for leapfrog. It does not preserve the symplectic structure but offers higher accuracy for smooth dynamics.

Parameters:

Name	Type	Description	Default
lattice	CoupledLattice	Current coupled lattice state	required
dt	float	Time step	required

Returns:

Type	Description
CoupledLattice	Updated CoupledLattice at t+dt

Note

For gauge fields, link updates use exponential map to preserve unitarity. Unlike leapfrog, this does NOT guarantee energy conservation.

Source code in jaxlatt/evolution/rk4.py

def coupled_rk4_step(
    lattice: CoupledLattice,
    dt: float,
) -> CoupledLattice:
    """
    Single RK4 step for coupled scalar-gauge field evolution (flat space).

    This is a 4th-order accurate integrator that requires 4 force evaluations
    per step, compared to 2 for leapfrog. It does not preserve the symplectic
    structure but offers higher accuracy for smooth dynamics.

    Args:
        lattice: Current coupled lattice state
        dt: Time step

    Returns:
        Updated CoupledLattice at t+dt

    Note:
        For gauge fields, link updates use exponential map to preserve unitarity.
        Unlike leapfrog, this does NOT guarantee energy conservation.
    """
    phi = lattice.phi
    pi = lattice.pi
    links = lattice.links
    E = lattice.E
    dx = lattice.dx
    m = lattice.m
    lambda_ = lattice.lambda_
    g = lattice.g

    # RK4 for the system:
    # dphi/dt = pi
    # dpi/dt = F_scalar(phi, links)
    # dE/dt = -F_gauge(links)
    # dU/dt = i*E*U (link evolution)

    # k1
    k1_phi = pi
    k1_pi = scalar_force(phi, links, m, lambda_, dx)
    k1_E = -gauge_force(links, dx, g)

    # k2
    phi_2 = phi + 0.5 * dt * k1_phi
    pi_2 = pi + 0.5 * dt * k1_pi
    links_2 = links * jnp.exp(0.5 * dt * 1j * E)
    links_2 = links_2 / jnp.abs(links_2)  # Ensure unitarity
    E_2 = E + 0.5 * dt * k1_E

    k2_phi = pi_2
    k2_pi = scalar_force(phi_2, links_2, m, lambda_, dx)
    k2_E = -gauge_force(links_2, dx, g)

    # k3
    phi_3 = phi + 0.5 * dt * k2_phi
    pi_3 = pi + 0.5 * dt * k2_pi
    links_3 = links * jnp.exp(0.5 * dt * 1j * E_2)
    links_3 = links_3 / jnp.abs(links_3)
    E_3 = E + 0.5 * dt * k2_E

    k3_phi = pi_3
    k3_pi = scalar_force(phi_3, links_3, m, lambda_, dx)
    k3_E = -gauge_force(links_3, dx, g)

    # k4
    phi_4 = phi + dt * k3_phi
    pi_4 = pi + dt * k3_pi
    links_4 = links * jnp.exp(dt * 1j * E_3)
    links_4 = links_4 / jnp.abs(links_4)
    E_4 = E + dt * k3_E

    k4_phi = pi_4
    k4_pi = scalar_force(phi_4, links_4, m, lambda_, dx)
    k4_E = -gauge_force(links_4, dx, g)

    # Combine with RK4 weights
    phi_new = phi + (dt / 6.0) * (k1_phi + 2 * k2_phi + 2 * k3_phi + k4_phi)
    pi_new = pi + (dt / 6.0) * (k1_pi + 2 * k2_pi + 2 * k3_pi + k4_pi)
    E_new = E + (dt / 6.0) * (k1_E + 2 * k2_E + 2 * k3_E + k4_E)

    # For links, use exponential of average E
    E_avg = (E + 2 * E_2 + 2 * E_3 + E_4) / 6.0
    links_new = links * jnp.exp(1j * dt * E_avg)
    links_new = links_new / jnp.abs(links_new)

    return lattice.replace(phi=phi_new, pi=pi_new, links=links_new, E=E_new)

gauge_rk4_step(lattice, dt)

Single RK4 step for pure gauge field evolution (flat space).

Parameters:

Name	Type	Description	Default
lattice	GaugeLattice	Current gauge lattice state	required
dt	float	Time step	required

Returns:

Type	Description
GaugeLattice	Updated GaugeLattice at t+dt

Source code in jaxlatt/evolution/rk4.py

def gauge_rk4_step(
    lattice: GaugeLattice,
    dt: float,
) -> GaugeLattice:
    """
    Single RK4 step for pure gauge field evolution (flat space).

    Args:
        lattice: Current gauge lattice state
        dt: Time step

    Returns:
        Updated GaugeLattice at t+dt
    """
    links = lattice.links
    E = lattice.E
    dx = lattice.dx
    g = lattice.g

    # k1
    k1_E = -gauge_force(links, dx, g)

    # k2
    links_2 = links * jnp.exp(0.5 * dt * 1j * E)
    links_2 = links_2 / jnp.abs(links_2)
    E_2 = E + 0.5 * dt * k1_E

    k2_E = -gauge_force(links_2, dx, g)

    # k3
    links_3 = links * jnp.exp(0.5 * dt * 1j * E_2)
    links_3 = links_3 / jnp.abs(links_3)
    E_3 = E + 0.5 * dt * k2_E

    k3_E = -gauge_force(links_3, dx, g)

    # k4
    links_4 = links * jnp.exp(dt * 1j * E_3)
    links_4 = links_4 / jnp.abs(links_4)
    E_4 = E + dt * k3_E

    k4_E = -gauge_force(links_4, dx, g)

    # Combine
    E_new = E + (dt / 6.0) * (k1_E + 2 * k2_E + 2 * k3_E + k4_E)

    # For links, use exponential of average E
    E_avg = (E + 2 * E_2 + 2 * E_3 + E_4) / 6.0
    links_new = links * jnp.exp(1j * dt * E_avg)
    links_new = links_new / jnp.abs(links_new)

    return lattice.replace(links=links_new, E=E_new)

evolve(lattice, potential, dt, steps, save_every=1, verbose=False, H=0.0, use_spectral=False)

Evolve a scalar field lattice for multiple time steps.

Parameters:

Name	Type	Description	Default
lattice	Lattice	Initial lattice state	required
potential	Callable[[Array], Array]	Potential function \(V(\phi)\)	required
dt	float	Time step	required
steps	int	Number of time steps	required
save_every	int	Save snapshot every N steps (default: 1)	1
verbose	bool	Print progress updates (default: False)	False
H	float	Hubble parameter for cosmological friction (default: 0.0, no friction)	0.0
use_spectral	bool	Use FFT-based spectral Laplacian (default: False)	False

Returns:

Type	Description
Array	Tuple of (times, snapshots) where:
list[Lattice]	times: Array of snapshot times
tuple[Array, list[Lattice]]	snapshots: List of Lattice objects at snapshot times

Example

from jaxlatt.core import ScalarPotential
from jaxlatt.utils import create_initial_lattice_1d

# Flat spacetime evolution
lattice = create_initial_lattice_1d(size=128, length=10.0)
potential = ScalarPotential.quadratic(m=1.0)
times, snapshots = evolve(lattice, potential, dt=0.01, steps=1000)

# Cosmological evolution with Hubble friction
times, snapshots = evolve(lattice, potential, dt=0.01, steps=1000, H=0.1)

Source code in jaxlatt/evolution/scalar.py

def evolve(
    lattice: Lattice,
    potential: Callable[[Array], Array],
    dt: float,
    steps: int,
    save_every: int = 1,
    verbose: bool = False,
    H: float = 0.0,
    use_spectral: bool = False,
) -> tuple[Array, list[Lattice]]:
    r"""
    Evolve a scalar field lattice for multiple time steps.

    Args:
        lattice: Initial lattice state
        potential: Potential function $V(\phi)$
        dt: Time step
        steps: Number of time steps
        save_every: Save snapshot every N steps (default: 1)
        verbose: Print progress updates (default: False)
        H: Hubble parameter for cosmological friction (default: 0.0, no friction)
        use_spectral: Use FFT-based spectral Laplacian (default: False)

    Returns:
        Tuple of (times, snapshots) where:
        - times: Array of snapshot times
        - snapshots: List of Lattice objects at snapshot times

    Example:
        ```python
        from jaxlatt.core import ScalarPotential
        from jaxlatt.utils import create_initial_lattice_1d

        # Flat spacetime evolution
        lattice = create_initial_lattice_1d(size=128, length=10.0)
        potential = ScalarPotential.quadratic(m=1.0)
        times, snapshots = evolve(lattice, potential, dt=0.01, steps=1000)

        # Cosmological evolution with Hubble friction
        times, snapshots = evolve(lattice, potential, dt=0.01, steps=1000, H=0.1)
        ```
    """

    def _step(_, lat: Lattice) -> Lattice:
        return evolve_step(lat, potential, dt, H=H, use_spectral=use_spectral)

    if verbose:
        # Python loop: debugging path, drives the progress bar. Not traceable.
        snapshots = []
        times = []
        current_lattice = lattice
        with _evolution_progress("Evolving field...", steps) as (progress, task):
            for step in range(steps):
                if step % save_every == 0:
                    snapshots.append(current_lattice)
                    times.append(step * dt)

                current_lattice = _step(step, current_lattice)
                progress.update(task, advance=1)

        if (steps - 1) % save_every != 0:
            snapshots.append(current_lattice)
            times.append(steps * dt)

        return jnp.array(times), snapshots

    # Fast path: lax.scan over segments, mirroring coupled_evolve.
    # Each segment emits the state *before* advancing, reproducing the Python
    # loop's convention of snapshotting at 0, save_every, 2*save_every, ...
    #
    # This is not a micro-optimisation. A Python loop unrolls `steps` copies of
    # the force kernel into the jaxpr, so tracing even a few hundred steps takes
    # minutes and vmap memory scales with `steps`.
    n_full = steps // save_every
    rem = steps % save_every

    def _segment(lat: Lattice, _):
        lat_next = lax.fori_loop(0, save_every, _step, lat)
        return lat_next, lat

    snapshots = []
    times = []
    current_lattice = lattice

    if n_full > 0:
        current_lattice, stacked = lax.scan(_segment, lattice, None, length=n_full)
        snapshots = [jax.tree_util.tree_map(lambda x, i=i: x[i], stacked) for i in range(n_full)]
        times = [i * save_every * dt for i in range(n_full)]

    if rem:
        # Trailing partial segment: emit, then advance the remaining steps.
        snapshots.append(current_lattice)
        times.append(n_full * save_every * dt)
        current_lattice = lax.fori_loop(0, rem, _step, current_lattice)

    # Save final state
    if (steps - 1) % save_every != 0:
        snapshots.append(current_lattice)
        times.append(steps * dt)

    return jnp.array(times), snapshots

evolve_step(lattice, potential, dt, H=0.0, use_spectral=False)

Perform a single leapfrog time step for a scalar field.

Uses the leapfrog (Verlet) scheme with optional Hubble friction:

1. \(\dot{\phi}_{n+1/2} = \dot{\phi}_n + (dt/2)\,F(\phi_n)\)
2. \(\phi_{n+1} = \phi_n + dt\,\dot{\phi}_{n+1/2}\)
3. \(\dot{\phi}_{n+1} = \dot{\phi}_{n+1/2} + (dt/2)\,F(\phi_{n+1})\)

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

For cosmological evolution with Hubble friction (\(H > 0\)), the equation becomes:

\[\ddot{\phi} + 2H\dot{\phi} = \nabla^2\phi - dV/d\phi\]

The friction term is treated semi-implicitly for numerical stability, applied once at the full-step velocity update:

• \(\dot{\phi}_{n+1/2} = \dot{\phi}_n + (dt/2)\,F_n\) (no friction at half-step)
• \(\dot{\phi}_{n+1} = [\dot{\phi}_{n+1/2} + (dt/2)\,F_{n+1}]\,/\,(1 + 2H\,dt)\)

Parameters:

Name	Type	Description	Default
lattice	Lattice	Current lattice state	required
potential	Callable[[Array], Array]	Potential function \(V(\phi)\)	required
dt	float	Time step	required
H	float	Hubble parameter (conformal). If H=0 (default), no friction is applied.	0.0
use_spectral	bool	If True, use FFT-based spectral Laplacian for exact \(k^2\) dispersion. Recommended for validation and periodic domains. Default: False (finite-difference).	False

Returns:

Type	Description
Lattice	New lattice state at time t + dt

Example

# Flat spacetime (no friction)
lattice_new = evolve_step(lattice, potential, dt=0.01)

# With Hubble friction
lattice_new = evolve_step(lattice, potential, dt=0.01, H=0.1)

# With spectral Laplacian for validation
lattice_new = evolve_step(lattice, potential, dt=0.01, H=0.1, use_spectral=True)

Source code in jaxlatt/evolution/scalar.py

def evolve_step(
    lattice: Lattice,
    potential: Callable[[Array], Array],
    dt: float,
    H: float = 0.0,
    use_spectral: bool = False,
) -> Lattice:
    r"""
    Perform a single leapfrog time step for a scalar field.

    Uses the leapfrog (Verlet) scheme with optional Hubble friction:

    1. $\dot{\phi}_{n+1/2} = \dot{\phi}_n + (dt/2)\,F(\phi_n)$
    2. $\phi_{n+1} = \phi_n + dt\,\dot{\phi}_{n+1/2}$
    3. $\dot{\phi}_{n+1} = \dot{\phi}_{n+1/2} + (dt/2)\,F(\phi_{n+1})$

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

    For cosmological evolution with Hubble friction ($H > 0$), the equation becomes:

    $$\ddot{\phi} + 2H\dot{\phi} = \nabla^2\phi - dV/d\phi$$

    The friction term is treated semi-implicitly for numerical stability, applied
    once at the full-step velocity update:

    - $\dot{\phi}_{n+1/2} = \dot{\phi}_n + (dt/2)\,F_n$  (no friction at half-step)
    - $\dot{\phi}_{n+1} = [\dot{\phi}_{n+1/2} + (dt/2)\,F_{n+1}]\,/\,(1 + 2H\,dt)$

    Args:
        lattice: Current lattice state
        potential: Potential function $V(\phi)$
        dt: Time step
        H: Hubble parameter (conformal). If H=0 (default), no friction is applied.
        use_spectral: If True, use FFT-based spectral Laplacian for exact $k^2$ dispersion.
            Recommended for validation and periodic domains. Default: False (finite-difference).

    Returns:
        New lattice state at time t + dt

    Example:
        ```python
        # Flat spacetime (no friction)
        lattice_new = evolve_step(lattice, potential, dt=0.01)

        # With Hubble friction
        lattice_new = evolve_step(lattice, potential, dt=0.01, H=0.1)

        # With spectral Laplacian for validation
        lattice_new = evolve_step(lattice, potential, dt=0.01, H=0.1, use_spectral=True)
        ```
    """
    # Current state
    field = lattice.field
    field_dot = lattice.field_dot

    # Compute force at current position
    force = compute_force(field, potential, lattice.dx, use_spectral=use_spectral)

    # Half-step velocity update (no friction)
    field_dot_half = field_dot + 0.5 * dt * force

    # Full-step position update
    field_new = field + dt * field_dot_half

    # Compute force at new position
    force_new = compute_force(field_new, potential, lattice.dx, use_spectral=use_spectral)

    # Full-step velocity update with semi-implicit Hubble friction (applied once)
    field_dot_new = field_dot_half + 0.5 * dt * force_new
    # jnp.where keeps H traceable (grad w.r.t. H, vmap over H). Mirrors the
    # idiom already used by FRWUniverse.H in core/cosmology/frw.py.
    field_dot_new = jnp.where(H > 0.0, field_dot_new / (1.0 + 2.0 * H * dt), field_dot_new)

    # Return new lattice state
    return Lattice(
        size=lattice.size,
        length=lattice.length,
        field=field_new,
        field_dot=field_dot_new,
    )