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

Fourier Integral Kernel model class.

This class implements the Fourier Integral Kernel network, which performs global convolution 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.

Initialization of the FourierIntegralKernel class.

Parameters:

  • input_numb_fields (int) – The number of input fields.

  • output_numb_fields (int) – The number of output fields.

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

  • dimensions (int) – The number of dimensions. It can be set to 1, 2, or 3. Default is 3.

  • padding (int) – The padding size. Default is 8.

  • padding_type (str) – The padding strategy. Default is constant.

  • inner_size (int) – The inner size. Default is 20.

  • n_layers (int) – The number of layers. Default is 2.

  • func (torch.nn.Module | list[torch.nn.Module]) – The activation function. If a list is passed, it must have the same length as n_layers. If a single function is passed, it is used for all layers, except for the last one. Default is torch.nn.Tanh.

  • layers (list[int]) – The list of the dimension of inner layers. If None, n_layers of dimension inner_size are used. Otherwise, it overrides the values passed to n_layers and inner_size. Default is None.

Raises:

  • RuntimeError – If the number of layers and functions are inconsistent.

  • RunTimeError – If the number of layers and modes are inconsistent.

forward(x)[source]#

Forward pass for the FourierIntegralKernel model.

Parameters:

x (torch.Tensor | LabelTensor) – The input tensor for performing the computation. Depending on the dimensions in the initialization, it expects a tensor with the following shapes: * 1D tensors: [batch, X, channels] * 2D tensors: [batch, X, Y, channels] * 3D tensors: [batch, X, Y, Z, channels]

Raises:

Warning – If a LabelTensor is passed as input.

Returns:

The output tensor.

Return type:

torch.Tensor