neuralop.models.FNO

class neuralop.models.FNO(*args, **kwargs)[source]

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.

Attributes:
n_modes

Methods

forward(x[, output_shape])

FNO's forward pass

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)[source]

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