neuralop.models.TFNO

class neuralop.models.TFNO(*args, **kwargs)

N-Dimensional Fourier Neural Operator. The FNO learns a mapping between spaces of functions discretized over regular grids using Fourier convolutions, as described in [1].

The key component of an FNO is its SpectralConv layer (see neuralop.layers.spectral_convolution), which is similar to a standard CNN conv layer but operates in the frequency domain.

For a deeper dive into the FNO architecture, refer to Fourier Neural Operators.

Parameters:
n_modesTuple[int]

number of modes to keep in Fourier Layer, along each dimension The dimensionality of the FNO is inferred from len(n_modes)

in_channelsint

Number of channels in input function

out_channelsint

Number of channels in output function

hidden_channelsint

width of the FNO (i.e. number of channels), by default 256

n_layersint, optional

Number of Fourier Layers, by default 4

Documentation for more advanced parameters is below.

Methods

forward(x[, output_shape])

FNO's forward pass
Other Parameters:
lifting_channel_ratioint, optional

ratio of lifting channels to hidden_channels, by default 2 The number of liting channels in the lifting block of the FNO is lifting_channel_ratio * hidden_channels (e.g. default 512)

projection_channel_ratioint, optional

ratio of projection channels to hidden_channels, by default 2 The number of projection channels in the projection block of the FNO is projection_channel_ratio * hidden_channels (e.g. default 512)

positional_embeddingUnion[str, nn.Module], optional

Positional embedding to apply to last channels of raw input before being passed through the FNO. Defaults to “grid”

  • If “grid”, appends a grid positional embedding with default settings to

the last channels of raw input. Assumes the inputs are discretized over a grid with entry [0,0,…] at the origin and side lengths of 1.

  • If an initialized GridEmbedding module, uses this module directly

See neuralop.embeddings.GridEmbeddingND for details.

  • If None, does nothing

non_linearitynn.Module, optional

Non-Linear activation function module to use, by default F.gelu

normstr {“ada_in”, “group_norm”, “instance_norm”}, optional

Normalization layer to use, by default None

complex_databool, optional

Whether data is complex-valued (default False) if True, initializes complex-valued modules.

channel_mlp_dropoutfloat, optional

dropout parameter for ChannelMLP in FNO Block, by default 0

channel_mlp_expansionfloat, optional

expansion parameter for ChannelMLP in FNO Block, by default 0.5

channel_mlp_skipstr {‘linear’, ‘identity’, ‘soft-gating’}, optional

Type of skip connection to use in channel-mixing mlp, by default ‘soft-gating’

fno_skipstr {‘linear’, ‘identity’, ‘soft-gating’}, optional

Type of skip connection to use in FNO layers, by default ‘linear’

resolution_scaling_factorUnion[Number, List[Number]], optional

layer-wise factor by which to scale the domain resolution of function, by default None

  • If a single number n, scales resolution by n at each layer

  • if a list of numbers [n_0, n_1,…] scales layer i’s resolution by n_i.

domain_paddingUnion[Number, List[Number]], optional

If not None, percentage of padding to use, by default None To vary the percentage of padding used along each input dimension, pass in a list of percentages e.g. [p1, p2, …, pN] such that p1 corresponds to the percentage of padding along dim 1, etc.

domain_padding_modestr {‘symmetric’, ‘one-sided’}, optional

How to perform domain padding, by default ‘one-sided’

fno_block_precisionstr {‘full’, ‘half’, ‘mixed’}, optional

precision mode in which to perform spectral convolution, by default “full”

stabilizerstr {‘tanh’} | None, optional

whether to use a tanh stabilizer in FNO block, by default None

Note: stabilizer greatly improves performance in the case fno_block_precision=’mixed’.

max_n_modesTuple[int] | None, optional

  • If not None, this allows to incrementally increase the number of

modes in Fourier domain during training. Has to verify n <= N for (n, m) in zip(max_n_modes, n_modes).

  • If None, all the n_modes are used.

This can be updated dynamically during training.

factorizationstr, optional

Tensor factorization of the FNO layer weights to use, by default None.

  • If None, a dense tensor parametrizes the Spectral convolutions

  • Otherwise, the specified tensor factorization is used.

rankfloat, optional

tensor rank to use in above factorization, by default 1.0

fixed_rank_modesbool, optional

Modes to not factorize, by default False

implementationstr {‘factorized’, ‘reconstructed’}, optional

  • If ‘factorized’, implements tensor contraction with the individual factors of the decomposition

  • If ‘reconstructed’, implements with the reconstructed full tensorized weight.

decomposition_kwargsdict, optional

extra kwargs for tensor decomposition (see tltorch.FactorizedTensor), by default dict()

separablebool, optional (DEACTIVATED)

if True, use a depthwise separable spectral convolution, by default False

preactivationbool, optional (DEACTIVATED)

whether to compute FNO forward pass with resnet-style preactivation, by default False

conv_modulenn.Module, optional

module to use for FNOBlock’s convolutions, by default SpectralConv

References

[1]

:

Li, Z. et al. “Fourier Neural Operator for Parametric Partial Differential

Equations” (2021). ICLR 2021, https://arxiv.org/pdf/2010.08895.

Examples

>>> from neuralop.models import FNO
>>> model = FNO(n_modes=(12,12), in_channels=1, out_channels=1, hidden_channels=64)
>>> model
FNO(
(positional_embedding): GridEmbeddingND()
(fno_blocks): FNOBlocks(
    (convs): SpectralConv(
    (weight): ModuleList(
        (0-3): 4 x DenseTensor(shape=torch.Size([64, 64, 12, 7]), rank=None)
    )
    )
        ... torch.nn.Module printout truncated ...
forward(x, output_shape=None, **kwargs)

FNO’s forward pass

  1. Applies optional positional encoding

  2. Sends inputs through a lifting layer to a high-dimensional latent

    space

  3. Applies optional domain padding to high-dimensional intermediate function representation

  4. Applies n_layers Fourier/FNO layers in sequence (SpectralConvolution + skip connections, nonlinearity)

  5. If domain padding was applied, domain padding is removed

  6. Projection of intermediate function representation to the output channels

Parameters:
xtensor

input tensor

output_shape{tuple, tuple list, None}, default is None

Gives the option of specifying the exact output shape for odd shaped inputs.

  • If None, don’t specify an output shape

  • If tuple, specifies the output-shape of the last FNO Block

  • If tuple list, specifies the exact output-shape of each FNO Block