Spectral Convolutions

Spectral Convolution¤

These layers evaluate the continuous integral operator in Fourier space. They perform the core global operations that make Fourier Neural Operators discretization-invariant.

`neojax.nn.spectral_conv.SpectralConvNd` ¤

General n-dimensional spectral convolution layer.

This layer computes the real n-dimensional forward FFT, truncates the higher frequency modes according to the specified modes, multiplies the remaining modes with learnable complex weights, and transforms the result back to the spatial domain using the inverse FFT.

Parameters:

Name	Type	Description	Default
`key`	`PRNGKeyArray`	PRNG key for weight initialization.	required
`in_channels`	`int`	Number of input channels.	required
`out_channels`	`int`	Number of output channels.	required
`modes`	`int \| Sequence[int]`	Number of Fourier modes to retain across each spatial dim. Can be an integer or sequence of integers.	required
`ranks`	`int \| Sequence[int] \| None`	Number of ranks to contract the spectral tensors to. If `ranks` is an Integer, the same number is used for all ranks. If not, should be a sequence of ranks. Default is None.	`None`
`init_std`	`Union[float, Literal['auto']]`	Standard deviation to use for weight initialization, by default 'auto'. If 'auto', uses (2 / (in_channels * out_channels)) ** 0.5.	`'auto'`
`enforce_hermitian_symmetry`	`bool`	Whether to enforce hermitian symmetry on the outputs before calling `irfftn`. Default is True. If set to False, cuFFT on GPU may cause line artifacts when calling irfftn.	`True`
`fft_norm`	`Optional[Literal['forward', 'backward', 'ortho']]`	FFT normalization. Can be `None`, `"backward"`, `"ortho"` or `"forward"`. Default is `"forward"`.	`'forward'`
`is_complex_data`	`bool`	Whether input data is complex valued. If True, uses full FFT. Default is False.	`False`
`resolution_scaling_factor`	`float \| int \| None`	Factor by which to scale the domain resolution of the function. Default is None, no scaling.	`None`
`factorization`	`Union[BaseTensor, Literal['tucker', 'cp', 'tt'], NoneType]`	Tensor factorization type. Can be a `BaseTensor` instance, a string ("tucker", "cp", "tt"), or None. Default is None.	`None`
`implementation`	`Literal['reconstructed', 'factorized']`	Weight reconstruction mode. Can be "reconstructed" or "factorized". Default is "factorized".	`'factorized'`
`separable`	`bool`	Whether to use separable implementation of contraction. If True, contracts factors of factorized tensor weight individually. Default is False.	`False`

Internal Attributes

These fields store the internal layers state (and weights).

weights (BaseTensor): Learnable complex weights for the retained Fourier modes.
in_channels (int): Number of input channels.
out_channels (int): Number of output channels.
modes (tuple[int, ...]): Number of Fourier modes to retain across each spatial dim.
fft_norm (str | None): FFT normalization.
is_complex_data (bool): Whether input data is complex valued.
enforce_hermitian_symmetry (bool): Whether to enforce hermitian symmetry on the outputs.
resolution_scaling_factor (float | int | None): Factor to scale domain resolution. Default is None.
implementation (str): Weight reconstruction mode.
separable (bool): Whether the convolution is separable.
num_corners (int): Number of corners of the n-dimensional FFT hypercube.

Examples:

import jax.numpy as jnp
import jax.random as jr
from neojax.nn import SpectralConvNd

key = jr.key(0)

# 1D Spectral Convolution
conv1d = SpectralConvNd(key, 3, 16, modes=16)
# or
conv1d = SpectralConvNd(key, 3, 16, modes=(16,))

# 2D Spectral Convolution
conv2d = SpectralConvNd(key, 3, 16, modes=(16, 16))

# 3D Spectral Convolution
conv3d = SpectralConvNd(key, 3, 16, modes=(16, 16, 16))

# Input shape: (channels, d1, ..., dN)
x = jnp.ones((3, 32, 32))
out = conv2d(x)

Implementation Details

Unlike the reference PyTorch neuraloperator library, this implementation does not shift the frequency spectrum (e.g., using fftshift / ifftshift). Instead, it directly indexes the positive and negative frequency modes from the boundaries of the unshifted FFT representation using a binary bitmask to select the corners of an n-dimensional hypercube. This avoids memory copy/transposition overhead at runtime and generalizes to arbitrary dimensions.

Limitation: Half & Mixed Precision

Unlike PyTorch, JAX does not natively support complex half-precision types (like complex32). Therefore, setting fno_block_precision to "half" or "mixed" is not supported, and all spectral operations run in full precision (complex64).

`call(x: Inexact[Array, 'in_c ...']) -> Inexact[Array, 'out_c ...']` ¤

Perform n-dimensional spectral convolution.

Parameters:

Name	Type	Description	Default
`x`	`Inexact[Array, 'in_c ...']`	Input signal.	required

Returns:

Type	Description
`Inexact[Array, 'out_c ...']`	Output signal.

`neojax.nn.geo_spectral_conv.GeoSpectralConvNd` ¤

N-dimensional geometry-aware continuous spectral convolution.

Wraps a standard SpectralConvNd and extends it to support continuous Fourier transforms on irregular meshes via direct DFT summation. When no coordinate meshes are provided at call time, it falls back to the standard fast grid-based spectral convolution.

Parameters:

Name	Type	Description	Default
`key`	`PRNGKeyArray`	PRNG key for weight initialization.	required
`in_channels`	`int`	Number of input channels.	required
`out_channels`	`int`	Number of output channels.	required
`modes`	`int \| Sequence[int]`	Number of Fourier modes to retain.	required
`ranks`	`int \| Sequence[int] \| None`	Number of factorization ranks.	`None`
`init_std`	`Union[float, Literal['auto']]`	Weight initialization standard deviation.	`'auto'`
`enforce_hermitian_symmetry`	`bool`	Enforce Hermitian symmetry for real outputs.	`True`
`fft_norm`	`Optional[Literal['forward', 'backward', 'ortho']]`	FFT normalization.	`'forward'`
`is_complex_data`	`bool`	If True, uses full complex FFT.	`False`
`resolution_scaling_factor`	`float \| int \| None`	Scale factor for grid resolution.	`None`
`factorization`	`Union[BaseTensor, Literal['tucker', 'cp', 'tt'], NoneType]`	Tensor factorization scheme.	`None`
`implementation`	`Literal['reconstructed', 'factorized']`	Weight reconstruction mode.	`'factorized'`
`separable`	`bool`	If True, uses separable contraction.	`False`

Internal Attributes

conv (SpectralConvNd): The underlying grid spectral convolution module whose weights and modes are reused for the continuous transform.

`call(u: Inexact[Array, 'in_c ...'], x_in: Float[Array, 'N d'] | None = None, x_out: Float[Array, 'M d'] | None = None) -> Inexact[Array, 'out_c ...']` ¤

Perform continuous or discrete spectral convolution.

Parameters:

Name	Type	Description	Default
`u`	`Inexact[Array, 'in_c ...']`	Input features on the grid or mesh.	required
`x_in`	`Float[Array, 'N d'] \| None`	Input spatial coordinates. If None, grid-based conv is used.	`None`
`x_out`	`Float[Array, 'M d'] \| None`	Output spatial coordinates. Defaults to `x_in`.	`None`

Returns:

Type	Description
`Inexact[Array, 'out_c ...']`	Output features after spectral convolution.

`get_corner_wavenumbers(corner_idx: int) -> Int[Array, 'modes_prod ndim']` ¤

Generates coordinate-aligned wavenumbers for a given corner slice of the Fourier space.

Parameters:

Name	Type	Description	Default
`corner_idx`	`int`	Integer index encoding which corner of the Fourier cube to use. The \(j\)-th bit of `corner_idx` controls the sign of the \(j\)-th frequency axis: `0` means non-negative frequencies, `1` means negative frequencies.	required

Returns:

Type	Description
`Int[Array, 'modes_prod ndim']`	Integer wavenumber matrix of shape `(modes_prod, ndim)` where `modes_prod = prod(modes)`.

Spectral Convolutions