Energy Functionals

functionals

Traceable (Array-returning) observables.

This module holds the implementations of the energy and stability diagnostics. jaxlatt.observables.energy and jaxlatt.observables.stability are the eager layer on top: their public functions are one-line float(...) wrappers around their namesakes here, so the physics is written exactly once.

The rule is simple: same name as in energy.py / stability.py, returns an Array. Use this module whenever an observable needs to appear inside jit, grad, vmap or scan; use the eager modules for notebooks and printing.

Only the functions that needed a traceable twin live here. Much of stability.py is already Array-returning and @jit-decorated (compute_mass_matrix, compute_curvature_at_origin, compute_linear_response, compute_fluctuation_spectrum_1d, mode_visibility); those stay where they are rather than being moved for symmetry's sake.

Do not jit these directly

Several functions take a potential_func: Callable, which is not a valid JAX type. Decorating them with @jit would fail at trace time. Wrap your outer loss in jax.jit instead — the same guidance the leapfrog kernels need for a different reason (donated buffers):

@jax.jit
def loss(phi):
    lat = dataclasses.replace(lattice, phi=phi)
    return functionals.compute_physical_energy_density(lat, a, dx, V)

grad_phi = jax.grad(loss)(lattice.phi)

Relation to operators.autodiff

jaxlatt.operators.autodiff has its own traceable energy functionals (scalar_hamiltonian, coupled_hamiltonian, compute_all_energy_components, ...). Those are a different convention and are not duplicates of these: they are flat-space, volume-integrated, use a finite-difference gauge-covariant gradient, and hardcode the quartic potential. The functions here are FRW-aware (explicit scale-factor scaling), volume-averaged, use an FFT gradient, and take an arbitrary potential callable. Pick by which convention you need.

EnergyHistory

Bases: Module

Fixed-size, traceable energy history — a valid pytree and scan carry.

The counterpart to EnergyTracker, which accumulates into Python lists and is therefore host-side by construction. This class preallocates its buffers and record returns a new instance, so it can be carried through lax.scan and differentiated through.

Attributes:

Name	Type	Description
times	Array	Recorded times, shape (capacity,).
components	dict[str, Array]	Per-component buffers, each of shape (capacity,).
index	Array	Number of slots written so far.

Example

history = EnergyHistory.empty(n_steps, ("kinetic", "gradient", "potential", "total"))

def step(carry, _):
    lattice, history = carry
    lattice = evolve_one_step(lattice)
    history = history.record(lattice.tau, energy_components_averaged(lattice, V))
    return (lattice, history), None

(lattice, history), _ = jax.lax.scan(step, (lattice, history), length=n_steps)
drift = history.conservation()["max_deviation"]

empty(capacity, keys=('kinetic', 'gradient', 'potential', 'total'), dtype=None) classmethod

Allocate an empty history.

Parameters:

Name	Type	Description	Default
capacity	int	Number of measurements to make room for. Must be a concrete Python int, since it sets the buffer shapes.	required
keys	tuple[str, ...]	Component names to track. Use the six physical-energy keys when recording compute_physical_energy_components output.	('kinetic', 'gradient', 'potential', 'total')
dtype	dtype | None	Buffer dtype; defaults to the current default float type, which honours jax_enable_x64.	None

Returns:

Type	Description
EnergyHistory	A history with zeroed buffers and index = 0.

Source code in jaxlatt/observables/functionals.py

@classmethod
def empty(
    cls,
    capacity: int,
    keys: tuple[str, ...] = ("kinetic", "gradient", "potential", "total"),
    dtype: jnp.dtype | None = None,
) -> "EnergyHistory":
    """Allocate an empty history.

    Args:
        capacity: Number of measurements to make room for. Must be a concrete
            Python int, since it sets the buffer shapes.
        keys: Component names to track. Use the six physical-energy keys when
            recording `compute_physical_energy_components` output.
        dtype: Buffer dtype; defaults to the current default float type, which
            honours `jax_enable_x64`.

    Returns:
        A history with zeroed buffers and `index = 0`.
    """
    dtype = jnp.result_type(float) if dtype is None else dtype
    zeros = jnp.zeros((capacity,), dtype=dtype)
    # The same zero buffer is shared across keys; arrays are immutable and
    # `record` writes through `.at[].set()`, which returns a new one.
    return cls(
        times=zeros,
        components=dict.fromkeys(keys, zeros),
        index=jnp.asarray(0),
    )

record(time, components)

Write one measurement and return the updated history.

Out-of-bounds writes are dropped by JAX's scatter semantics, so recording more than capacity measurements silently keeps the first capacity.

Parameters:

Name	Type	Description	Default
time	Array	Current simulation time.	required
components	dict[str, Array]	Energy components, with the same keys the history was allocated with.	required

Returns:

Type	Description
EnergyHistory	A new EnergyHistory; the receiver is unchanged.

Source code in jaxlatt/observables/functionals.py

def record(self, time: Array, components: dict[str, Array]) -> "EnergyHistory":
    """Write one measurement and return the updated history.

    Out-of-bounds writes are dropped by JAX's scatter semantics, so recording
    more than `capacity` measurements silently keeps the first `capacity`.

    Args:
        time: Current simulation time.
        components: Energy components, with the same keys the history was
            allocated with.

    Returns:
        A new `EnergyHistory`; the receiver is unchanged.
    """
    i = self.index
    return EnergyHistory(
        times=self.times.at[i].set(time),
        components={k: v.at[i].set(components[k]) for k, v in self.components.items()},
        index=i + 1,
    )

conservation()

Conservation metrics over the recorded 'total' series.

Returns:

Type	Description
dict[str, Array]	The dictionary described in
dict[str, Array]	energy_conservation.

Source code in jaxlatt/observables/functionals.py

def conservation(self) -> dict[str, Array]:
    """Conservation metrics over the recorded `'total'` series.

    Returns:
        The dictionary described in
        [`energy_conservation`][jaxlatt.observables.functionals.energy_conservation].
    """
    return energy_conservation(self.components["total"])

kinetic_energy_averaged(lattice)

Volume-averaged kinetic energy density, as an Array.

Parameters:

Name	Type	Description	Default
lattice	Lattice	Lattice object with field_dot.	required

Returns:

Type	Description
Array	\(\langle \rho_\mathrm{kin} \rangle = \frac{1}{2}\langle \dot{\phi}^2 \rangle\)

Source code in jaxlatt/observables/functionals.py

def kinetic_energy_averaged(lattice: Lattice) -> Array:
    r"""Volume-averaged kinetic energy density, as an `Array`.

    Args:
        lattice: Lattice object with `field_dot`.

    Returns:
        $\langle \rho_\mathrm{kin} \rangle = \frac{1}{2}\langle \dot{\phi}^2 \rangle$
    """
    return _kinetic_energy_averaged_impl(lattice.field_dot)

gradient_energy_averaged(lattice)

Volume-averaged gradient energy density, as an Array.

Uses FFT-based gradient computation for spectral accuracy.

Parameters:

Name	Type	Description	Default
lattice	Lattice	Lattice object with field.	required

Returns:

Type	Description
Array	\(\langle \rho_\mathrm{grad} \rangle = \frac{1}{2}\langle (\nabla\phi)^2 \rangle\)

Raises:

Type	Description
ValueError	If the field is not 1D, 2D or 3D. The check is on ndim, which is static under tracing.

Source code in jaxlatt/observables/functionals.py

def gradient_energy_averaged(lattice: Lattice) -> Array:
    r"""Volume-averaged gradient energy density, as an `Array`.

    Uses FFT-based gradient computation for spectral accuracy.

    Args:
        lattice: Lattice object with `field`.

    Returns:
        $\langle \rho_\mathrm{grad} \rangle = \frac{1}{2}\langle (\nabla\phi)^2 \rangle$

    Raises:
        ValueError: If the field is not 1D, 2D or 3D. The check is on `ndim`,
            which is static under tracing.
    """
    if lattice.field.ndim not in (1, 2, 3):
        raise ValueError(f"Unsupported dimensionality: {lattice.field.ndim}")
    dx_per_dim = tuple(L / N for L, N in zip(lattice.length, lattice.size))
    return _gradient_energy_averaged_impl(lattice.field, dx_per_dim)

potential_energy_averaged(lattice, potential_func)

Volume-averaged potential energy density, as an Array.

Parameters:

Name	Type	Description	Default
lattice	Lattice	Lattice object with field.	required
potential_func	Callable	Potential function \(V(\phi) \to V\).	required

Returns:

Type	Description
Array	\(\langle \rho_\mathrm{pot} \rangle = \langle V(\phi) \rangle\)

Source code in jaxlatt/observables/functionals.py

def potential_energy_averaged(lattice: Lattice, potential_func: Callable) -> Array:
    r"""Volume-averaged potential energy density, as an `Array`.

    Args:
        lattice: Lattice object with `field`.
        potential_func: Potential function $V(\phi) \to V$.

    Returns:
        $\langle \rho_\mathrm{pot} \rangle = \langle V(\phi) \rangle$
    """
    V_field = potential_func(lattice.field)
    return _potential_energy_averaged_impl(V_field)

total_energy_averaged(lattice, potential_func)

Total volume-averaged energy density, as an Array.

Parameters:

Name	Type	Description	Default
lattice	Lattice	Lattice object.	required
potential_func	Callable	Potential function \(V(\phi) \to V\).	required

Returns:

Type	Description
Array	\(E_\mathrm{kin} + E_\mathrm{grad} + E_\mathrm{pot}\)

Source code in jaxlatt/observables/functionals.py

def total_energy_averaged(lattice: Lattice, potential_func: Callable) -> Array:
    r"""Total volume-averaged energy density, as an `Array`.

    Args:
        lattice: Lattice object.
        potential_func: Potential function $V(\phi) \to V$.

    Returns:
        $E_\mathrm{kin} + E_\mathrm{grad} + E_\mathrm{pot}$
    """
    return (
        kinetic_energy_averaged(lattice)
        + gradient_energy_averaged(lattice)
        + potential_energy_averaged(lattice, potential_func)
    )

energy_components_averaged(lattice, potential_func)

All volume-averaged energy components, as Arrays.

Parameters:

Name	Type	Description	Default
lattice	Lattice	Lattice object.	required
potential_func	Callable	Potential function \(V(\phi) \to V\).	required

Returns:

Type	Description
dict[str, Array]	Dictionary with keys 'kinetic', 'gradient', 'potential', 'total'.

Source code in jaxlatt/observables/functionals.py

def energy_components_averaged(lattice: Lattice, potential_func: Callable) -> dict[str, Array]:
    r"""All volume-averaged energy components, as `Array`s.

    Args:
        lattice: Lattice object.
        potential_func: Potential function $V(\phi) \to V$.

    Returns:
        Dictionary with keys `'kinetic'`, `'gradient'`, `'potential'`, `'total'`.
    """
    E_kin = kinetic_energy_averaged(lattice)
    E_grad = gradient_energy_averaged(lattice)
    E_pot = potential_energy_averaged(lattice, potential_func)

    return {
        "kinetic": E_kin,
        "gradient": E_grad,
        "potential": E_pot,
        "total": E_kin + E_grad + E_pot,
    }

kinetic_energy_integrated(lattice)

Kinetic energy integrated over the box, as an Array.

Parameters:

Name	Type	Description	Default
lattice	Lattice	Lattice state containing field_dot and geometry.	required

Returns:

Type	Description
Array	Total kinetic energy.

Source code in jaxlatt/observables/functionals.py

def kinetic_energy_integrated(lattice: Lattice) -> Array:
    """Kinetic energy integrated over the box, as an `Array`.

    Args:
        lattice: Lattice state containing `field_dot` and geometry.

    Returns:
        Total kinetic energy.
    """
    return kinetic_energy_averaged(lattice) * lattice.volume

potential_energy_integrated(lattice, potential)

Potential energy integrated over the box, as an Array.

Parameters:

Name	Type	Description	Default
lattice	Lattice	Lattice state containing scalar field values.	required
potential	Callable	Potential function V(phi).	required

Returns:

Type	Description
Array	Total potential energy.

Source code in jaxlatt/observables/functionals.py

def potential_energy_integrated(lattice: Lattice, potential: Callable) -> Array:
    """Potential energy integrated over the box, as an `Array`.

    Args:
        lattice: Lattice state containing scalar field values.
        potential: Potential function `V(phi)`.

    Returns:
        Total potential energy.
    """
    return potential_energy_averaged(lattice, potential) * lattice.volume

total_energy_integrated(lattice, potential)

Total energy integrated over the box, as an Array.

Parameters:

Name	Type	Description	Default
lattice	Lattice	Lattice state.	required
potential	Callable	Potential function V(phi).	required

Returns:

Type	Description
Array	Total energy.

Source code in jaxlatt/observables/functionals.py

def total_energy_integrated(lattice: Lattice, potential: Callable) -> Array:
    """Total energy integrated over the box, as an `Array`.

    Args:
        lattice: Lattice state.
        potential: Potential function `V(phi)`.

    Returns:
        Total energy.
    """
    return total_energy_averaged(lattice, potential) * lattice.volume

energy_components_integrated(lattice, potential)

All integrated energy components, as Arrays.

Parameters:

Name	Type	Description	Default
lattice	Lattice	Lattice state.	required
potential	Callable	Potential function V(phi).	required

Returns:

Type	Description
dict[str, Array]	Dictionary with integrated kinetic, gradient, potential, and total.

Source code in jaxlatt/observables/functionals.py

def energy_components_integrated(lattice: Lattice, potential: Callable) -> dict[str, Array]:
    """All integrated energy components, as `Array`s.

    Args:
        lattice: Lattice state.
        potential: Potential function `V(phi)`.

    Returns:
        Dictionary with integrated `kinetic`, `gradient`, `potential`, and `total`.
    """
    volume = lattice.volume
    return {k: v * volume for k, v in energy_components_averaged(lattice, potential).items()}

compute_physical_scalar_gradient_energy(field, scale_factor, dx)

Compute physical gradient energy density with scale factor scaling.

In expanding universe (conformal time):

\[\rho_\mathrm{grad}^\mathrm{phys} = \frac{(\nabla\phi)^2}{2a^2}\]

where \(\nabla\) is the comoving gradient operator.

Parameters:

Name	Type	Description	Default
field	Array	Scalar field configuration	required
scale_factor	float	FLRW scale factor \(a(\tau)\)	required
dx	float	Lattice spacing (comoving)	required

Returns:

Type	Description
Array	Physical gradient energy density (averaged over lattice)

Note

Uses FFT-based gradient computation for spectral accuracy. For expanding universe, this should be called with comoving coordinates.

Source code in jaxlatt/observables/functionals.py

@jit
def compute_physical_scalar_gradient_energy(
    field: Array,
    scale_factor: float,
    dx: float,
) -> Array:
    r"""
    Compute physical gradient energy density with scale factor scaling.

    In expanding universe (conformal time):

    $$\rho_\mathrm{grad}^\mathrm{phys} = \frac{(\nabla\phi)^2}{2a^2}$$

    where $\nabla$ is the comoving gradient operator.

    Args:
        field: Scalar field configuration
        scale_factor: FLRW scale factor $a(\tau)$
        dx: Lattice spacing (comoving)

    Returns:
        Physical gradient energy density (averaged over lattice)

    Note:
        Uses FFT-based gradient computation for spectral accuracy.
        For expanding universe, this should be called with comoving coordinates.
    """
    field_k = jnp.fft.fftn(field)
    grad_sq_k = _k_squared(field.shape, dx) * jnp.abs(field_k) ** 2
    return 0.5 * jnp.mean(grad_sq_k) / (field.size * scale_factor**2)

compute_physical_scalar_kinetic_energy(field_dot, scale_factor)

Compute physical kinetic energy density with scale factor scaling.

field_dot is the conformal time derivative \(\phi' = d\phi/d\tau\), which is what the integrators store as \(\pi\) (see docs/api/units.md). Since \(\dot\phi = \phi'/a\):

\[\rho_\mathrm{kin}^\mathrm{phys} = \frac{(\phi')^2}{2a^2} = \frac{\dot{\phi}^2}{2}\]

Parameters:

Name	Type	Description	Default
field_dot	Array	Time derivative of field (conformal time)	required
scale_factor	float	FLRW scale factor \(a(\tau)\)	required

Returns:

Type	Description
Array	Physical kinetic energy density (averaged)

Note

Convention depends on how field_dot is stored. Here we assume it's the conformal time derivative, giving \(a^{-2}\) scaling.

Source code in jaxlatt/observables/functionals.py

@jit
def compute_physical_scalar_kinetic_energy(
    field_dot: Array,
    scale_factor: float,
) -> Array:
    r"""
    Compute physical kinetic energy density with scale factor scaling.

    `field_dot` is the conformal time derivative $\phi' = d\phi/d\tau$, which is what
    the integrators store as $\pi$ (see `docs/api/units.md`). Since $\dot\phi = \phi'/a$:

    $$\rho_\mathrm{kin}^\mathrm{phys} = \frac{(\phi')^2}{2a^2} = \frac{\dot{\phi}^2}{2}$$

    Args:
        field_dot: Time derivative of field (conformal time)
        scale_factor: FLRW scale factor $a(\tau)$

    Returns:
        Physical kinetic energy density (averaged)

    Note:
        Convention depends on how field_dot is stored. Here we assume
        it's the conformal time derivative, giving $a^{-2}$ scaling.
    """
    a_squared = scale_factor * scale_factor
    kinetic_comoving = 0.5 * jnp.mean(field_dot**2)

    return kinetic_comoving / a_squared

compute_physical_potential_energy(field, potential_func)

Compute physical potential energy density.

Potential energy has no scale factor dependence:

\[\rho_\mathrm{pot}^\mathrm{phys} = V(\phi)\]

Parameters:

Name	Type	Description	Default
field	Array	Scalar field configuration	required
potential_func	Callable[[Array], ndarray]	Potential function \(V(\phi)\)	required

Returns:

Type	Description
Array	Physical potential energy density (averaged)

Note

Deliberately not @jit-decorated: potential_func is not a valid JAX type, so decorating this would fail at trace time. Wrap your outer loss in jax.jit instead.

Source code in jaxlatt/observables/functionals.py

def compute_physical_potential_energy(
    field: Array,
    potential_func: Callable[[Array], jnp.ndarray],
) -> Array:
    r"""
    Compute physical potential energy density.

    Potential energy has no scale factor dependence:

    $$\rho_\mathrm{pot}^\mathrm{phys} = V(\phi)$$

    Args:
        field: Scalar field configuration
        potential_func: Potential function $V(\phi)$

    Returns:
        Physical potential energy density (averaged)

    Note:
        Deliberately **not** `@jit`-decorated: `potential_func` is not a valid JAX
        type, so decorating this would fail at trace time. Wrap your outer loss in
        `jax.jit` instead.
    """
    V_field = potential_func(field)
    return jnp.mean(V_field)

compute_physical_gauge_electric_energy(E_field, scale_factor)

Compute physical electric field energy density.

Electric field energy in expanding universe:

\[\rho_E^\mathrm{phys} = \frac{E^2}{2a^4}\]

where \(E\) is the comoving electric field.

Parameters:

Name	Type	Description	Default
E_field	Array	Electric field components (3, N, N, N)	required
scale_factor	float	FLRW scale factor \(a(\tau)\)	required

Returns:

Type	Description
Array	Physical electric energy density (averaged)

Source code in jaxlatt/observables/functionals.py

@jit
def compute_physical_gauge_electric_energy(
    E_field: Array,
    scale_factor: float,
) -> Array:
    r"""
    Compute physical electric field energy density.

    Electric field energy in expanding universe:

    $$\rho_E^\mathrm{phys} = \frac{E^2}{2a^4}$$

    where $E$ is the comoving electric field.

    Args:
        E_field: Electric field components (3, N, N, N)
        scale_factor: FLRW scale factor $a(\tau)$

    Returns:
        Physical electric energy density (averaged)
    """
    a_4 = scale_factor**4
    E_squared = jnp.sum(E_field**2, axis=0)  # Sum over field components
    electric_comoving = 0.5 * jnp.mean(E_squared)

    return electric_comoving / a_4

compute_physical_gauge_magnetic_energy(links, scale_factor, dx, g)

Compute physical magnetic field energy density.

Magnetic field energy in expanding universe:

\[\rho_B^\mathrm{phys} = \frac{B^2}{2a^4}\]

where \(B\) is extracted from plaquettes.

Parameters:

Name	Type	Description	Default
links	Array	Gauge link variables (3, N, N, N)	required
scale_factor	float	FLRW scale factor \(a(\tau)\)	required
dx	float	Lattice spacing (comoving)	required
g	float	Gauge coupling	required

Returns:

Type	Description
Array	Physical magnetic energy density (averaged)

Note

Imports magnetic_field from operators.gauge to avoid circular dependency.

Source code in jaxlatt/observables/functionals.py

@jit
def compute_physical_gauge_magnetic_energy(
    links: Array,
    scale_factor: float,
    dx: float,
    g: float,
) -> Array:
    r"""
    Compute physical magnetic field energy density.

    Magnetic field energy in expanding universe:

    $$\rho_B^\mathrm{phys} = \frac{B^2}{2a^4}$$

    where $B$ is extracted from plaquettes.

    Args:
        links: Gauge link variables (3, N, N, N)
        scale_factor: FLRW scale factor $a(\tau)$
        dx: Lattice spacing (comoving)
        g: Gauge coupling

    Returns:
        Physical magnetic energy density (averaged)

    Note:
        Imports `magnetic_field` from `operators.gauge` to avoid circular dependency.
    """
    from jaxlatt.operators.gauge import magnetic_field

    B = magnetic_field(links, dx, g)
    a_4 = scale_factor**4
    B_squared = jnp.sum(B**2, axis=0)
    magnetic_comoving = 0.5 * jnp.mean(B_squared)

    return magnetic_comoving / a_4

compute_physical_energy_density(lattice, scale_factor, dx, potential_func, g=None)

Total physical energy density for the Friedmann equation, as an Array.

This is the source term \(\rho_\mathrm{total}\) in:

\[H^2 = \frac{8\pi G}{3} \rho_\mathrm{total}\]

Parameters:

Name	Type	Description	Default
lattice	CoupledLattice	CoupledLattice with scalar and gauge fields	required
scale_factor	float	Current scale factor \(a(\tau)\)	required
dx	float	Lattice spacing (comoving)	required
potential_func	Callable[[Array], ndarray]	Scalar potential \(V(\phi)\)	required
g	float | None	Optional gauge coupling override. Defaults to lattice.g, which is a static field and so carries no gradient; pass it explicitly to differentiate with respect to the coupling.	None

Returns:

Type	Description
Array	Total physical energy density \(\rho_\mathrm{phys}\) (averaged over lattice)

Source code in jaxlatt/observables/functionals.py

def compute_physical_energy_density(
    lattice: "CoupledLattice",
    scale_factor: float,
    dx: float,
    potential_func: Callable[[Array], jnp.ndarray],
    g: float | None = None,
) -> Array:
    r"""
    Total physical energy density for the Friedmann equation, as an `Array`.

    This is the source term $\rho_\mathrm{total}$ in:

    $$H^2 = \frac{8\pi G}{3} \rho_\mathrm{total}$$

    Args:
        lattice: CoupledLattice with scalar and gauge fields
        scale_factor: Current scale factor $a(\tau)$
        dx: Lattice spacing (comoving)
        potential_func: Scalar potential $V(\phi)$
        g: Optional gauge coupling override. Defaults to `lattice.g`, which is a
            static field and so carries no gradient; pass it explicitly to
            differentiate with respect to the coupling.

    Returns:
        Total physical energy density $\rho_\mathrm{phys}$ (averaged over lattice)
    """
    rho_total = sum(_physical_energy_components(lattice, scale_factor, dx, potential_func, g))
    return _real(rho_total)

compute_physical_energy_components(lattice, scale_factor, dx, potential_func, g=None)

All physical energy density components separately, as Arrays.

Parameters:

Name	Type	Description	Default
lattice	CoupledLattice	CoupledLattice state	required
scale_factor	float	Current scale factor \(a(\tau)\)	required
dx	float	Lattice spacing (comoving)	required
potential_func	Callable[[Array], ndarray]	Scalar potential \(V(\phi)\)	required
g	float | None	Optional gauge coupling override; see compute_physical_energy_density.	None

Returns:

Type	Description
dict[str, Array]	Dictionary with keys 'kinetic', 'gradient', 'potential', 'electric',
dict[str, Array]	'magnetic' and 'total'.

Source code in jaxlatt/observables/functionals.py

def compute_physical_energy_components(
    lattice: "CoupledLattice",
    scale_factor: float,
    dx: float,
    potential_func: Callable[[Array], jnp.ndarray],
    g: float | None = None,
) -> dict[str, Array]:
    r"""
    All physical energy density components separately, as `Array`s.

    Args:
        lattice: CoupledLattice state
        scale_factor: Current scale factor $a(\tau)$
        dx: Lattice spacing (comoving)
        potential_func: Scalar potential $V(\phi)$
        g: Optional gauge coupling override; see
            [`compute_physical_energy_density`][jaxlatt.observables.functionals.compute_physical_energy_density].

    Returns:
        Dictionary with keys `'kinetic'`, `'gradient'`, `'potential'`, `'electric'`,
        `'magnetic'` and `'total'`.
    """
    components = _physical_energy_components(lattice, scale_factor, dx, potential_func, g)
    rho_kinetic, rho_gradient, rho_potential, rho_electric, rho_magnetic = components

    return {
        "kinetic": _real(rho_kinetic),
        "gradient": _real(rho_gradient),
        "potential": _real(rho_potential),
        "electric": _real(rho_electric),
        "magnetic": _real(rho_magnetic),
        "total": _real(sum(components)),
    }

compute_comoving_energy_density(lattice, dx, potential_func, g=None)

Comoving energy density (no scale factor corrections), as an Array.

This is what standard (non-expanding) simulations compute.

Parameters:

Name	Type	Description	Default
lattice	CoupledLattice	CoupledLattice state	required
dx	float	Lattice spacing	required
potential_func	Callable[[Array], ndarray]	Scalar potential \(V(\phi)\)	required
g	float | None	Optional gauge coupling override.	None

Returns:

Type	Description
Array	Comoving energy density (what you'd get with \(a=1\))

Source code in jaxlatt/observables/functionals.py

def compute_comoving_energy_density(
    lattice: "CoupledLattice",
    dx: float,
    potential_func: Callable[[Array], jnp.ndarray],
    g: float | None = None,
) -> Array:
    r"""
    Comoving energy density (no scale factor corrections), as an `Array`.

    This is what standard (non-expanding) simulations compute.

    Args:
        lattice: CoupledLattice state
        dx: Lattice spacing
        potential_func: Scalar potential $V(\phi)$
        g: Optional gauge coupling override.

    Returns:
        Comoving energy density (what you'd get with $a=1$)
    """
    return compute_physical_energy_density(
        lattice, scale_factor=1.0, dx=dx, potential_func=potential_func, g=g
    )

kinetic_energy(field_dot, dV)

Total kinetic energy \(E_\mathrm{kin} = \frac{1}{2} \int (\partial\phi/\partial t)^2\,dV\).

Parameters:

Name	Type	Description	Default
field_dot	Array	Time derivative of the field	required
dV	float	Volume element (grid cell volume)	required

Returns:

Type	Description
Array	Total kinetic energy (integrated, not averaged)

Source code in jaxlatt/observables/functionals.py

@jit
def kinetic_energy(field_dot: Array, dV: float) -> Array:
    r"""
    Total kinetic energy $E_\mathrm{kin} = \frac{1}{2} \int (\partial\phi/\partial t)^2\,dV$.

    Args:
        field_dot: Time derivative of the field
        dV: Volume element (grid cell volume)

    Returns:
        Total kinetic energy (integrated, not averaged)
    """
    return 0.5 * jnp.sum(field_dot**2) * dV

potential_energy(field, potential, dV)

Total potential energy \(E_\mathrm{pot} = \int V(\phi)\,dV\).

Parameters:

Name	Type	Description	Default
field	Array	Field configuration	required
potential	Callable[[Array], ndarray]	Potential function \(V(\phi)\)	required
dV	float	Volume element	required

Returns:

Type	Description
Array	Total potential energy (integrated, not averaged)

Source code in jaxlatt/observables/functionals.py

def potential_energy(field: Array, potential: Callable[[Array], jnp.ndarray], dV: float) -> Array:
    r"""
    Total potential energy $E_\mathrm{pot} = \int V(\phi)\,dV$.

    Args:
        field: Field configuration
        potential: Potential function $V(\phi)$
        dV: Volume element

    Returns:
        Total potential energy (integrated, not averaged)
    """
    return jnp.sum(potential(field)) * dV

compute_energy(lattice, potential)

Total energy of a field configuration (volume-integrated), as an Array.

Parameters:

Name	Type	Description	Default
lattice	Lattice	Lattice object containing field and field_dot	required
potential	Callable[[Array], ndarray]	Potential function \(V(\phi)\)	required

Returns:

Type	Description
Array	Total energy (integrated)

Source code in jaxlatt/observables/functionals.py

def compute_energy(lattice: Lattice, potential: Callable[[Array], jnp.ndarray]) -> Array:
    r"""
    Total energy of a field configuration (volume-integrated), as an `Array`.

    Args:
        lattice: Lattice object containing field and field_dot
        potential: Potential function $V(\phi)$

    Returns:
        Total energy (integrated)
    """
    return total_energy_integrated(lattice, potential)

compute_energy_components(lattice, potential)

All energy components separately (volume-integrated), as Arrays.

Parameters:

Name	Type	Description	Default
lattice	Lattice	Lattice object containing field and field_dot	required
potential	Callable[[Array], ndarray]	Potential function \(V(\phi)\)	required

Returns:

Type	Description
dict[str, Array]	Dictionary with keys 'kinetic', 'gradient', 'potential', 'total'.

Source code in jaxlatt/observables/functionals.py

def compute_energy_components(
    lattice: Lattice, potential: Callable[[Array], jnp.ndarray]
) -> dict[str, Array]:
    r"""
    All energy components separately (volume-integrated), as `Array`s.

    Args:
        lattice: Lattice object containing field and field_dot
        potential: Potential function $V(\phi)$

    Returns:
        Dictionary with keys `'kinetic'`, `'gradient'`, `'potential'`, `'total'`.
    """
    return energy_components_integrated(lattice, potential)

energy_density(field, field_dot, potential, dx)

Local energy density at each lattice point.

Parameters:

Name	Type	Description	Default
field	Array	Field configuration	required
field_dot	Array	Time derivative of field	required
potential	Callable[[Array], ndarray]	Potential function	required
dx	float	Lattice spacing	required

Returns:

Type	Description
Array	Energy density at each lattice point

Raises:

Type	Description
ValueError	If the field is not 1D, 2D or 3D.

Source code in jaxlatt/observables/functionals.py

def energy_density(
    field: Array,
    field_dot: Array,
    potential: Callable[[Array], jnp.ndarray],
    dx: float,
) -> Array:
    """
    Local energy density at each lattice point.

    Args:
        field: Field configuration
        field_dot: Time derivative of field
        potential: Potential function
        dx: Lattice spacing

    Returns:
        Energy density at each lattice point

    Raises:
        ValueError: If the field is not 1D, 2D or 3D.
    """
    rho_kin = 0.5 * field_dot**2
    rho_pot = potential(field)

    if field.ndim == 1:
        grad = gradient_1d(field, dx)
        rho_grad = 0.5 * grad**2
    elif field.ndim == 2:
        grad_x, grad_y = gradient_2d(field, dx)
        rho_grad = 0.5 * (grad_x**2 + grad_y**2)
    elif field.ndim == 3:
        gx, gy, gz = gradient_3d(field, dx)
        rho_grad = 0.5 * (gx**2 + gy**2 + gz**2)
    else:
        raise ValueError("energy_density only supports 1D, 2D, 3D")

    return rho_kin + rho_grad + rho_pot

energy_conservation(E_tot)

Conservation metrics for a stacked total-energy series.

This is the traceable form of EnergyTracker.energy_conservation, and the one an energy-drift objective wants: it takes the stacked series that the evolvers already return, so nothing has to round-trip through a Python list.

Parameters:

Name	Type	Description	Default
E_tot	Array	1D array of total energy at each recorded step.	required

Returns:

Type	Description
dict[str, Array]	Dictionary with 'initial', 'final', 'mean', 'std',
dict[str, Array]	'relative_change' and 'max_deviation'.

Example

drift = functionals.energy_conservation(E_series)["max_deviation"]
grad_dt = jax.grad(lambda dt: run_and_measure(dt))(dt0)

Source code in jaxlatt/observables/functionals.py

@jit
def energy_conservation(E_tot: Array) -> dict[str, Array]:
    """Conservation metrics for a stacked total-energy series.

    This is the traceable form of
    [`EnergyTracker.energy_conservation`][jaxlatt.observables.energy.EnergyTracker.energy_conservation],
    and the one an energy-drift objective wants: it takes the stacked series that
    the evolvers already return, so nothing has to round-trip through a Python list.

    Args:
        E_tot: 1D array of total energy at each recorded step.

    Returns:
        Dictionary with `'initial'`, `'final'`, `'mean'`, `'std'`,
        `'relative_change'` and `'max_deviation'`.

    Example:
        ```python
        drift = functionals.energy_conservation(E_series)["max_deviation"]
        grad_dt = jax.grad(lambda dt: run_and_measure(dt))(dt0)
        ```
    """
    E_initial = E_tot[0]
    E_final = E_tot[-1]
    E_mean = jnp.mean(E_tot)

    return {
        "initial": E_initial,
        "final": E_final,
        "mean": E_mean,
        "std": jnp.std(E_tot),
        "relative_change": (E_final - E_initial) / E_initial,
        "max_deviation": jnp.max(jnp.abs(E_tot - E_mean)) / E_mean,
    }

compute_mass_matrix(phi, m, lambda_)

Effective mass matrix \(\partial^2 V/\partial\phi\,\partial\phi^*\).

Convention: \(m_\text{eff}^2 = \text{sign}(m)\,m^2 + \lambda|\phi|^2\).

• \(m > 0\): quartic potential \(V = m^2/2\,|\phi|^2 + \lambda/4\,|\phi|^4\). Curvature at origin is \(+m^2\) (stable vacuum at \(\phi=0\)).
• \(m < 0\): double-well potential with tachyonic mass. Curvature at origin is \(-m^2\) (unstable vacuum, SSB).

Parameters:

Name	Type	Description	Default
phi	Array	Scalar field (N, N, N) complex.	required
m	float	Signed mass parameter. Positive → stable; negative → tachyonic at origin.	required
lambda_	float	Self-coupling.	required

Returns:

Type	Description
Array	Effective mass squared at each point (N, N, N) real.

Source code in jaxlatt/observables/functionals.py

@jit
def compute_mass_matrix(phi: Array, m: float, lambda_: float) -> Array:
    r"""Effective mass matrix $\partial^2 V/\partial\phi\,\partial\phi^*$.

    Convention: $m_\text{eff}^2 = \text{sign}(m)\,m^2 + \lambda|\phi|^2$.

    - $m > 0$: quartic potential $V = m^2/2\,|\phi|^2 + \lambda/4\,|\phi|^4$.
      Curvature at origin is $+m^2$ (stable vacuum at $\phi=0$).
    - $m < 0$: double-well potential with tachyonic mass.
      Curvature at origin is $-m^2$ (unstable vacuum, SSB).

    Args:
        phi: Scalar field (N, N, N) complex.
        m: Signed mass parameter. Positive → stable; negative → tachyonic at origin.
        lambda_: Self-coupling.

    Returns:
        Effective mass squared at each point (N, N, N) real.
    """
    phi_sq = jnp.abs(phi) ** 2
    # sign(m) * m² + λ|φ|²: when m < 0 this gives negative curvature at φ=0
    # (tachyonic / symmetry-breaking). m * abs(m) = sign(m) * m² for all real m.
    return m * jnp.abs(m) + lambda_ * phi_sq

compute_curvature_at_origin(m, lambda_)

Compute \(\partial^2 V/\partial\phi^2|_{\phi=0} = \text{sign}(m)\,m^2\).

• If negative: symmetry broken (\(\phi = 0\) unstable).
• If positive: symmetric phase (\(\phi = 0\) stable).

Parameters:

Name	Type	Description	Default
m	float	Mass parameter.	required
lambda_	float	Coupling (unused; kept for interface symmetry).	required

Returns:

Type	Description
Array	Curvature at the origin.

Source code in jaxlatt/observables/functionals.py

@jit
def compute_curvature_at_origin(m: float, lambda_: float) -> Array:
    r"""Compute $\partial^2 V/\partial\phi^2|_{\phi=0} = \text{sign}(m)\,m^2$.

    - If negative: symmetry broken ($\phi = 0$ unstable).
    - If positive: symmetric phase ($\phi = 0$ stable).

    Args:
        m: Mass parameter.
        lambda_: Coupling (unused; kept for interface symmetry).

    Returns:
        Curvature at the origin.
    """
    return m * jnp.abs(m)

detect_tachyonic_modes(phi, m, lambda_)

Detect tachyonic (unstable) modes where \(m_\mathrm{eff}^2 < 0\).

Parameters:

Name	Type	Description	Default
phi	Array	Scalar field configuration.	required
m	float	Mass parameter.	required
lambda_	float	Self-coupling.	required

Returns:

Type	Description
dict[str, Array]	Dictionary with counts and fractions of unstable sites, extrema of
dict[str, Array]	effective mass-squared, and a boolean instability flag, all as Array.

Source code in jaxlatt/observables/functionals.py

def detect_tachyonic_modes(phi: Array, m: float, lambda_: float) -> dict[str, Array]:
    """Detect tachyonic (unstable) modes where $m_\\mathrm{eff}^2 < 0$.

    Args:
        phi: Scalar field configuration.
        m: Mass parameter.
        lambda_: Self-coupling.

    Returns:
        Dictionary with counts and fractions of unstable sites, extrema of
        effective mass-squared, and a boolean instability flag, all as `Array`.
    """
    m_eff_sq = compute_mass_matrix(phi, m, lambda_)

    is_tachyonic = m_eff_sq < 0
    min_m_sq = jnp.min(m_eff_sq)

    return {
        "n_tachyonic_sites": jnp.sum(is_tachyonic),
        "fraction_tachyonic": jnp.mean(is_tachyonic),
        "min_mass_squared": min_m_sq,
        "max_mass_squared": jnp.max(m_eff_sq),
        "has_instability": min_m_sq < 0,
    }

compute_vev_estimate(m, lambda_)

Estimate the vacuum expectation value for symmetry breaking.

\(|\langle\phi\rangle| = \sqrt{m^2/\lambda}\) when \(m < 0\) and \(\lambda > 0\), and zero otherwise.

The guard matters for gradients as well as values: \(\lambda\) is sanitised before the division, so the unused branch never evaluates \(m^2/0\).

Parameters:

Name	Type	Description	Default
m	float	Mass parameter (negative for symmetry breaking).	required
lambda_	float	Self-coupling.	required

Returns:

Type	Description
Array	Estimated \(|\langle\phi\rangle|\).

Source code in jaxlatt/observables/functionals.py

def compute_vev_estimate(m: float, lambda_: float) -> Array:
    r"""Estimate the vacuum expectation value for symmetry breaking.

    $|\langle\phi\rangle| = \sqrt{m^2/\lambda}$ when $m < 0$ and $\lambda > 0$,
    and zero otherwise.

    The guard matters for gradients as well as values: $\lambda$ is sanitised
    before the division, so the unused branch never evaluates $m^2/0$.

    Args:
        m: Mass parameter (negative for symmetry breaking).
        lambda_: Self-coupling.

    Returns:
        Estimated $|\langle\phi\rangle|$.
    """
    broken = (m < 0) & (lambda_ > 0)
    # Two sanitisations, not one. Guarding only the division still leaves
    # sqrt(0) in the dead branch at m = 0, whose derivative is infinite, and the
    # `where` transposes into both branches — so d(vev)/dm came back NaN there.
    # Feeding the dead branch 1.0 keeps sqrt'(1) = 1/2 finite.
    lambda_safe = jnp.where(broken, lambda_, 1.0)
    ratio = jnp.where(broken, m**2 / lambda_safe, 1.0)
    return jnp.where(broken, jnp.sqrt(ratio), 0.0)

check_symmetry_breaking(phi, m, lambda_)

Check whether the field has a non-zero VEV (symmetry breaking).

Parameters:

Name	Type	Description	Default
phi	Array	Scalar field configuration.	required
m	float	Mass parameter.	required
lambda_	float	Self-coupling.	required

Returns:

Type	Description
dict[str, Array]	Dictionary with measured field statistics, expected VEV, curvature at the
dict[str, Array]	origin, and a symmetry-breaking indicator, all as Array.

Source code in jaxlatt/observables/functionals.py

def check_symmetry_breaking(phi: Array, m: float, lambda_: float) -> dict[str, Array]:
    """Check whether the field has a non-zero VEV (symmetry breaking).

    Args:
        phi: Scalar field configuration.
        m: Mass parameter.
        lambda_: Self-coupling.

    Returns:
        Dictionary with measured field statistics, expected VEV, curvature at the
        origin, and a symmetry-breaking indicator, all as `Array`.
    """
    phi_mean = jnp.mean(phi)
    phi_rms = jnp.sqrt(jnp.mean(jnp.abs(phi) ** 2))

    vev_expected = compute_vev_estimate(m, lambda_)
    curvature = compute_curvature_at_origin(m, lambda_)

    return {
        "mean_field_magnitude": jnp.abs(phi_mean),
        "rms_field": phi_rms,
        "expected_vev": vev_expected,
        "curvature_at_origin": curvature,
        "is_symmetry_broken": curvature < 0,
        "vev_ratio": phi_rms / (vev_expected + 1e-10),
    }

estimate_effective_mass_from_spectrum(k_values, power_spectrum, dx)

Estimate the effective mass from a fluctuation spectrum.

For a free theory \(P(k) \approx 1/(k^2 + m_\mathrm{eff}^2)\), so the low-\(k\) plateau fixes \(m_\mathrm{eff}\).

The eager version selected the low-\(k\) modes with boolean-mask indexing, whose output shape depends on the values rather than the shapes — untraceable. Since k_values is ascending, the first retained mode is simply element 0, so the mask is only needed for the "enough modes to fit" test, which becomes a jnp.where. The two agree exactly: if element 0 is not below the cutoff then no element is, and both return zero.

Parameters:

Name	Type	Description	Default
k_values	Array	Momentum values, ascending.	required
power_spectrum	Array	\(P(k)\).	required
dx	float	Lattice spacing.	required

Returns:

Type	Description
Array	Effective mass estimate.

Source code in jaxlatt/observables/functionals.py

def estimate_effective_mass_from_spectrum(
    k_values: Array,
    power_spectrum: Array,
    dx: float,
) -> Array:
    r"""Estimate the effective mass from a fluctuation spectrum.

    For a free theory $P(k) \approx 1/(k^2 + m_\mathrm{eff}^2)$, so the low-$k$
    plateau fixes $m_\mathrm{eff}$.

    The eager version selected the low-$k$ modes with boolean-mask indexing, whose
    output shape depends on the values rather than the shapes — untraceable. Since
    `k_values` is ascending, the first retained mode is simply element 0, so the
    mask is only needed for the "enough modes to fit" test, which becomes a
    `jnp.where`. The two agree exactly: if element 0 is not below the cutoff then
    no element is, and both return zero.

    Args:
        k_values: Momentum values, ascending.
        power_spectrum: $P(k)$.
        dx: Lattice spacing.

    Returns:
        Effective mass estimate.
    """
    k_max = jnp.pi / (4 * dx)
    enough_modes = jnp.sum(k_values < k_max) > 2

    m_eff_sq = 1.0 / (power_spectrum[0] + 1e-10)
    return jnp.where(enough_modes, jnp.sqrt(jnp.maximum(m_eff_sq, 0)), 0.0)

estimate_bounce_action(phi_false, phi_true, m, lambda_)

Estimate the bounce action for tunnelling \(\phi_\mathrm{false} \to \phi_\mathrm{true}\).

\[S_\mathrm{bounce} \sim (\Delta V)^{-2}\,(\text{barrier height})^4\]

with tunnelling rate \(\Gamma \sim \exp(-S_\mathrm{bounce})\). A crude estimate: the barrier is taken at \(\phi = 0\) when \(m < 0\), and absent otherwise.

Parameters:

Name	Type	Description	Default
phi_false	float	False vacuum field value.	required
phi_true	float	True vacuum field value.	required
m	float	Mass.	required
lambda_	float	Coupling.	required

Returns:

Type	Description
Array	Bounce action estimate.

Source code in jaxlatt/observables/functionals.py

def estimate_bounce_action(
    phi_false: float,
    phi_true: float,
    m: float,
    lambda_: float,
) -> Array:
    r"""Estimate the bounce action for tunnelling $\phi_\mathrm{false} \to \phi_\mathrm{true}$.

    $$S_\mathrm{bounce} \sim (\Delta V)^{-2}\,(\text{barrier height})^4$$

    with tunnelling rate $\Gamma \sim \exp(-S_\mathrm{bounce})$. A crude estimate:
    the barrier is taken at $\phi = 0$ when $m < 0$, and absent otherwise.

    Args:
        phi_false: False vacuum field value.
        phi_true: True vacuum field value.
        m: Mass.
        lambda_: Coupling.

    Returns:
        Bounce action estimate.
    """
    V_false = scalar_potential_energy(phi_false, m, lambda_)
    V_true = scalar_potential_energy(phi_true, m, lambda_)
    Delta_V = jnp.abs(V_true - V_false)

    # Barrier at phi = 0 for the symmetry-broken case, none otherwise.
    barrier_height = jnp.where(m < 0, jnp.abs(0.0 - V_false), 0.0)

    # The +1e-10 keeps the quotient finite even at Delta_V = 0, so the unused
    # branch is harmless and a single `where` is safe here.
    return jnp.where(
        Delta_V > 1e-10,
        barrier_height**4 / (Delta_V**2 + 1e-10),
        1e10,  # degenerate vacua: tunnelling is maximally suppressed
    )

generate_stability_report(phi, m, lambda_, dx)

Generate a comprehensive stability analysis report.

Parameters:

Name	Type	Description	Default
phi	Array	Scalar field configuration.	required
m	float	Mass parameter.	required
lambda_	float	Self-coupling.	required
dx	float	Lattice spacing.	required

Returns:

Type	Description
dict[str, Array]	Dictionary combining mass-matrix statistics, tachyonic-mode diagnostics,
dict[str, Array]	and symmetry-breaking indicators, all as Array.

Source code in jaxlatt/observables/functionals.py

def generate_stability_report(
    phi: Array,
    m: float,
    lambda_: float,
    dx: float,
) -> dict[str, Array]:
    """Generate a comprehensive stability analysis report.

    Args:
        phi: Scalar field configuration.
        m: Mass parameter.
        lambda_: Self-coupling.
        dx: Lattice spacing.

    Returns:
        Dictionary combining mass-matrix statistics, tachyonic-mode diagnostics,
        and symmetry-breaking indicators, all as `Array`.
    """
    m_eff_sq = compute_mass_matrix(phi, m, lambda_)
    tach_info = detect_tachyonic_modes(phi, m, lambda_)
    ssb_info = check_symmetry_breaking(phi, m, lambda_)

    return {
        "min_mass_sq": jnp.min(m_eff_sq),
        "max_mass_sq": jnp.max(m_eff_sq),
        "mean_mass_sq": jnp.mean(m_eff_sq),
        "has_tachyonic_modes": tach_info["has_instability"],
        "fraction_unstable": tach_info["fraction_tachyonic"],
        "field_rms": ssb_info["rms_field"],
        "symmetry_broken": ssb_info["is_symmetry_broken"],
        "curvature_at_origin": ssb_info["curvature_at_origin"],
    }

compute_constraint_violation_norm(phi, pi, E, links, dx, scale_factor=1.0)

Norms of the Gauss-constraint violation, as Arrays.

Parameters:

Name	Type	Description	Default
phi	Array	Scalar field.	required
pi	Array	Conjugate momentum.	required
E	Array	Electric field.	required
links	Array	Gauge links.	required
dx	float	Comoving lattice spacing.	required
scale_factor	float	Scale factor \(a\). The FRW Gauss law is \(\nabla\cdot E + a^2\operatorname{Im}(\phi^*\pi)/dx = 0\), so this must be supplied for expanding runs; the default of 1.0 is the flat case.	1.0

Returns:

Type	Description
dict[str, Array]	Dictionary with 'L1_norm', 'L2_norm', 'Linf_norm' and
dict[str, Array]	'total_violation'.

Source code in jaxlatt/observables/functionals.py

def compute_constraint_violation_norm(
    phi: Array,
    pi: Array,
    E: Array,
    links: Array,
    dx: float,
    scale_factor: float = 1.0,
) -> dict[str, Array]:
    r"""Norms of the Gauss-constraint violation, as `Array`s.

    Args:
        phi: Scalar field.
        pi: Conjugate momentum.
        E: Electric field.
        links: Gauge links.
        dx: Comoving lattice spacing.
        scale_factor: Scale factor $a$. The FRW Gauss law is
            $\nabla\cdot E + a^2\operatorname{Im}(\phi^*\pi)/dx = 0$, so this must
            be supplied for expanding runs; the default of 1.0 is the flat case.

    Returns:
        Dictionary with `'L1_norm'`, `'L2_norm'`, `'Linf_norm'` and
        `'total_violation'`.
    """
    return _constraint_violation_norm_impl(phi, pi, E, links, dx, scale_factor)

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

Gauss-constraint drift across a time step, as Arrays.

Parameters:

Name	Type	Description	Default
phi_old	Array	Scalar field at the previous time.	required
pi_old	Array	Conjugate momentum at the previous time.	required
E_old	Array	Electric field at the previous time.	required
links_old	Array	Gauge links at the previous time.	required
phi_new	Array	Scalar field at the current time.	required
pi_new	Array	Conjugate momentum at the current time.	required
E_new	Array	Electric field at the current time.	required
links_new	Array	Gauge links at the current time.	required
dx	float	Comoving lattice spacing.	required
scale_factor	float	Scale factor \(a\); the default of 1.0 is the flat case.	1.0

Returns:

Type	Description
dict[str, Array]	Dictionary with 'violation_before', 'violation_after', 'drift' and
dict[str, Array]	'drift_rate'.

Source code in jaxlatt/observables/functionals.py

def check_constraint_preservation(
    phi_old: Array,
    pi_old: Array,
    E_old: Array,
    links_old: Array,
    phi_new: Array,
    pi_new: Array,
    E_new: Array,
    links_new: Array,
    dx: float,
    scale_factor: float = 1.0,
) -> dict[str, Array]:
    r"""Gauss-constraint drift across a time step, as `Array`s.

    Args:
        phi_old: Scalar field at the previous time.
        pi_old: Conjugate momentum at the previous time.
        E_old: Electric field at the previous time.
        links_old: Gauge links at the previous time.
        phi_new: Scalar field at the current time.
        pi_new: Conjugate momentum at the current time.
        E_new: Electric field at the current time.
        links_new: Gauge links at the current time.
        dx: Comoving lattice spacing.
        scale_factor: Scale factor $a$; the default of 1.0 is the flat case.

    Returns:
        Dictionary with `'violation_before'`, `'violation_after'`, `'drift'` and
        `'drift_rate'`.
    """
    return _constraint_preservation_impl(
        phi_old, pi_old, E_old, links_old, phi_new, pi_new, E_new, links_new, dx, scale_factor
    )