neuralop.layers.spectral_convolution.SpectralConv

class neuralop.layers.spectral_convolution.SpectralConv(in_channels, out_channels, n_modes, max_n_modes=None, bias=True, n_layers=1, separable=False, output_scaling_factor: int | float | List[float | int] | None = None, fno_block_precision='full', rank=0.5, factorization=None, implementation='reconstructed', fixed_rank_modes=False, joint_factorization=False, decomposition_kwargs: dict | None = None, init_std='auto', fft_norm='backward', device=None, dtype=None)[source]

Generic N-Dimensional Fourier Neural Operator

Parameters:
in_channelsint, optional

Number of input channels

out_channelsint, optional

Number of output channels

max_n_modesNone or int tuple, default is None

Number of modes to use for contraction in Fourier domain during training.

Warning

We take care of the redundancy in the Fourier modes, therefore, for an input of size I_1, …, I_N, please provide modes M_K that are I_1 < M_K <= I_N We will automatically keep the right amount of modes: specifically, for the last mode only, if you specify M_N modes we will use M_N // 2 + 1 modes as the real FFT is redundant along that last dimension.

Note

Provided modes should be even integers. odd numbers will be rounded to the closest even number.

This can be updated dynamically during training.

max_n_modesint tuple or None, default is None
  • If not None, maximum number of modes to keep in Fourier Layer, along each dim

    The number of modes (n_modes) cannot be increased beyond that.

  • If None, all the n_modes are used.

separablebool, default is True
init_stdfloat or ‘auto’, default is ‘auto’

std to use for the init

n_layersint, optional

Number of Fourier Layers, by default 4

factorizationstr or None, {‘tucker’, ‘cp’, ‘tt’}, default is None

If None, a single dense weight is learned for the FNO. Otherwise, that weight, used for the contraction in the Fourier domain is learned in factorized form. In that case, factorization is the tensor factorization of the parameters weight used.

joint_factorizationbool, optional

Whether all the Fourier Layers should be parametrized by a single tensor (vs one per layer), by default False Ignored if factorization is None

rankfloat or rank, optional

Rank of the tensor factorization of the Fourier weights, by default 1.0 Ignored if factorization is None

fixed_rank_modesbool, optional

Modes to not factorize, by default False Ignored if factorization is None

fft_normstr, optional

by default ‘forward’

implementation{‘factorized’, ‘reconstructed’}, optional, default is ‘factorized’

If factorization is not None, forward mode to use:: * reconstructed : the full weight tensor is reconstructed from the

factorization and used for the forward pass

  • factorized : the input is directly contracted with the factors of the decomposition

Ignored if factorization is None

decomposition_kwargsdict, optional, default is {}

Optionaly additional parameters to pass to the tensor decomposition Ignored if factorization is None

Attributes:
n_modes

Methods

forward(x[, indices, output_shape])

Generic forward pass for the Factorized Spectral Conv

get_conv(indices)

Returns a sub-convolutional layer from the joint parametrize main-convolution

transform(x[, layer_index, output_shape])

Transforms an input x for a skip connection, by default just an identity map
transform(x, layer_index=0, output_shape=None)[source]

Transforms an input x for a skip connection, by default just an identity map

If your function transforms the input then you should also implement this transform method so the skip connection can also work.

Typical usecases are:

  • Your upsample or downsample the input in the Spectral conv: the skip connection has to be similarly scaled. This allows you to deal with it however you want (e.g. avoid aliasing)

  • You perform a change of basis in your Spectral Conv, again, this needs to be applied to the skip connection too.

forward(x: Tensor, indices=0, output_shape: Tuple[int] | None = None)[source]

Generic forward pass for the Factorized Spectral Conv

Parameters:
xtorch.Tensor

input activation of size (batch_size, channels, d1, …, dN)

indicesint, default is 0

if joint_factorization, index of the layers for n_layers > 1

Returns:
tensorized_spectral_conv(x)
get_conv(indices)[source]

Returns a sub-convolutional layer from the joint parametrize main-convolution

The parametrization of sub-convolutional layers is shared with the main one.