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.

spectrum_asymmetry(scaling, real_shape[, ...])

Tool: measure violation of the conjugate-pair precondition of a spectrum.

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)[source]

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.

Momenta are folded into the first Brillouin zone: mode \(n\) of an axis of length \(L\) carries momentum \(2\pi n/L\) for \(n \le L/2\) and \(2\pi (n - L)/L\) above that, so momenta run over \((-\pi, \pi]\) and modes \(n\) and \(L - n\) are related by \(k \to -k\). Folding is what makes \(|k|^2\) the momentum of that mode (mode \(L-1\) is the slowest non-zero mode, not the fastest). With reduced=True the last axis only runs to \(L/2\) and is unaffected; folding matters on every other axis, i.e. in two or more space dimensions.

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 raw non-negative integer indices into the FFT grid instead of momenta. These are array indices, not momenta, and are deliberately not folded; for integer-valued momenta use fft_momenta(shape, unit=True) folded by the caller, or FourierMeta.ks_full, which does exactly that.

Return type:

Array

Returns:

Momentum grid array with shape (*spatial_shape, spatial_rank). For continuum: momenta in units of 2π/L, folded to \((-\pi, \pi]\). For lattice: \(2\sin(k/2)\) of those momenta, odd in \(k\).

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)
bijx.fourier.spectrum_asymmetry(scaling, real_shape, channel_dim=0)[source]

Tool: measure violation of the conjugate-pair precondition of a spectrum.

A diagonal scaling in rFFT space defines an invertible map on real fields (with log-Jacobian \(\sum_k w_k \log|s_k|\)) only if two conditions hold:

  1. Entries whose conjugate partner is also stored in the rFFT grid are consistent, \(s[\text{copy\_to}] == s[\text{copy\_from}]^*\).

  2. The spectrum is real at self-conjugate modes (mr & ~mi), where the field itself carries no imaginary degree of freedom.

All other entries are unconstrained. See also SpectrumScaling.

This is not checked inside the bijection (it cannot raise under jit); call this function to assert the precondition in tests / debugging.

Parameters:

  • scaling – Spectrum with leading axes matching the rFFT shape of real_shape, optionally followed by channel axes.

  • real_shape – Shape of the real-space (spatial) data.

  • channel_dim – Number of channel dimensions of scaling.

Returns:

Maximum absolute violation of either condition; the larger of $max |s[text{copy_to}] - s[text{copy_from}]^*|$ and $max |operatorname{Im} s|$ over the self-conjugate modes.

Any spectrum built from fft_momenta() satisfies both conditions, since those momenta are folded and real; this function is for hand-built or externally supplied spectra.

Example

>>> k = fft_momenta((6, 6))
>>> spectrum = jnp.exp(-0.1 * jnp.sum(k**2, axis=-1))
>>> assert spectrum_asymmetry(spectrum, (6, 6)) < 1e-12
>>> asymmetric = spectrum.at[1, 0].multiply(1.8)  # breaks one pair
>>> assert spectrum_asymmetry(asymmetric, (6, 6)) > 0
class bijx.fourier.FourierMeta[source]

Bases: Pytree

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 (Array) – Boolean mask for real (independent) Fourier modes.

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

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

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

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

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

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

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

Note

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

Note

ks_full is \(|k|^2\) in units of \((2\pi/L)^2\), from indices folded into the first Brillouin zone, so conjugate partners (copy_from/ copy_to) always share a \(|k|^2\) class. A spectrum parametrised per class through unique_idc/unique_unfold therefore satisfies the conjugate-pair precondition of SpectrumScaling by construction, learnable or not.

shape_info: ShapeInfo
mr: Array
mi: Array
copy_from: Array
copy_to: Array
ks_full: Array
ks_reduced: Array
unique_idc: Array
unique_unfold: Array
replace(**changes)[source]

Create new config with specified parameters replaced.

classmethod create(real_shape, channel_dim=0)[source]

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)[source]
class bijx.fourier.FFTRep[source]

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[source]

Bases: Pytree

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
replace(**changes)[source]

Create new config with specified parameters replaced.

classmethod from_real(x, real_shape, to=None, channel_dim=0)[source]
Parameters:

to (FFTRep | None)

to(rep)[source]
Parameters:

rep (FFTRep | None)

static rfft_to_real(rfft, meta)[source]
static real_to_rfft(x, meta)[source]
static complex_to_rfft(xk, meta)[source]
static rfft_to_complex(rfft, meta)[source]
static rdof_to_complex(rdof, meta)[source]
static complex_to_rdof(xk, meta)[source]