Spectral Operators

spectral

FFT-based spatial derivative operators for lattice fields.

This module implements spatial derivatives using spectral methods (FFT), which are efficient and accurate for periodic boundary conditions.

Supported dimensions: 1D, 2D, and 3D.

Nyquist convention

First derivatives drop the Nyquist plane of every even-length axis; second derivatives keep it. The asymmetry is required, not stylistic.

At \(p = N/2\) the momentum has no well-defined sign: \(-N/2\) and \(+N/2\) are the same grid index, so \(k(-k) = -k(k)\) cannot hold there. The first-derivative symbol \(ik\) is odd in \(k\), so keeping that unpaired mode breaks Hermiticity and makes the derivative of a real field complex. The Laplacian symbol \(-k^2\) is even, so Hermiticity already holds and the Nyquist mode is legitimately resolved. Odd \(N\) has no exact Nyquist and is untouched.

Dtype contract

These operators preserve the input dtype: a real field yields a real result, a complex field yields a complex one. Do not reinstate an unconditional .real — that silently truncates complex fields such as CoupledLattice.phi. For real input it is the Nyquist zeroing above, not the cast, that keeps the result real.

laplacian_1d(field, dx)

Compute the Laplacian of a 1D field using FFT.

For periodic boundary conditions, the Laplacian in Fourier space is:

\[\nabla^2\phi(x) = \mathrm{IFFT}[-k^2 \cdot \mathrm{FFT}(\phi(x))]\]

The Nyquist mode is retained: \(-k^2\) is even in \(k\), so Hermiticity holds and the mode is legitimately resolved (see the module docstring).

Parameters:

Name	Type	Description	Default
field	Array	1D array of field values	required
dx	float	Lattice spacing	required

Returns:

Type	Description
Array	Laplacian of the field, in the input's dtype

Source code in jaxlatt/operators/spectral.py

@jit
def laplacian_1d(field: Array, dx: float) -> Array:
    """
    Compute the Laplacian of a 1D field using FFT.

    For periodic boundary conditions, the Laplacian in Fourier space is:

    $$\\nabla^2\\phi(x) = \\mathrm{IFFT}[-k^2 \\cdot \\mathrm{FFT}(\\phi(x))]$$

    The Nyquist mode is retained: $-k^2$ is even in $k$, so Hermiticity holds
    and the mode is legitimately resolved (see the module docstring).

    Args:
        field: 1D array of field values
        dx: Lattice spacing

    Returns:
        Laplacian of the field, in the input's dtype
    """
    n = field.shape[0]
    # Compute k-space coordinates
    k = jnp.fft.fftfreq(n, d=dx) * 2 * jnp.pi

    # Transform to Fourier space
    field_k = jnp.fft.fft(field)

    # Apply Laplacian operator: -k²
    laplacian_k = -(k**2) * field_k

    # Transform back to real space
    laplacian = _preserve_dtype(jnp.fft.ifft(laplacian_k), field)

    return laplacian

laplacian_2d(field, dx)

Compute the Laplacian of a 2D field using FFT.

For periodic boundary conditions:

\[\nabla^2\phi(x,y) = \mathrm{IFFT}[-(k_x^2 + k_y^2) \cdot \mathrm{FFT}(\phi(x,y))]\]

Parameters:

Name	Type	Description	Default
field	Array	2D array of field values	required
dx	float	Lattice spacing (assumed equal in both directions)	required

Returns:

Type	Description
Array	Laplacian of the field, in the input's dtype

Source code in jaxlatt/operators/spectral.py

@jit
def laplacian_2d(field: Array, dx: float) -> Array:
    """
    Compute the Laplacian of a 2D field using FFT.

    For periodic boundary conditions:

    $$\\nabla^2\\phi(x,y) = \\mathrm{IFFT}[-(k_x^2 + k_y^2) \\cdot \\mathrm{FFT}(\\phi(x,y))]$$

    Args:
        field: 2D array of field values
        dx: Lattice spacing (assumed equal in both directions)

    Returns:
        Laplacian of the field, in the input's dtype
    """
    result = jnp.fft.ifft2(-_k_squared(field.shape, dx) * jnp.fft.fft2(field))
    return _preserve_dtype(result, field)

laplacian_3d(field, dx)

Compute the Laplacian of a 3D field using FFT.

\[\nabla^2\phi(x,y,z) = \mathrm{IFFTN}[-(k_x^2 + k_y^2 + k_z^2) \cdot \mathrm{FFTN}(\phi)]\]

Parameters:

Name	Type	Description	Default
field	Array	3D array of field values	required
dx	float	Lattice spacing (assumed equal in all directions)	required

Returns:

Type	Description
Array	Laplacian of the field, in the input's dtype

Source code in jaxlatt/operators/spectral.py

@jit
def laplacian_3d(field: Array, dx: float) -> Array:
    """
    Compute the Laplacian of a 3D field using FFT.

    $$\\nabla^2\\phi(x,y,z) = \\mathrm{IFFTN}[-(k_x^2 + k_y^2 + k_z^2) \\cdot \\mathrm{FFTN}(\\phi)]$$

    Args:
        field: 3D array of field values
        dx: Lattice spacing (assumed equal in all directions)

    Returns:
        Laplacian of the field, in the input's dtype
    """
    result = jnp.fft.ifftn(-_k_squared(field.shape, dx) * jnp.fft.fftn(field))
    return _preserve_dtype(result, field)

laplacian(field, dx)

Compute the Laplacian of a field (supports 1D, 2D, 3D).

Parameters:

Name	Type	Description	Default
field	Array	Array of field values	required
dx	float	Lattice spacing	required

Returns:

Type	Description
Array	Laplacian of the field

Source code in jaxlatt/operators/spectral.py

def laplacian(field: Array, dx: float) -> Array:
    """
    Compute the Laplacian of a field (supports 1D, 2D, 3D).

    Args:
        field: Array of field values
        dx: Lattice spacing

    Returns:
        Laplacian of the field
    """
    if field.ndim == 1:
        return laplacian_1d(field, dx)
    elif field.ndim == 2:
        return laplacian_2d(field, dx)
    elif field.ndim == 3:
        return laplacian_3d(field, dx)
    else:
        raise ValueError(f"laplacian only supports 1D, 2D, 3D fields, got {field.ndim}D")

gradient_1d(field, dx)

Compute the gradient of a 1D field using FFT.

The Nyquist mode of an even-length axis is dropped, because \(ik\) is odd in \(k\) (see the module docstring).

Parameters:

Name	Type	Description	Default
field	Array	1D array of field values	required
dx	float	Lattice spacing	required

Returns:

Type	Description
Array	Gradient of the field, in the input's dtype

Source code in jaxlatt/operators/spectral.py

@jit
def gradient_1d(field: Array, dx: float) -> Array:
    """
    Compute the gradient of a 1D field using FFT.

    The Nyquist mode of an even-length axis is dropped, because $ik$ is odd in
    $k$ (see the module docstring).

    Args:
        field: 1D array of field values
        dx: Lattice spacing

    Returns:
        Gradient of the field, in the input's dtype
    """
    (k,) = _k_vectors(field.shape, dx, drop_nyquist=True)

    field_k = jnp.fft.fft(field)
    gradient_k = 1j * k * field_k
    gradient = _preserve_dtype(jnp.fft.ifft(gradient_k), field)

    return gradient

gradient_2d(field, dx)

Compute the gradient of a 2D field using FFT.

The Nyquist plane of each even-length axis is dropped, because \(ik\) is odd in \(k\) (see the module docstring).

Parameters:

Name	Type	Description	Default
field	Array	2D array of field values	required
dx	float	Lattice spacing	required

Returns:

Type	Description
tuple[Array, Array]	Tuple of (gradient_x, gradient_y), in the input's dtype

Source code in jaxlatt/operators/spectral.py

@jit
def gradient_2d(field: Array, dx: float) -> tuple[Array, Array]:
    """
    Compute the gradient of a 2D field using FFT.

    The Nyquist plane of each even-length axis is dropped, because $ik$ is odd
    in $k$ (see the module docstring).

    Args:
        field: 2D array of field values
        dx: Lattice spacing

    Returns:
        Tuple of (gradient_x, gradient_y), in the input's dtype
    """
    KX, KY = _k_vectors(field.shape, dx, drop_nyquist=True)
    field_k = jnp.fft.fft2(field)
    return (
        _preserve_dtype(jnp.fft.ifft2(1j * KX * field_k), field),
        _preserve_dtype(jnp.fft.ifft2(1j * KY * field_k), field),
    )

gradient_energy_1d(field, dx)

Compute the gradient energy \(\frac{1}{2}\int(\nabla\phi)^2 \, dV\) for a 1D field.

Parameters:

Name	Type	Description	Default
field	Array	1D array of field values	required
dx	float	Lattice spacing	required

Returns:

Type	Description
Array	Total gradient energy (always real, including for a complex field)

Source code in jaxlatt/operators/spectral.py

@jit
def gradient_energy_1d(field: Array, dx: float) -> Array:
    """
    Compute the gradient energy $\\frac{1}{2}\\int(\\nabla\\phi)^2 \\, dV$ for a 1D field.

    Args:
        field: 1D array of field values
        dx: Lattice spacing

    Returns:
        Total gradient energy (always real, including for a complex field)
    """
    grad = gradient_1d(field, dx)
    return 0.5 * jnp.sum(jnp.abs(grad) ** 2) * dx

gradient_energy_2d(field, dx)

Compute the gradient energy \(\frac{1}{2}\int(\nabla\phi)^2 \, dV\) for a 2D field.

Parameters:

Name	Type	Description	Default
field	Array	2D array of field values	required
dx	float	Lattice spacing	required

Returns:

Type	Description
Array	Total gradient energy (always real, including for a complex field)

Source code in jaxlatt/operators/spectral.py

@jit
def gradient_energy_2d(field: Array, dx: float) -> Array:
    """
    Compute the gradient energy $\\frac{1}{2}\\int(\\nabla\\phi)^2 \\, dV$ for a 2D field.

    Args:
        field: 2D array of field values
        dx: Lattice spacing

    Returns:
        Total gradient energy (always real, including for a complex field)
    """
    grad_x, grad_y = gradient_2d(field, dx)
    return 0.5 * jnp.sum(jnp.abs(grad_x) ** 2 + jnp.abs(grad_y) ** 2) * dx**2

gradient_3d(field, dx)

Compute the gradient of a 3D field using FFT.

The Nyquist plane of each even-length axis is dropped, because \(ik\) is odd in \(k\) (see the module docstring).

Parameters:

Name	Type	Description	Default
field	Array	3D array of field values	required
dx	float	Lattice spacing	required

Returns:

Type	Description
Array	Tuple \((\partial\phi/\partial x, \partial\phi/\partial y, \partial\phi/\partial z)\),
Array	in the input's dtype

Source code in jaxlatt/operators/spectral.py

@jit
def gradient_3d(field: Array, dx: float) -> tuple[Array, Array, Array]:
    """
    Compute the gradient of a 3D field using FFT.

    The Nyquist plane of each even-length axis is dropped, because $ik$ is odd
    in $k$ (see the module docstring).

    Args:
        field: 3D array of field values
        dx: Lattice spacing

    Returns:
        Tuple $(\\partial\\phi/\\partial x, \\partial\\phi/\\partial y, \\partial\\phi/\\partial z)$,
        in the input's dtype
    """
    KX, KY, KZ = _k_vectors(field.shape, dx, drop_nyquist=True)
    field_k = jnp.fft.fftn(field)
    return (
        _preserve_dtype(jnp.fft.ifftn(1j * KX * field_k), field),
        _preserve_dtype(jnp.fft.ifftn(1j * KY * field_k), field),
        _preserve_dtype(jnp.fft.ifftn(1j * KZ * field_k), field),
    )

gradient_energy_3d(field, dx)

Compute 3D gradient energy using FFT-based derivatives.

Parameters:

Name	Type	Description	Default
field	Array	Three-dimensional scalar field.	required
dx	float	Lattice spacing.	required

Returns:

Type	Description
Array	Total gradient energy including the 1/2 prefactor. Always real,
Array	including for a complex field.

Source code in jaxlatt/operators/spectral.py

@jit
def gradient_energy_3d(field: Array, dx: float) -> Array:
    """Compute 3D gradient energy using FFT-based derivatives.

    Args:
        field: Three-dimensional scalar field.
        dx: Lattice spacing.

    Returns:
        Total gradient energy including the `1/2` prefactor. Always real,
        including for a complex field.
    """
    gx, gy, gz = gradient_3d(field, dx)
    squared = jnp.abs(gx) ** 2 + jnp.abs(gy) ** 2 + jnp.abs(gz) ** 2
    return 0.5 * jnp.sum(squared) * dx**3

gradient_energy(field, dx)

Compute the gradient energy \(\frac{1}{2}\int(\nabla\phi)^2 \, dV\) (supports 1D, 2D, 3D).

Parameters:

Name	Type	Description	Default
field	Array	Array of field values	required
dx	float	Lattice spacing	required

Returns:

Type	Description
Array	Total gradient energy including 1/2 factor

Source code in jaxlatt/operators/spectral.py

def gradient_energy(field: Array, dx: float) -> Array:
    """
    Compute the gradient energy $\\frac{1}{2}\\int(\\nabla\\phi)^2 \\, dV$ (supports 1D, 2D, 3D).

    Args:
        field: Array of field values
        dx: Lattice spacing

    Returns:
        Total gradient energy including 1/2 factor
    """
    if field.ndim == 1:
        return gradient_energy_1d(field, dx)
    elif field.ndim == 2:
        return gradient_energy_2d(field, dx)
    elif field.ndim == 3:
        return gradient_energy_3d(field, dx)
    else:
        raise ValueError("gradient_energy only supports 1D, 2D, 3D fields")