Fourier Neural Operator Block#

class FourierBlock1D(input_numb_fields, output_numb_fields, n_modes, activation=<class 'torch.nn.modules.activation.Tanh'>)[source]#

Bases: Module

The inner block of the Fourier Neural Operator for 1-dimensional input tensors.

The module computes the spectral convolution of the input with a linear kernel in the fourier space, and then it maps the input back to the physical space. The output is then added to a Linear tranformation of the input in the physical space. Finally an activation function is applied to the output.

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 FourierBlock1D class.

Parameters:

  • input_numb_fields (int) – The number of channels for the input.

  • output_numb_fields (int) – The number of channels for the output.

  • n_modes (list[int] | tuple[int]) – The number of modes to select for each dimension. It must be at most equal to \(\floor(Nx/2)+1\).

  • activation (torch.nn.Module) – The activation function. Default is torch.nn.Tanh.

forward(x)[source]#

Forward pass of the block. It performs a spectral convolution and a linear transformation of the input. Then, it sums the results.

Parameters:

x (torch.Tensor) – The input tensor for performing the computation.

Returns:

The output tensor.

Return type:

torch.Tensor

class FourierBlock2D(input_numb_fields, output_numb_fields, n_modes, activation=<class 'torch.nn.modules.activation.Tanh'>)[source]#

Bases: Module

The inner block of the Fourier Neural Operator for 2-dimensional input tensors.

The module computes the spectral convolution of the input with a linear kernel in the fourier space, and then it maps the input back to the physical space. The output is then added to a Linear tranformation of the input in the physical space. Finally an activation function is applied to the output.

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 FourierBlock2D class.

Parameters:

  • input_numb_fields (int) – The number of channels for the input.

  • output_numb_fields (int) – The number of channels for the output.

  • n_modes (list[int] | tuple[int]) – The number of modes to select for each dimension. It must be at most equal to \(\floor(Nx/2)+1\), \(\floor(Ny/2)+1\).

  • activation (torch.nn.Module) – The activation function. Default is torch.nn.Tanh.

forward(x)[source]#

Forward pass of the block. It performs a spectral convolution and a linear transformation of the input. Then, it sums the results.

Parameters:

x (torch.Tensor) – The input tensor for performing the computation.

Returns:

The output tensor.

Return type:

torch.Tensor

class FourierBlock3D(input_numb_fields, output_numb_fields, n_modes, activation=<class 'torch.nn.modules.activation.Tanh'>)[source]#

Bases: Module

The inner block of the Fourier Neural Operator for 3-dimensional input tensors.

The module computes the spectral convolution of the input with a linear kernel in the fourier space, and then it maps the input back to the physical space. The output is then added to a Linear tranformation of the input in the physical space. Finally an activation function is applied to the output.

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 FourierBlock3D class.

Parameters:

  • input_numb_fields (int) – The number of channels for the input.

  • output_numb_fields (int) – The number of channels for the output.

  • n_modes (list[int] | tuple[int]) – The number of modes to select for each dimension. It must be at most equal to \(\floor(Nx/2)+1\), \(\floor(Ny/2)+1\), \(\floor(Nz/2)+1\).

  • activation (torch.nn.Module) – The activation function. Default is torch.nn.Tanh.

forward(x)[source]#

Forward pass of the block. It performs a spectral convolution and a linear transformation of the input. Then, it sums the results.

Parameters:

x (torch.Tensor) – The input tensor for performing the computation.

Returns:

The output tensor.

Return type:

torch.Tensor