List of Fourier Layers

All Fourier layers are types of convolutions, although they do not have a direct counterpart in torch.

Formula

Using the notation from Li et al. [40], the Fourier layer is defined as

\[FL(x) = \underbrace{\mathcal{F}^{-1}( R ( \mathcal{F}(x) ) )}_\text{Spectral Convolution} + \underbrace{W(x)}_\text{Convolution}\]

The spectral convolution \(\mathcal{F}^{-1}( R ( \mathcal{F}(x) ) )\) is computed by UQpy’s sml.functional.spectral_conv function and \(W(x)\) is computed by torch.nn.functional.conv. The Fourier1d layer calls spectral_conv1d and conv1d and the higher dimensional Fourier layers call the appropriate higher-dimensional functions.

The forward Fourier transform \(\mathcal{F}\) and its inverse \(\mathcal{F}^{-1}\) are computed by torch.fft. The linear transformation \(R\) is the learnable parameter weight_spectral. weight_spectral contains real numbers (torch.float) that are cast to complex (torch.cfloat) with 0 in the imaginary component for compatibility with the spectral convolutions computed by the Scientific Machine Learning functional submodule. The convolution \(W\) is computed by the appropriate convolution from torch.nn.functional using the learnable weights weight_conv and optional bias bias_conv.

Fourier1d

class Fourier1d(width, modes, bias=True, device=None)[source]

A 1d Fourier layer to compute \(\mathcal{F}^{-1} (R (\mathcal{F}x)) + W(x)\)

Parameters:

  • width (int) – Number of neurons in the layer and channels in the spectral convolution

  • modes (int) – Number of Fourier modes to keep, at most \(\lfloor L / 2 \rfloor + 1\)

  • bias (bool) – If True, adds a learnable bias to the convolution. Default: True

Note

This class does not accept the dtype argument since Fourier layers require real and complex tensors as described by the attributes.

Shape:

  • Input: \((N, \text{width}, L)\)

  • Output: \((N, \text{width}, L)\)

Attributes:

  • weight_spectral (torch.nn.Parameter): The learnable weights of the spectral convolution of shape \((\text{width}, \text{width}, \text{modes})\) with complex entries. The initial values of these weights are sampled from \(\mathcal{U}(-\sqrt{k}, \sqrt{k})\) where \(k = \frac{1}{\text{width}}\).

  • weight_conv (torch.nn.Parameter): The learnable weights of the convolution of shape \((\text{width}, \text{width}, \text{kernel_size})\) with real entries. The \(\text{kernel_size}=1\). The initial values of these weights are sampled from \(\mathcal{U}(-\sqrt{k}, \sqrt{k})\) where \(k = \frac{1}{\text{width}}\).

  • bias_conv (torch.nn.Parameter): The learnable bias of the convolution of shape \((\text{width})\) with real entries. If bias is True, then the initial values of these weights are sampled from \(\mathcal{U}(-\sqrt{k}, \sqrt{k})\) where \(k = \frac{1}{\text{width}}\).

The kernel for the convolution is fixed as \(\text{kernel_size} = 1\).

Example:

>>> length = 128
>>> modes = (length // 2) + 1
>>> width = 8
>>> f = sml.Fourier1d(width, modes)
>>> input = torch.rand(2, width, length)
>>> output = f(input)
forward(x)[source]

Compute \(\mathcal{F}^{-1} (R (\mathcal{F}x)) + W(x)\)

Parameters:

x (Tensor) – Tensor of shape \((N, \text{width}, L)\)

Return type:

Tensor

Returns:

Tensor of shape \((N, \text{width}, L)\)

Fourier2d

class Fourier2d(width, modes, bias=True, device=None)[source]

A 2d Fourier layer to compute \(\mathcal{F}^{-1} (R (\mathcal{F}x)) + W(x)\)

Parameters:

  • width (int) – Number of neurons in the layer and channels in the spectral convolution

  • modes (tuple[int, int]) – Tuple of Fourier modes to keep. At most \((\lfloor H / 2 \rfloor + 1, \lfloor W / 2 \rfloor + 1)\)

  • bias (bool) – If True, adds a learnable bias to the convolution. Default: True

Shape:

  • Input: \((N, \text{width}, H, W)\)

  • Output: \((N, \text{width}, H, W)\)

Attributes:

  • weight_spectral (torch.nn.Parameter): The learnable weights for the spectral convolution of shape \((2, \text{width}, \text{width}, \text{modes[0]}, \text{modes[1]})\). The initial values of these weights are sampled from \(\mathcal{U}(-\sqrt{k}, \sqrt{k})\) where \(k = \frac{1}{\text{width}}\).

  • weight_conv (torch.nn.Parameter): The learnable weights of the convolution of shape \((\text{width}, \text{width}, \text{kernel_size[0]}, \text{kernel_size[1]})\) with real entries. The \(\text{kernel_size} = (1, 1)\). The initial values of these weights are sampled from \(\mathcal{U}(-\sqrt{k}, \sqrt{k})\) where \(k = \frac{1}{\text{width}}\).

  • bias_conv (torch.nn.Parameter): The learnable bias of the convolution of shape \((\text{width})\) with real entries. If bias is True, then the initial values of these weights are sampled from \(\mathcal{U}(-\sqrt{k}, \sqrt{k})\) where \(k = \frac{1}{\text{width}}\).

The kernel for the convolution is fixed as \(\text{kernel_size}=(1, 1)\).

Example:

>>> h, w = 64, 128
>>> modes = (h//2 + 1, w//2 + 1)
>>> width = 5
>>> f = sml.Fourier2d(width, modes)
>>> input = torch.rand(2, width, h, w)
>>> output = f(input)
forward(x)[source]

Compute \(\mathcal{F}^{-1} (R (\mathcal{F}x)) + W(x)\)

Parameters:

x (Tensor) – Tensor of shape \((N, \text{width}, H, W)\)

Return type:

Tensor

Returns:

Tensor of shape \((N, \text{width}, H, W)\)

Fourier3d

class Fourier3d(width, modes, bias=True, device=None)[source]

A 3d Fourier layer to compute \(\mathcal{F}^{-1} (R (\mathcal{F}x)) + W(x)\)

Parameters:

  • width (int) – Number of neurons in the layer and channels in the spectral convolution

  • modes (tuple[int, int, int]) – Tuple of Fourier modes to keep. At most \((\lfloor D / 2 \rfloor + 1, \lfloor H / 2 \rfloor + 1, \lfloor W / 2 \rfloor + 1)\)

  • bias (bool) – If True, adds a learnable bias to the convolution. Default: True

Shape:

  • Input: \((N, \text{width}, D, H, W)\)

  • Output: \((N, \text{width}, D, H, W)\)

Attributes:

  • weight_spectral (torch.nn.Parameter): The learnable weights for the spectral convolution of shape \((4, \text{width}, \text{width}, \text{modes[0]}, \text{modes[1]}, \text{modes[2]})\) with complex entries. The initial values of these weights are sampled from \(\mathcal{U}(-\sqrt{k}, \sqrt{k})\) where \(k = \frac{1}{\text{width}}\).

  • weight_conv (torch.nn.Parameter): The learnable weights of the convolution of shape \((\text{width}, \text{width}, \text{kernel_size[0]}, \text{kernel_size[1]}, \text{kernel_size[2]})\) with real entries. The \(\text{kernel_size} = (1, 1, 1)\). The initial values of these weights are sampled from \(\mathcal{U}(-\sqrt{k}, \sqrt{k})\) where \(k = \frac{1}{\text{width}}\).

  • bias_conv (torch.nn.Parameter): The learnable bias of the convolution of shape \((\text{width})\) with real entries. If bias is True, then the initial values of these weights are sampled from \(\mathcal{U}(-\sqrt{k}, \sqrt{k})\) where \(k = \frac{1}{\text{width}}\).

The kernel for the convolution is fixed as \(\text{kernel_size}=(1, 1, 1)\).

Example:

>>> d, h, w = 64, 128, 256
>>> modes = (d//2 + 1, h//2 + 1, w//2 + 1)
>>> width = 5
>>> f = sml.Fourier3d(width, modes)
>>> input = torch.rand(2, width, d, h, w)
>>> output = f(input)
forward(x)[source]

Compute \(\mathcal{F}^{-1} (R (\mathcal{F}x)) + W(x)\)

Parameters:

x (Tensor) – Tensor of shape \((N, \text{width}, D, H, W)\)

Return type:

Tensor

Returns:

Tensor of shape \((N, \text{width}, D, H, W)\)