credit.skebs

credit.skebs#

Attributes#

`logger`
`supported_models`

Classes#

`BackscatterFCNN`	Base class for all neural network modules.
`BackscatterFCNNWide`	Base class for all neural network modules.
`BackscatterCNN`	Base class for all neural network modules.
`BackscatterUnet`	Base class for all neural network modules.
`BackscatterFixedCol`	Base class for all neural network modules.
`BackscatterPrescribed`	Base class for all neural network modules.
`SKEBS`	Stochastic Kinetic Energy Backscatter Torch module.

Functions#

`concat_for_inplace_ops`(y_orig, y_inplace_slice, ...)	alternate way to concat tensors along first dim,
`load_premade_encoder_model`(model_conf)

Module Contents#

credit.skebs.logger#

credit.skebs.concat_for_inplace_ops(y_orig, y_inplace_slice, ind_start, ind_end)#: alternate way to concat tensors along first dim, given a set of indices to replace that are contiguous

class credit.skebs.BackscatterFCNN(in_channels, levels)#

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.

in_channels#

levels#

fc1#

relu1#

fc2#

relu2#

forward(x)#

class credit.skebs.BackscatterFCNNWide(in_channels, levels)#

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.

in_channels#

levels#

fc1#

relu1#

fc2#

relu2#

fc3#

relu3#

fc4#

relu4#

forward(x)#

class credit.skebs.BackscatterCNN(in_channels, levels, nlat, nlon)#

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.

nlat#

nlon#

in_channels#

levels#

pad_lon#

pad_lat#

conv#

sigmoid#

pad(x)#

unpad(x)#

forward(x)#

credit.skebs.supported_models#

credit.skebs.load_premade_encoder_model(model_conf)#

class credit.skebs.BackscatterUnet(in_channels, levels, nlat, nlon, architecture, padding)#

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.

nlat#

nlon#

in_channels#

levels#

pad#

relu#

model#

forward(x)#

class credit.skebs.BackscatterFixedCol(levels)#

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.

backscatter_array#

forward(x)#

class credit.skebs.BackscatterPrescribed(nlat, nlon, levels, std_path, sigma_max)#

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.

nlat#

nlon#

backscatter_array#

forward(x)#

class credit.skebs.SKEBS(post_conf)#

Bases: torch.nn.Module

Stochastic Kinetic Energy Backscatter Torch module.

Parameters:: post_conf – dictionary with config options for PostBlock. if post_conf is not specified in config, defaults are set in the parser

post_conf#

retain_graph#

nlon#

nlat#

channels#

levels#

surface_channels#

output_only_channels#

input_only_channels#

frames#

forecast_len#

valid_forecast_len#

multistep#

lmax#

mmax#

grid#

U_inds#

V_inds#

T_inds#

Q_inds#

sp_index#

static_inds#

cos_lat#

state_trans = None#

eps = 1e-12#

alpha_init#

zero_out_levels_top_of_model#

tropics_only_dissipation#

spec_coef_is_initialized = False#

freeze_pattern_weights#

use_statics#

relu1#

dissipation_scaling_coefficient#

dissipation_type#

train_alpha#

is_training = False#

iteration = 0#

write_rollout_debug_files#

write_train_debug_files#

write_every#

debug_save_loc#

iteration_stop#

initialize_sht()#: Initialize spherical harmonics and inverse spherical harmonics transformations for both scalar and vector fields.

initialize_skebs_parameters()#: initialize the trainable parameters of skebs (everything excluding the backscatter)

initialize_filters()#: initialize the spectral filters for the backscatter prediction, and the spectral pattern

clip_parameters()#: clip the trainable parameters so that they are always physical

initialize_pattern(y_pred)#: initialize the random red noise pattern. in Berner et al

m is zonal wavenumber -> mmax n is total wavenumber -> lmax

cycle_pattern(spec_coef)#: cycle the random red noise pattern that is temporally correlated in Berner et al. m is zonal wavenumber -> mmax n is total wavenumber -> lmax

forward(x)#: the inverse sht operation requires float32 or greater

spec2grid(uspec)#: spatial data from spectral coefficients

getuv(vrtdivspec)#: compute wind vector from spectral coeffs of vorticity and divergence

getgrad(chispec)#

Compute vector gradient on grid given complex spectral coefficients.

Parameters:: chispec – rank 1 or 2 or 3 tensor complex array with shape (ntrunc+1)*(ntrunc+2)/2 or ((ntrunc+1)*(ntrunc+2)/2,nt) containing complex spherical harmonic coefficients (where ntrunc is the triangular truncation limit and nt is the number of spectral arrays to be transformed). If chispec is rank 1, nt is assumed to be 1.
Returns:: C{B{uchi, vchi}} - rank 2 or 3 numpy float32 arrays containing gridded zonal and meridional components of the vector gradient. Shapes are either (nlat,nlon) or (nlat,nlon,nt).