credit.skebs ============ .. py:module:: credit.skebs Attributes ---------- .. autoapisummary:: credit.skebs.logger credit.skebs.supported_models Classes ------- .. autoapisummary:: credit.skebs.BackscatterFCNN credit.skebs.BackscatterFCNNWide credit.skebs.BackscatterCNN credit.skebs.BackscatterUnet credit.skebs.BackscatterFixedCol credit.skebs.BackscatterPrescribed credit.skebs.SKEBS Functions --------- .. autoapisummary:: credit.skebs.concat_for_inplace_ops credit.skebs.load_premade_encoder_model Module Contents --------------- .. py:data:: logger .. py:function:: 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 .. py:class:: BackscatterFCNN(in_channels, levels) Bases: :py:obj:`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 :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool .. py:attribute:: in_channels .. py:attribute:: levels .. py:attribute:: fc1 .. py:attribute:: relu1 .. py:attribute:: fc2 .. py:attribute:: relu2 .. py:method:: forward(x) .. py:class:: BackscatterFCNNWide(in_channels, levels) Bases: :py:obj:`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 :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool .. py:attribute:: in_channels .. py:attribute:: levels .. py:attribute:: fc1 .. py:attribute:: relu1 .. py:attribute:: fc2 .. py:attribute:: relu2 .. py:attribute:: fc3 .. py:attribute:: relu3 .. py:attribute:: fc4 .. py:attribute:: relu4 .. py:method:: forward(x) .. py:class:: BackscatterCNN(in_channels, levels, nlat, nlon) Bases: :py:obj:`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 :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool .. py:attribute:: nlat .. py:attribute:: nlon .. py:attribute:: in_channels .. py:attribute:: levels .. py:attribute:: pad_lon .. py:attribute:: pad_lat .. py:attribute:: conv .. py:attribute:: sigmoid .. py:method:: pad(x) .. py:method:: unpad(x) .. py:method:: forward(x) .. py:data:: supported_models .. py:function:: load_premade_encoder_model(model_conf) .. py:class:: BackscatterUnet(in_channels, levels, nlat, nlon, architecture, padding) Bases: :py:obj:`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 :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool .. py:attribute:: nlat .. py:attribute:: nlon .. py:attribute:: in_channels .. py:attribute:: levels .. py:attribute:: pad .. py:attribute:: relu .. py:attribute:: model .. py:method:: forward(x) .. py:class:: BackscatterFixedCol(levels) Bases: :py:obj:`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 :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool .. py:attribute:: backscatter_array .. py:method:: forward(x) .. py:class:: BackscatterPrescribed(nlat, nlon, levels, std_path, sigma_max) Bases: :py:obj:`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 :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool .. py:attribute:: nlat .. py:attribute:: nlon .. py:attribute:: backscatter_array .. py:method:: forward(x) .. py:class:: SKEBS(post_conf) Bases: :py:obj:`torch.nn.Module` Stochastic Kinetic Energy Backscatter Torch module. :param post_conf: dictionary with config options for PostBlock. if post_conf is not specified in config, defaults are set in the parser .. py:attribute:: post_conf .. py:attribute:: retain_graph .. py:attribute:: nlon .. py:attribute:: nlat .. py:attribute:: channels .. py:attribute:: levels .. py:attribute:: surface_channels .. py:attribute:: output_only_channels .. py:attribute:: input_only_channels .. py:attribute:: frames .. py:attribute:: forecast_len .. py:attribute:: valid_forecast_len .. py:attribute:: multistep .. py:attribute:: lmax .. py:attribute:: mmax .. py:attribute:: grid .. py:attribute:: U_inds .. py:attribute:: V_inds .. py:attribute:: T_inds .. py:attribute:: Q_inds .. py:attribute:: sp_index .. py:attribute:: static_inds .. py:attribute:: cos_lat .. py:attribute:: state_trans :value: None .. py:attribute:: eps :value: 1e-12 .. py:attribute:: alpha_init .. py:attribute:: zero_out_levels_top_of_model .. py:attribute:: tropics_only_dissipation .. py:attribute:: spec_coef_is_initialized :value: False .. py:attribute:: freeze_pattern_weights .. py:attribute:: use_statics .. py:attribute:: relu1 .. py:attribute:: dissipation_scaling_coefficient .. py:attribute:: dissipation_type .. py:attribute:: train_alpha .. py:attribute:: is_training :value: False .. py:attribute:: iteration :value: 0 .. py:attribute:: write_rollout_debug_files .. py:attribute:: write_train_debug_files .. py:attribute:: write_every .. py:attribute:: debug_save_loc .. py:attribute:: iteration_stop .. py:method:: initialize_sht() Initialize spherical harmonics and inverse spherical harmonics transformations for both scalar and vector fields. .. py:method:: initialize_skebs_parameters() initialize the trainable parameters of skebs (everything excluding the backscatter) .. py:method:: initialize_filters() initialize the spectral filters for the backscatter prediction, and the spectral pattern .. py:method:: clip_parameters() clip the trainable parameters so that they are always physical .. py:method:: initialize_pattern(y_pred) initialize the random red noise pattern. in Berner et al m is zonal wavenumber -> mmax n is total wavenumber -> lmax .. py:method:: 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 .. py:method:: forward(x) the inverse sht operation requires float32 or greater .. py:method:: spec2grid(uspec) spatial data from spectral coefficients .. py:method:: getuv(vrtdivspec) compute wind vector from spectral coeffs of vorticity and divergence .. py:method:: getgrad(chispec) Compute vector gradient on grid given complex spectral coefficients. :param 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).