FNO#

class FNO(lifting_net, projecting_net, n_modes, dimensions=3, padding=8, padding_type='constant', inner_size=20, n_layers=2, func=<class 'torch.nn.modules.activation.Tanh'>, layers=None)[source]#

Bases: KernelNeuralOperator

The PINA implementation of Fourier Neural Operator network.

Fourier Neural Operator (FNO) is a general architecture for learning Operators. Unlike traditional machine learning methods FNO is designed to map entire functions to other functions. It can be trained with Supervised learning strategies. FNO does global convolution by performing the operation on the Fourier space.

See also

Original reference: Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., Bhattacharya, K., Stuart, A., & Anandkumar, A. (2020). Fourier neural operator for parametric partial differential equations. DOI: arXiv preprint arXiv:2010.08895.

Parameters:

  • lifting_net (torch.nn.Module) – The neural network for lifting the input.

  • projecting_net (torch.nn.Module) – The neural network for projecting the output.

  • n_modes (int | list[int]) – Number of modes.

  • dimensions (int) – Number of dimensions (1, 2, or 3).

  • padding (int) – Padding size, defaults to 8.

  • padding_type (str) – Type of padding, defaults to constant.

  • inner_size (int) – Inner size, defaults to 20.

  • n_layers (int) – Number of layers, defaults to 2.

  • func (torch.nn.Module) – Activation function, defaults to nn.Tanh.

  • layers (list[int]) – List of layer sizes, defaults to None.

forward(x)[source]#

Forward computation for Fourier Neural Operator. It performs a lifting of the input by the lifting_net. Then different layers of Fourier Blocks are applied. Finally the output is projected to the final dimensionality by the projecting_net.

Parameters:

x (torch.Tensor) –

The input tensor for fourier block, depending on dimension in the initialization. In particular it is expected:

  • 1D tensors: [batch, X, channels]

  • 2D tensors: [batch, X, Y, channels]

  • 3D tensors: [batch, X, Y, Z, channels]

Returns:

The output tensor obtained from FNO.

Return type:

torch.Tensor