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 transformation 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.

Example:
>>> import torch
>>> from pina.model.block import FourierBlock1D
>>> block = FourierBlock1D(
...     input_numb_fields=2, output_numb_fields=2, n_modes=16
... )
>>> x = torch.randn(10, 2, 50)
>>> out = block(x)

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 transformation 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.

Example:
>>> import torch
>>> from pina.model.block import FourierBlock2D
>>> block = FourierBlock2D(
...     input_numb_fields=2, output_numb_fields=2, n_modes=[8, 8]
... )
>>> x = torch.randn(10, 2, 50, 50)
>>> out = block(x)

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 transformation 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.

Example:
>>> import torch
>>> from pina.model.block import FourierBlock3D
>>> block = FourierBlock3D(
...     input_numb_fields=2, output_numb_fields=2, n_modes=[4, 4, 4]
... )
>>> x = torch.randn(10, 2, 20, 20, 20)
>>> out = block(x)

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