bijx.SpectrumScaling

class bijx.SpectrumScaling

Bases: ApplyBijection

Diagonal scaling transformation in Fourier space.

Applies element-wise scaling to the Fourier transform of real-valued fields, implementing diagonal transformations in momentum space. This is particularly useful for implementing free field theories and spectral preconditioning.

Type: \(\mathbb{R}^{H \times W \times C} \to \mathbb{R}^{H \times W \times C}\) Transform: \(\mathcal{F}^{-1}[s(\mathbf{k}) \mathcal{F}[\mathbf{x}]]\)

The scaling factors correspond to momentum-dependent transformations, with the log-Jacobian computed from FFT multiplicities to handle real FFT symmetries.

Parameters:

• scaling (Array | Variable) – Scaling factors with shape matching rFFT output. If not an nnx.Variable/nnx.Param, by default treated as constant.

• channel_dim (int) – Number of channel dimensions.

• space_dim (int | None) – Number of spatial dimensions. If None, inferred from the rank of scaling, which assumes a single spectrum shared by all channels. Must be given explicitly for a per-channel spectrum (and whenever scaling is None).

Note

The spatial part of the scaling array must have the same shape as the output of jnp.fft.rfftn over the space axes. With channel_dim > 0 the scaling may either have rank space_dim (one spectrum shared by all channels) or rank space_dim + channel_dim (per-channel spectra).

Note

Those rFFT entries whose conjugate partner is also stored in the rFFT grid are constrained, and they must satisfy s[copy_to] == conj(s[copy_from]) for the index pairs of FourierMeta.create(space_shape) (plain equality for a real spectrum). For a real spectrum the condition is vacuous in one dimension (no such pairs exist). It is separate from the other requirement that the spectrum is symmetric under \(k \to -k\) everywhere.

Violating it fails silently: mr + mi vanishes on every copy_to entry, so the log-Jacobian weights assume the symmetry rather than check it. The map then stops being invertible and the reported log-density change is wrong. Use spectrum_asymmetry() to check the precondition.

Any real function of fft_momenta() satisfies this automatically, since those momenta are folded into the first Brillouin zone and are therefore related by \(k \to -k\) across each conjugate pair, as is a spectrum parametrised per \(|k|^2\) class through FourierMeta.unique_unfold.

A complex spectrum carries a second condition: it must be real at self-conjugate modes (mr & ~mi), where the field carries no imaginary degree of freedom. See spectrum_asymmetry().

Example

>>> # Create momentum-dependent scaling
>>> k = fft_momenta((8, 8))
>>> scaling = jnp.exp(-0.1 * jnp.sum(k**2, axis=-1))
>>> bijection = SpectrumScaling(scaling)
>>> y, log_det = bijection.forward(phi, log_density)

__init__(scaling, channel_dim=0, space_dim=None)

Parameters:

• scaling (Array | Variable)

• channel_dim (int)

• space_dim (int | None)

Methods

apply(x, log_density[, reverse])	Unified transformation method.
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.

Attributes

scaling

property scaling

apply(x, log_density, reverse=False, **kwargs)

Unified transformation method.

Parameters:

• x – Input data of any pytree structure.

• log_density – Log density values corresponding to the input.

• reverse – If True, apply reverse transformation; if False, forward.

• **kwargs – Additional transformation-specific arguments.

Returns:

Tuple of (transformed_data, updated_log_density).

Raises:

NotImplementedError – Must be implemented by subclasses.