FourierIntegralKernel#

class FourierIntegralKernel(input_numb_fields, output_numb_fields, 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: Module

Implementation of Fourier Integral Kernel network.

This class implements the Fourier Integral Kernel network, which is a PINA implementation of Fourier Neural Operator kernel network. It performs global convolution by operating in 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:

  • input_numb_fields (int) – Number of input fields.

  • output_numb_fields (int) – Number of output fields.

  • 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 the kernels convolution.

Return type:

torch.Tensor