Fourier Neural Operators#

class physicsnemo.models.fno.fno.FNO(*args, **kwargs)[source]#

Bases: Module

Fourier neural operator (FNO) model.

The FNO architecture supports options for 1D, 2D, 3D and 4D fields which can be controlled using the dimension parameter.

Parameters:

in_channels (int) – Number of input channels.
out_channels (int) – Number of output channels.
decoder_layers (int, optional, default=1) – Number of decoder layers.
decoder_layer_size (int, optional, default=32) – Number of neurons in decoder layers.
decoder_activation_fn (str, optional, default="silu") – Activation function for decoder.
dimension (int, optional, default=2) – Model dimensionality (supports 1, 2, 3, 4).
latent_channels (int, optional, default=32) – Latent features size in spectral convolutions.
num_fno_layers (int, optional, default=4) – Number of spectral convolutional layers.
num_fno_modes (Union[int, List[int]], optional, default=16) – Number of Fourier modes kept in spectral convolutions.
padding (int, optional, default=8) – Domain padding for spectral convolutions.
padding_type (str, optional, default="constant") – Type of padding for spectral convolutions.
activation_fn (str, optional, default="gelu") – Activation function.
coord_features (bool, optional, default=True) – Use coordinate grid as additional feature map.

Forward:

x (torch.Tensor) – Input tensor. Shape depends on dimension:

1D: \((B, C_{in}, L)\) where \(L\) is sequence length
2D: \((B, C_{in}, H, W)\)
3D: \((B, C_{in}, D, H, W)\)
4D: \((B, C_{in}, X, Y, Z, T)\)

Outputs:

torch.Tensor – Output tensor with same spatial dimensions as input:

1D: \((B, C_{out}, L)\)
2D: \((B, C_{out}, H, W)\)
3D: \((B, C_{out}, D, H, W)\)
4D: \((B, C_{out}, X, Y, Z, T)\)

Examples

>>> import torch
>>> import physicsnemo
>>> # Define a 2D FNO model
>>> model = physicsnemo.models.fno.FNO(
...     in_channels=4,
...     out_channels=3,
...     decoder_layers=2,
...     decoder_layer_size=32,
...     dimension=2,
...     latent_channels=32,
...     num_fno_layers=2,
...     padding=0,
... )
>>> input = torch.randn(32, 4, 32, 32)  # (N, C, H, W)
>>> output = model(input)
>>> output.size()
torch.Size([32, 3, 32, 32])

See also

Fourier Neural Operator (FNO) :: Original FNO paper.

class physicsnemo.models.fno.fno.FNO1DEncoder(*args, **kwargs)[source]#

Bases: Module

1D Spectral encoder for FNO.

This encoder applies a lifting network followed by spectral convolution layers in the Fourier domain for 1D input data.

Parameters:

in_channels (int, optional, default=1) – Number of input channels.
num_fno_layers (int, optional, default=4) – Number of spectral convolutional layers.
fno_layer_size (int, optional, default=32) – Latent features size in spectral convolutions.
num_fno_modes (Union[int, List[int]], optional, default=16) – Number of Fourier modes kept in spectral convolutions.
padding (Union[int, List[int]], optional, default=8) – Domain padding for spectral convolutions.
padding_type (str, optional, default="constant") – Type of padding for spectral convolutions.
activation_fn (nn.Module, optional, default=nn.GELU()) – Activation function.
coord_features (bool, optional, default=True) – Use coordinate grid as additional feature map.

Forward:

x (torch.Tensor) – Input tensor of shape \((B, C_{in}, L)\) where \(B\) is batch size, \(C_{in}\) is the number of input channels, and \(L\) is the sequence length (spatial dimension).

Outputs:

torch.Tensor – Output tensor of shape \((B, C_{latent}, L)\) where \(C_{latent}\) is fno_layer_size.

Examples

>>> import torch
>>> encoder = FNO1DEncoder(in_channels=3, fno_layer_size=32, num_fno_modes=8)
>>> x = torch.randn(4, 3, 64)
>>> output = encoder(x)
>>> output.shape
torch.Size([4, 32, 64])

grid_to_points( value: Tensor, ) → Tuple[Tensor, List[int]][source]#

Convert from grid-based (image) to point-based representation.

Parameters:: value (Tensor) – Grid tensor of shape \((B, C, L)\).
Returns:: Tuple of (flattened tensor, original shape).
Return type:: Tuple[Tensor, List[int]]

points_to_grid( value: Tensor, shape: List[int], ) → Tensor[source]#

Convert from point-based to grid-based (image) representation.

Parameters:

value (Tensor) – Point tensor of shape \((B \times X, C)\).
shape (List[int]) – Original grid shape as [B, C, L].

Returns:

Grid tensor of shape \((B, C, L)\).

Return type:

Tensor

class physicsnemo.models.fno.fno.FNO2DEncoder(*args, **kwargs)[source]#

Bases: Module

2D Spectral encoder for FNO.

This encoder applies a lifting network followed by spectral convolution layers in the Fourier domain for 2D input data.

Parameters:

in_channels (int, optional, default=1) – Number of input channels.
num_fno_layers (int, optional, default=4) – Number of spectral convolutional layers.
fno_layer_size (int, optional, default=32) – Latent features size in spectral convolutions.
num_fno_modes (Union[int, List[int]], optional, default=16) – Number of Fourier modes kept in spectral convolutions.
padding (Union[int, List[int]], optional, default=8) – Domain padding for spectral convolutions.
padding_type (str, optional, default="constant") – Type of padding for spectral convolutions.
activation_fn (nn.Module, optional, default=nn.GELU()) – Activation function.
coord_features (bool, optional, default=True) – Use coordinate grid as additional feature map.

Forward:

x (torch.Tensor) – Input tensor of shape \((B, C_{in}, H, W)\) where \(B\) is batch size, \(C_{in}\) is the number of input channels, and \(H, W\) are spatial dimensions.

Outputs:

torch.Tensor – Output tensor of shape \((B, C_{latent}, H, W)\) where \(C_{latent}\) is fno_layer_size.

Examples

>>> import torch
>>> encoder = FNO2DEncoder(in_channels=3, fno_layer_size=32, num_fno_modes=8)
>>> x = torch.randn(4, 3, 32, 32)
>>> output = encoder(x)
>>> output.shape
torch.Size([4, 32, 32, 32])

grid_to_points( value: Tensor, ) → Tuple[Tensor, List[int]][source]#

Convert from grid-based (image) to point-based representation.

Parameters:: value (Tensor) – Grid tensor of shape \((B, C, H, W)\).
Returns:: Tuple of (flattened tensor, original shape).
Return type:: Tuple[Tensor, List[int]]

points_to_grid( value: Tensor, shape: List[int], ) → Tensor[source]#

Convert from point-based to grid-based (image) representation.

Parameters:

value (Tensor) – Point tensor of shape \((B \times H \times W, C)\).
shape (List[int]) – Original grid shape as [B, C, H, W].

Returns:

Grid tensor of shape \((B, C, H, W)\).

Return type:

Tensor

class physicsnemo.models.fno.fno.FNO3DEncoder(*args, **kwargs)[source]#

Bases: Module

3D Spectral encoder for FNO.

This encoder applies a lifting network followed by spectral convolution layers in the Fourier domain for 3D input data.

Parameters:

in_channels (int, optional, default=1) – Number of input channels.
num_fno_layers (int, optional, default=4) – Number of spectral convolutional layers.
fno_layer_size (int, optional, default=32) – Latent features size in spectral convolutions.
num_fno_modes (Union[int, List[int]], optional, default=16) – Number of Fourier modes kept in spectral convolutions.
padding (Union[int, List[int]], optional, default=8) – Domain padding for spectral convolutions.
padding_type (str, optional, default="constant") – Type of padding for spectral convolutions.
activation_fn (nn.Module, optional, default=nn.GELU()) – Activation function.
coord_features (bool, optional, default=True) – Use coordinate grid as additional feature map.

Forward:

x (torch.Tensor) – Input tensor of shape \((B, C_{in}, D, H, W)\) where \(B\) is batch size, \(C_{in}\) is the number of input channels, and \(D, H, W\) are spatial dimensions.

Outputs:

torch.Tensor – Output tensor of shape \((B, C_{latent}, D, H, W)\) where \(C_{latent}\) is fno_layer_size.

Examples

>>> import torch
>>> encoder = FNO3DEncoder(in_channels=3, fno_layer_size=32, num_fno_modes=8)
>>> x = torch.randn(4, 3, 16, 16, 16)
>>> output = encoder(x)
>>> output.shape
torch.Size([4, 32, 16, 16, 16])

grid_to_points( value: Tensor, ) → Tuple[Tensor, List[int]][source]#

Convert from grid-based (image) to point-based representation.

Parameters:: value (Tensor) – Grid tensor of shape \((B, C, D, H, W)\).
Returns:: Tuple of (flattened tensor, original shape).
Return type:: Tuple[Tensor, List[int]]

points_to_grid( value: Tensor, shape: List[int], ) → Tensor[source]#

Convert from point-based to grid-based (image) representation.

Parameters:

value (Tensor) – Point tensor of shape \((B \times D \times H \times W, C)\).
shape (List[int]) – Original grid shape as [B, C, D, H, W].

Returns:

Grid tensor of shape \((B, C, D, H, W)\).

Return type:

Tensor

class physicsnemo.models.fno.fno.FNO4DEncoder(*args, **kwargs)[source]#

Bases: Module

4D Spectral encoder for FNO.

This encoder applies a lifting network followed by spectral convolution layers in the Fourier domain for 4D input data (3D spatial + time).

Parameters:

in_channels (int, optional, default=1) – Number of input channels.
num_fno_layers (int, optional, default=4) – Number of spectral convolutional layers.
fno_layer_size (int, optional, default=32) – Latent features size in spectral convolutions.
num_fno_modes (Union[int, List[int]], optional, default=16) – Number of Fourier modes kept in spectral convolutions.
padding (Union[int, List[int]], optional, default=8) – Domain padding for spectral convolutions.
padding_type (str, optional, default="constant") – Type of padding for spectral convolutions.
activation_fn (nn.Module, optional, default=nn.GELU()) – Activation function.
coord_features (bool, optional, default=True) – Use coordinate grid as additional feature map.

Forward:

x (torch.Tensor) – Input tensor of shape \((B, C_{in}, X, Y, Z, T)\) where \(B\) is batch size, \(C_{in}\) is the number of input channels, and \(X, Y, Z, T\) are spatial and temporal dimensions.

Outputs:

torch.Tensor – Output tensor of shape \((B, C_{latent}, X, Y, Z, T)\) where \(C_{latent}\) is fno_layer_size.

Examples

>>> import torch
>>> encoder = FNO4DEncoder(in_channels=3, fno_layer_size=32, num_fno_modes=4)
>>> x = torch.randn(2, 3, 8, 8, 8, 8)
>>> output = encoder(x)
>>> output.shape
torch.Size([2, 32, 8, 8, 8, 8])

grid_to_points( value: Tensor, ) → Tuple[Tensor, List[int]][source]#

Convert from grid-based (image) to point-based representation.

Parameters:: value (Tensor) – Grid tensor of shape \((B, C, X, Y, Z, T)\).
Returns:: Tuple of (flattened tensor, original shape).
Return type:: Tuple[Tensor, List[int]]

points_to_grid( value: Tensor, shape: List[int], ) → Tensor[source]#

Convert from point-based to grid-based (image) representation.

Parameters:

value (Tensor) – Point tensor of shape \((B \times X \times Y \times Z \times T, C)\).
shape (List[int]) – Original grid shape as [B, C, X, Y, Z, T].

Returns:

Grid tensor of shape \((B, C, X, Y, Z, T)\).

Return type:

Tensor

class physicsnemo.models.afno.afno.AFNO(*args, **kwargs)[source]#

Bases: Module

Adaptive Fourier neural operator (AFNO) model.

AFNO is a model that is designed for 2D images only. It combines patch embedding with spectral convolution blocks in the Fourier domain.

Parameters:

inp_shape (List[int]) – Input image dimensions as [height, width].
in_channels (int) – Number of input channels.
out_channels (int) – Number of output channels.
patch_size (List[int], optional, default=[16, 16]) – Size of image patches as [patch_height, patch_width].
embed_dim (int, optional, default=256) – Embedded channel size.
depth (int, optional, default=4) – Number of AFNO layers.
mlp_ratio (float, optional, default=4.0) – Ratio of layer MLP latent variable size to input feature size.
drop_rate (float, optional, default=0.0) – Drop out rate in layer MLPs.
num_blocks (int, optional, default=16) – Number of blocks in the block-diag frequency weight matrices.
sparsity_threshold (float, optional, default=0.01) – Sparsity threshold (softshrink) of spectral features.
hard_thresholding_fraction (float, optional, default=1.0) – Threshold for limiting number of modes used, in range [0, 1].

Forward:

x (torch.Tensor) – Input tensor of shape \((B, C_{in}, H, W)\) where \(B\) is batch size, \(C_{in}\) is the number of input channels, and \(H, W\) are spatial dimensions matching inp_shape.

Outputs:

torch.Tensor – Output tensor of shape \((B, C_{out}, H, W)\) where \(C_{out}\) is out_channels.

Examples

>>> import torch
>>> import physicsnemo
>>> model = physicsnemo.models.afno.AFNO(
...     inp_shape=[32, 32],
...     in_channels=2,
...     out_channels=1,
...     patch_size=(8, 8),
...     embed_dim=16,
...     depth=2,
...     num_blocks=2,
... )
>>> input = torch.randn(32, 2, 32, 32)  # (N, C, H, W)
>>> output = model(input)
>>> output.size()
torch.Size([32, 1, 32, 32])

See also

DistributedAFNO :: Distributed (model-parallel) AFNO for multi-GPU training.
Adaptive Fourier Neural Operator (AFNO) :: Original AFNO paper.

forward( x: Float[Tensor, 'B C_in H W'], ) → Float[Tensor, 'B C_out H W'][source]#: Forward pass of the AFNO model.

class physicsnemo.models.afno.afno.AFNO2DLayer(*args, **kwargs)[source]#

Bases: Module

AFNO spectral convolution layer.

This layer performs spectral mixing using block-diagonal weight matrices in the Fourier domain with soft shrinkage for sparsity.

Parameters:

hidden_size (int) – Feature dimensionality.
num_blocks (int, optional, default=8) – Number of blocks used in the block diagonal weight matrix.
sparsity_threshold (float, optional, default=0.01) – Sparsity threshold (softshrink) of spectral features.
hard_thresholding_fraction (float, optional, default=1) – Threshold for limiting number of modes used, in range [0, 1].
hidden_size_factor (int, optional, default=1) – Factor to increase spectral features by after weight multiplication.

Forward:

x (torch.Tensor) – Input tensor of shape \((B, H, W, C)\) where \(B\) is batch size, \(H, W\) are spatial dimensions, and \(C\) is hidden_size.

Outputs:

torch.Tensor – Output tensor of shape \((B, H, W, C)\).

Examples

>>> import torch
>>> layer = AFNO2DLayer(hidden_size=64, num_blocks=8)
>>> x = torch.randn(4, 32, 32, 64)
>>> output = layer(x)
>>> output.shape
torch.Size([4, 32, 32, 64])

class physicsnemo.models.afno.afno.AFNOMlp(*args, **kwargs)[source]#

Bases: Module

Fully-connected Multi-layer perception used inside AFNO.

Parameters:

in_features (int) – Input feature size.
latent_features (int) – Latent feature size.
out_features (int) – Output feature size.
activation_fn (nn.Module, optional, default=nn.GELU()) – Activation function.
drop (float, optional, default=0.0) – Drop out rate.

Forward:

x (torch.Tensor) – Input tensor of shape \((*, D_{in})\) where \(D_{in}\) is in_features.

Outputs:

torch.Tensor – Output tensor of shape \((*, D_{out})\) where \(D_{out}\) is out_features.

Examples

>>> import torch
>>> mlp = AFNOMlp(in_features=64, latent_features=128, out_features=64)
>>> x = torch.randn(4, 32, 32, 64)
>>> output = mlp(x)
>>> output.shape
torch.Size([4, 32, 32, 64])

class physicsnemo.models.afno.afno.Block(*args, **kwargs)[source]#

Bases: Module

AFNO block consisting of spectral convolution and MLP.

Parameters:

embed_dim (int) – Embedded feature dimensionality.
num_blocks (int, optional, default=8) – Number of blocks used in the block diagonal weight matrix.
mlp_ratio (float, optional, default=4.0) – Ratio of MLP latent variable size to input feature size.
drop (float, optional, default=0.0) – Drop out rate in MLP.
activation_fn (nn.Module, optional, default=nn.GELU()) – Activation function used in MLP.
norm_layer (nn.Module, optional, default=nn.LayerNorm) – Normalization function.
double_skip (bool, optional, default=True) – Whether to use double skip connections.
sparsity_threshold (float, optional, default=0.01) – Sparsity threshold (softshrink) of spectral features.
hard_thresholding_fraction (float, optional, default=1.0) – Threshold for limiting number of modes used, in range [0, 1].

Forward:

x (torch.Tensor) – Input tensor of shape \((B, H, W, C)\) where \(B\) is batch size, \(H, W\) are spatial dimensions, and \(C\) is embed_dim.

Outputs:

torch.Tensor – Output tensor of shape \((B, H, W, C)\).

Examples

>>> import torch
>>> from physicsnemo.models.afno.afno import Block
>>> block = Block(embed_dim=64, num_blocks=8)
>>> x = torch.randn(2, 8, 8, 64)  # (B, H, W, C)
>>> out = block(x)
>>> out.shape
torch.Size([2, 8, 8, 64])

physicsnemo.models.afno.afno.PatchEmbed#: alias of AFNOPatchEmbed

class physicsnemo.models.afno.modafno.ModAFNO(*args, **kwargs)[source]#

Bases: Module

Modulated Adaptive Fourier neural operator (ModAFNO) model.

ModAFNO extends AFNO with modulation capabilities for conditioning on auxiliary inputs (e.g., time, parameters).

Parameters:

inp_shape (List[int]) – Input image dimensions as [height, width].
in_channels (int, optional, default=155) – Number of input channels.
out_channels (int, optional, default=73) – Number of output channels.
embed_model (dict, optional) – Dictionary of arguments to pass to the ModEmbedNet embedding model.
patch_size (List[int], optional, default=[2, 2]) – Size of image patches as [patch_height, patch_width].
embed_dim (int, optional, default=512) – Embedded channel size.
mod_dim (int, optional, default=64) – Modulation input dimensionality.
modulate_filter (bool, optional, default=True) – Whether to compute the modulation for the FFT filter.
modulate_mlp (bool, optional, default=True) – Whether to compute the modulation for the MLP.
scale_shift_mode (Literal["complex", "real"], optional, default="complex") – If "complex", compute the scale-shift operation using complex operations. If "real", use real operations.
depth (int, optional, default=12) – Number of AFNO layers.
mlp_ratio (float, optional, default=2.0) – Ratio of layer MLP latent variable size to input feature size.
drop_rate (float, optional, default=0.0) – Drop out rate in layer MLPs.
num_blocks (int, optional, default=1) – Number of blocks in the block-diag frequency weight matrices.
sparsity_threshold (float, optional, default=0.01) – Sparsity threshold (softshrink) of spectral features.
hard_thresholding_fraction (float, optional, default=1.0) – Threshold for limiting number of modes used, in range [0, 1].

Forward:

x (torch.Tensor) – Input tensor of shape \((B, C_{in}, H, W)\).
mod (torch.Tensor) – Modulation input of shape \((B, 1)\) or \((B,)\).

Outputs:

torch.Tensor – Output tensor of shape \((B, C_{out}, H, W)\).

Examples

>>> import torch
>>> from physicsnemo.models.afno import ModAFNO
>>> model = ModAFNO(
...     inp_shape=[32, 32],
...     in_channels=2,
...     out_channels=1,
...     patch_size=(8, 8),
...     embed_dim=16,
...     depth=2,
...     num_blocks=2,
... )
>>> input = torch.randn(32, 2, 32, 32)  # (N, C, H, W)
>>> time = torch.full((32, 1), 0.5)
>>> output = model(input, time)
>>> output.size()
torch.Size([32, 1, 32, 32])