Fourier Layers#

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

Bases: Module

Fourier block implementation for three dimensional input tensor. The combination of Fourier blocks make up the Fourier Neural Operator

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.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]#

Forward computation for Fourier Block. It performs a spectral convolution and a linear transformation of the input and sum the results.

Parameters:

x (torch.Tensor) – The input tensor for fourier block, expect of size [batch, input_numb_fields, x].

Returns:

The output tensor obtained from the fourier block of size [batch, output_numb_fields, x].

Return type:

torch.Tensor

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

Bases: Module

Fourier block implementation for two dimensional input tensor. The combination of Fourier blocks make up the Fourier Neural Operator

See also

Original reference: Li, Zongyi, et al. Fourier neural operator for parametric partial differential equations. arXiv preprint arXiv:2010.08895 (2020) <https://arxiv.org/abs/2010.08895.pdf>`_.

PINA implementation of Fourier block two dimensions. 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.

The block expects an input of size [batch, input_numb_fields, Nx, Ny] and returns an output of size [batch, output_numb_fields, Nx, Ny].

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 | tuple) – Number of modes to select for each dimension. It must be at most equal to the floor(Nx/2)+1 and floor(Ny/2)+1.

  • activation (torch.nn.Module) – The activation function.

forward(x)[source]#

Forward computation for Fourier Block. It performs a spectral convolution and a linear transformation of the input and sum the results.

Parameters:

x (torch.Tensor) – The input tensor for fourier block, expect of size [batch, input_numb_fields, x, y].

Returns:

The output tensor obtained from the fourier block of size [batch, output_numb_fields, x, y, z].

Return type:

torch.Tensor

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

Bases: Module

Fourier block implementation for three dimensional input tensor. The combination of Fourier blocks make up the Fourier Neural Operator

See also

Original reference: Li, Zongyi, et al. Fourier neural operator for parametric partial differential equations. arXiv preprint arXiv:2010.08895 (2020) <https://arxiv.org/abs/2010.08895.pdf>`_.

PINA implementation of Fourier block three dimensions. 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.

The block expects an input of size [batch, input_numb_fields, Nx, Ny, Nz] and returns an output of size [batch, output_numb_fields, Nx, Ny, Nz].

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 | tuple) – Number of modes to select for each dimension. It must be at most equal to the floor(Nx/2)+1, floor(Ny/2)+1 and floor(Nz/2)+1.

  • activation (torch.nn.Module) – The activation function.

forward(x)[source]#

Forward computation for Fourier Block. It performs a spectral convolution and a linear transformation of the input and sum the results.

Parameters:

x (torch.Tensor) – The input tensor for fourier block, expect of size [batch, input_numb_fields, x, y, z].

Returns:

The output tensor obtained from the fourier block of size [batch, output_numb_fields, x, y, z].

Return type:

torch.Tensor