Lattice

lattice

Core lattice structures for scalar and gauge fields.

This module defines the fundamental lattice data structures:

• Lattice: Simple scalar field (1D, 2D, 3D)
• CoupledLattice: Scalar + U(1) gauge (Abelian Higgs model)
• GaugeLattice: Pure U(1) gauge field

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()),
    )

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