Weather / Climate Models#

class physicsnemo.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

>>> model = physicsnemo.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]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class physicsnemo.models.dlwp_healpix.HEALPixRecUNet.HEALPixRecUNet(*args, **kwargs)[source]#

Bases: Module

Deep Learning Weather Prediction (DLWP) recurrent UNet model on the HEALPix mesh.

forward( inputs: Sequence, output_only_last=False, ) → Tensor[source]#

Forward pass of the HEALPixUnet

Parameters:

inputs (Sequence) – Inputs to the model, of the form [prognostics|TISR|constants] [B, F, T, C, H, W] is the format for prognostics and TISR [F, C, H, W] is the format for constants
output_only_last (bool, optional) – If only the last dimension of the outputs should be returned

Returns:

th.Tensor

Return type:

Predicted outputs

property integration_steps#: Number of integration steps

reset()[source]#: Resets the state of the network

class physicsnemo.models.graphcast.graph_cast_net.GraphCastNet(*args, **kwargs)[source]#

Bases: Module

GraphCast network architecture

Parameters:

multimesh_level (int, optional) – Level of the latent mesh, by default 6
multimesh (bool, optional) – If the latent mesh is a multimesh, by default True If True, the latent mesh includes the nodes corresponding to the specified mesh_level`and incorporates the edges from all mesh levels ranging from level 0 up to and including `mesh_level.
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_type (str, optional) – The type of processor used in this model. Available options are ‘MessagePassing’, and ‘GraphTransformer’, which correspond to the processors in GraphCast and GenCast, respectively. By default ‘MessagePassing’.
khop_neighbors (int, optional) – Number of khop neighbors used in the GraphTransformer. This option is ignored if ‘MessagePassing’ processor is used. By default 0.
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 [“TELayerNorm”, “LayerNorm”]. Use “TELayerNorm” for optimal performance. 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_concat_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.
partition_size (int, default=1) – Number of process groups across which graphs are distributed. If equal to 1, the model is run in a normal Single-GPU configuration.
partition_group_name (str, default=None) – Name of process group across which graphs are distributed. If partition_size is set to 1, the model is run in a normal Single-GPU configuration and the specification of a process group is not necessary. If partitition_size > 1, passing no process group name leads to a parallelism across the default process group. Otherwise, the group size of a process group is expected to match partition_size.
use_lat_lon_partitioning (bool, default=False) – flag to specify whether all graphs (grid-to-mesh, mesh, mesh-to-grid) are partitioned based on lat-lon-coordinates of nodes or based on IDs.
expect_partitioned_input (bool, default=False) – Flag indicating whether the model expects the input to be already partitioned. This can be helpful e.g. in multi-step rollouts to avoid aggregating the output just to distribute it in the next step again.
global_features_on_rank_0 (bool, default=False) – Flag indicating whether the model expects the input to be present in its “global” form only on group_rank 0. During the input preparation phase, the model will take care of scattering the input accordingly onto all ranks of the process group across which the graph is partitioned. Note that only either this flag or expect_partitioned_input can be set at a time.
produce_aggregated_output (bool, default=True) – Flag indicating whether the model produces the aggregated output on each rank of the procress group across which the graph is distributed or whether the output is kept distributed. This can be helpful e.g. in multi-step rollouts to avoid aggregating the output just to distribute it in the next step again.
produce_aggregated_output_on_all_ranks (bool, default=True) – Flag indicating - if produce_aggregated_output is True - whether the model produces the aggregated output on each rank of the process group across which the group is distributed or only on group_rank 0. This can be helpful for computing the loss using global targets only on a single rank which can avoid either having to distribute the computation of a loss function.
graph_backend (str, default="pyg") – Backend to use for the graph. Available options are “dgl” and “pyg”.

Note

Based on these papers:

“GraphCast: Learning skillful medium-range global weather forecasting”
https://arxiv.org/abs/2212.12794
“Forecasting Global Weather with Graph Neural Networks”
https://arxiv.org/abs/2202.07575
“Learning Mesh-Based Simulation with Graph Networks”
https://arxiv.org/abs/2010.03409
“MultiScale MeshGraphNets”
https://arxiv.org/abs/2210.00612
“GenCast: Diffusion-based ensemble forecasting for medium-range weather”
https://arxiv.org/abs/2312.15796

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]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

prepare_input( invar: Tensor, expect_partitioned_input: bool, global_features_on_rank_0: bool, ) → Tensor[source]#

Prepares the input to the model in the required shape.

Parameters:

invar (Tensor) – Input in the shape [N, C, H, W].
expect_partitioned_input (bool) – flag indicating whether input is partioned according to graph partitioning scheme
global_features_on_rank_0 (bool) – Flag indicating whether input is in its “global” form only on group_rank 0 which requires a scatter operation beforehand. Note that only either this flag or expect_partitioned_input can be set at a time.

Returns:

Reshaped input.

Return type:

Tensor

prepare_output( outvar: Tensor, produce_aggregated_output: bool, produce_aggregated_output_on_all_ranks: bool = True, ) → 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.
produce_aggregated_output (bool) – flag indicating whether output is gathered onto each rank or kept distributed
produce_aggregated_output_on_all_ranks (bool) – flag indicating whether output is gatherered on each rank or only gathered at group_rank 0, True by default and only valid if produce_aggregated_output is set.

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,

) → Self[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

physicsnemo.models.graphcast.graph_cast_net.get_lat_lon_partition_separators(partition_size: int)[source]#

Utility Function to get separation intervals for lat-lon grid for partition_sizes of interest.

Parameters:: partition_size (int) – size of graph partition

class physicsnemo.models.fengwu.fengwu.Fengwu(*args, **kwargs)[source]#

Bases: Module

FengWu PyTorch impl of: FengWu: Pushing the Skillful Global Medium-range Weather Forecast beyond 10 Days Lead - https://arxiv.org/pdf/2304.02948.pdf

Parameters:

img_size – Image size(Lat, Lon). Default: (721,1440)
pressure_level – Number of pressure_level. Default: 37
embed_dim (int) – Patch embedding dimension. Default: 192
patch_size (tuple[int]) – Patch token size. Default: (4,4)
num_heads (tuple[int]) – Number of attention heads in different layers.
window_size (tuple[int]) – Window size.

forward(x)[source]#

Parameters:

surface (torch.Tensor) – 2D n_lat=721, n_lon=1440, chans=4.
z (torch.Tensor) – 2D n_lat=721, n_lon=1440, chans=37.
r (torch.Tensor) – 2D n_lat=721, n_lon=1440, chans=37.
u (torch.Tensor) – 2D n_lat=721, n_lon=1440, chans=37.
v (torch.Tensor) – 2D n_lat=721, n_lon=1440, chans=37.
t (torch.Tensor) – 2D n_lat=721, n_lon=1440, chans=37.

prepare_input(surface, z, r, u, v, t)[source]#: Prepares the input to the model in the required shape. :param surface: 2D n_lat=721, n_lon=1440, chans=4. :type surface: torch.Tensor :param z: 2D n_lat=721, n_lon=1440, chans=37. :type z: torch.Tensor :param r: 2D n_lat=721, n_lon=1440, chans=37. :type r: torch.Tensor :param u: 2D n_lat=721, n_lon=1440, chans=37. :type u: torch.Tensor :param v: 2D n_lat=721, n_lon=1440, chans=37. :type v: torch.Tensor :param t: 2D n_lat=721, n_lon=1440, chans=37. :type t: torch.Tensor

class physicsnemo.models.pangu.pangu.Pangu(*args, **kwargs)[source]#

Bases: Module

Pangu A PyTorch impl of: Pangu-Weather: A 3D High-Resolution Model for Fast and Accurate Global Weather Forecast - https://arxiv.org/abs/2211.02556

Parameters:

img_size (tuple[int]) – Image size [Lat, Lon].
patch_size (tuple[int]) – Patch token size [Lat, Lon].
embed_dim (int) – Patch embedding dimension. Default: 192
num_heads (tuple[int]) – Number of attention heads in different layers.
window_size (tuple[int]) – Window size.

forward(x)[source]#

Parameters:: x (torch.Tensor) – [batch, 4+3+5*13, lat, lon]

prepare_input(surface, surface_mask, upper_air)[source]#: Prepares the input to the model in the required shape. :param surface: 2D n_lat=721, n_lon=1440, chans=4. :type surface: torch.Tensor :param surface_mask: 2D n_lat=721, n_lon=1440, chans=3. :type surface_mask: torch.Tensor :param upper_air: 3D n_pl=13, n_lat=721, n_lon=1440, chans=5. :type upper_air: torch.Tensor

class physicsnemo.models.swinvrnn.swinvrnn.SwinRNN(*args, **kwargs)[source]#

Bases: Module

Implementation of SwinRNN https://arxiv.org/abs/2205.13158 :param img_size: Image size [T, Lat, Lon]. :type img_size: Sequence[int], optional :param patch_size: Patch token size [T, Lat, Lon]. :type patch_size: Sequence[int], optional :param in_chans: number of input channels. :type in_chans: int, optional :param out_chans: number of output channels. :type out_chans: int, optional :param embed_dim: number of embed channels. :type embed_dim: int, optional :param num_groups: number of groups to separate the channels into. :type num_groups: Sequence[int] | int, optional :param num_heads: Number of attention heads. :type num_heads: int, optional :param window_size: Local window size. :type window_size: int | tuple[int], optional

forward(x: Tensor)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.