Energy

energy

Energy diagnostics for field dynamics.

This module provides functions to compute and track various energy components:

• Kinetic energy: \(\frac{1}{2} \langle \dot{\phi}^2 \rangle\)
• Gradient energy: \(\frac{1}{2} \langle (\nabla \phi)^2 \rangle\)
• Potential energy: \(\langle V(\phi) \rangle\)
• Total energy and conservation checking

Both volume-averaged and volume-integrated functions are provided.

For expanding universe simulations, provides physical energy densities with proper scale factor scaling for Friedmann equation integration.

Inspired by CosmoLattice's energy measurement framework.

EnergyTracker()

Track energy evolution and conservation over time.

This class stores energy measurements at each timestep and provides methods to analyze energy conservation.

Attributes:

Name	Type	Description
times		List of measurement times
energies		List of energy component dictionaries

Example

from jaxlatt.potentials import quadratic_potential

V = quadratic_potential(m=1.0)
tracker = EnergyTracker()

# During evolution
for t, snapshot in zip(times, snapshots):
    tracker.add_measurement(t, snapshot, V)

# Check conservation
conservation = tracker.energy_conservation()
print(f"Energy drift: {conservation['relative_change']:.2e}")

Initialize an empty energy tracker.

Creates storage for simulation times and energy-component snapshots.

Source code in jaxlatt/observables/energy.py

def __init__(self):
    """Initialize an empty energy tracker.

    Creates storage for simulation times and energy-component snapshots.
    """
    self.times = []
    self.energies = []

add_measurement(time, lattice, potential_func)

Add energy measurement at given time.

Parameters:

Name	Type	Description	Default
time	float	Current simulation time	required
lattice	Lattice	Current lattice state	required
potential_func	Callable	Potential function	required

Source code in jaxlatt/observables/energy.py

def add_measurement(self, time: float, lattice: Lattice, potential_func: Callable) -> None:
    """Add energy measurement at given time.

    Args:
        time: Current simulation time
        lattice: Current lattice state
        potential_func: Potential function
    """
    self.times.append(float(time))
    self.energies.append(energy_components_averaged(lattice, potential_func))

get_component(component)

Get time series of a specific energy component.

Parameters:

Name	Type	Description	Default
component	str	Energy component name ('kinetic', 'gradient', 'potential', or 'total')	required

Returns:

Type	Description
tuple[ndarray, ndarray]	Tuple of (times, energy_values) as numpy arrays

Example

times, E_kin = tracker.get_component('kinetic')
plt.plot(times, E_kin, label='Kinetic')

Source code in jaxlatt/observables/energy.py

def get_component(self, component: str) -> tuple[jnp.ndarray, jnp.ndarray]:
    """Get time series of a specific energy component.

    Args:
        component: Energy component name ('kinetic', 'gradient',
                  'potential', or 'total')

    Returns:
        Tuple of (times, energy_values) as numpy arrays

    Example:
        ```python
        times, E_kin = tracker.get_component('kinetic')
        plt.plot(times, E_kin, label='Kinetic')
        ```
    """
    if not self.energies:
        raise ValueError("No measurements recorded yet")

    values = [e[component] for e in self.energies]
    return jnp.array(self.times), jnp.array(values)

energy_conservation()

Analyze energy conservation.

Computes various metrics to assess how well energy is conserved during the simulation.

Returns:

Type	Description
dict[str, float]	Dictionary with: - 'initial': Initial total energy - 'final': Final total energy - 'mean': Mean total energy - 'std': Standard deviation - 'relative_change': (E_final - E_initial) / E_initial - 'max_deviation': max|E - E_mean| / E_mean

Example

conservation = tracker.energy_conservation()
if abs(conservation['relative_change']) > 1e-6:
    print("WARNING: Significant energy drift detected!")

Source code in jaxlatt/observables/energy.py

def energy_conservation(self) -> dict[str, float]:
    """Analyze energy conservation.

    Computes various metrics to assess how well energy is conserved
    during the simulation.

    Returns:
        Dictionary with:
            - 'initial': Initial total energy
            - 'final': Final total energy
            - 'mean': Mean total energy
            - 'std': Standard deviation
            - 'relative_change': (E_final - E_initial) / E_initial
            - 'max_deviation': max|E - E_mean| / E_mean

    Example:
        ```python
        conservation = tracker.energy_conservation()
        if abs(conservation['relative_change']) > 1e-6:
            print("WARNING: Significant energy drift detected!")
        ```
    """
    if len(self.energies) < 2:
        raise ValueError("Need at least 2 measurements for conservation check")

    _, E_tot = self.get_component("total")

    E_initial = E_tot[0]
    E_final = E_tot[-1]
    E_mean = jnp.mean(E_tot)
    E_std = jnp.std(E_tot)

    relative_change = float((E_final - E_initial) / E_initial)
    max_deviation = float(jnp.max(jnp.abs(E_tot - E_mean)) / E_mean)

    return {
        "initial": float(E_initial),
        "final": float(E_final),
        "mean": float(E_mean),
        "std": float(E_std),
        "relative_change": relative_change,
        "max_deviation": max_deviation,
    }

summary()

Generate human-readable summary of energy tracking.

Returns:

Type	Description
str	Formatted string with energy statistics

Example

print(tracker.summary())

Source code in jaxlatt/observables/energy.py

def summary(self) -> str:
    """Generate human-readable summary of energy tracking.

    Returns:
        Formatted string with energy statistics

    Example:
        ```python
        print(tracker.summary())
        ```
    """
    if not self.energies:
        return "No measurements recorded"

    conservation = self.energy_conservation()

    summary_lines = [
        "=" * 60,
        "ENERGY CONSERVATION SUMMARY",
        "=" * 60,
        f"Time range: {self.times[0]:.2f} to {self.times[-1]:.2f}",
        f"Measurements: {len(self.times)}",
        "",
        "Total Energy:",
        f"  Initial:  {conservation['initial']:.10e}",
        f"  Final:    {conservation['final']:.10e}",
        f"  Mean:     {conservation['mean']:.10e}",
        f"  Std Dev:  {conservation['std']:.10e}",
        "",
        "Conservation Metrics:",
        f"  Relative change: {conservation['relative_change']:+.6e}",
        f"  Max deviation:   {conservation['max_deviation']:+.6e}",
        "",
    ]

    # Add final energy breakdown
    final = self.energies[-1]
    summary_lines.extend(
        [
            "Final Energy Breakdown:",
            f"  Kinetic:   {final['kinetic']:.10e}  ({final['kinetic'] / final['total'] * 100:.2f}%)",
            f"  Gradient:  {final['gradient']:.10e}  ({final['gradient'] / final['total'] * 100:.2f}%)",
            f"  Potential: {final['potential']:.10e}  ({final['potential'] / final['total'] * 100:.2f}%)",
            "=" * 60,
        ]
    )

    return "\n".join(summary_lines)

to_dict()

Export all data as dictionary for saving.

Returns:

Type	Description
dict	Dictionary with 'times' and 'energies' arrays

Example

import json
data = tracker.to_dict()
with open('energies.json', 'w') as f:
    json.dump(data, f)

Source code in jaxlatt/observables/energy.py

def to_dict(self) -> dict:
    """Export all data as dictionary for saving.

    Returns:
        Dictionary with 'times' and 'energies' arrays

    Example:
        ```python
        import json
        data = tracker.to_dict()
        with open('energies.json', 'w') as f:
            json.dump(data, f)
        ```
    """
    return {"times": self.times, "energies": self.energies}

kinetic_energy_averaged(lattice)

Compute volume-averaged kinetic energy density.

Kinetic energy density: \(\rho_\mathrm{kin} = \frac{1}{2}\dot{\phi}^2\)

Parameters:

Name	Type	Description	Default
lattice	Lattice	Lattice object with field_dot	required

Returns:

Type	Description
float	Volume-averaged kinetic energy: \(\langle \rho_\mathrm{kin} \rangle = \frac{1}{2}\langle \dot{\phi}^2 \rangle\)

Example

E_kin = kinetic_energy_averaged(lattice)
print(f"Kinetic energy: {E_kin:.6e}")

Source code in jaxlatt/observables/energy.py

def kinetic_energy_averaged(lattice: Lattice) -> float:
    r"""Compute volume-averaged kinetic energy density.

    Kinetic energy density: $\rho_\mathrm{kin} = \frac{1}{2}\dot{\phi}^2$

    Args:
        lattice: Lattice object with field_dot

    Returns:
        Volume-averaged kinetic energy: $\langle \rho_\mathrm{kin} \rangle = \frac{1}{2}\langle \dot{\phi}^2 \rangle$

    Example:
        ```python
        E_kin = kinetic_energy_averaged(lattice)
        print(f"Kinetic energy: {E_kin:.6e}")
        ```
    """
    return float(_kinetic_energy_averaged_impl(lattice.field_dot))

gradient_energy_averaged(lattice)

Compute volume-averaged gradient energy density.

Gradient energy density: \(\rho_\mathrm{grad} = \frac{1}{2}(\nabla\phi)^2\)

Uses FFT-based gradient computation for accuracy.

Parameters:

Name	Type	Description	Default
lattice	Lattice	Lattice object with field	required

Returns:

Type	Description
float	Volume-averaged gradient energy: \(\langle \rho_\mathrm{grad} \rangle = \frac{1}{2}\langle (\nabla\phi)^2 \rangle\)

Example

E_grad = gradient_energy_averaged(lattice)
print(f"Gradient energy: {E_grad:.6e}")

Source code in jaxlatt/observables/energy.py

def gradient_energy_averaged(lattice: Lattice) -> float:
    r"""Compute volume-averaged gradient energy density.

    Gradient energy density: $\rho_\mathrm{grad} = \frac{1}{2}(\nabla\phi)^2$

    Uses FFT-based gradient computation for accuracy.

    Args:
        lattice: Lattice object with field

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

    Example:
        ```python
        E_grad = gradient_energy_averaged(lattice)
        print(f"Gradient energy: {E_grad:.6e}")
        ```
    """
    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 float(_gradient_energy_averaged_impl(lattice.field, dx_per_dim))

potential_energy_averaged(lattice, potential_func)

Compute volume-averaged potential energy density.

Potential energy density: \(\rho_\mathrm{pot} = V(\phi)\)

Parameters:

Name	Type	Description	Default
lattice	Lattice	Lattice object with field	required
potential_func	Callable	Potential function \(V(\phi) \to V\); should be JIT-compiled.	required

Returns:

Type	Description
float	Volume-averaged potential energy: \(\langle \rho_\mathrm{pot} \rangle = \langle V(\phi) \rangle\)

Example

from jaxlatt.potentials import quadratic_potential
V = quadratic_potential(m=1.0)
E_pot = potential_energy_averaged(lattice, V)
print(f"Potential energy: {E_pot:.6e}")

Source code in jaxlatt/observables/energy.py

def potential_energy_averaged(lattice: Lattice, potential_func: Callable) -> float:
    r"""Compute volume-averaged potential energy density.

    Potential energy density: $\rho_\mathrm{pot} = V(\phi)$

    Args:
        lattice: Lattice object with field
        potential_func: Potential function $V(\phi) \to V$; should be JIT-compiled.

    Returns:
        Volume-averaged potential energy: $\langle \rho_\mathrm{pot} \rangle = \langle V(\phi) \rangle$

    Example:
        ```python
        from jaxlatt.potentials import quadratic_potential
        V = quadratic_potential(m=1.0)
        E_pot = potential_energy_averaged(lattice, V)
        print(f"Potential energy: {E_pot:.6e}")
        ```
    """
    # Evaluate potential at each grid point (potential_func should be JIT-compiled)
    V_field = potential_func(lattice.field)
    return float(_potential_energy_averaged_impl(V_field))

total_energy_averaged(lattice, potential_func)

Compute total volume-averaged energy density.

Total energy: \(E_\mathrm{tot} = E_\mathrm{kin} + E_\mathrm{grad} + E_\mathrm{pot}\)

Parameters:

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

Returns:

Type	Description
float	Total volume-averaged energy density

Example

from jaxlatt.potentials import quadratic_potential
V = quadratic_potential(m=1.0)
E_tot = total_energy_averaged(lattice, V)
print(f"Total energy: {E_tot:.6e}")

Source code in jaxlatt/observables/energy.py

def total_energy_averaged(lattice: Lattice, potential_func: Callable) -> float:
    r"""Compute total volume-averaged energy density.

    Total energy: $E_\mathrm{tot} = E_\mathrm{kin} + E_\mathrm{grad} + E_\mathrm{pot}$

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

    Returns:
        Total volume-averaged energy density

    Example:
        ```python
        from jaxlatt.potentials import quadratic_potential
        V = quadratic_potential(m=1.0)
        E_tot = total_energy_averaged(lattice, V)
        print(f"Total energy: {E_tot:.6e}")
        ```
    """
    E_kin = kinetic_energy_averaged(lattice)
    E_grad = gradient_energy_averaged(lattice)
    E_pot = potential_energy_averaged(lattice, potential_func)

    return E_kin + E_grad + E_pot

energy_components_averaged(lattice, potential_func)

Compute all energy components.

Returns a dictionary with all energy contributions for easy tracking.

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, float]	Dictionary with keys:
dict[str, float]	'kinetic': Kinetic energy \(\langle \dot{\phi}^2/2 \rangle\)
dict[str, float]	'gradient': Gradient energy \(\langle (\nabla\phi)^2/2 \rangle\)
dict[str, float]	'potential': Potential energy \(\langle V(\phi) \rangle\)
dict[str, float]	'total': Total energy

Example

from jaxlatt.potentials import quadratic_potential
V = quadratic_potential(m=1.0)
energies = energy_components_averaged(lattice, V)
for name, value in energies.items():
    print(f"{name:12s}: {value:.6e}")

Source code in jaxlatt/observables/energy.py

def energy_components_averaged(lattice: Lattice, potential_func: Callable) -> dict[str, float]:
    r"""Compute all energy components.

    Returns a dictionary with all energy contributions for easy tracking.

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

    Returns:
        Dictionary with keys:

        - `'kinetic'`: Kinetic energy $\langle \dot{\phi}^2/2 \rangle$
        - `'gradient'`: Gradient energy $\langle (\nabla\phi)^2/2 \rangle$
        - `'potential'`: Potential energy $\langle V(\phi) \rangle$
        - `'total'`: Total energy

    Example:
        ```python
        from jaxlatt.potentials import quadratic_potential
        V = quadratic_potential(m=1.0)
        energies = energy_components_averaged(lattice, V)
        for name, value in energies.items():
            print(f"{name:12s}: {value:.6e}")
        ```
    """
    E_kin = kinetic_energy_averaged(lattice)
    E_grad = gradient_energy_averaged(lattice)
    E_pot = potential_energy_averaged(lattice, potential_func)
    E_tot = E_kin + E_grad + E_pot

    return {
        "kinetic": float(E_kin),
        "gradient": float(E_grad),
        "potential": float(E_pot),
        "total": float(E_tot),
    }

kinetic_energy_integrated(lattice)

Compute integrated kinetic energy.

Parameters:

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

Returns:

Type	Description
float	Total kinetic energy, integrated over the full simulation volume.

Source code in jaxlatt/observables/energy.py

def kinetic_energy_integrated(lattice: Lattice) -> float:
    """Compute integrated kinetic energy.

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

    Returns:
        Total kinetic energy, integrated over the full simulation volume.
    """
    return kinetic_energy_averaged(lattice) * lattice.volume

potential_energy_integrated(lattice, potential)

Compute integrated potential energy.

Parameters:

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

Returns:

Type	Description
float	Total potential energy, integrated over the full simulation volume.

Source code in jaxlatt/observables/energy.py

def potential_energy_integrated(lattice: Lattice, potential: Callable) -> float:
    """Compute integrated potential energy.

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

    Returns:
        Total potential energy, integrated over the full simulation volume.
    """
    return potential_energy_averaged(lattice, potential) * lattice.volume

energy_components_integrated(lattice, potential)

Compute all integrated energy components.

Parameters:

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

Returns:

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

Source code in jaxlatt/observables/energy.py

def energy_components_integrated(lattice: Lattice, potential: Callable) -> dict[str, float]:
    """Compute all integrated energy components.

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

    Returns:
        Dictionary with integrated `kinetic`, `gradient`, `potential`, and `total`.
    """
    # Get volume-averaged components
    components = energy_components_averaged(lattice, potential)

    # Convert to integrated quantities
    volume = lattice.volume
    return {
        "kinetic": components["kinetic"] * volume,
        "gradient": components["gradient"] * volume,
        "potential": components["potential"] * volume,
        "total": components["total"] * volume,
    }

total_energy_integrated(lattice, potential)

Compute integrated total energy.

Parameters:

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

Returns:

Type	Description
float	Total energy integrated over the full simulation volume.

Source code in jaxlatt/observables/energy.py

def total_energy_integrated(lattice: Lattice, potential: Callable) -> float:
    """Compute integrated total energy.

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

    Returns:
        Total energy integrated over the full simulation volume.
    """
    return total_energy_averaged(lattice, potential) * lattice.volume

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/energy.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.

In conformal time with canonical momentum \(\pi = a^3 \dot{\phi}\):

\[\rho_\mathrm{kin}^\mathrm{phys} = \frac{\pi^2}{2a^6} = \frac{\dot{\phi}^2}{2}\]

For proper implementation, we assume field_dot stores the conformal time derivative (\(\phi'\) in conformal time), so:

\[\rho_\mathrm{kin}^\mathrm{phys} = \frac{(\phi')^2}{2a^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/energy.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.

    In conformal time with canonical momentum $\pi = a^3 \dot{\phi}$:

    $$\rho_\mathrm{kin}^\mathrm{phys} = \frac{\pi^2}{2a^6} = \frac{\dot{\phi}^2}{2}$$

    For proper implementation, we assume `field_dot` stores the conformal
    time derivative ($\phi'$ in conformal time), so:

    $$\rho_\mathrm{kin}^\mathrm{phys} = \frac{(\phi')^2}{2a^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)

Source code in jaxlatt/observables/energy.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)
    """
    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/energy.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/energy.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)

Compute total physical energy density for Friedmann equation.

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

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

Includes all energy components with correct scale factor scaling:

• Scalar kinetic: \(|\pi|^2/(2a^6)\) where \(\pi\) is canonical momentum
• Scalar gradient: \((\nabla\phi)^2/(2a^2)\)
• Scalar potential: \(V(\phi)\)
• Electric field: \(E^2/(2a^4)\)
• Magnetic field: \(B^2/(2a^4)\)

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

Returns:

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

Note

For scalar-only simulations, pass lattice with dummy gauge fields or use scalar-specific functions above.

Source code in jaxlatt/observables/energy.py

def compute_physical_energy_density(
    lattice: "CoupledLattice",
    scale_factor: float,
    dx: float,
    potential_func: Callable[[Array], jnp.ndarray],
) -> float:
    r"""
    Compute total physical energy density for Friedmann equation.

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

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

    Includes all energy components with correct scale factor scaling:

    - Scalar kinetic: $|\pi|^2/(2a^6)$ where $\pi$ is canonical momentum
    - Scalar gradient: $(\nabla\phi)^2/(2a^2)$
    - Scalar potential: $V(\phi)$
    - Electric field: $E^2/(2a^4)$
    - Magnetic field: $B^2/(2a^4)$

    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)$

    Returns:
        Total physical energy density $\rho_\mathrm{phys}$ (averaged over lattice)

    Note:
        For scalar-only simulations, pass lattice with dummy gauge fields
        or use scalar-specific functions above.
    """
    rho_total = sum(_physical_energy_components(lattice, scale_factor, dx, potential_func))
    return float(rho_total.real if jnp.iscomplexobj(rho_total) else rho_total)

compute_physical_energy_components(lattice, scale_factor, dx, potential_func)

Compute all physical energy density components separately.

Useful for diagnostics and tracking how energy is distributed between different field components during expansion.

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

Returns:

Type	Description
dict[str, float]	Dictionary with keys:
dict[str, float]	'kinetic': Scalar kinetic energy density
dict[str, float]	'gradient': Scalar gradient energy density
dict[str, float]	'potential': Scalar potential energy density
dict[str, float]	'electric': Electric field energy density
dict[str, float]	'magnetic': Magnetic field energy density
dict[str, float]	'total': Total physical energy density

Example

components = compute_physical_energy_components(lattice, a=2.0, dx=0.5, V)
for name, value in components.items():
    print(f"{name:12s}: {value:.6e}")

Source code in jaxlatt/observables/energy.py

def compute_physical_energy_components(
    lattice: "CoupledLattice",
    scale_factor: float,
    dx: float,
    potential_func: Callable[[Array], jnp.ndarray],
) -> dict[str, float]:
    r"""
    Compute all physical energy density components separately.

    Useful for diagnostics and tracking how energy is distributed
    between different field components during expansion.

    Args:
        lattice: CoupledLattice state
        scale_factor: Current scale factor $a(\tau)$
        dx: Lattice spacing (comoving)
        potential_func: Scalar potential $V(\phi)$

    Returns:
        Dictionary with keys:

        - `'kinetic'`: Scalar kinetic energy density
        - `'gradient'`: Scalar gradient energy density
        - `'potential'`: Scalar potential energy density
        - `'electric'`: Electric field energy density
        - `'magnetic'`: Magnetic field energy density
        - `'total'`: Total physical energy density

    Example:
        ```python
        components = compute_physical_energy_components(lattice, a=2.0, dx=0.5, V)
        for name, value in components.items():
            print(f"{name:12s}: {value:.6e}")
        ```
    """
    components = _physical_energy_components(lattice, scale_factor, dx, potential_func)
    rho_kinetic, rho_gradient, rho_potential, rho_electric, rho_magnetic = components
    rho_total = sum(components)

    def _f(x: Array) -> float:
        return float(x.real if jnp.iscomplexobj(x) else x)

    return {
        "kinetic": _f(rho_kinetic),
        "gradient": _f(rho_gradient),
        "potential": _f(rho_potential),
        "electric": _f(rho_electric),
        "magnetic": _f(rho_magnetic),
        "total": _f(rho_total),
    }

compute_comoving_energy_density(lattice, dx, potential_func)

Compute comoving energy density (no scale factor corrections).

This is what standard (non-expanding) simulations compute. Useful for comparison and debugging.

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

Returns:

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

Source code in jaxlatt/observables/energy.py

def compute_comoving_energy_density(
    lattice: "CoupledLattice",
    dx: float,
    potential_func: Callable[[Array], jnp.ndarray],
) -> float:
    r"""
    Compute comoving energy density (no scale factor corrections).

    This is what standard (non-expanding) simulations compute.
    Useful for comparison and debugging.

    Args:
        lattice: CoupledLattice state
        dx: Lattice spacing
        potential_func: Scalar potential $V(\phi)$

    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
    )

kinetic_energy(field_dot, dV)

Compute 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
float	Total kinetic energy (integrated, not averaged)

Source code in jaxlatt/observables/energy.py

def kinetic_energy(field_dot: Array, dV: float) -> float:
    r"""
    Compute 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 float(0.5 * jnp.sum(field_dot**2) * dV)

potential_energy(field, potential, dV)

Compute 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
float	Total potential energy (integrated, not averaged)

Source code in jaxlatt/observables/energy.py

def potential_energy(field: Array, potential: Callable[[Array], jnp.ndarray], dV: float) -> float:
    r"""
    Compute 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)
    """
    V = potential(field)
    return float(jnp.sum(V) * dV)

compute_energy_components(lattice, potential)

Compute all energy components separately (volume-integrated).

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, float]	Dictionary with keys 'kinetic', 'gradient', 'potential', 'total'

Source code in jaxlatt/observables/energy.py

def compute_energy_components(
    lattice: Lattice, potential: Callable[[Array], jnp.ndarray]
) -> dict[str, float]:
    r"""
    Compute all energy components separately (volume-integrated).

    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)

compute_energy(lattice, potential)

Compute total energy of a field configuration (volume-integrated).

\[E_\mathrm{total} = E_\mathrm{kin} + E_\mathrm{grad} + E_\mathrm{pot} = \int \left[\frac{1}{2}\left(\frac{\partial\phi}{\partial t}\right)^2 + \frac{1}{2}(\nabla\phi)^2 + V(\phi)\right] dV\]

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
float	Total energy (integrated)

Source code in jaxlatt/observables/energy.py

def compute_energy(lattice: Lattice, potential: Callable[[Array], jnp.ndarray]) -> float:
    r"""
    Compute total energy of a field configuration (volume-integrated).

    $$E_\mathrm{total} = E_\mathrm{kin} + E_\mathrm{grad} + E_\mathrm{pot} = \int \left[\frac{1}{2}\left(\frac{\partial\phi}{\partial t}\right)^2 + \frac{1}{2}(\nabla\phi)^2 + V(\phi)\right] dV$$

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

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

energy_density(field, field_dot, potential, dx)

Compute 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

Source code in jaxlatt/observables/energy.py

def energy_density(
    field: Array,
    field_dot: Array,
    potential: Callable[[Array], jnp.ndarray],
    dx: float,
) -> Array:
    """
    Compute 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
    """
    # Kinetic contribution
    rho_kin = 0.5 * field_dot**2

    # Potential contribution
    rho_pot = potential(field)

    # Gradient contribution
    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