Spectra

spectra

Power spectrum calculations for scalar fields on lattices.

This module provides functions to compute power spectra \(P(k)\) from field configurations in 1D, 2D, and 3D. The power spectrum measures the distribution of power across different wavenumbers \(k\), which is fundamental for understanding field fluctuations and correlations.

Physical interpretation:

• \(P(k)\) gives the mean square field amplitude at wavenumber \(k\)
• For Gaussian random fields: \(\langle \phi(k) \phi^*(k') \rangle = (2\pi)^d \delta(k - k') P(k)\)
• In expanding universes: \(P(k,t)\) tracks growth of structure

Usage

from jaxlatt.observables import power_spectrum_3d, bin_power_spectrum
field = ...  # 3D field configuration
k_values, P_k = power_spectrum_3d(field, L)
k_binned, P_binned, counts = bin_power_spectrum(k_values, P_k, num_bins=50)

power_spectrum_1d(field, L)

Compute 1D power spectrum \(P(k)\) from a field configuration.

The power spectrum is defined as:

\[P(k) = (dx)^2 |\tilde{\phi}(k)|^2\]

where \(\tilde{\phi}(k)\) is the DFT of the field and \(dx = L/N\) is the lattice spacing. This normalization ensures Parseval's theorem: \(\int (dk/2\pi)\,P(k) = \int dx\,|\phi|^2\).

Parameters:

Name	Type	Description	Default
field	Array	1D array of field values, shape (N,)	required
L	float	Physical length of the lattice	required

Returns:

Type	Description
Array	Tuple of (k_values, P_k) where:
Array	k_values: Array of wavenumbers (positive frequencies only)
tuple[Array, Array]	P_k: Power spectrum values at each k

Example

field = jnp.array([...])  # 1D field
k, P = power_spectrum_1d(field, L=10.0)
# Plot: plt.loglog(k, P)

Source code in jaxlatt/observables/spectra.py

def power_spectrum_1d(field: Array, L: float) -> tuple[Array, Array]:
    r"""
    Compute 1D power spectrum $P(k)$ from a field configuration.

    The power spectrum is defined as:

    $$P(k) = (dx)^2 |\tilde{\phi}(k)|^2$$

    where $\tilde{\phi}(k)$ is the DFT of the field and $dx = L/N$ is the lattice spacing.
    This normalization ensures Parseval's theorem: $\int (dk/2\pi)\,P(k) = \int dx\,|\phi|^2$.

    Args:
        field: 1D array of field values, shape (N,)
        L: Physical length of the lattice

    Returns:
        Tuple of (k_values, P_k) where:
        - k_values: Array of wavenumbers (positive frequencies only)
        - P_k: Power spectrum values at each k

    Example:
        ```python
        field = jnp.array([...])  # 1D field
        k, P = power_spectrum_1d(field, L=10.0)
        # Plot: plt.loglog(k, P)
        ```
    """
    N = field.shape[0]

    # FFT of field
    field_k = jnp.fft.rfft(field)

    # Wavenumbers (positive frequencies only for rfft)
    k_values = jnp.fft.rfftfreq(N, d=L / N) * 2 * jnp.pi

    # Power spectrum: (dx)^2 |phi_k|^2 with dx = L/N
    # Factor of 2 for k>0 to account for negative frequencies
    # (except Nyquist for even N, which has no negative frequency partner)
    dx = L / N
    P_k = jnp.abs(field_k) ** 2 * dx**2

    # Double for positive frequencies (exclude k=0 and Nyquist if present)
    if N % 2 == 0:
        # Even N: indices 0, 1, ..., N/2-1, N/2 (Nyquist)
        # Double 1 to N/2-1, but not 0 or N/2
        P_k = P_k.at[1:-1].multiply(2)
    else:
        # Odd N: no Nyquist, double all k>0
        P_k = P_k.at[1:].multiply(2)

    return k_values, P_k

power_spectrum_2d(field, L)

Compute 2D power spectrum P(k_x, k_y) from a field configuration.

Returns the full 2D spectrum in k-space. For radially averaged spectrum, use bin_power_spectrum() on the output.

Parameters:

Name	Type	Description	Default
field	Array	2D array of field values, shape (N, N)	required
L	float	Physical length of the lattice (assumed square)	required

Returns:

Type	Description
Array	Tuple of (k_x, k_y, P_k) where:
Array	k_x: 2D array of x-wavenumbers
Array	k_y: 2D array of y-wavenumbers
tuple[Array, Array, Array]	P_k: 2D power spectrum values

Example

field = jnp.array([...])  # 2D field
kx, ky, P = power_spectrum_2d(field, L=20.0)
# Radial average:
k_mag = jnp.sqrt(kx**2 + ky**2)
k_binned, P_binned, counts = bin_power_spectrum(
    k_mag.flatten(), P.flatten(), num_bins=50
)

Source code in jaxlatt/observables/spectra.py

def power_spectrum_2d(field: Array, L: float) -> tuple[Array, Array, Array]:
    """
    Compute 2D power spectrum P(k_x, k_y) from a field configuration.

    Returns the full 2D spectrum in k-space. For radially averaged spectrum,
    use bin_power_spectrum() on the output.

    Args:
        field: 2D array of field values, shape (N, N)
        L: Physical length of the lattice (assumed square)

    Returns:
        Tuple of (k_x, k_y, P_k) where:
        - k_x: 2D array of x-wavenumbers
        - k_y: 2D array of y-wavenumbers
        - P_k: 2D power spectrum values

    Example:
        ```python
        field = jnp.array([...])  # 2D field
        kx, ky, P = power_spectrum_2d(field, L=20.0)
        # Radial average:
        k_mag = jnp.sqrt(kx**2 + ky**2)
        k_binned, P_binned, counts = bin_power_spectrum(
            k_mag.flatten(), P.flatten(), num_bins=50
        )
        ```
    """
    N = field.shape[0]

    # FFT of field
    field_k = jnp.fft.fft2(field)

    # Wavenumbers
    kx_1d = jnp.fft.fftfreq(N, d=L / N) * 2 * jnp.pi
    ky_1d = jnp.fft.fftfreq(N, d=L / N) * 2 * jnp.pi
    k_x, k_y = jnp.meshgrid(kx_1d, ky_1d, indexing="ij")

    # Power spectrum: (dx)^4 |phi_k|^2 with dx = L/N
    # This normalization ensures Parseval's theorem
    dx = L / N
    P_k = jnp.abs(field_k) ** 2 * dx**4

    return k_x, k_y, P_k

power_spectrum_3d(field, L)

Compute 3D power spectrum P(k_x, k_y, k_z) from a field configuration.

Returns the full 3D spectrum in k-space. For spherically averaged spectrum, use bin_power_spectrum() on the output.

Parameters:

Name	Type	Description	Default
field	Array	3D array of field values, shape (N, N, N)	required
L	float	Physical length of the lattice (assumed cubic)	required

Returns:

Type	Description
Array	Tuple of (k_x, k_y, k_z, P_k) where:
Array	k_x, k_y, k_z: 3D arrays of wavenumber components
Array	P_k: 3D power spectrum values

Example

field = jnp.array([...])  # 3D field
kx, ky, kz, P = power_spectrum_3d(field, L=20.0)
# Spherical average:
k_mag = jnp.sqrt(kx**2 + ky**2 + kz**2)
k_binned, P_binned, counts = bin_power_spectrum(
    k_mag.flatten(), P.flatten(), num_bins=50
)

Source code in jaxlatt/observables/spectra.py

def power_spectrum_3d(field: Array, L: float) -> tuple[Array, Array, Array, Array]:
    """
    Compute 3D power spectrum P(k_x, k_y, k_z) from a field configuration.

    Returns the full 3D spectrum in k-space. For spherically averaged spectrum,
    use bin_power_spectrum() on the output.

    Args:
        field: 3D array of field values, shape (N, N, N)
        L: Physical length of the lattice (assumed cubic)

    Returns:
        Tuple of (k_x, k_y, k_z, P_k) where:
        - k_x, k_y, k_z: 3D arrays of wavenumber components
        - P_k: 3D power spectrum values

    Example:
        ```python
        field = jnp.array([...])  # 3D field
        kx, ky, kz, P = power_spectrum_3d(field, L=20.0)
        # Spherical average:
        k_mag = jnp.sqrt(kx**2 + ky**2 + kz**2)
        k_binned, P_binned, counts = bin_power_spectrum(
            k_mag.flatten(), P.flatten(), num_bins=50
        )
        ```
    """
    N = field.shape[0]

    # FFT of field
    field_k = jnp.fft.fftn(field)

    # Wavenumbers
    kx_1d = jnp.fft.fftfreq(N, d=L / N) * 2 * jnp.pi
    ky_1d = jnp.fft.fftfreq(N, d=L / N) * 2 * jnp.pi
    kz_1d = jnp.fft.fftfreq(N, d=L / N) * 2 * jnp.pi
    k_x, k_y, k_z = jnp.meshgrid(kx_1d, ky_1d, kz_1d, indexing="ij")

    # Power spectrum: (dx)^6 |phi_k|^2 with dx = L/N
    # This normalization ensures Parseval's theorem
    dx = L / N
    P_k = jnp.abs(field_k) ** 2 * dx**6

    return k_x, k_y, k_z, P_k

power_spectrum(field, L)

Compute power spectrum automatically based on field dimensionality.

Dispatches to the appropriate 1D/2D/3D function based on field shape.

Parameters:

Name	Type	Description	Default
field	Array	Field array (1D, 2D, or 3D)	required
L	float | tuple[float, ...]	Physical length (scalar or tuple for anisotropic boxes)	required

Returns:

Type	Description
tuple	Tuple of (wavenumbers..., P_k) - format depends on dimensionality

Example

field = jnp.array([...])  # Any dimension
result = power_spectrum(field, L=20.0)

Source code in jaxlatt/observables/spectra.py

def power_spectrum(field: Array, L: float | tuple[float, ...]) -> tuple:
    """
    Compute power spectrum automatically based on field dimensionality.

    Dispatches to the appropriate 1D/2D/3D function based on field shape.

    Args:
        field: Field array (1D, 2D, or 3D)
        L: Physical length (scalar or tuple for anisotropic boxes)

    Returns:
        Tuple of (wavenumbers..., P_k) - format depends on dimensionality

    Example:
        ```python
        field = jnp.array([...])  # Any dimension
        result = power_spectrum(field, L=20.0)
        ```
    """
    ndim = len(field.shape)

    # Convert L to scalar if needed
    if isinstance(L, (tuple, list)):
        if len(set(L)) > 1:
            raise ValueError("Anisotropic boxes not yet supported - use equal lengths")
        L_scalar = L[0]
    else:
        L_scalar = L

    if ndim == 1:
        return power_spectrum_1d(field, L_scalar)
    elif ndim == 2:
        return power_spectrum_2d(field, L_scalar)
    elif ndim == 3:
        return power_spectrum_3d(field, L_scalar)
    else:
        raise ValueError(f"Field must be 1D, 2D, or 3D, got {ndim}D")

make_log_bin_edges(k_min, k_max, num_bins)

Logarithmically spaced bin edges, shape (num_bins + 1,).

Compute once outside a traced region and pass to :func:bin_power_spectrum_static, which needs the edges to be static so its output shape is known at trace time.

Source code in jaxlatt/observables/spectra.py

def make_log_bin_edges(k_min: float, k_max: float, num_bins: int) -> Array:
    """Logarithmically spaced bin edges, shape ``(num_bins + 1,)``.

    Compute once outside a traced region and pass to
    :func:`bin_power_spectrum_static`, which needs the edges to be static so
    its output shape is known at trace time.
    """
    return jnp.logspace(jnp.log10(k_min), jnp.log10(k_max), num_bins + 1)

make_linear_bin_edges(k_min, k_max, num_bins)

Linearly spaced bin edges, shape (num_bins + 1,).

Source code in jaxlatt/observables/spectra.py

def make_linear_bin_edges(k_min: float, k_max: float, num_bins: int) -> Array:
    """Linearly spaced bin edges, shape ``(num_bins + 1,)``."""
    return jnp.linspace(k_min, k_max, num_bins + 1)

bin_power_spectrum_static(k_values, P_values, bin_edges)

Radially bin a power spectrum with a static output shape.

Unlike :func:bin_power_spectrum, this is a pure JAX function: it is jittable, vmappable and differentiable, because it always returns exactly num_bins values instead of dropping empty bins. Empty bins are reported with counts == 0 and a value of zero; mask on counts downstream.

Parameters:

Name	Type	Description	Default
k_values	Array	Wavenumber magnitudes, any shape.	required
P_values	Array	Power values, same shape as k_values.	required
bin_edges	Array	Monotonic edges of shape (num_bins + 1,), e.g. from :func:make_log_bin_edges.	required

Returns:

Type	Description
tuple[Array, Array, Array]	Tuple (k_binned, P_binned, counts), each of shape (num_bins,).

Notes

• The k = 0 mode and any non-finite entries are excluded by weight, not by indexing, so shapes stay static.
• Modes outside [bin_edges[0], bin_edges[-1]] are excluded. A mode exactly equal to bin_edges[-1] falls in the last bin (this differs from :func:bin_power_spectrum, whose last bin is half-open and therefore drops the largest mode).

Source code in jaxlatt/observables/spectra.py

def bin_power_spectrum_static(
    k_values: Array,
    P_values: Array,
    bin_edges: Array,
) -> tuple[Array, Array, Array]:
    """Radially bin a power spectrum with a **static output shape**.

    Unlike :func:`bin_power_spectrum`, this is a pure JAX function: it is
    jittable, vmappable and differentiable, because it always returns exactly
    ``num_bins`` values instead of dropping empty bins. Empty bins are reported
    with ``counts == 0`` and a value of zero; mask on ``counts`` downstream.

    Args:
        k_values: Wavenumber magnitudes, any shape.
        P_values: Power values, same shape as ``k_values``.
        bin_edges: Monotonic edges of shape ``(num_bins + 1,)``, e.g. from
            :func:`make_log_bin_edges`.

    Returns:
        Tuple ``(k_binned, P_binned, counts)``, each of shape ``(num_bins,)``.

    Notes:
        - The ``k = 0`` mode and any non-finite entries are excluded by weight,
          not by indexing, so shapes stay static.
        - Modes outside ``[bin_edges[0], bin_edges[-1]]`` are excluded. A mode
          exactly equal to ``bin_edges[-1]`` falls in the last bin (this
          differs from :func:`bin_power_spectrum`, whose last bin is half-open
          and therefore drops the largest mode).
    """
    k = jnp.ravel(k_values)
    P = jnp.ravel(P_values)
    n_bins = bin_edges.shape[0] - 1

    idx = jnp.digitize(k, bin_edges) - 1
    # Include the exact right edge in the final bin.
    idx = jnp.where(k == bin_edges[-1], n_bins - 1, idx)

    valid = (k > 0) & jnp.isfinite(k) & jnp.isfinite(P) & (idx >= 0) & (idx < n_bins)
    # Clip only *after* folding validity into the weight, so out-of-range
    # modes are dropped rather than piled into the first/last bin.
    idx = jnp.clip(idx, 0, n_bins - 1)
    w = valid.astype(P.dtype)

    counts = jnp.zeros(n_bins, P.dtype).at[idx].add(w)
    P_sum = jnp.zeros(n_bins, P.dtype).at[idx].add(P * w)
    k_sum = jnp.zeros(n_bins, k.dtype).at[idx].add(k * jnp.real(w).astype(k.dtype))

    denom = jnp.maximum(counts, 1)
    return k_sum / jnp.real(denom).astype(k.dtype), P_sum / denom, counts

bin_power_spectrum(k_values, P_values, num_bins=50, k_min=None, k_max=None, log_bins=True)

Bin power spectrum values into radial/spherical shells.

This function takes the raw k-space power spectrum and averages it into bins of |k| to produce a 1D radially-averaged spectrum P(k).

Eager convenience wrapper around :func:bin_power_spectrum_static: it picks bin edges from the data and drops empty bins, which makes the output shape value-dependent and therefore not traceable. Use :func:bin_power_spectrum_static inside jit/vmap/grad.

Parameters:

Name	Type	Description	Default
k_values	Array	Flattened array of wavenumber magnitudes |k|	required
P_values	Array	Flattened array of power spectrum values	required
num_bins	int	Number of bins for averaging	50
k_min	float | None	Minimum k value (default: smallest non-zero k)	None
k_max	float | None	Maximum k value (default: maximum k)	None
log_bins	bool	Use logarithmic binning (default: True, recommended for wide k-range)	True

Returns:

Type	Description
ndarray	Tuple of (k_binned, P_binned, counts) where:
ndarray	k_binned: Center wavenumber of each bin
ndarray	P_binned: Average power in each bin
tuple[ndarray, ndarray, ndarray]	counts: Number of k-modes in each bin

Example

# 3D case
kx, ky, kz, P = power_spectrum_3d(field, L=20.0)
k_mag = jnp.sqrt(kx**2 + ky**2 + kz**2)
k_binned, P_binned, counts = bin_power_spectrum(
    k_mag.flatten(), P.flatten(), num_bins=50
)
plt.loglog(k_binned, P_binned)

Notes

• Excludes k=0 mode (DC component)
• Log binning is recommended for fields with power-law spectra
• Empty bins are removed from output

Source code in jaxlatt/observables/spectra.py

def bin_power_spectrum(
    k_values: Array,
    P_values: Array,
    num_bins: int = 50,
    k_min: float | None = None,
    k_max: float | None = None,
    log_bins: bool = True,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Bin power spectrum values into radial/spherical shells.

    This function takes the raw k-space power spectrum and averages it into
    bins of |k| to produce a 1D radially-averaged spectrum P(k).

    Eager convenience wrapper around :func:`bin_power_spectrum_static`: it
    picks bin edges from the data and drops empty bins, which makes the output
    shape value-dependent and therefore **not** traceable. Use
    :func:`bin_power_spectrum_static` inside ``jit``/``vmap``/``grad``.

    Args:
        k_values: Flattened array of wavenumber magnitudes |k|
        P_values: Flattened array of power spectrum values
        num_bins: Number of bins for averaging
        k_min: Minimum k value (default: smallest non-zero k)
        k_max: Maximum k value (default: maximum k)
        log_bins: Use logarithmic binning (default: True, recommended for wide k-range)

    Returns:
        Tuple of (k_binned, P_binned, counts) where:
        - k_binned: Center wavenumber of each bin
        - P_binned: Average power in each bin
        - counts: Number of k-modes in each bin

    Example:
        ```python
        # 3D case
        kx, ky, kz, P = power_spectrum_3d(field, L=20.0)
        k_mag = jnp.sqrt(kx**2 + ky**2 + kz**2)
        k_binned, P_binned, counts = bin_power_spectrum(
            k_mag.flatten(), P.flatten(), num_bins=50
        )
        plt.loglog(k_binned, P_binned)
        ```

    Notes:
        - Excludes k=0 mode (DC component)
        - Log binning is recommended for fields with power-law spectra
        - Empty bins are removed from output
    """
    k_vals = np.asarray(k_values).ravel()
    P_vals = np.asarray(P_values).ravel()

    finite = (k_vals > 0) & np.isfinite(k_vals) & np.isfinite(P_vals)
    if not finite.any():
        return np.array([]), np.array([]), np.array([])

    lo = float(k_vals[finite].min()) if k_min is None else float(k_min)
    hi = float(k_vals[finite].max()) if k_max is None else float(k_max)
    if log_bins and lo <= 0:
        lo = float(k_vals[finite].min())

    edges = (
        make_log_bin_edges(lo, hi, num_bins)
        if log_bins
        else make_linear_bin_edges(lo, hi, num_bins)
    )
    k_b, P_b, counts = bin_power_spectrum_static(k_vals, P_vals, edges)

    counts = np.asarray(counts)
    keep = counts > 0
    return (
        np.asarray(k_b)[keep],
        np.asarray(P_b)[keep],
        counts[keep].astype(int),
    )

power_spectrum_with_binning(field, L, num_bins=50, log_bins=True)

Convenience function: compute and bin power spectrum in one call.

Automatically handles 1D/2D/3D fields and returns radially/spherically averaged power spectrum.

Parameters:

Name	Type	Description	Default
field	Array	Field array (1D, 2D, or 3D)	required
L	float	Physical length of the lattice	required
num_bins	int	Number of bins for averaging	50
log_bins	bool	Use logarithmic binning	True

Returns:

Type	Description
tuple[ndarray, ndarray, ndarray]	Tuple of (k, P, counts) - binned power spectrum

Example

field = jnp.array([...])  # 3D field
k, P, counts = power_spectrum_with_binning(field, L=20.0, num_bins=50)
plt.loglog(k, P)
plt.xlabel('k')
plt.ylabel('P(k)')

Source code in jaxlatt/observables/spectra.py

def power_spectrum_with_binning(
    field: Array,
    L: float,
    num_bins: int = 50,
    log_bins: bool = True,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Convenience function: compute and bin power spectrum in one call.

    Automatically handles 1D/2D/3D fields and returns radially/spherically
    averaged power spectrum.

    Args:
        field: Field array (1D, 2D, or 3D)
        L: Physical length of the lattice
        num_bins: Number of bins for averaging
        log_bins: Use logarithmic binning

    Returns:
        Tuple of (k, P, counts) - binned power spectrum

    Example:
        ```python
        field = jnp.array([...])  # 3D field
        k, P, counts = power_spectrum_with_binning(field, L=20.0, num_bins=50)
        plt.loglog(k, P)
        plt.xlabel('k')
        plt.ylabel('P(k)')
        ```
    """
    ndim = len(field.shape)

    if ndim == 1:
        k_values, P_k = power_spectrum_1d(field, L)
        return bin_power_spectrum(k_values, P_k, num_bins, log_bins=log_bins)

    elif ndim == 2:
        k_x, k_y, P_k = power_spectrum_2d(field, L)
        k_mag = jnp.sqrt(k_x**2 + k_y**2)
        return bin_power_spectrum(k_mag.flatten(), P_k.flatten(), num_bins, log_bins=log_bins)

    elif ndim == 3:
        k_x, k_y, k_z, P_k = power_spectrum_3d(field, L)
        k_mag = jnp.sqrt(k_x**2 + k_y**2 + k_z**2)
        return bin_power_spectrum(k_mag.flatten(), P_k.flatten(), num_bins, log_bins=log_bins)

    else:
        raise ValueError(f"Field must be 1D, 2D, or 3D, got {ndim}D")