bijx.fourier

Fourier transform utilities for lattice field theory and physics applications.

This module provides comprehensive utilities for working with Fourier transforms of real-valued fields based on the FFT implementation in JAX.

Functions

fft_momenta(shape[, reduced, lattice, unit])

Generate momentum grid for Fourier transforms.

Classes

FourierMeta	Metadata for handling real FFT constraints and symmetries.
FFTRep	Enumeration of different Fourier data representations.
FourierData	Multi-representation container for Fourier data.

bijx.fourier.fft_momenta(shape, reduced=True, lattice=False, unit=False)

Generate momentum grid for Fourier transforms.

Creates momentum coordinate arrays suitable for physics applications, supporting both continuum and lattice formulations. Handles the reduced form appropriate for real FFTs with Hermitian symmetry.

Parameters:

• shape (tuple[int, ...]) – Spatial grid dimensions.

• reduced (bool) – If True, use reduced form for real FFT (Hermitian symmetry).

• lattice (bool) – If True, use lattice momenta; otherwise continuum momenta.

• unit (bool) – If True, return integer indices instead of momentum values.

Return type:

Array

Returns:

Momentum grid array with shape (*spatial_shape, spatial_rank). For continuum: momenta in units of 2π/L. For lattice: momenta appropriate for lattice derivatives.

Example

>>> # Continuum momenta for 2D lattice
>>> k = fft_momenta((64, 64), lattice=False)
>>> k_squared = jnp.sum(k**2, axis=-1)  # |k|²
>>> # Lattice momenta for finite difference operators
>>> k_lat = fft_momenta((64, 64), lattice=True)

class bijx.fourier.FourierMeta

Bases: object

Metadata for handling real FFT constraints and symmetries.

Encapsulates all the bookkeeping needed to work with real-valued Fourier transforms, including Hermitian symmetry constraints, multiplicities for log-Jacobian computation, and indexing for different representations.

The metadata handles the reduction from full complex FFT to the independent real degrees of freedom.

Parameters:

• shape_info (ShapeInfo) – Shape information for spatial and channel dimensions.

• mr (ndarray) – Boolean mask for real (independent) Fourier modes.

• mi (ndarray) – Boolean mask for imaginary (independent) Fourier modes.

• copy_from (ndarray) – Indices of modes that are copied due to Hermitian symmetry.

• copy_to (ndarray) – Target indices for Hermitian symmetry copying.

• ks_full (ndarray) – Full momentum magnitude squared values.

• ks_reduced (ndarray) – Reduced momentum magnitude squared values.

• unique_idc (ndarray) – Indices of unique momentum magnitudes.

• unique_unfold (ndarray) – Mapping from reduced to unique momentum magnitudes.

Note

This class is created automatically by FourierMeta.create() and should usually not be instantiated directly.

shape_info: ShapeInfo

mr: ndarray

mi: ndarray

copy_from: ndarray

copy_to: ndarray

ks_full: ndarray

ks_reduced: ndarray

unique_idc: ndarray

unique_unfold: ndarray

classmethod create(real_shape, channel_dim=0)

Create FourierMeta for given real-space shape.

Parameters:

• real_shape – Shape of real-space data.

• channel_dim – Number of channel dimensions.

Returns:

FourierMeta instance with all symmetry constraints computed.

property real_shape

property have_imag

property channel_slices

property idc_rfft_independent

property idc_have_imag

property idc_copy_from

property idc_copy_to

get_complex_dtype(real_data)

replace(**updates)

Returns a new object replacing the specified fields with new values.

class bijx.fourier.FFTRep

Bases: IntEnum

Enumeration of different Fourier data representations.

Defines the various ways to represent Fourier data for real-valued fields, each with different trade-offs in terms of memory usage, computational efficiency, and mathematical convenience.

Values:

real_space: Original real-space field data. rfft: Raw output from real FFT (includes redundant information). comp_complex: Independent complex Fourier components only. comp_real: All independent real degrees of freedom as a single array.

Note

The comp_real representation packs both real and imaginary parts of independent modes into a single real-valued array, maximizing compatibility with standard bijection layers.

real_space = 0

rfft = 1

comp_complex = 2

comp_real = 3

class bijx.fourier.FourierData

Bases: object

Multi-representation container for Fourier data.

Provides a unified interface for working with Fourier data in different representations, with automatic conversion between formats. This enables seamless switching between representations based on computational needs.

The container maintains the data, its current representation type, and the associated metadata needed for conversions. All conversions preserve the underlying mathematical content while changing the format.

Parameters:

• data (Array) – The actual data array in the current representation.

• rep (FFTRep) – Current representation type (FFTRep enum).

• meta (FourierMeta) – FourierMeta containing symmetry and indexing information.

Example

>>> # Create from real-space data
>>> fd = FourierData.from_real(x, (64, 64))
>>> # Convert to complex components
>>> fd_complex = fd.to(FFTRep.comp_complex)
>>> # Convert to real degrees of freedom
>>> fd_real = fd.to(FFTRep.comp_real)

data: Array

rep: FFTRep

meta: FourierMeta

classmethod from_real(x, real_shape, to=None, channel_dim=0)

Parameters:

to (FFTRep | None)

to(rep)

Parameters:

rep (FFTRep | None)

static rfft_to_real(rfft, meta)

static real_to_rfft(x, meta)

static complex_to_rfft(xk, meta)

static rfft_to_complex(rfft, meta)

static rdof_to_complex(rdof, meta)

static complex_to_rdof(xk, meta)

replace(**updates)

Returns a new object replacing the specified fields with new values.