Spectral Convolution Block#

class SpectralConvBlock1D(input_numb_fields, output_numb_fields, n_modes)[source]#

Bases: Module

Spectral Convolution Block for one-dimensional tensors.

This class 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 block expects an input of size [batch, input_numb_fields, N] and returns an output of size [batch, output_numb_fields, N].

Initialization of the SpectralConvBlock1D 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 (int) – The number of modes to select for each dimension. It must be at most equal to \(\floor(Nx/2)+1\).

forward(x)[source]#

Forward pass.

Parameters:

x (torch.Tensor) – The input tensor. Expected of size [batch, input_numb_fields, N].

Returns:

The input tensor. Expected of size [batch, output_numb_fields, N].

Return type:

torch.Tensor

class SpectralConvBlock2D(input_numb_fields, output_numb_fields, n_modes)[source]#

Bases: Module

Spectral Convolution Block for two-dimensional tensors.

This class 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 block expects an input of size [batch, input_numb_fields, Nx, Ny] and returns an output of size [batch, output_numb_fields, Nx, Ny].

Initialization of the SpectralConvBlock2D 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\).

Raises:

  • ValueError – If the number of modes is not consistent.

  • ValueError – If the number of modes is not a list or tuple.

forward(x)[source]#

Forward pass.

Parameters:

x (torch.Tensor) – The input tensor. Expected of size [batch, input_numb_fields, Nx, Ny].

Returns:

The input tensor. Expected of size [batch, output_numb_fields, Nx, Ny].

Return type:

torch.Tensor

class SpectralConvBlock3D(input_numb_fields, output_numb_fields, n_modes)[source]#

Bases: Module

Spectral Convolution Block for three-dimensional tensors.

This class 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 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].

Initialization of the SpectralConvBlock3D 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\).

Raises:

  • ValueError – If the number of modes is not consistent.

  • ValueError – If the number of modes is not a list or tuple.

forward(x)[source]#

Forward pass.

Parameters:

x (torch.Tensor) – The input tensor. Expected of size [batch, input_numb_fields, Nx, Ny, Nz].

Returns:

The input tensor. Expected of size [batch, output_numb_fields, Nx, Ny, Nz].

Return type:

torch.Tensor