credit.models.swin_wrf

credit.models.swin_wrf#

Attributes#

logger

Classes#

`CubeEmbedding`
`DownBlock`	Base class for all neural network modules.
`UpBlock`	Base class for all neural network modules.
`UTransformer`	U-Transformer
`WRFTransformer`

Functions#

`apply_spectral_norm`(model)	add spectral norm to all the conv and linear layers
`get_pad3d`(input_resolution, window_size)	Estimate the size of padding based on the given window size and the original input size.
`get_pad2d`(input_resolution, window_size)

Module Contents#

credit.models.swin_wrf.logger#

credit.models.swin_wrf.apply_spectral_norm(model)#: add spectral norm to all the conv and linear layers

credit.models.swin_wrf.get_pad3d(input_resolution, window_size)#

Estimate the size of padding based on the given window size and the original input size.

Parameters:

input_resolution (tuple[int]) – (Pl, Lat, Lon)
window_size (tuple[int]) – (Pl, Lat, Lon)

Returns:

(padding_left, padding_right, padding_top, padding_bottom, padding_front, padding_back)

Return type:

padding (tuple[int])

credit.models.swin_wrf.get_pad2d(input_resolution, window_size)#

Parameters:

input_resolution (tuple[int]) – Lat, Lon
window_size (tuple[int]) – Lat, Lon

Returns:

(padding_left, padding_right, padding_top, padding_bottom)

Return type:

padding (tuple[int])

class credit.models.swin_wrf.CubeEmbedding(img_size, patch_size, in_chans, embed_dim, norm_layer=nn.LayerNorm)#

Bases: torch.nn.Module

Parameters:

img_size – T, Lat, Lon
patch_size – T, Lat, Lon

img_size#

patches_resolution#

embed_dim#

proj#

forward(x: torch.Tensor)#

class credit.models.swin_wrf.DownBlock(in_chans: int, out_chans: int, num_groups: int, num_residuals: int = 2)#

Bases: torch.nn.Module

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:: training (bool) – Boolean represents whether this module is in training or evaluation mode.

conv#

b#

forward(x)#

class credit.models.swin_wrf.UpBlock(in_chans, out_chans, num_groups, num_residuals=2)#

Bases: torch.nn.Module

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:: training (bool) – Boolean represents whether this module is in training or evaluation mode.

conv#

b#

forward(x)#

class credit.models.swin_wrf.UTransformer(embed_dim, num_groups, input_resolution, num_heads, window_size, depth, drop_path)#

Bases: torch.nn.Module

U-Transformer :param embed_dim: Patch embedding dimension. :type embed_dim: int :param num_groups: number of groups to separate the channels into. :type num_groups: int | tuple[int] :param input_resolution: Lat, Lon. :type input_resolution: tuple[int] :param num_heads: Number of attention heads in different layers. :type num_heads: int :param window_size: Window size. :type window_size: int | tuple[int] :param depth: Number of blocks. :type depth: int

padding#

pad#

down#

layer#

up#

forward(x)#

class credit.models.swin_wrf.WRFTransformer(param_interior, param_outside, time_encode_dim=12, num_groups=32, num_heads=8, depth=48, window_size=7, use_spectral_norm=True, interp=True, drop_path=0, padding_conf=None, post_conf=None, **kwargs)#

Bases: credit.models.base_model.BaseModel

Parameters:

img_size (Sequence[int], optional) – T, Lat, Lon.
patch_size (Sequence[int], optional) – T, Lat, Lon.
in_chans (int, optional) – number of input channels.
out_chans (int, optional) – number of output channels.
dim (int, optional) – number of embed channels.
num_groups (Sequence[int] | int, optional) – number of groups to separate the channels into.
num_heads (int, optional) – Number of attention heads.
window_size (int | tuple[int], optional) – Local window size.

time_encode = 12#

use_interp = True#

use_spectral_norm = True#

use_padding#

use_post_block#

cube_embedding_inside#

cube_embedding_outside#

total_dim#

u_transformer#

fc#

patch_size#

input_resolution#

out_chans#

img_size#

film#

_match_spatial(src: torch.Tensor, ref: torch.Tensor)#

forward(x: torch.Tensor, x_outside: torch.Tensor, x_extra: torch.Tensor)#