Core Module

core

Core lattice structures and cosmology.

FRWUniverse

Bases: Module

FLRW universe parameters in conformal time.

Attributes:

Name	Type	Description
a	Array	Scale factor (dimensionless, typically normalized to \(a(\tau_0) = 1\))
adot	Array	Conformal time derivative \(da/d\tau\) (same units as a/time)
tau	Array	Conformal time coordinate \(\tau\)
M_pl	float	Reduced Planck mass in simulation units (sets gravity strength)

Derived quantities:

• \(H = a'/a\): Conformal Hubble parameter
• \(\rho_{\mathrm{crit}} = 3M_{\mathrm{pl}}^2 H^2\): Critical density

H property

Compute the conformal Hubble parameter.

Returns:

Type	Description
float	Value of \(H = a'/a\) for the current universe state.

rho_crit property

Compute the critical density.

Returns:

Type	Description
float	Critical density computed from 3 * H**2 / (8 * pi * G).

copy()

Create a copy of the universe state.

Returns:

Type	Description
	New FRWUniverse with identical values.

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

def copy(self):
    """Create a copy of the universe state.

    Returns:
        New `FRWUniverse` with identical values.
    """
    return dataclasses.replace(self)

replace(**kwargs)

Create a new universe with selected fields replaced.

Parameters:

Name	Type	Description	Default
**kwargs		Fields to override (a, adot, tau).	{}

Returns:

Type	Description
FRWUniverse	New FRWUniverse with replaced values.

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

def replace(self, **kwargs) -> "FRWUniverse":
    """Create a new universe with selected fields replaced.

    Args:
        **kwargs: Fields to override (a, adot, tau).

    Returns:
        New `FRWUniverse` with replaced values.
    """
    return dataclasses.replace(self, **kwargs)

CoupledLattice

Bases: Module

Combined scalar and gauge field lattice for Abelian Higgs model.

Attributes:

Name	Type	Description
phi	Array	Complex scalar field with shape (Nx, Ny, Nz).
pi	Array	Conjugate momentum to phi with shape (Nx, Ny, Nz).
links	Array	U(1) gauge links with shape (3, Nx, Ny, Nz).
E	Array	Electric field components with shape (3, Nx, Ny, Nz).
m	float	Scalar mass parameter.
lambda_	float	Scalar self-coupling.
g	float	Gauge coupling.
dx	float	Lattice spacing.
size	tuple[int, int, int]	Lattice dimensions (Nx, Ny, Nz).
length	tuple[float, float, float]	Physical box lengths (Lx, Ly, Lz).

volume property

Physical volume of the simulation box.

dV property

Volume element (cell volume).

copy()

Create a deep copy of the coupled lattice.

Returns:

Type	Description
CoupledLattice	New CoupledLattice instance with independent array buffers.

Source code in jaxlatt/core/lattice.py

def copy(self) -> "CoupledLattice":
    """Create a deep copy of the coupled lattice.

    Returns:
        New `CoupledLattice` instance with independent array buffers.
    """
    return dataclasses.replace(
        self,
        phi=jnp.copy(self.phi),
        pi=jnp.copy(self.pi),
        links=jnp.copy(self.links),
        E=jnp.copy(self.E),
    )

update(phi, pi, links, E)

Create a new CoupledLattice with updated dynamical fields.

Parameters:

Name	Type	Description	Default
phi	Array	New scalar field.	required
pi	Array	New conjugate momentum.	required
links	Array	New gauge links.	required
E	Array	New electric field.	required

Returns:

Type	Description
CoupledLattice	New CoupledLattice with static parameters (m, lambda_, g, dx, size, length) preserved.

Source code in jaxlatt/core/lattice.py

def update(self, phi: Array, pi: Array, links: Array, E: Array) -> "CoupledLattice":
    """Create a new CoupledLattice with updated dynamical fields.

    Args:
        phi: New scalar field.
        pi: New conjugate momentum.
        links: New gauge links.
        E: New electric field.

    Returns:
        New `CoupledLattice` with static parameters (m, lambda_, g, dx, size, length) preserved.
    """
    return dataclasses.replace(self, phi=phi, pi=pi, links=links, E=E)

replace(**kwargs)

Create a new lattice with selected attributes replaced.

Parameters:

Name	Type	Description	Default
**kwargs		Field values or parameters to override.	{}

Returns:

Type	Description
CoupledLattice	New CoupledLattice containing replaced attributes.

Source code in jaxlatt/core/lattice.py

def replace(self, **kwargs) -> "CoupledLattice":
    """Create a new lattice with selected attributes replaced.

    Args:
        **kwargs: Field values or parameters to override.

    Returns:
        New `CoupledLattice` containing replaced attributes.
    """
    return dataclasses.replace(self, **kwargs)

GaugeLattice(size, length, g=1.0, links=None, E=None, a=1.0, H=0.0)

Bases: Module

U(1) gauge fields on a 3D periodic lattice.

Uses compact formulation with link variables \(U_i \in \mathrm{U}(1)\).

Attributes:

Name	Type	Description
links	Array	Complex link variables (3, N, N, N)
E	Array	Electric field components (3, N, N, N)
dx	float	Lattice spacing
size	tuple[int, int, int]	Lattice dimensions
length	tuple[float, float, float]	Physical box size
g	float	Gauge coupling
a	Array	Scale factor (cosmology)
H	Array	Conformal Hubble parameter (cosmology)

Initialize a gauge lattice.

Parameters:

Name	Type	Description	Default
size	tuple[int, int, int]	Lattice dimensions (Nx, Ny, Nz).	required
length	float	Physical box length (assumed isotropic).	required
g	float	Gauge coupling.	1.0
links	Array | None	Optional initial link variables.	None
E	Array | None	Optional initial electric field.	None
a	float	Cosmological scale factor.	1.0
H	float	Conformal Hubble parameter.	0.0

Source code in jaxlatt/core/lattice.py

def __init__(
    self,
    size: tuple[int, int, int],
    length: float,
    g: float = 1.0,
    links: Array | None = None,
    E: Array | None = None,
    a: float = 1.0,
    H: float = 0.0,
):
    """Initialize a gauge lattice.

    Args:
        size: Lattice dimensions `(Nx, Ny, Nz)`.
        length: Physical box length (assumed isotropic).
        g: Gauge coupling.
        links: Optional initial link variables.
        E: Optional initial electric field.
        a: Cosmological scale factor.
        H: Conformal Hubble parameter.
    """
    self.size = size
    self.length = (length, length, length)
    self.dx = length / size[0]
    self.g = g
    self.a = jnp.asarray(a)
    self.H = jnp.asarray(H)

    if links is None:
        self.links = jnp.ones((3, *size), dtype=jnp.complex64)
    else:
        self.links = links.astype(jnp.result_type(complex))

    if E is None:
        self.E = jnp.zeros((3, *size), dtype=jnp.float32)
    else:
        self.E = E.astype(jnp.result_type(float))

volume property

Compute the total physical lattice volume.

Returns:

Type	Description
float	Total box volume.

dV property

Compute the volume element for one lattice site.

Returns:

Type	Description
float	Cell volume dx**3.

copy()

Create a deep copy of the gauge lattice state.

Returns:

Type	Description
GaugeLattice	New GaugeLattice instance with independent array buffers.

Source code in jaxlatt/core/lattice.py

def copy(self) -> "GaugeLattice":
    """Create a deep copy of the gauge lattice state.

    Returns:
        New `GaugeLattice` instance with independent array buffers.
    """
    return eqx.tree_at(
        lambda t: (t.links, t.E),
        self,
        (jnp.copy(self.links), jnp.copy(self.E)),
    )

update(links, E, a=None, H=None)

Create a new gauge lattice with updated dynamical variables.

Parameters:

Name	Type	Description	Default
links	Array	Updated link variables.	required
E	Array	Updated electric field.	required
a	Array | None	Optional updated scale factor.	None
H	Array | None	Optional updated Hubble parameter.	None

Returns:

Type	Description
GaugeLattice	New GaugeLattice instance with updated values.

Source code in jaxlatt/core/lattice.py

def update(
    self,
    links: Array,
    E: Array,
    a: Array | None = None,
    H: Array | None = None,
) -> "GaugeLattice":
    """Create a new gauge lattice with updated dynamical variables.

    Args:
        links: Updated link variables.
        E: Updated electric field.
        a: Optional updated scale factor.
        H: Optional updated Hubble parameter.

    Returns:
        New `GaugeLattice` instance with updated values.
    """
    result = eqx.tree_at(lambda t: (t.links, t.E), self, (links, E))
    if a is not None:
        result = eqx.tree_at(lambda t: t.a, result, jnp.asarray(a))
    if H is not None:
        result = eqx.tree_at(lambda t: t.H, result, jnp.asarray(H))
    return result

replace(**kwargs)

Create a new lattice with selected dynamic attributes replaced.

Supports replacing links, E, a, and H. To change static fields (g, size, length) construct a new GaugeLattice directly.

Parameters:

Name	Type	Description	Default
**kwargs		Dynamic field values to override.	{}

Returns:

Type	Description
GaugeLattice	New GaugeLattice with replaced attributes.

Source code in jaxlatt/core/lattice.py

def replace(self, **kwargs) -> "GaugeLattice":
    """Create a new lattice with selected dynamic attributes replaced.

    Supports replacing ``links``, ``E``, ``a``, and ``H``.
    To change static fields (``g``, ``size``, ``length``) construct a new
    ``GaugeLattice`` directly.

    Args:
        **kwargs: Dynamic field values to override.

    Returns:
        New `GaugeLattice` with replaced attributes.
    """
    result = self
    for key, value in kwargs.items():
        result = eqx.tree_at(lambda t, k=key: getattr(t, k), result, value)
    return result

Lattice(size, length, field=None, field_dot=None)

Bases: Module

Represents a scalar field on a periodic lattice (supports 1D, 2D, 3D).

Attributes:

Name	Type	Description
field	Array	Scalar field values (shape matches size tuple)
field_dot	Array	Time derivative of the field (same shape)
dx	float	Lattice spacing (assumed uniform, taken from first dimension)
size	tuple[int, ...]	Tuple of lattice grid counts per dimension (always a tuple internally)
length	tuple[float, ...]	Tuple of physical lengths per dimension (always a tuple internally)
ndim	int	Number of spatial dimensions

Initialize a lattice with given size and physical length.

Parameters:

Name	Type	Description	Default
size	int | tuple[int, ...]	Number of grid points (int for 1D, tuple for 2D)	required
length	float | tuple[float, ...]	Physical length (float for 1D, tuple for 2D)	required
field	Array | None	Initial field values (if None, initialized to zeros)	None
field_dot	Array | None	Initial time derivative (if None, initialized to zeros)	None

Source code in jaxlatt/core/lattice.py

def __init__(
    self,
    size: int | tuple[int, ...],
    length: float | tuple[float, ...],
    field: Array | None = None,
    field_dot: Array | None = None,
):
    """
    Initialize a lattice with given size and physical length.

    Args:
        size: Number of grid points (int for 1D, tuple for 2D)
        length: Physical length (float for 1D, tuple for 2D)
        field: Initial field values (if None, initialized to zeros)
        field_dot: Initial time derivative (if None, initialized to zeros)
    """
    # Handle size and length for 1D and 2D
    if isinstance(size, int):
        self.size = (size,)
        self.ndim = 1
    else:
        self.size = size
        self.ndim = len(size)

    if isinstance(length, (int, float)):
        # Replicate scalar length across all spatial dimensions
        if isinstance(size, int):
            self.length = (float(length),)
        else:
            self.length = tuple(float(length) for _ in range(len(size)))
    else:
        # If a tuple was provided but length dimensionality doesn't match size, replicate or raise
        if isinstance(size, tuple) and len(length) == 1 and len(size) > 1:
            self.length = tuple(float(length[0]) for _ in range(len(size)))
        elif isinstance(size, tuple) and len(length) != len(size):
            raise ValueError(
                f"Length tuple dimensionality {len(length)} does not match size {len(size)}"
            )
        else:
            self.length = tuple(float(val) for val in length)

    # Compute lattice spacing
    # Assume uniform spacing; use first dimension for dx
    self.dx = self.length[0] / self.size[0]

    # Initialize fields
    if field is None:
        self.field = jnp.zeros(self.size)
    else:
        self.field = field

    if field_dot is None:
        self.field_dot = jnp.zeros(self.size)
    else:
        self.field_dot = field_dot

volume property

Compute the total physical lattice volume.

Returns:

Type	Description
float	Product of all box lengths.

dV property

Compute the volume element for one grid cell.

Returns:

Type	Description
float	Physical volume associated with a single lattice site.

copy()

Create a deep copy of the lattice state.

Uses type(self) so a subclass copies to its own type rather than being silently downcast to Lattice.

Returns:

Type	Description
Lattice	New instance of the same class, with independent array buffers.

Source code in jaxlatt/core/lattice.py

def copy(self) -> "Lattice":
    """Create a deep copy of the lattice state.

    Uses ``type(self)`` so a subclass copies to its own type rather than
    being silently downcast to `Lattice`.

    Returns:
        New instance of the same class, with independent array buffers.
    """
    return type(self)(
        size=self.size,
        length=self.length,
        field=jnp.copy(self.field),
        field_dot=jnp.copy(self.field_dot),
    )

update(field, field_dot)

Create a new lattice with updated field values.

Uses type(self), so a subclass keeps its own type.

Parameters:

Name	Type	Description	Default
field	Array	New field values	required
field_dot	Array	New field time derivative	required

Returns:

Type	Description
Lattice	New instance of the same class, with updated values

Source code in jaxlatt/core/lattice.py

def update(self, field: Array, field_dot: Array) -> "Lattice":
    """
    Create a new lattice with updated field values.

    Uses ``type(self)``, so a subclass keeps its own type.

    Args:
        field: New field values
        field_dot: New field time derivative

    Returns:
        New instance of the same class, with updated values
    """
    return type(self)(size=self.size, length=self.length, field=field, field_dot=field_dot)

replace(**kwargs)

Create a new lattice with selected dynamic attributes replaced.

Supports the dynamic fields field and field_dot. Static fields (size, length, dx, ndim) are part of the pytree's treedef rather than its leaves, so replacing one raises; construct a new lattice directly instead.

CoupledLattice uses dataclasses.replace for this, which cannot work here: Lattice takes a custom __init__ while dx and ndim are init-eligible dataclass fields, so dataclasses.replace would try to pass them and raise TypeError. GaugeLattice has the same constraint and resolves it the same way, via eqx.tree_at.

Parameters:

Name	Type	Description	Default
**kwargs	Array	Dynamic field values to override.	{}

Returns:

Type	Description
Lattice	New instance of the same class with the given attributes replaced.

Source code in jaxlatt/core/lattice.py

def replace(self, **kwargs: Array) -> "Lattice":
    """Create a new lattice with selected dynamic attributes replaced.

    Supports the dynamic fields ``field`` and ``field_dot``. Static fields
    (``size``, ``length``, ``dx``, ``ndim``) are part of the pytree's
    treedef rather than its leaves, so replacing one raises; construct a new
    lattice directly instead.

    `CoupledLattice` uses ``dataclasses.replace`` for this, which cannot
    work here: `Lattice` takes a custom ``__init__`` while ``dx`` and
    ``ndim`` are init-eligible dataclass fields, so ``dataclasses.replace``
    would try to pass them and raise ``TypeError``. `GaugeLattice` has the
    same constraint and resolves it the same way, via ``eqx.tree_at``.

    Args:
        **kwargs: Dynamic field values to override.

    Returns:
        New instance of the same class with the given attributes replaced.
    """
    if not kwargs:
        return self
    return eqx.tree_at(
        lambda s: tuple(getattr(s, k) for k in kwargs),
        self,
        tuple(kwargs.values()),
    )

ScalarPotential(name, params) dataclass

Unified potential representation with autodiff-derived forces.

This class encapsulates a potential function \(V(\phi)\) along with its parameters, providing automatic computation of forces via JAX autodiff. The frozen dataclass ensures hashability, enabling efficient JIT compilation caching.

Attributes:

Name	Type	Description
name	str	Identifier for the potential type (e.g., "quadratic", "quartic")
params	tuple	Tuple of parameters (must be hashable for caching)

Example

V = ScalarPotential.quartic(m=1.0, lambda_=0.1)
energy_density = V(field)            # V(phi) at each point
total_energy = V.total_energy(field) # sum of V(phi)
force = V.force(field)               # F = -dV/dphi (autodiff)

__call__(field)

Evaluate potential energy density \(V(\phi)\) at each lattice point.

Parameters:

Name	Type	Description	Default
field	Array	Scalar field configuration (real or complex)	required

Returns:

Type	Description
Array	Potential energy density array (same shape as field)

Source code in jaxlatt/core/potentials.py

def __call__(self, field: Array) -> Array:
    """
    Evaluate potential energy density $V(\\phi)$ at each lattice point.

    Args:
        field: Scalar field configuration (real or complex)

    Returns:
        Potential energy density array (same shape as field)
    """
    return _get_potential_fn(self)(field)

total_energy(field)

Compute total potential energy \(\sum V(\phi)\).

Parameters:

Name	Type	Description	Default
field	Array	Scalar field configuration	required

Returns:

Type	Description
Array	Total potential energy (scalar)

Source code in jaxlatt/core/potentials.py

def total_energy(self, field: Array) -> Array:
    """
    Compute total potential energy $\\sum V(\\phi)$.

    Args:
        field: Scalar field configuration

    Returns:
        Total potential energy (scalar)
    """
    return jnp.sum(self(field)).real

force(field)

Compute force \(F = -dV/d\phi\) using autodiff.

For complex fields, computes the Wirtinger derivative \(-dV/d\phi^*\).

Parameters:

Name	Type	Description	Default
field	Array	Scalar field configuration	required

Returns:

Type	Description
Array	Force array (same shape as field)

Source code in jaxlatt/core/potentials.py

def force(self, field: Array) -> Array:
    """
    Compute force $F = -dV/d\\phi$ using autodiff.

    For complex fields, computes the Wirtinger derivative $-dV/d\\phi^*$.

    Args:
        field: Scalar field configuration

    Returns:
        Force array (same shape as field)
    """
    return _get_cached_force(self)(field)

quadratic(m=1.0) staticmethod

Create quadratic potential \(V(\phi) = \frac{1}{2} m^2 |\phi|^2\).

This is the simplest non-trivial potential, commonly used in chaotic inflation models and free field theory.

Parameters:

Name	Type	Description	Default
m	float	Mass parameter	1.0

Returns:

Type	Description
ScalarPotential	ScalarPotential instance

Source code in jaxlatt/core/potentials.py

@staticmethod
def quadratic(m: float = 1.0) -> "ScalarPotential":
    """
    Create quadratic potential $V(\\phi) = \\frac{1}{2} m^2 |\\phi|^2$.

    This is the simplest non-trivial potential, commonly used in
    chaotic inflation models and free field theory.

    Args:
        m: Mass parameter

    Returns:
        `ScalarPotential` instance
    """
    return ScalarPotential(name="quadratic", params=(m,))

quartic(m, lambda_) staticmethod

Create quartic potential \(V(\phi) = \frac{1}{2} m^2 |\phi|^2 + \frac{1}{4} \lambda |\phi|^4\).

Standard scalar field potential with mass term and self-interaction. Used in Higgs-like models and preheating studies.

Parameters:

Name	Type	Description	Default
m	float	Mass parameter	required
lambda_	float	Quartic self-coupling	required

Returns:

Type	Description
ScalarPotential	ScalarPotential instance

Source code in jaxlatt/core/potentials.py

@staticmethod
def quartic(m: float, lambda_: float) -> "ScalarPotential":
    """
    Create quartic potential $V(\\phi) = \\frac{1}{2} m^2 |\\phi|^2 + \\frac{1}{4} \\lambda |\\phi|^4$.

    Standard scalar field potential with mass term and self-interaction.
    Used in Higgs-like models and preheating studies.

    Args:
        m: Mass parameter
        lambda_: Quartic self-coupling

    Returns:
        `ScalarPotential` instance
    """
    return ScalarPotential(name="quartic", params=(m, lambda_))

double_well(mu2=1.0, lam=1.0) staticmethod

Create double-well potential \(V(\phi) = -\frac{1}{2} \mu^2 \phi^2 + \frac{1}{4} \lambda \phi^4\).

Has minima at \(\phi = \pm\sqrt{\mu^2/\lambda}\), used for symmetry breaking studies.

Parameters:

Name	Type	Description	Default
mu2	float	Negative mass squared coefficient	1.0
lam	float	Quartic coupling	1.0

Returns:

Type	Description
ScalarPotential	ScalarPotential instance

Source code in jaxlatt/core/potentials.py

@staticmethod
def double_well(mu2: float = 1.0, lam: float = 1.0) -> "ScalarPotential":
    """
    Create double-well potential $V(\\phi) = -\\frac{1}{2} \\mu^2 \\phi^2 + \\frac{1}{4} \\lambda \\phi^4$.

    Has minima at $\\phi = \\pm\\sqrt{\\mu^2/\\lambda}$, used for symmetry breaking studies.

    Args:
        mu2: Negative mass squared coefficient
        lam: Quartic coupling

    Returns:
        `ScalarPotential` instance
    """
    return ScalarPotential(name="double_well", params=(mu2, lam))

mexican_hat(lam=1.0, v=1.0) staticmethod

Create Mexican hat potential \(V(\phi) = \lambda (|\phi|^2 - v^2)^2\).

Classic symmetry-breaking potential with circular minimum at \(|\phi| = v\).

Parameters:

Name	Type	Description	Default
lam	float	Coupling constant	1.0
v	float	Vacuum expectation value	1.0

Returns:

Type	Description
ScalarPotential	ScalarPotential instance

Source code in jaxlatt/core/potentials.py

@staticmethod
def mexican_hat(lam: float = 1.0, v: float = 1.0) -> "ScalarPotential":
    """
    Create Mexican hat potential $V(\\phi) = \\lambda (|\\phi|^2 - v^2)^2$.

    Classic symmetry-breaking potential with circular minimum at $|\\phi| = v$.

    Args:
        lam: Coupling constant
        v: Vacuum expectation value

    Returns:
        `ScalarPotential` instance
    """
    return ScalarPotential(name="mexican_hat", params=(lam, v))

from_function(potential_fn, name='custom') staticmethod

Create ScalarPotential from an arbitrary function.

Note: Custom functions may not cache as efficiently since they are identified by object id rather than parameters.

Parameters:

Name	Type	Description	Default
potential_fn	PotentialFunction	Function \(V(\phi) \to\) energy density	required
name	str	Identifier for this potential	'custom'

Returns:

Type	Description
ScalarPotential	ScalarPotential instance

Source code in jaxlatt/core/potentials.py

@staticmethod
def from_function(potential_fn: PotentialFunction, name: str = "custom") -> "ScalarPotential":
    """
    Create `ScalarPotential` from an arbitrary function.

    Note: Custom functions may not cache as efficiently since they
    are identified by object id rather than parameters.

    Args:
        potential_fn: Function $V(\\phi) \\to$ energy density
        name: Identifier for this potential

    Returns:
        `ScalarPotential` instance
    """
    # Use id of function as unique identifier
    return ScalarPotential(name=name, params=(id(potential_fn), potential_fn))

RealScalarLattice(size, length, field=None, field_dot=None)

Bases: Lattice

Real scalar field on a periodic lattice (1D / 2D / 3D).

Semantically equivalent to :class:Lattice, but named for a standalone real scalar rather than the gauge-coupled sector's field.

Attributes:

Name	Type	Description
field	Array	chi(x, tau) -- real scalar field values.
field_dot	Array	chi'(x, tau) -- conformal-time velocity.
dx	float	Lattice spacing (uniform in all directions).
size	tuple[int, ...]	Tuple of grid counts per dimension.
length	tuple[float, ...]	Tuple of physical box lengths.
ndim	int	Number of spatial dimensions.

Source code in jaxlatt/core/lattice.py

def __init__(
    self,
    size: int | tuple[int, ...],
    length: float | tuple[float, ...],
    field: Array | None = None,
    field_dot: Array | None = None,
):
    """
    Initialize a lattice with given size and physical length.

    Args:
        size: Number of grid points (int for 1D, tuple for 2D)
        length: Physical length (float for 1D, tuple for 2D)
        field: Initial field values (if None, initialized to zeros)
        field_dot: Initial time derivative (if None, initialized to zeros)
    """
    # Handle size and length for 1D and 2D
    if isinstance(size, int):
        self.size = (size,)
        self.ndim = 1
    else:
        self.size = size
        self.ndim = len(size)

    if isinstance(length, (int, float)):
        # Replicate scalar length across all spatial dimensions
        if isinstance(size, int):
            self.length = (float(length),)
        else:
            self.length = tuple(float(length) for _ in range(len(size)))
    else:
        # If a tuple was provided but length dimensionality doesn't match size, replicate or raise
        if isinstance(size, tuple) and len(length) == 1 and len(size) > 1:
            self.length = tuple(float(length[0]) for _ in range(len(size)))
        elif isinstance(size, tuple) and len(length) != len(size):
            raise ValueError(
                f"Length tuple dimensionality {len(length)} does not match size {len(size)}"
            )
        else:
            self.length = tuple(float(val) for val in length)

    # Compute lattice spacing
    # Assume uniform spacing; use first dimension for dx
    self.dx = self.length[0] / self.size[0]

    # Initialize fields
    if field is None:
        self.field = jnp.zeros(self.size)
    else:
        self.field = field

    if field_dot is None:
        self.field_dot = jnp.zeros(self.size)
    else:
        self.field_dot = field_dot

create_frw_universe(a_initial=1.0, rho_initial=1.0, M_pl=1.0, *, w=1.0 / 3.0)

Initialize a generic FRW universe for a perfect fluid with equation of state \(p = w \rho\).

The first Friedmann equation in conformal time,

\[\left(\frac{a'}{a}\right)^2 = \frac{8\pi G}{3} \rho a^2,\]

fixes the initial \(a'\) from \(\rho\) and \(a\) alone — \(w\) does not enter the initial conditions. Pass the matching pressure function to coupled_evolve_expanding to realize the intended dynamics:

• \(w = 1/3\) (radiation): use make_radiation_pressure
• \(w = 0\) (matter): use make_matter_pressure
• arbitrary: use make_eos_pressure(w)

Parameters:

Name	Type	Description	Default
a_initial	float	Initial scale factor.	1.0
rho_initial	float	Initial energy density.	1.0
M_pl	float	Reduced Planck mass.	1.0
w	float	Equation-of-state parameter \(p/\rho\). Does not affect the returned object; provided for call-site documentation.	1.0 / 3.0

Returns:

Type	Description
FRWUniverse	Initialized FRWUniverse.

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

def create_frw_universe(
    a_initial: float = 1.0,
    rho_initial: float = 1.0,
    M_pl: float = 1.0,
    *,
    w: float = 1.0 / 3.0,
) -> FRWUniverse:
    r"""
    Initialize a generic FRW universe for a perfect fluid with equation of state $p = w \rho$.

    The first Friedmann equation in conformal time,

    $$\left(\frac{a'}{a}\right)^2 = \frac{8\pi G}{3} \rho a^2,$$

    fixes the initial $a'$ from $\rho$ and $a$ alone — $w$ does not enter
    the initial conditions. Pass the matching pressure function to
    ``coupled_evolve_expanding`` to realize the intended dynamics:

    - $w = 1/3$ (radiation): use ``make_radiation_pressure``
    - $w = 0$ (matter): use ``make_matter_pressure``
    - arbitrary: use ``make_eos_pressure(w)``

    Args:
        a_initial: Initial scale factor.
        rho_initial: Initial energy density.
        M_pl: Reduced Planck mass.
        w: Equation-of-state parameter $p/\rho$. Does not affect the returned
            object; provided for call-site documentation.

    Returns:
        Initialized ``FRWUniverse``.
    """
    G = 1.0 / (M_pl * M_pl)
    H_conformal_squared = (8.0 * jnp.pi * G / 3.0) * rho_initial * (a_initial**2)
    adot_initial = a_initial * jnp.sqrt(H_conformal_squared)
    return FRWUniverse(
        a=jnp.asarray(a_initial),
        adot=jnp.asarray(adot_initial),
        tau=jnp.asarray(0.0),
        M_pl=M_pl,
    )

create_matter_universe(a_initial=1.0, rho_initial=1.0, M_pl=1.0)

Initialize FRW universe in matter-dominated era (\(w = 0\)).

See create_frw_universe for full documentation. Pass make_matter_pressure to coupled_evolve_expanding for matter dynamics.

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

def create_matter_universe(
    a_initial: float = 1.0,
    rho_initial: float = 1.0,
    M_pl: float = 1.0,
) -> FRWUniverse:
    """Initialize FRW universe in matter-dominated era ($w = 0$).

    See `create_frw_universe` for full documentation.
    Pass ``make_matter_pressure`` to ``coupled_evolve_expanding`` for matter dynamics.
    """
    return create_frw_universe(a_initial, rho_initial, M_pl, w=0.0)

create_radiation_universe(a_initial=1.0, rho_initial=1.0, M_pl=1.0)

Initialize FRW universe in radiation-dominated era (\(w = 1/3\)).

See create_frw_universe for full documentation. Pass make_radiation_pressure to coupled_evolve_expanding for radiation dynamics.

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

def create_radiation_universe(
    a_initial: float = 1.0,
    rho_initial: float = 1.0,
    M_pl: float = 1.0,
) -> FRWUniverse:
    """Initialize FRW universe in radiation-dominated era ($w = 1/3$).

    See `create_frw_universe` for full documentation.
    Pass ``make_radiation_pressure`` to ``coupled_evolve_expanding`` for radiation dynamics.
    """
    return create_frw_universe(a_initial, rho_initial, M_pl, w=1.0 / 3.0)

friedmann_acceleration(rho, p, M_pl, a)

Compute conformal acceleration \(a''\) from Friedmann equation.

In conformal time \(\tau\) (related to cosmic time \(t\) by \(dt = a(\tau)d\tau\)), the second Friedmann equation is:

\[a'' = \frac{4\pi G}{3} a^3 (\rho - 3p)\]

where:

• \(a'' = d^2a/d\tau^2\)
• \(G = 1/M_{\mathrm{pl}}^2\) (gravitational constant)
• \(\rho\) = energy density
• \(p\) = pressure

Special cases:

• Radiation (\(p = \rho/3\)): \(a'' = 0\) \(\Rightarrow\) \(a(\tau)\) linear in \(\tau\)
• Matter (\(p = 0\)): \(a'' = \frac{4\pi G}{3} a^3 \rho\) \(\Rightarrow\) \(a(\tau) \propto \tau^2\)

References

• Dodelson & Schmidt, "Modern Cosmology" (2nd ed.), Eq. (2.26)
• Mukhanov, "Physical Foundations of Cosmology", Eq. (2.18)
• Baumann, "Cosmology", Eq. (2.35)

Parameters:

Name	Type	Description	Default
rho	float	Total energy density	required
p	float	Total pressure	required
M_pl	float	Reduced Planck mass (default: 1.0 in natural units)	required
a	float	Current scale factor	required

Returns:

Type	Description
Array	Conformal acceleration \(a''\) (second derivative with respect to \(\tau\))

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

@jit
def friedmann_acceleration(rho: float, p: float, M_pl: float, a: float) -> Array:
    """
    Compute conformal acceleration $a''$ from Friedmann equation.

    In conformal time $\\tau$ (related to cosmic time $t$ by $dt = a(\\tau)d\\tau$),
    the second Friedmann equation is:

    $$a'' = \\frac{4\\pi G}{3} a^3 (\\rho - 3p)$$

    where:

    - $a'' = d^2a/d\\tau^2$
    - $G = 1/M_{\\mathrm{pl}}^2$ (gravitational constant)
    - $\\rho$ = energy density
    - $p$ = pressure

    Special cases:

    - Radiation ($p = \\rho/3$): $a'' = 0$ $\\Rightarrow$ $a(\\tau)$ linear in $\\tau$
    - Matter ($p = 0$): $a'' = \\frac{4\\pi G}{3} a^3 \\rho$ $\\Rightarrow$ $a(\\tau) \\propto \\tau^2$

    References:
        - Dodelson & Schmidt, "Modern Cosmology" (2nd ed.), Eq. (2.26)
        - Mukhanov, "Physical Foundations of Cosmology", Eq. (2.18)
        - Baumann, "Cosmology", Eq. (2.35)

    Args:
        rho: Total energy density
        p: Total pressure
        M_pl: Reduced Planck mass (default: 1.0 in natural units)
        a: Current scale factor

    Returns:
        Conformal acceleration $a''$ (second derivative with respect to $\\tau$)
    """
    G = 1.0 / (M_pl * M_pl)
    addot = (4.0 * jnp.pi * G / 3.0) * (rho - 3.0 * p) * (a**3)
    return addot

friedmann_step_leapfrog(universe, rho_func, pressure_func, dt)

Advance scale factor using leapfrog integration (symplectic).

Leapfrog for \(a(\tau)\):

\[a'(\tau + d\tau/2) = a'(\tau) + \frac{d\tau}{2} a''(\tau)\]

\[a(\tau + d\tau) = a(\tau) + d\tau \, a'(\tau + d\tau/2)\]

\[a'(\tau + d\tau) = a'(\tau + d\tau/2) + \frac{d\tau}{2} a''(\tau + d\tau)\]

This preserves symplectic structure if rho_func and pressure_func are Hamiltonian.

Parameters:

Name	Type	Description	Default
universe	FRWUniverse	Current FRWUniverse state	required
rho_func	Callable[[float], float]	Function computing \(\rho(a)\)	required
pressure_func	Callable[[float], float]	Function computing \(p(a)\)	required
dt	float	Conformal timestep \(d\tau\)	required

Returns:

Type	Description
FRWUniverse	Updated FRWUniverse

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

def friedmann_step_leapfrog(
    universe: FRWUniverse,
    rho_func: Callable[[float], float],
    pressure_func: Callable[[float], float],
    dt: float,
) -> FRWUniverse:
    """
    Advance scale factor using leapfrog integration (symplectic).

    Leapfrog for $a(\\tau)$:

    $$a'(\\tau + d\\tau/2) = a'(\\tau) + \\frac{d\\tau}{2} a''(\\tau)$$

    $$a(\\tau + d\\tau) = a(\\tau) + d\\tau \\, a'(\\tau + d\\tau/2)$$

    $$a'(\\tau + d\\tau) = a'(\\tau + d\\tau/2) + \\frac{d\\tau}{2} a''(\\tau + d\\tau)$$

    This preserves symplectic structure if `rho_func` and `pressure_func` are Hamiltonian.

    Args:
        universe: Current FRWUniverse state
        rho_func: Function computing $\\rho(a)$
        pressure_func: Function computing $p(a)$
        dt: Conformal timestep $d\\tau$

    Returns:
        Updated `FRWUniverse`
    """
    a = universe.a
    adot = universe.adot
    tau = universe.tau
    M_pl = universe.M_pl

    # Half-step velocity
    rho = rho_func(a)
    p = pressure_func(a)
    addot = friedmann_acceleration(rho, p, M_pl, a)
    adot_half = adot + 0.5 * dt * addot

    # Full step position
    a_new = a + dt * adot_half

    # Half-step velocity (second half)
    rho_new = rho_func(a_new)
    p_new = pressure_func(a_new)
    addot_new = friedmann_acceleration(rho_new, p_new, M_pl, a_new)
    adot_new = adot_half + 0.5 * dt * addot_new

    tau_new = tau + dt

    return FRWUniverse(a=a_new, adot=adot_new, tau=tau_new, M_pl=M_pl)

friedmann_step_predictor_corrector(universe, rho_func, pressure_func, dt)

Advance scale factor by one timestep using predictor-corrector.

This uses a 2nd-order accurate predictor-corrector scheme:

1. Predictor: Euler step to estimate \(a(\tau + d\tau)\)
2. Corrector: Trapezoidal rule using predicted value

The energy density \(\rho\) and pressure \(p\) are computed from field values via rho_func and pressure_func, which may depend on \(a\) (for rescaled fields).

Parameters:

Name	Type	Description	Default
universe	FRWUniverse	Current FRWUniverse state	required
rho_func	Callable[[float], float]	Function computing total energy density from scale factor. Signature: rho_func(a) -> rho	required
pressure_func	Callable[[float], float]	Function computing total pressure from scale factor. Signature: pressure_func(a) -> p	required
dt	float	Conformal timestep \(d\tau\)	required

Returns:

Type	Description
FRWUniverse	Updated FRWUniverse at \(\tau + d\tau\)

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

def friedmann_step_predictor_corrector(
    universe: FRWUniverse,
    rho_func: Callable[[float], float],
    pressure_func: Callable[[float], float],
    dt: float,
) -> FRWUniverse:
    """
    Advance scale factor by one timestep using predictor-corrector.

    This uses a 2nd-order accurate predictor-corrector scheme:

    1. Predictor: Euler step to estimate $a(\\tau + d\\tau)$
    2. Corrector: Trapezoidal rule using predicted value

    The energy density $\\rho$ and pressure $p$ are computed from field values via
    `rho_func` and `pressure_func`, which may depend on $a$ (for rescaled fields).

    Args:
        universe: Current FRWUniverse state
        rho_func: Function computing total energy density from scale factor.
                  Signature: ``rho_func(a) -> rho``
        pressure_func: Function computing total pressure from scale factor.
                       Signature: ``pressure_func(a) -> p``
        dt: Conformal timestep $d\\tau$

    Returns:
        Updated `FRWUniverse` at $\\tau + d\\tau$
    """

    def rho_p_callable(u: FRWUniverse):
        return rho_func(u.a), pressure_func(u.a)

    return friedmann_step_predictor_corrector_from_state(universe, rho_p_callable, dt)

friedmann_step_predictor_corrector_from_state(universe, rho_p_callable, dt)

Advance scale factor using predictor-corrector with state-dependent ρ, p.

More general than :func:friedmann_step_predictor_corrector: instead of separate rho_func(a) and pressure_func(a) callables, it accepts a single rho_p_callable(universe) → (rho, p). Because the callable receives the whole universe rather than just a, it can close over lattice field state, which is what makes self-consistent field-sourced expansion possible: the fields set ρ and p, which drive a(τ), which in turn evolves the fields.

friedmann_step_predictor_corrector is a thin wrapper around this function, so the two share one integrator and cannot drift apart.

The predictor-corrector scheme is identical to :func:friedmann_step_predictor_corrector; the difference is purely in how ρ and p are evaluated.

Parameters:

Name	Type	Description	Default
universe	FRWUniverse	Current FRWUniverse state.	required
rho_p_callable	Callable[[FRWUniverse], tuple]	Callable (FRWUniverse) → (rho, p).	required
dt	float	Conformal timestep.	required

Returns:

Name	Type	Description
Updated	FRWUniverse	class:FRWUniverse at τ + dt.

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

def friedmann_step_predictor_corrector_from_state(
    universe: FRWUniverse,
    rho_p_callable: Callable[["FRWUniverse"], tuple],
    dt: float,
) -> FRWUniverse:
    """Advance scale factor using predictor-corrector with state-dependent ρ, p.

    More general than :func:`friedmann_step_predictor_corrector`: instead of
    separate ``rho_func(a)`` and ``pressure_func(a)`` callables, it accepts a
    single ``rho_p_callable(universe) → (rho, p)``. Because the callable
    receives the whole universe rather than just ``a``, it can close over
    lattice field state, which is what makes self-consistent field-sourced
    expansion possible: the fields set ρ and p, which drive a(τ), which in turn
    evolves the fields.

    ``friedmann_step_predictor_corrector`` is a thin wrapper around this
    function, so the two share one integrator and cannot drift apart.

    The predictor-corrector scheme is identical to
    :func:`friedmann_step_predictor_corrector`; the difference is purely in
    how ρ and p are evaluated.

    Args:
        universe: Current FRWUniverse state.
        rho_p_callable: Callable ``(FRWUniverse) → (rho, p)``.
        dt: Conformal timestep.

    Returns:
        Updated :class:`FRWUniverse` at τ + dt.
    """
    a = universe.a
    adot = universe.adot
    tau = universe.tau
    M_pl = universe.M_pl

    rho_current, p_current = rho_p_callable(universe)
    addot_current = friedmann_acceleration(rho_current, p_current, M_pl, a)

    a_pred = jnp.maximum(a + dt * adot, 1e-12)
    adot_pred = adot + dt * addot_current

    universe_pred = FRWUniverse(a=a_pred, adot=adot_pred, tau=tau, M_pl=M_pl)
    rho_pred, p_pred = rho_p_callable(universe_pred)
    addot_pred = friedmann_acceleration(rho_pred, p_pred, M_pl, a_pred)

    a_new = jnp.maximum(a + dt * 0.5 * (adot + adot_pred), 1e-12)
    adot_new = adot + dt * 0.5 * (addot_current + addot_pred)
    tau_new = tau + dt
    return FRWUniverse(a=a_new, adot=adot_new, tau=tau_new, M_pl=M_pl)

make_eos_pressure(w, rho_func)

Create pressure function for equation of state \(p = w\rho\).

Parameters:

Name	Type	Description	Default
w	float	Equation of state parameter	required
rho_func	Callable[[float], float]	Energy density function \(\rho(a)\)	required

Returns:

Type	Description
	Function \(p(a) = w\rho(a)\)

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

def make_eos_pressure(w: float, rho_func: Callable[[float], float]):
    """
    Create pressure function for equation of state $p = w\\rho$.

    Args:
        w: Equation of state parameter
        rho_func: Energy density function $\\rho(a)$

    Returns:
        Function $p(a) = w\\rho(a)$
    """

    def pressure_func(a: float) -> float:
        return w * rho_func(a)

    return pressure_func

make_matter_pressure(a=None)

Create matter pressure function \(p = 0\).

Parameters:

Name	Type	Description	Default
a	float	Scale factor (unused, for API consistency)	None

Returns:

Type	Description
	Function \(p(a) = 0\) for matter

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

def make_matter_pressure(a: float = None):
    """
    Create matter pressure function $p = 0$.

    Args:
        a: Scale factor (unused, for API consistency)

    Returns:
        Function $p(a) = 0$ for matter
    """

    def pressure_func(a):
        return 0.0 * jnp.asarray(a)

    return pressure_func

make_radiation_pressure(rho_initial, a_initial)

Create radiation pressure function \(p = \rho/3\) with \(\rho \propto a^{-4}\).

Parameters:

Name	Type	Description	Default
rho_initial	float	Initial energy density	required
a_initial	float	Initial scale factor	required

Returns:

Type	Description
	Function \(p(a)\) for radiation

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

def make_radiation_pressure(rho_initial: float, a_initial: float):
    """
    Create radiation pressure function $p = \\rho/3$ with $\\rho \\propto a^{-4}$.

    Args:
        rho_initial: Initial energy density
        a_initial: Initial scale factor

    Returns:
        Function $p(a)$ for radiation
    """

    def pressure_func(a: float) -> float:
        rho = rho_initial * (a_initial / a) ** 4
        return rho / 3.0

    return pressure_func

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

matter_scaling_check(a_initial, a_final, rho_initial, rho_final)

Check if energy density scales as matter (\(\rho \propto a^{-3}\)).

Parameters:

Name	Type	Description	Default
a_initial	float	Initial scale factor.	required
a_final	float	Final scale factor.	required
rho_initial	float	Initial energy density.	required
rho_final	float	Final energy density.	required

Returns:

Type	Description
float	Tuple (scaling_exponent, relative_error) where scaling_exponent is the
float	measured exponent \(n\) in \(\rho \propto a^{-n}\) and relative_error is
tuple[float, float]	the deviation from \(n = 3\).

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

def matter_scaling_check(
    a_initial: float,
    a_final: float,
    rho_initial: float,
    rho_final: float,
) -> tuple[float, float]:
    """
    Check if energy density scales as matter ($\\rho \\propto a^{-3}$).

    Args:
        a_initial: Initial scale factor.
        a_final: Final scale factor.
        rho_initial: Initial energy density.
        rho_final: Final energy density.

    Returns:
        Tuple `(scaling_exponent, relative_error)` where `scaling_exponent` is the
        measured exponent $n$ in $\\rho \\propto a^{-n}$ and `relative_error` is
        the deviation from $n = 3$.
    """
    return _scaling_check(a_initial, a_final, rho_initial, rho_final, 3.0)

radiation_scaling_check(a_initial, a_final, rho_initial, rho_final)

Check if energy density scales as radiation (\(\rho \propto a^{-4}\)).

For pure radiation:

\[\rho(a) = \rho_0 \left(\frac{a_0}{a}\right)^4\]

Parameters:

Name	Type	Description	Default
a_initial	float	Initial scale factor.	required
a_final	float	Final scale factor.	required
rho_initial	float	Initial energy density.	required
rho_final	float	Final energy density.	required

Returns:

Type	Description
float	Tuple (scaling_exponent, relative_error) where scaling_exponent is the
float	measured exponent \(n\) in \(\rho \propto a^{-n}\) and relative_error is
tuple[float, float]	the deviation from \(n = 4\).

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

def radiation_scaling_check(
    a_initial: float,
    a_final: float,
    rho_initial: float,
    rho_final: float,
) -> tuple[float, float]:
    """
    Check if energy density scales as radiation ($\\rho \\propto a^{-4}$).

    For pure radiation:

    $$\\rho(a) = \\rho_0 \\left(\\frac{a_0}{a}\\right)^4$$

    Args:
        a_initial: Initial scale factor.
        a_final: Final scale factor.
        rho_initial: Initial energy density.
        rho_final: Final energy density.

    Returns:
        Tuple `(scaling_exponent, relative_error)` where `scaling_exponent` is the
        measured exponent $n$ in $\\rho \\propto a^{-n}$ and `relative_error` is
        the deviation from $n = 4$.
    """
    return _scaling_check(a_initial, a_final, rho_initial, rho_final, 4.0)

create_higgs_vev_lattice(size, length, m=1.0, lambda_=1.0, g=1.0, vev_amplitude=None)

Create a coupled lattice with Higgs field at vacuum expectation value.

For the symmetry breaking potential \(V = -\frac{m^2}{2}|\phi|^2 + \frac{\lambda}{4}|\phi|^4\), the VEV is \(|\phi| = \sqrt{m^2/\lambda}\).

Parameters:

Name	Type	Description	Default
size	tuple[int, int, int]	Lattice dimensions (Nx, Ny, Nz).	required
length	float | tuple[float, float, float]	Physical box size.	required
m	float	Scalar mass parameter.	1.0
lambda_	float	Scalar self-coupling.	1.0
g	float	Gauge coupling.	1.0
vev_amplitude	float	Optional VEV amplitude. If None, uses \(\sqrt{m^2/\lambda}\).	None

Returns:

Type	Description
CoupledLattice	CoupledLattice with Higgs field initialized near the VEV.

Source code in jaxlatt/core/fields.py

def create_higgs_vev_lattice(
    size: tuple[int, int, int],
    length: float | tuple[float, float, float],
    m: float = 1.0,
    lambda_: float = 1.0,
    g: float = 1.0,
    vev_amplitude: float = None,
) -> CoupledLattice:
    """
    Create a coupled lattice with Higgs field at vacuum expectation value.

    For the symmetry breaking potential $V = -\\frac{m^2}{2}|\\phi|^2 + \\frac{\\lambda}{4}|\\phi|^4$,
    the VEV is $|\\phi| = \\sqrt{m^2/\\lambda}$.

    Args:
        size: Lattice dimensions `(Nx, Ny, Nz)`.
        length: Physical box size.
        m: Scalar mass parameter.
        lambda_: Scalar self-coupling.
        g: Gauge coupling.
        vev_amplitude: Optional VEV amplitude. If `None`, uses $\\sqrt{m^2/\\lambda}$.

    Returns:
        `CoupledLattice` with Higgs field initialized near the VEV.
    """
    if isinstance(length, (int, float)):
        length = (length, length, length)

    dx = length[0] / size[0]

    # Vacuum expectation value
    if vev_amplitude is None:
        if lambda_ > 0 and m > 0:
            vev_amplitude = jnp.sqrt(m**2 / lambda_)
        else:
            vev_amplitude = 0.0

    # Constant field at VEV (real for simplicity)
    phi = jnp.full(size, vev_amplitude, dtype=jnp.complex64)
    pi = jnp.zeros(size, dtype=jnp.complex64)
    links = jnp.ones((3,) + size, dtype=jnp.complex64)
    E = jnp.zeros((3,) + size, dtype=jnp.float32)

    return CoupledLattice(
        phi=phi,
        pi=pi,
        links=links,
        E=E,
        m=m,
        lambda_=lambda_,
        g=g,
        dx=dx,
        size=size,
        length=length,
    )

create_random_coupled_lattice(key, size, length, m=1.0, lambda_=1.0, g=1.0, amplitude=0.1)

Create a coupled lattice with random initial conditions.

Parameters:

Name	Type	Description	Default
key	PRNGKey	Random number generator key.	required
size	tuple[int, int, int]	Lattice dimensions (Nx, Ny, Nz).	required
length	float | tuple[float, float, float]	Physical box size.	required
m	float	Scalar mass.	1.0
lambda_	float	Scalar self-coupling.	1.0
g	float	Gauge coupling.	1.0
amplitude	float	Amplitude of random perturbations.	0.1

Returns:

Type	Description
CoupledLattice	CoupledLattice with random initial fields.

Source code in jaxlatt/core/fields.py

def create_random_coupled_lattice(
    key: random.PRNGKey,
    size: tuple[int, int, int],
    length: float | tuple[float, float, float],
    m: float = 1.0,
    lambda_: float = 1.0,
    g: float = 1.0,
    amplitude: float = 0.1,
) -> CoupledLattice:
    """
    Create a coupled lattice with random initial conditions.

    Args:
        key: Random number generator key.
        size: Lattice dimensions `(Nx, Ny, Nz)`.
        length: Physical box size.
        m: Scalar mass.
        lambda_: Scalar self-coupling.
        g: Gauge coupling.
        amplitude: Amplitude of random perturbations.

    Returns:
        `CoupledLattice` with random initial fields.
    """
    if isinstance(length, (int, float)):
        length = (length, length, length)

    dx = length[0] / size[0]

    # Split random key
    k1, k2, k3, k4, k5, k6 = random.split(key, 6)

    # Random scalar field (complex)
    phi_re = amplitude * random.normal(k1, size, dtype=jnp.float32)
    phi_im = amplitude * random.normal(k2, size, dtype=jnp.float32)
    phi = phi_re + 1j * phi_im

    # Random conjugate momentum (complex)
    pi_re = amplitude * random.normal(k3, size, dtype=jnp.float32)
    pi_im = amplitude * random.normal(k4, size, dtype=jnp.float32)
    pi = pi_re + 1j * pi_im

    # Random gauge links (phases on unit circle)
    theta = amplitude * random.normal(k5, (3,) + size, dtype=jnp.float32)
    links = jnp.exp(1j * theta)

    # Random electric field
    E = amplitude * random.normal(k6, (3,) + size, dtype=jnp.float32)

    return CoupledLattice(
        phi=phi,
        pi=pi,
        links=links,
        E=E,
        m=m,
        lambda_=lambda_,
        g=g,
        dx=dx,
        size=size,
        length=length,
    )

create_random_gauge_lattice(key, size, length, g=1.0, amplitude=0.1)

Create a gauge lattice with small random perturbations.

Initializes links with small random phases and electric field with small random values, suitable for testing dynamics.

Parameters:

Name	Type	Description	Default
key	PRNGKey	JAX random key	required
size	tuple[int, int, int]	Grid dimensions	required
length	float	Physical box size	required
g	float	Gauge coupling	1.0
amplitude	float	Amplitude of random fluctuations	0.1

Returns:

Type	Description
GaugeLattice	GaugeLattice with random initial conditions

Source code in jaxlatt/core/fields.py

def create_random_gauge_lattice(
    key: random.PRNGKey,
    size: tuple[int, int, int],
    length: float,
    g: float = 1.0,
    amplitude: float = 0.1,
) -> GaugeLattice:
    """
    Create a gauge lattice with small random perturbations.

    Initializes links with small random phases and electric field with
    small random values, suitable for testing dynamics.

    Args:
        key: JAX random key
        size: Grid dimensions
        length: Physical box size
        g: Gauge coupling
        amplitude: Amplitude of random fluctuations

    Returns:
        GaugeLattice with random initial conditions
    """
    key1, key2 = random.split(key)

    # Random phases for links (small perturbation around identity)
    phases = amplitude * random.normal(key1, shape=(3, *size))
    links = jnp.exp(1j * phases).astype(jnp.result_type(complex))

    # Random electric field
    E = amplitude * random.normal(key2, shape=(3, *size)).astype(jnp.result_type(float))

    return GaugeLattice(size=size, length=length, g=g, links=links, E=E)

create_random_real_scalar_lattice(key, size, length, amplitude=0.001)

Initialise chi with small random perturbations, chi' = 0.

Parameters:

Name	Type	Description	Default
key	Array	JAX PRNG key.	required
size	tuple[int, ...]	Grid dimensions.	required
length	float	Physical side length.	required
amplitude	float	Standard deviation of the Gaussian perturbation.	0.001

Returns:

Type	Description
RealScalarLattice	class:RealScalarLattice with random field and zero field_dot.

Source code in jaxlatt/core/fields.py

def create_random_real_scalar_lattice(
    key: Array,
    size: tuple[int, ...],
    length: float,
    amplitude: float = 1e-3,
) -> RealScalarLattice:
    """Initialise chi with small random perturbations, chi' = 0.

    Args:
        key: JAX PRNG key.
        size: Grid dimensions.
        length: Physical side length.
        amplitude: Standard deviation of the Gaussian perturbation.

    Returns:
        :class:`RealScalarLattice` with random ``field`` and zero ``field_dot``.
    """
    # astype matches the gauge factories above: random.normal's default dtype
    # ignores jax_enable_x64, so without this the field is float32 even in a
    # run that has explicitly asked for double precision.
    field = amplitude * random.normal(key, shape=size).astype(jnp.result_type(float))
    return RealScalarLattice(size=size, length=length, field=field)

create_vacuum_coupled_lattice(size, length, m=1.0, lambda_=1.0, g=1.0)

Create a coupled lattice in vacuum state.

Vacuum: \(\phi = 0, \pi = 0, U = 1, E = 0\)

Parameters:

Name	Type	Description	Default
size	tuple[int, int, int]	Lattice dimensions (Nx, Ny, Nz).	required
length	float | tuple[float, float, float]	Physical box size.	required
m	float	Scalar mass.	1.0
lambda_	float	Scalar self-coupling.	1.0
g	float	Gauge coupling.	1.0

Returns:

Type	Description
CoupledLattice	CoupledLattice in vacuum state.

Source code in jaxlatt/core/fields.py

def create_vacuum_coupled_lattice(
    size: tuple[int, int, int],
    length: float | tuple[float, float, float],
    m: float = 1.0,
    lambda_: float = 1.0,
    g: float = 1.0,
) -> CoupledLattice:
    """
    Create a coupled lattice in vacuum state.

    Vacuum: $\\phi = 0, \\pi = 0, U = 1, E = 0$

    Args:
        size: Lattice dimensions `(Nx, Ny, Nz)`.
        length: Physical box size.
        m: Scalar mass.
        lambda_: Scalar self-coupling.
        g: Gauge coupling.

    Returns:
        `CoupledLattice` in vacuum state.
    """
    if isinstance(length, (int, float)):
        length = (length, length, length)

    dx = length[0] / size[0]

    phi = jnp.zeros(size, dtype=jnp.complex64)
    pi = jnp.zeros(size, dtype=jnp.complex64)
    links = jnp.ones((3,) + size, dtype=jnp.complex64)
    E = jnp.zeros((3,) + size, dtype=jnp.float32)

    return CoupledLattice(
        phi=phi,
        pi=pi,
        links=links,
        E=E,
        m=m,
        lambda_=lambda_,
        g=g,
        dx=dx,
        size=size,
        length=length,
    )

create_vacuum_gauge_lattice(size, length, g=1.0)

Create a gauge lattice in vacuum state (links=1, E=0).

Parameters:

Name	Type	Description	Default
size	tuple[int, int, int]	Grid dimensions	required
length	float	Physical box size	required
g	float	Gauge coupling	1.0

Returns:

Type	Description
GaugeLattice	GaugeLattice in vacuum configuration

Source code in jaxlatt/core/fields.py

def create_vacuum_gauge_lattice(
    size: tuple[int, int, int],
    length: float,
    g: float = 1.0,
) -> GaugeLattice:
    """
    Create a gauge lattice in vacuum state (links=1, E=0).

    Args:
        size: Grid dimensions
        length: Physical box size
        g: Gauge coupling

    Returns:
        GaugeLattice in vacuum configuration
    """
    return GaugeLattice(size=size, length=length, g=g)

create_vacuum_real_scalar_lattice(size, length)

Initialise chi = 0, chi' = 0 (vacuum) on a periodic lattice.

Parameters:

Name	Type	Description	Default
size	tuple[int, ...]	Grid dimensions, e.g. (32, 32, 32) for a cubic 3-D box.	required
length	float	Physical side length (same for all dimensions).	required

Returns:

Type	Description
RealScalarLattice	Zero-initialised :class:RealScalarLattice.

Source code in jaxlatt/core/fields.py

def create_vacuum_real_scalar_lattice(
    size: tuple[int, ...],
    length: float,
) -> RealScalarLattice:
    """Initialise chi = 0, chi' = 0 (vacuum) on a periodic lattice.

    Args:
        size: Grid dimensions, e.g. ``(32, 32, 32)`` for a cubic 3-D box.
        length: Physical side length (same for all dimensions).

    Returns:
        Zero-initialised :class:`RealScalarLattice`.
    """
    return RealScalarLattice(size=size, length=length)