Modulus Models
NVIDIA Docs Hub  NVIDIA Modulus  NVIDIA Modulus Core v0.3.0  Modulus Models
Core v0.3.0

Fully Connected Network

class modulus.models.mlp.fully_connected.FullyConnected(*args, **kwargs)[source]

Bases: Module

A densely-connected MLP architecture

Parameters

  • in_features (int, optional) – Size of input features, by default 512

  • layer_size (int, optional) – Size of every hidden layer, by default 512

  • out_features (int, optional) – Size of output features, by default 512

  • num_layers (int, optional) – Number of hidden layers, by default 6

  • activation_fn (Union[str, List[str]], optional) – Activation function to use, by default ‘silu’

  • skip_connections (bool, optional) – Add skip connections every 2 hidden layers, by default False

  • adaptive_activations (bool, optional) – Use an adaptive activation function, by default False

  • weight_norm (bool, optional) – Use weight norm on fully connected layers, by default False

Example

Copy
Copied!
            

            
>>> model = modulus.models.mlp.FullyConnected(in_features=32, out_features=64)
>>> input = torch.randn(128, 32)
>>> output = model(input)
>>> output.size()
torch.Size([128, 64])

forward(x: Tensor) → Tensor[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
class modulus.models.mlp.fully_connected.MetaData(name: str = 'FullyConnected', jit: bool = True, cuda_graphs: bool = True, amp: bool = True, amp_cpu: bool = None, amp_gpu: bool = None, torch_fx: bool = True, onnx: bool = True, onnx_gpu: bool = None, onnx_cpu: bool = None, onnx_runtime: bool = True, trt: bool = False, var_dim: int = -1, func_torch: bool = True, auto_grad: bool = True)[source]

Bases: ModelMetaData

Fourier Neural Operators

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

Bases: Module

Fourier neural operator (FNO) model.

Note

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) – Number of decoder layers, by default 1

  • decoder_layer_size (int, optional) – Number of neurons in decoder layers, by default 32

  • decoder_activation_fn (str, optional) – Activation function for decoder, by default “silu”

  • dimension (int) – Model dimensionality (supports 1, 2, 3).

  • latent_channels (int, optional) – Latent features size in spectral convolutions, by default 32

  • num_fno_layers (int, optional) – Number of spectral convolutional layers, by default 4

  • num_fno_modes (Union[int, List[int]], optional) – Number of Fourier modes kept in spectral convolutions, by default 16

  • padding (int, optional) – Domain padding for spectral convolutions, by default 8

  • padding_type (str, optional) – Type of padding for spectral convolutions, by default “constant”

  • activation_fn (str, optional) – Activation function, by default “gelu”

  • coord_features (bool, optional) – Use coordinate grid as additional feature map, by default True

Example

Copy
Copied!
            

            
>>> # define the 2d FNO model
>>> model = modulus.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])

Note

Reference: Li, Zongyi, et al. “Fourier neural operator for parametric partial differential equations.” arXiv preprint arXiv:2010.08895 (2020).

forward(x: Tensor) → Tensor[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
class modulus.models.fno.fno.FNO1DEncoder(in_channels: int = 1, num_fno_layers: int = 4, fno_layer_size: int = 32, num_fno_modes: Union[int, List[int]] = 16, padding: Union[int, List[int]] = 8, padding_type: str = 'constant', activation_fn: Module = GELU(approximate='none'), coord_features: bool = True)[source]

Bases: Module

1D Spectral encoder for FNO

Parameters

  • in_channels (int, optional) – Number of input channels, by default 1

  • num_fno_layers (int, optional) – Number of spectral convolutional layers, by default 4

  • fno_layer_size (int, optional) – Latent features size in spectral convolutions, by default 32

  • num_fno_modes (Union[int, List[int]], optional) – Number of Fourier modes kept in spectral convolutions, by default 16

  • padding (Union[int, List[int]], optional) – Domain padding for spectral convolutions, by default 8

  • padding_type (str, optional) – Type of padding for spectral convolutions, by default “constant”

  • activation_fn (nn.Module, optional) – Activation function, by default nn.GELU

  • coord_features (bool, optional) – Use coordinate grid as additional feature map, by default True
forward(x: Tensor) → Tensor[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
meshgrid(shape: List[int], device: device) → Tensor[source]

Creates 1D meshgrid feature

Parameters

  • shape (List[int]) – Tensor shape

  • device (torch.device) – Device model is on
Returns

Meshgrid tensor
Return type

Tensor
class modulus.models.fno.fno.FNO2DEncoder(in_channels: int = 1, num_fno_layers: int = 4, fno_layer_size: int = 32, num_fno_modes: Union[int, List[int]] = 16, padding: Union[int, List[int]] = 8, padding_type: str = 'constant', activation_fn: Module = GELU(approximate='none'), coord_features: bool = True)[source]

Bases: Module

2D Spectral encoder for FNO

Parameters

  • in_channels (int, optional) – Number of input channels, by default 1

  • num_fno_layers (int, optional) – Number of spectral convolutional layers, by default 4

  • fno_layer_size (int, optional) – Latent features size in spectral convolutions, by default 32

  • num_fno_modes (Union[int, List[int]], optional) – Number of Fourier modes kept in spectral convolutions, by default 16

  • padding (Union[int, List[int]], optional) – Domain padding for spectral convolutions, by default 8

  • padding_type (str, optional) – Type of padding for spectral convolutions, by default “constant”

  • activation_fn (nn.Module, optional) – Activation function, by default nn.GELU

  • coord_features (bool, optional) – Use coordinate grid as additional feature map, by default True
forward(x: Tensor) → Tensor[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
meshgrid(shape: List[int], device: device) → Tensor[source]

Creates 2D meshgrid feature

Parameters

  • shape (List[int]) – Tensor shape

  • device (torch.device) – Device model is on
Returns

Meshgrid tensor
Return type

Tensor
class modulus.models.fno.fno.FNO3DEncoder(in_channels: int = 1, num_fno_layers: int = 4, fno_layer_size: int = 32, num_fno_modes: Union[int, List[int]] = 16, padding: Union[int, List[int]] = 8, padding_type: str = 'constant', activation_fn: Module = GELU(approximate='none'), coord_features: bool = True)[source]

Bases: Module

3D Spectral encoder for FNO

Parameters

  • in_channels (int, optional) – Number of input channels, by default 1

  • num_fno_layers (int, optional) – Number of spectral convolutional layers, by default 4

  • fno_layer_size (int, optional) – Latent features size in spectral convolutions, by default 32

  • num_fno_modes (Union[int, List[int]], optional) – Number of Fourier modes kept in spectral convolutions, by default 16

  • padding (Union[int, List[int]], optional) – Domain padding for spectral convolutions, by default 8

  • padding_type (str, optional) – Type of padding for spectral convolutions, by default “constant”

  • activation_fn (nn.Module, optional) – Activation function, by default nn.GELU

  • coord_features (bool, optional) – Use coordinate grid as additional feature map, by default True
forward(x: Tensor) → Tensor[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
meshgrid(shape: List[int], device: device) → Tensor[source]

Creates 3D meshgrid feature

Parameters

  • shape (List[int]) – Tensor shape

  • device (torch.device) – Device model is on
Returns

Meshgrid tensor
Return type

Tensor
class modulus.models.fno.fno.FNO4DEncoder(in_channels: int = 1, num_fno_layers: int = 4, fno_layer_size: int = 32, num_fno_modes: Union[int, List[int]] = 16, padding: Union[int, List[int]] = 8, padding_type: str = 'constant', activation_fn: Module = GELU(approximate='none'), coord_features: bool = True)[source]

Bases: Module

4D Spectral encoder for FNO

Parameters

  • in_channels (int, optional) – Number of input channels, by default 1

  • num_fno_layers (int, optional) – Number of spectral convolutional layers, by default 4

  • fno_layer_size (int, optional) – Latent features size in spectral convolutions, by default 32

  • num_fno_modes (Union[int, List[int]], optional) – Number of Fourier modes kept in spectral convolutions, by default 16

  • padding (Union[int, List[int]], optional) – Domain padding for spectral convolutions, by default 8

  • padding_type (str, optional) – Type of padding for spectral convolutions, by default “constant”

  • activation_fn (nn.Module, optional) – Activation function, by default nn.GELU

  • coord_features (bool, optional) – Use coordinate grid as additional feature map, by default True
forward(x: Tensor) → Tensor[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
meshgrid(shape: List[int], device: device) → Tensor[source]

Creates 4D meshgrid feature

Parameters

  • shape (List[int]) – Tensor shape

  • device (torch.device) – Device model is on
Returns

Meshgrid tensor
Return type

Tensor
class modulus.models.fno.fno.MetaData(name: str = 'FourierNeuralOperator', jit: bool = True, cuda_graphs: bool = True, amp: bool = False, amp_cpu: bool = None, amp_gpu: bool = None, torch_fx: bool = False, onnx: bool = False, onnx_gpu: bool = False, onnx_cpu: bool = False, onnx_runtime: bool = False, trt: bool = False, var_dim: int = 1, func_torch: bool = False, auto_grad: bool = False)[source]

Bases: ModelMetaData
class modulus.models.afno.afno.AFNO(*args, **kwargs)[source]

Bases: Module

Adaptive Fourier neural operator (AFNO) model.

Note

AFNO is a model that is designed for 2D images only.

Parameters

  • img_size (Tuple[int, int]) – Input image dimensions (height, width)

  • in_channels (int) – Number of input channels

  • out_channels (int) – Number of output channels

  • patch_size (Tuple[int, int], optional) – Size of image patches, by default (16, 16)

  • embed_dim (int, optional) – Embedded channel size, by default 256

  • depth (int, optional) – Number of AFNO layers, by default 4

  • mlp_ratio (float, optional) – Ratio of layer MLP latent variable size to input feature size, by default 4.0

  • drop_rate (float, optional) – Drop out rate in layer MLPs, by default 0.0

  • num_blocks (int, optional) – Number of blocks in the block-diag frequency weight matrices, by default 16

  • sparsity_threshold (float, optional) – Sparsity threshold (softshrink) of spectral features, by default 0.01

  • hard_thresholding_fraction (float, optional) – Threshold for limiting number of modes used [0,1], by default 1

Example

Copy
Copied!
            

            
>>> model = modulus.models.afno.AFNO(
...     img_size=(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])

Note

Reference: Guibas, John, et al. “Adaptive fourier neural operators: Efficient token mixers for transformers.” arXiv preprint arXiv:2111.13587 (2021).

forward(x: Tensor) → Tensor[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
forward_features(x: Tensor) → Tensor[source]

Forward pass of core AFNO
class modulus.models.afno.afno.AFNO2DLayer(hidden_size: int, num_blocks: int = 8, sparsity_threshold: float = 0.01, hard_thresholding_fraction: float = 1, hidden_size_factor: int = 1)[source]

Bases: Module

AFNO spectral convolution layer

Parameters

  • hidden_size (int) – Feature dimensionality

  • num_blocks (int, optional) – Number of blocks used in the block diagonal weight matrix, by default 8

  • sparsity_threshold (float, optional) – Sparsity threshold (softshrink) of spectral features, by default 0.01

  • hard_thresholding_fraction (float, optional) – Threshold for limiting number of modes used [0,1], by default 1

  • hidden_size_factor (int, optional) – Factor to increase spectral features by after weight multiplication, by default 1
forward(x: Tensor) → Tensor[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
class modulus.models.afno.afno.AFNOMlp(in_features: int, latent_features: int, out_features: int, activation_fn: Module = GELU(approximate='none'), drop: float = 0.0)[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) – Activation function, by default nn.GELU

  • drop (float, optional) – Drop out rate, by default 0.0
forward(x: Tensor) → Tensor[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
class modulus.models.afno.afno.Block(embed_dim: int, num_blocks: int = 8, mlp_ratio: float = 4.0, drop: float = 0.0, activation_fn: ~torch.nn.modules.module.Module = GELU(approximate='none'), norm_layer: ~torch.nn.modules.module.Module = <class 'torch.nn.modules.normalization.LayerNorm'>, double_skip: bool = True, sparsity_threshold: float = 0.01, hard_thresholding_fraction: float = 1.0)[source]

Bases: Module

AFNO block, spectral convolution and MLP

Parameters

  • embed_dim (int) – Embedded feature dimensionality

  • num_blocks (int, optional) – Number of blocks used in the block diagonal weight matrix, by default 8

  • mlp_ratio (float, optional) – Ratio of MLP latent variable size to input feature size, by default 4.0

  • drop (float, optional) – Drop out rate in MLP, by default 0.0

  • activation_fn (nn.Module, optional) – Activation function used in MLP, by default nn.GELU

  • norm_layer (nn.Module, optional) – Normalization function, by default nn.LayerNorm

  • double_skip (bool, optional) – Residual, by default True

  • sparsity_threshold (float, optional) – Sparsity threshold (softshrink) of spectral features, by default 0.01

  • hard_thresholding_fraction (float, optional) – Threshold for limiting number of modes used [0,1], by default 1
forward(x: Tensor) → Tensor[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
class modulus.models.afno.afno.MetaData(name: str = 'AFNO', jit: bool = False, cuda_graphs: bool = True, amp: bool = True, amp_cpu: bool = None, amp_gpu: bool = None, torch_fx: bool = False, onnx: bool = False, onnx_gpu: bool = True, onnx_cpu: bool = False, onnx_runtime: bool = True, trt: bool = False, var_dim: int = 1, func_torch: bool = False, auto_grad: bool = False)[source]

Bases: ModelMetaData
class modulus.models.afno.afno.PatchEmbed(img_size: Tuple[int, int], in_channels: int, patch_size: Tuple[int, int] = (16, 16), embed_dim: int = 256)[source]

Bases: Module

Patch embedding layer

Converts 2D patch into a 1D vector for input to AFNO

Parameters

  • img_size (Tuple[int, int]) – Input image dimensions (height, width)

  • in_channels (int) – Number of input channels

  • patch_size (Tuple[int, int], optional) – Size of image patches, by default (16, 16)

  • embed_dim (int, optional) – Embedded channel size, by default 256
forward(x: Tensor) → Tensor[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.

Graph Neural Networks

class modulus.models.meshgraphnet.meshgraphnet.MeshGraphNet(*args, **kwargs)[source]

Bases: Module

MeshGraphNet network architecture

Parameters

  • input_dim_nodes (int) – Number of node features

  • input_dim_edges (int) – Number of edge features

  • output_dim (int) – Number of outputs

  • processor_size (int, optional) – Number of message passing blocks, by default 15

  • num_layers_node_processor (int, optional) – Number of MLP layers for processing nodes in each message passing block, by default 2

  • num_layers_edge_processor (int, optional) – Number of MLP layers for processing edge features in each message passing block, by default 2

  • hidden_dim_processor (int, optional) – Hidden layer size for the message passing blocks, by default 128

  • hidden_dim_node_encoder (int, optional) – Hidden layer size for the node feature encoder, by default 128

  • num_layers_node_encoder (int, optional) – Number of MLP layers for the node feature encoder, by default 2

  • hidden_dim_edge_encoder (int, optional) – Hidden layer size for the edge feature encoder, by default 128

  • num_layers_edge_encoder (int, optional) – Number of MLP layers for the edge feature encoder, by default 2

  • hidden_dim_node_decoder (int, optional) – Hidden layer size for the node feature decoder, by default 128

  • num_layers_node_decoder (int, optional) – Number of MLP layers for the node feature decoder, by default 2

  • aggregation (str, optional) – Message aggregation type, by default “sum”

  • do_conat_trick (: bool, default=False) – Whether to replace concat+MLP with MLP+idx+sum

  • num_processor_checkpoint_segments (int, optional) – Number of processor segments for gradient checkpointing, by default 0 (checkpointing disabled)

Example

Copy
Copied!
            

            
>>> model = modulus.models.meshgraphnet.MeshGraphNet(
...         input_dim_nodes=4,
...         input_dim_edges=3,
...         output_dim=2,
...     )
>>> graph = dgl.rand_graph(10, 5)
>>> node_features = torch.randn(10, 4)
>>> edge_features = torch.randn(5, 3)
>>> output = model(node_features, edge_features, graph)
>>> output.size()
torch.Size([10, 2])

Note

Reference: Pfaff, Tobias, et al. “Learning mesh-based simulation with graph networks.” arXiv preprint arXiv:2010.03409 (2020).

forward(node_features: Tensor, edge_features: Tensor, graph: Union[DGLGraph, List[DGLGraph]]) → Tensor[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
class modulus.models.meshgraphnet.meshgraphnet.MeshGraphNetProcessor(processor_size: int = 15, input_dim_node: int = 128, input_dim_edge: int = 128, num_layers_node: int = 2, num_layers_edge: int = 2, aggregation: str = 'sum', norm_type: str = 'LayerNorm', activation_fn: Module = ReLU(), do_concat_trick: bool = False, num_processor_checkpoint_segments: int = 0)[source]

Bases: Module

MeshGraphNet processor block

forward(node_features: Tensor, edge_features: Tensor, graph: Union[DGLGraph, List[DGLGraph], CuGraphCSC]) → Tensor[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
run_function(segment_start: int, segment_end: int) → Callable[[Tensor, Tensor, Union[DGLGraph, List[DGLGraph]]], Tuple[Tensor, Tensor]][source]

Custom forward for gradient checkpointing

Parameters

  • segment_start (int) – Layer index as start of the segment

  • segment_end (int) – Layer index as end of the segment
Returns

Custom forward function
Return type

Callable
set_checkpoint_segments(checkpoint_segments: int)[source]

Set the number of checkpoint segments

Parameters

checkpoint_segments (int) – number of checkpoint segments
Raises

ValueError – if the number of processor layers is not a multiple of the number of checkpoint segments
class modulus.models.meshgraphnet.meshgraphnet.MetaData(name: str = 'MeshGraphNet', jit: bool = False, cuda_graphs: bool = False, amp: bool = False, amp_cpu: bool = False, amp_gpu: bool = True, torch_fx: bool = False, onnx: bool = False, onnx_gpu: bool = None, onnx_cpu: bool = None, onnx_runtime: bool = False, trt: bool = False, var_dim: int = -1, func_torch: bool = True, auto_grad: bool = True)[source]

Bases: ModelMetaData
class modulus.models.graphcast.graph_cast_net.GraphCastNet(*args, **kwargs)[source]

Bases: Module

GraphCast network architecture

Parameters

  • meshgraph_path (str) – Path to the meshgraph file. If not provided, the meshgraph will be created using PyMesh.

  • static_dataset_path (str) – Path to the static dataset file.

  • input_res (Tuple[int, int]) – Input resolution of the latitude-longitude grid

  • input_dim_grid_nodes (int, optional) – Input dimensionality of the grid node features, by default 474

  • input_dim_mesh_nodes (int, optional) – Input dimensionality of the mesh node features, by default 3

  • input_dim_edges (int, optional) – Input dimensionality of the edge features, by default 4

  • output_dim_grid_nodes (int, optional) – Final output dimensionality of the edge features, by default 227

  • processor_layers (int, optional) – Number of processor layers, by default 16

  • hidden_layers (int, optional) – Number of hiddel layers, by default 1

  • hidden_dim (int, optional) – Number of neurons in each hidden layer, by default 512

  • aggregation (str, optional) – Message passing aggregation method (“sum”, “mean”), by default “sum”

  • activation_fn (str, optional) – Type of activation function, by default “silu”

  • norm_type (str, optional) – Normalization type, by default “LayerNorm”

  • use_cugraphops_encoder (bool, default=False) – Flag to select cugraphops kernels in encoder

  • use_cugraphops_processor (bool, default=False) – Flag to select cugraphops kernels in the processor

  • use_cugraphops_decoder (bool, default=False) – Flag to select cugraphops kernels in the decoder

  • do_conat_trick (: bool, default=False) – Whether to replace concat+MLP with MLP+idx+sum

  • recompute_activation (bool, optional) – Flag for recomputing activation in backward to save memory, by default False. Currently, only SiLU is supported.
Note

Based on these papers: - “GraphCast: Learning skillful medium-range global weather forecasting”

https://arxiv.org/abs/2212.12794

custom_forward(grid_nfeat: Tensor) → Tensor[source]

GraphCast forward method with support for gradient checkpointing.

Parameters

grid_nfeat (Tensor) – Node features of the latitude-longitude graph.
Returns

grid_nfeat_finale – Predicted node features of the latitude-longitude graph.
Return type

Tensor
decoder_forward(mesh_efeat_processed: Tensor, mesh_nfeat_processed: Tensor, grid_nfeat_encoded: Tensor) → Tensor[source]

Forward method for the last layer of the processor, the decoder, and the final MLP.

Parameters

  • mesh_efeat_processed (Tensor) – Multimesh edge features processed by the processor.

  • mesh_nfeat_processed (Tensor) – Multi-mesh node features processed by the processor.

  • grid_nfeat_encoded (Tensor) – The encoded node features for the latitude-longitude grid.
Returns

grid_nfeat_finale – The final node features for the latitude-longitude grid.
Return type

Tensor
encoder_forward(grid_nfeat: Tensor) → Tensor[source]

Forward method for the embedder, encoder, and the first of the processor.

Parameters

grid_nfeat (Tensor) – Node features for the latitude-longitude grid.
Returns

  • mesh_efeat_processed (Tensor) – Processed edge features for the multimesh.

  • mesh_nfeat_processed (Tensor) – Processed node features for the multimesh.

  • grid_nfeat_encoded (Tensor) – Encoded node features for the latitude-longitude grid.
forward(grid_nfeat: Tensor) → Tensor[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
prepare_input(invar: Tensor) → Tensor[source]

Prepares the input to the model in the required shape.

Parameters

invar (Tensor) – Input in the shape [N, C, H, W].
Returns

Reshaped input.
Return type

Tensor
prepare_output(outvar: Tensor) → Tensor[source]

Prepares the output of the model in the shape [N, C, H, W].

Parameters

outvar (Tensor) – Output of the final MLP of the model.
Returns

The reshaped output of the model.
Return type

Tensor
set_checkpoint_decoder(checkpoint_flag: bool)[source]

Sets checkpoint function for the last layer of the processor, the decoder, and the final MLP.

This function returns the appropriate checkpoint function based on the provided checkpoint_flag flag. If checkpoint_flag is True, the function returns the checkpoint function from PyTorch’s torch.utils.checkpoint. Otherwise, it returns an identity function that simply passes the inputs through the given layer.

Parameters

checkpoint_flag (bool) – Whether to use checkpointing for gradient computation. Checkpointing can reduce memory usage during backpropagation at the cost of increased computation time.
Returns

The selected checkpoint function to use for gradient computation.
Return type

Callable
set_checkpoint_encoder(checkpoint_flag: bool)[source]

Sets checkpoint function for the embedder, encoder, and the first of the processor.

This function returns the appropriate checkpoint function based on the provided checkpoint_flag flag. If checkpoint_flag is True, the function returns the checkpoint function from PyTorch’s torch.utils.checkpoint. Otherwise, it returns an identity function that simply passes the inputs through the given layer.

Parameters

checkpoint_flag (bool) – Whether to use checkpointing for gradient computation. Checkpointing can reduce memory usage during backpropagation at the cost of increased computation time.
Returns

The selected checkpoint function to use for gradient computation.
Return type

Callable
set_checkpoint_model(checkpoint_flag: bool)[source]

Sets checkpoint function for the entire model.

This function returns the appropriate checkpoint function based on the provided checkpoint_flag flag. If checkpoint_flag is True, the function returns the checkpoint function from PyTorch’s torch.utils.checkpoint. In this case, all the other gradient checkpoitings will be disabled. Otherwise, it returns an identity function that simply passes the inputs through the given layer.

Parameters

checkpoint_flag (bool) – Whether to use checkpointing for gradient computation. Checkpointing can reduce memory usage during backpropagation at the cost of increased computation time.
Returns

The selected checkpoint function to use for gradient computation.
Return type

Callable
set_checkpoint_processor(checkpoint_segments: int)[source]

Sets checkpoint function for the processor excluding the first and last layers.

This function returns the appropriate checkpoint function based on the provided checkpoint_segments flag. If checkpoint_segments is positive,

the function returns the checkpoint function from PyTorch’s

torch.utils.checkpoint, with number of checkpointing segments equal to checkpoint_segments. Otherwise, it returns an identity function that simply passes the inputs through the given layer.

Parameters

checkpoint_segments (int) – Number of checkpointing segments for gradient computation. Checkpointing can reduce memory usage during backpropagation at the cost of increased computation time.
Returns

The selected checkpoint function to use for gradient computation.
Return type

Callable
to(*args: Any, **kwargs: Any) → GraphCastNet[source]

Moves the object to the specified device, dtype, or format. This method moves the object and its underlying graph and graph features to the specified device, dtype, or format, and returns the updated object.

Parameters

  • *args (Any) – Positional arguments to be passed to the torch._C._nn._parse_to function.

  • **kwargs (Any) – Keyword arguments to be passed to the torch._C._nn._parse_to function.
Returns

The updated object after moving to the specified device, dtype, or format.
Return type

GraphCastNet
class modulus.models.graphcast.graph_cast_net.MetaData(name: str = 'GraphCastNet', jit: bool = False, cuda_graphs: bool = False, amp: bool = False, amp_cpu: bool = False, amp_gpu: bool = True, torch_fx: bool = False, onnx: bool = False, onnx_gpu: bool = None, onnx_cpu: bool = None, onnx_runtime: bool = False, trt: bool = False, var_dim: int = -1, func_torch: bool = False, auto_grad: bool = False)[source]

Bases: ModelMetaData

Pix2Pix Net

class modulus.models.pix2pix.pix2pix.MetaData(name: str = 'Pix2Pix', jit: bool = True, cuda_graphs: bool = True, amp: bool = False, amp_cpu: bool = False, amp_gpu: bool = True, torch_fx: bool = False, onnx: bool = True, onnx_gpu: bool = None, onnx_cpu: bool = None, onnx_runtime: bool = False, trt: bool = False, var_dim: int = 1, func_torch: bool = True, auto_grad: bool = True)[source]

Bases: ModelMetaData
class modulus.models.pix2pix.pix2pix.Pix2Pix(*args, **kwargs)[source]

Bases: Module

Convolutional encoder-decoder based on pix2pix generator models.

Note

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

Parameters

  • in_channels (int) – Number of input channels

  • out_channels (Union[int, Any], optional) – Number of outout channels

  • dimension (int) – Model dimensionality (supports 1, 2, 3).

  • conv_layer_size (int, optional) – Latent channel size after first convolution, by default 64

  • n_downsampling (int, optional) – Number of downsampling blocks, by default 3

  • n_upsampling (int, optional) – Number of upsampling blocks, by default 3

  • n_blocks (int, optional) – Number of residual blocks in middle of model, by default 3

  • activation_fn (str, optional) – Activation function, by default “relu”

  • batch_norm (bool, optional) – Batch normalization, by default False

  • padding_type (str, optional) – Padding type (‘reflect’, ‘replicate’ or ‘zero’), by default “reflect”

Example

Copy
Copied!
            

            
>>> #2D convolutional encoder decoder
>>> model = modulus.models.pix2pix.Pix2Pix(
... in_channels=1,
... out_channels=2,
... dimension=2,
... conv_layer_size=4)
>>> input = torch.randn(4, 1, 32, 32) #(N, C, H, W)
>>> output = model(input)
>>> output.size()
torch.Size([4, 2, 32, 32])

Note

Reference: Isola, Phillip, et al. “Image-To-Image translation with conditional adversarial networks” Conference on Computer Vision and Pattern Recognition, 2017. https://arxiv.org/abs/1611.07004

Reference: Wang, Ting-Chun, et al. “High-Resolution image synthesis and semantic manipulation with conditional GANs” Conference on Computer Vision and Pattern Recognition, 2018. https://arxiv.org/abs/1711.11585

Note

Based on the implementation: https://github.com/NVIDIA/pix2pixHD

forward(input: Tensor) → Tensor[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
class modulus.models.pix2pix.pix2pix.ResnetBlock(dimension: int, channels: int, padding_type: str = 'reflect', activation: Module = ReLU(), use_batch_norm: bool = False, use_dropout: bool = False)[source]

Bases: Module

A simple ResNet block

Parameters

  • dimension (int) – Model dimensionality (supports 1, 2, 3).

  • channels (int) – Number of feature channels

  • padding_type (str, optional) – Padding type (‘reflect’, ‘replicate’ or ‘zero’), by default “reflect”

  • activation (nn.Module, optional) – Activation function, by default nn.ReLU()

  • use_batch_norm (bool, optional) – Batch normalization, by default False
forward(x: Tensor) → Tensor[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.

Recurrent Neural Networks

class modulus.models.rnn.rnn_one2many.MetaData(name: str = 'One2ManyRNN', jit: bool = False, cuda_graphs: bool = False, amp: bool = True, amp_cpu: bool = None, amp_gpu: bool = None, torch_fx: bool = True, onnx: bool = False, onnx_gpu: bool = None, onnx_cpu: bool = None, onnx_runtime: bool = False, trt: bool = False, var_dim: int = -1, func_torch: bool = False, auto_grad: bool = False)[source]

Bases: ModelMetaData
class modulus.models.rnn.rnn_one2many.One2ManyRNN(*args, **kwargs)[source]

Bases: Module

A RNN model with encoder/decoder for 2d/3d problems that provides predictions based on single initial condition.

Parameters

  • input_channels (int) – Number of channels in the input

  • dimension (int, optional) – Spatial dimension of the input. Only 2d and 3d are supported, by default 2

  • nr_latent_channels (int, optional) – Channels for encoding/decoding, by default 512

  • nr_residual_blocks (int, optional) – Number of residual blocks, by default 2

  • activation_fn (str, optional) – Activation function to use, by default “relu”

  • nr_downsamples (int, optional) – Number of downsamples, by default 2

  • nr_tsteps (int, optional) – Time steps to predict, by default 32

Example

Copy
Copied!
            

            
>>> model = modulus.models.rnn.One2ManyRNN(
... input_channels=6,
... dimension=2,
... nr_latent_channels=32,
... activation_fn="relu",
... nr_downsamples=2,
... nr_tsteps=16,
... )
>>> input = invar = torch.randn(4, 6, 1, 16, 16) # [N, C, T, H, W]
>>> output = model(input)
>>> output.size()
torch.Size([4, 6, 16, 16, 16])

forward(x: Tensor) → Tensor[source]

Forward pass

Parameters

x (Tensor) – Expects a tensor of size [N, C, 1, H, W] for 2D or [N, C, 1, D, H, W] for 3D Where, N is the batch size, C is the number of channels, 1 is the number of input timesteps and D, H, W are spatial dimensions.
Returns

Size [N, C, T, H, W] for 2D or [N, C, T, D, H, W] for 3D. Where, T is the number of timesteps being predicted.
Return type

Tensor
class modulus.models.rnn.rnn_seq2seq.MetaData(name: str = 'Seq2SeqRNN', jit: bool = False, cuda_graphs: bool = False, amp: bool = True, amp_cpu: bool = None, amp_gpu: bool = None, torch_fx: bool = True, onnx: bool = False, onnx_gpu: bool = None, onnx_cpu: bool = None, onnx_runtime: bool = False, trt: bool = False, var_dim: int = -1, func_torch: bool = False, auto_grad: bool = False)[source]

Bases: ModelMetaData
class modulus.models.rnn.rnn_seq2seq.Seq2SeqRNN(*args, **kwargs)[source]

Bases: Module

A RNN model with encoder/decoder for 2d/3d problems. Given input 0 to t-1, predicts signal t to t + nr_tsteps

Parameters

  • input_channels (int) – Number of channels in the input

  • dimension (int, optional) – Spatial dimension of the input. Only 2d and 3d are supported, by default 2

  • nr_latent_channels (int, optional) – Channels for encoding/decoding, by default 512

  • nr_residual_blocks (int, optional) – Number of residual blocks, by default 2

  • activation_fn (str, optional) – Activation function to use, by default “relu”

  • nr_downsamples (int, optional) – Number of downsamples, by default 2

  • nr_tsteps (int, optional) – Time steps to predict, by default 32

Example

Copy
Copied!
            

            
>>> model = modulus.models.rnn.Seq2SeqRNN(
... input_channels=6,
... dimension=2,
... nr_latent_channels=32,
... activation_fn="relu",
... nr_downsamples=2,
... nr_tsteps=16,
... )
>>> input = invar = torch.randn(4, 6, 16, 16, 16) # [N, C, T, H, W]
>>> output = model(input)
>>> output.size()
torch.Size([4, 6, 16, 16, 16])

forward(x: Tensor) → Tensor[source]

Forward pass

Parameters

x (Tensor) – Expects a tensor of size [N, C, T, H, W] for 2D or [N, C, T, D, H, W] for 3D Where, N is the batch size, C is the number of channels, T is the number of input timesteps and D, H, W are spatial dimensions. Currently, this requires input time steps to be same as predicted time steps.
Returns

Size [N, C, T, H, W] for 2D or [N, C, T, D, H, W] for 3D. Where, T is the number of timesteps being predicted.
Return type

Tensor

Super Resolution Network

class modulus.models.srrn.super_res_net.ConvolutionalBlock3d(in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, batch_norm: bool = False, activation_fn: Module = Identity())[source]

Bases: Module

3D convolutional block

Parameters

  • in_channels (int) – Input channels

  • out_channels (int) – Output channels

  • kernel_size (int) – Kernel size

  • stride (int, optional) – Convolutional stride, by default 1

  • batch_norm (bool, optional) – Use batchnorm, by default False
forward(input: Tensor) → Tensor[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
class modulus.models.srrn.super_res_net.MetaData(name: str = 'SuperResolution', jit: bool = True, cuda_graphs: bool = False, amp: bool = False, amp_cpu: bool = False, amp_gpu: bool = False, torch_fx: bool = False, onnx: bool = True, onnx_gpu: bool = None, onnx_cpu: bool = None, onnx_runtime: bool = False, trt: bool = False, var_dim: int = 1, func_torch: bool = True, auto_grad: bool = True)[source]

Bases: ModelMetaData
class modulus.models.srrn.super_res_net.PixelShuffle3d(scale: int)[source]

Bases: Module

3D pixel-shuffle operation

Parameters

scale (int) – Factor to downscale channel count by
Note

Reference: http://www.multisilicon.com/blog/a25332339.html

forward(input: Tensor) → Tensor[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
class modulus.models.srrn.super_res_net.ResidualConvBlock3d(n_layers: int = 1, kernel_size: int = 3, conv_layer_size: int = 64, activation_fn: Module = Identity())[source]

Bases: Module

3D ResNet block

Parameters

  • n_layers (int, optional) – Number of convolutional layers, by default 1

  • kernel_size (int, optional) – Kernel size, by default 3

  • conv_layer_size (int, optional) – Latent channel size, by default 64

  • activation_fn (nn.Module, optional) – Activation function, by default nn.Identity()
forward(input: Tensor) → Tensor[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
class modulus.models.srrn.super_res_net.SRResNet(*args, **kwargs)[source]

Bases: Module

3D convolutional super-resolution network

Parameters

  • in_channels (int) – Number of input channels

  • out_channels (int) – Number of outout channels

  • large_kernel_size (int, optional) – convolutional kernel size for first and last convolution, by default 7

  • small_kernel_size (int, optional) – convolutional kernel size for internal convolutions, by default 3

  • conv_layer_size (int, optional) – Latent channel size, by default 32

  • n_resid_blocks (int, optional) – Number of residual blocks before , by default 8

  • scaling_factor (int, optional) – Scaling factor to increase the output feature size compared to the input (2, 4, or 8), by default 8

  • activation_fn (str, optional) – Activation function, by default “prelu”

Example

Copy
Copied!
            

            
>>> #3D convolutional encoder decoder
>>> model = modulus.models.srrn.SRResNet(
... in_channels=1,
... out_channels=2,
... conv_layer_size=4,
... scaling_factor=2)
>>> input = torch.randn(4, 1, 8, 8, 8) #(N, C, D, H, W)
>>> output = model(input)
>>> output.size()
torch.Size([4, 2, 16, 16, 16])

Note

Based on the implementation: https://github.com/sgrvinod/a-PyTorch-Tutorial-to-Super-Resolution

forward(in_vars: Tensor) → Tensor[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
class modulus.models.srrn.super_res_net.SubPixel_ConvolutionalBlock3d(kernel_size: int = 3, conv_layer_size: int = 64, scaling_factor: int = 2)[source]

Bases: Module

Convolutional block with Pixel Shuffle operation

Parameters

  • kernel_size (int, optional) – Kernel size, by default 3

  • conv_layer_size (int, optional) – Latent channel size, by default 64

  • scaling_factor (int, optional) – Pixel shuffle scaling factor, by default 2
forward(input: Tensor) → Tensor[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.

DLWP Model

class modulus.models.dlwp.dlwp.DLWP(*args, **kwargs)[source]

Bases: Module

A Convolutional model for Deep Learning Weather Prediction that works on Cubed-sphere grids.

This model expects the input to be of shape [N, C, 6, Res, Res]

Parameters

  • nr_input_channels (int) – Number of channels in the input

  • nr_output_channels (int) – Number of channels in the output

  • nr_initial_channels (int) – Number of channels in the initial convolution. This governs the overall channels in the model.

  • activation_fn (str) – Activation function for the convolutions

  • depth (int) – Depth for the U-Net

  • clamp_activation (Tuple of ints, floats or None) – The min and max value used for torch.clamp()

Example

Copy
Copied!
            

            
>>> model = modulus.models.dlwp.DLWP(
... nr_input_channels=2,
... nr_output_channels=4,
... )
>>> input = torch.randn(4, 2, 6, 64, 64) # [N, C, F, Res, Res]
>>> output = model(input)
>>> output.size()
torch.Size([4, 4, 6, 64, 64])

Note
Reference: Weyn, Jonathan A., et al. “Sub‐seasonal forecasting with a large ensemble

of deep‐learning weather prediction models.” Journal of Advances in Modeling Earth Systems 13.7 (2021): e2021MS002502.

forward(cubed_sphere_input)[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
class modulus.models.dlwp.dlwp.MetaData(name: str = 'DLWP', jit: bool = False, cuda_graphs: bool = True, amp: bool = False, amp_cpu: bool = True, amp_gpu: bool = True, torch_fx: bool = False, onnx: bool = False, onnx_gpu: bool = None, onnx_cpu: bool = None, onnx_runtime: bool = False, trt: bool = False, var_dim: int = 1, func_torch: bool = False, auto_grad: bool = False)[source]

Bases: ModelMetaData

Spherical Harmonics Fourier Neural Operator

class modulus.models.sfno.sfnonet.FourierNeuralOperatorBlock(forward_transform, inverse_transform, embed_dim, filter_type='linear', operator_type='diagonal', mlp_ratio=2.0, drop_rate=0.0, drop_path=0.0, act_layer='gelu', norm_layer=(<class 'torch.nn.modules.normalization.LayerNorm'>, <class 'torch.nn.modules.normalization.LayerNorm'>), sparsity_threshold=0.0, use_complex_kernels=True, rank=1.0, factorization=None, separable=False, inner_skip='linear', outer_skip=None, use_mlp=False, comm_feature_inp_name=None, comm_feature_hidden_name=None, complex_network=True, complex_activation='real', spectral_layers=1, checkpointing=0)[source]

Bases: Module

Fourier Neural Operator Block

forward(x)[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
class modulus.models.sfno.sfnonet.MetaData(name: str = 'SFNO', jit: bool = False, cuda_graphs: bool = True, amp: bool = False, amp_cpu: bool = True, amp_gpu: bool = True, torch_fx: bool = False, onnx: bool = False, onnx_gpu: bool = None, onnx_cpu: bool = None, onnx_runtime: bool = False, trt: bool = False, var_dim: int = -1, func_torch: bool = False, auto_grad: bool = False)[source]

Bases: ModelMetaData
class modulus.models.sfno.sfnonet.SpectralFilterLayer(forward_transform, inverse_transform, embed_dim, filter_type='linear', operator_type='diagonal', sparsity_threshold=0.0, use_complex_kernels=True, hidden_size_factor=1, rank=1.0, factorization=None, separable=False, complex_network=True, complex_activation='real', spectral_layers=1, drop_rate=0.0)[source]

Bases: Module

Spectral filter layer

forward(x)[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
class modulus.models.sfno.sfnonet.SphericalFourierNeuralOperatorNet(*args, **kwargs)[source]

Bases: Module

Spherical Fourier Neural Operator Network

Parameters

  • params (dict) – Dictionary of parameters

  • spectral_transform (str, optional) – Type of spectral transformation to use, by default “sht”

  • grid (str, optional) – Type of grid to use, by default “legendre-gauss”

  • filter_type (str, optional) – Type of filter to use (‘linear’, ‘non-linear’), by default “non-linear”

  • operator_type (str, optional) – Type of operator to use (‘diaginal’, ‘dhconv’), by default “diagonal”

  • inp_shape (tuple, optional) – Shape of the input channels, by default (721, 1440)

  • scale_factor (int, optional) – Scale factor to use, by default 16

  • in_chans (int, optional) – Number of input channels, by default 2

  • out_chans (int, optional) – Number of output channels, by default 2

  • embed_dim (int, optional) – Dimension of the embeddings, by default 256

  • num_layers (int, optional) – Number of layers in the network, by default 12

  • repeat_layers (int, optional) – Number of times to repeat the layers, by default 1

  • use_mlp (int, optional) – Whether to use MLP, by default True

  • mlp_ratio (int, optional) – Ratio of MLP to use, by default 2.0

  • activation_function (str, optional) – Activation function to use, by default “gelu”

  • encoder_layers (int, optional) – Number of layers in the encoder, by default 1

  • pos_embed (str, optional) – Type of positional embedding to use, by default “direct”

  • drop_rate (float, optional) – Dropout rate, by default 0.0

  • drop_path_rate (float, optional) – Dropout path rate, by default 0.0

  • sparsity_threshold (float, optional) – Threshold for sparsity, by default 0.0

  • normalization_layer (str, optional) – Type of normalization layer to use (“layer_norm”, “instance_norm”, “none”), by default “instance_norm”

  • max_modes (Any, optional) – Maximum modes to use, by default None

  • hard_thresholding_fraction (float, optional) – Fraction of hard thresholding to apply, by default 1.0

  • use_complex_kernels (bool, optional) – Whether to use complex kernels, by default True

  • big_skip (bool, optional) – Whether to use big skip connections, by default True

  • rank (float, optional) – Rank of the approximation, by default 1.0

  • factorization (Any, optional) – Type of factorization to use, by default None

  • separable (bool, optional) – Whether to use separable convolutions, by default False

  • complex_network (bool, optional) – Whether to use a complex network architecture, by default True

  • complex_activation (str, optional) – Type of complex activation function to use, by default “real”

  • spectral_layers (int, optional) – Number of spectral layers, by default 3

  • output_transform (bool, optional) – Whether to use an output transform, by default False

  • checkpointing (int, optional) – Number of checkpointing segments, by default 0

  • Example –

  • -------- –

  • SFNO (>>> from modulus.models.sfno.sfnonet import SphericalFourierNeuralOperatorNet as) –

  • SFNO( (>>> model =) –

  • params={} (...) –

:param : :param … inp_shape=(8: :param 16): :param : :param … scale_factor=4: :param : :param … in_chans=2: :param : :param … out_chans=2: :param : :param … embed_dim=16: :param : :param … num_layers=2: :param : :param … encoder_layers=1: :param : :param … spectral_layers=2: :param : :param … use_mlp=True: :param ): :param >>> model(torch.randn(1: :param 2: :param 8: :param 16)).shape: :param torch.Size([1: :param 2: :param 8: :param 16]):

forward(x)[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
no_weight_decay()[source]

Helper
class modulus.models.sfno.activations.ComplexActivation(activation, mode='cartesian', bias_shape=None)[source]

Bases: Module

A module implementing complex-valued activation functions. The module supports different modes of operation, depending on how the complex numbers are treated for the activation function: - “cartesian”: the activation function is applied separately to the

real and imaginary parts of the complex input.

  • “modulus”: the activation function is applied to the modulus of

    the complex input, after adding a learnable bias.

  • any other mode: the complex input is returned as-is (identity operation).

forward(z: Tensor) → Tensor[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
class modulus.models.sfno.activations.ComplexReLU(negative_slope=0.0, mode='real', bias_shape=None, scale=1.0)[source]

Bases: Module

Complex-valued variants of the ReLU activation function

forward(z: Tensor) → Tensor[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
modulus.models.sfno.factorizations.get_contract_fun(weight, implementation='reconstructed', separable=False, operator_type='diagonal', complex=True)[source]

Generic ND implementation of Fourier Spectral Conv contraction

Parameters

  • weight (tensorly-torch's FactorizedTensor) –

  • implementation ({'reconstructed', 'factorized'}, default is 'reconstructed') – whether to reconstruct the weight and do a forward pass (reconstructed) or contract directly the factors of the factorized weight with the input (factorized)
Returns

function
Return type

(x, weight) -> x * weight in Fourier space
class modulus.models.sfno.layers.DropPath(drop_prob=None)[source]

Bases: Module

Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).

forward(x)[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
class modulus.models.sfno.layers.EncoderDecoder(num_layers, input_dim, output_dim, hidden_dim, act)[source]

Bases: Module

Basic Encoder/Decoder

forward(x)[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
class modulus.models.sfno.layers.InverseRealFFT2(nlat, nlon, lmax=None, mmax=None)[source]

Bases: Module

Helper routine to wrap FFT similarly to the SHT

forward(x)[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
class modulus.models.sfno.layers.MLP(in_features, hidden_features=None, out_features=None, act_layer='gelu', output_bias=True, drop_rate=0.0, checkpointing=0, **kwargs)[source]

Bases: Module

Basic CNN with support for gradient checkpointing

checkpoint_forward(x)[source]

Forward method with support for gradient checkpointing
forward(x)[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
class modulus.models.sfno.layers.PatchEmbed(img_size=(224, 224), patch_size=(16, 16), in_chans=3, embed_dim=768)[source]

Bases: Module

Divides the input image into patches and embeds them into a specified dimension using a convolutional layer.

forward(x)[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
class modulus.models.sfno.layers.RealFFT2(nlat, nlon, lmax=None, mmax=None)[source]

Bases: Module

Helper routine to wrap FFT similarly to the SHT

forward(x)[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
class modulus.models.sfno.layers.SpectralAttention2d(forward_transform, inverse_transform, embed_dim, sparsity_threshold=0.0, hidden_size_factor=2, use_complex_network=True, use_complex_kernels=False, complex_activation='real', bias=False, spectral_layers=1, drop_rate=0.0)[source]

Bases: Module

2d Spectral Attention layer

forward(x)[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
forward_mlp(xr)[source]

forward method for the MLP part of the network
class modulus.models.sfno.layers.SpectralAttentionS2(forward_transform, inverse_transform, embed_dim, sparsity_threshold=0.0, hidden_size_factor=2, use_complex_network=True, complex_activation='real', bias=False, spectral_layers=1, drop_rate=0.0)[source]

Bases: Module

geometrical Spectral Attention layer

forward(x)[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
forward_mlp(xr)[source]

forward method for the MLP part of the network
class modulus.models.sfno.layers.SpectralConv2d(forward_transform, inverse_transform, in_channels, out_channels, scale='auto', hard_thresholding_fraction=1, compression=None, rank=0, bias=False)[source]

Bases: Module

Spectral Convolution as utilized in

forward(x)[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
class modulus.models.sfno.s2convolutions.SpectralAttentionS2(forward_transform, inverse_transform, embed_dim, operator_type='diagonal', sparsity_threshold=0.0, hidden_size_factor=2, complex_activation='real', scale='auto', bias=False, spectral_layers=1, drop_rate=0.0)[source]

Bases: Module

Spherical non-linear FNO layer

forward(x)[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
forward_mlp(x)[source]

forward pass of the MLP
class modulus.models.sfno.s2convolutions.SpectralConvS2(forward_transform, inverse_transform, in_channels, out_channels, scale='auto', operator_type='diagonal', rank=0.2, factorization=None, separable=False, decomposition_kwargs={}, bias=False, use_tensorly=True)[source]

Bases: Module

Spectral Convolution according to Driscoll & Healy. Designed for convolutions on the two-sphere S2 using the Spherical Harmonic Transforms in torch-harmonics, but supports convolutions on the periodic domain via the RealFFT2 and InverseRealFFT2 wrappers.

forward(x)[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.
modulus.models.sfno.initialization.trunc_normal_(tensor, mean=0.0, std=1.0, a=-2.0, b=2.0)[source]

Fills the input Tensor with values drawn from a truncated normal distribution. The values are effectively drawn from the normal distribution \(\mathcal{N}(\text{mean}, \text{std}^2)\) with values outside \([a, b]\) redrawn until they are within the bounds. The method used for generating the random values works best when \(a \leq \text{mean} \leq b\). Args: tensor: an n-dimensional torch.Tensor mean: the mean of the normal distribution std: the standard deviation of the normal distribution a: the minimum cutoff value b: the maximum cutoff value
class modulus.models.sfno.preprocessor.Preprocessor2D(params)[source]

Bases: Module

Preprocessing methods to flatten image history, add static features, and convert the data format from NCHW to NHWC.

add_static_features(x)[source]

Adds static features to the input
append_channels(x, xc)[source]

Appends channels
append_history(x1, x2, step)[source]

Appends history to the main input. Without history, just returns the second tensor (x2).
append_unpredicted_features(inp)[source]

Appends features not predicted by the model (such as zenith angle) from the input
cache_unpredicted_features(x, y=None, xz=None, yz=None)[source]

Caches features not predicted by the model (such as zenith angle)
expand_history(x, nhist)[source]

Expand history from flattened data
flatten_history(x)[source]

Flatten input so that history is included as part of channels
history_compute_stats(x)[source]

Compute stats from history timesteps
history_denormalize(xn, target=False)[source]

Denormalize history
history_normalize(x, target=False)[source]

Normalize history
remove_static_features(x)[source]

Removes static features from the input only remove if something was added in the first place
remove_unpredicted_features(inp)[source]

Removes features not predicted by the model (such as zenith angle) from the input
modulus.models.sfno.preprocessor.get_preprocessor(params)[source]

Returns the preprocessor module
Previous Welcome to Modulus Core’s documentation!
Next Modulus Datapipes
© Copyright 2023, NVIDIA Modulus Team. Last updated on Jan 25, 2024.