bijx.ToFourierData

class bijx.ToFourierData

Bases: Bijection

Bijection for converting between real and Fourier data representations.

Provides a bijective transformation between real-space arrays and their various Fourier representations (FFT output, independent components, etc.). This is useful for working with different Fourier data formats within flows.

Type: \(\mathbb{R}^{\text{real_shape}} \leftrightarrow \text{FourierData}\) Transform: Format conversion between real and Fourier representations

Parameters:

• real_shape – Shape of the real-space data.

• rep – Target Fourier representation (FFTRep enum value).

• channel_dim – Number of channel dimensions.

• unpack – Whether to unpack FourierData to raw array.

Note

The Fourier transform can be understood as a rotation, and is therefore volume-preserving. The log-density remains unchanged.

Example

>>> # Convert to complex Fourier components
>>> converter = ToFourierData((32, 32), rep=FFTRep.comp_complex)
>>> fourier_data, log_det = converter.forward(real_data, log_density)
>>> # log_det == 0 (volume-preserving)

__init__(real_shape, rep=None, channel_dim=0, unpack=False)

Methods

forward(x, log_density, **kwargs)	Apply forward transformation.
invert()	Create an inverted version of this bijection.
reverse(x, log_density, **kwargs)	Apply reverse (inverse) transformation.

forward(x, log_density, **kwargs)

Apply forward transformation.

Transforms input through the bijection and updates log-density according to the change of variables formula.

For convenience Bijection() gives the default identity bijection.

Parameters:

• x – Input data of any pytree structure.

• log_density – Log density values corresponding to the input.

• **kwargs – Additional transformation-specific arguments.

Returns:

Tuple of (transformed_data, updated_log_density) where the log-density incorporates the log absolute determinant of the transformation Jacobian.

reverse(x, log_density, **kwargs)

Apply reverse (inverse) transformation.

Transforms input through the inverse bijection and updates log-density accordingly.

Parameters:

• x – Input data of any pytree structure.

• log_density – Log density values corresponding to the input.

• **kwargs – Additional transformation-specific arguments.

Returns:

Tuple of (inverse_transformed_data, updated_log_density) where the log-density change has the opposite sign compared to forward().