credit.postblock

credit.postblock#

Content:

  • PostBlock

  • TracerFixer

  • GlobalMassFixer

  • GlobalWaterFixer

  • GlobalEnergyFixer

Attributes#

PI

logger

Classes#

PostBlock

Base class for all neural network modules.

TracerFixer

This module fixes tracer values by replacing their values to a given threshold

GlobalMassFixer

This module applies global mass conservation fixes for both dry air and water budget.

GlobalWaterFixer

Base class for all neural network modules.

GlobalEnergyFixer

This module applys global energy conservation fixes. The output ensures that the global sum

Functions#

concat_fix(y_pred, q_pred_correct, q_ind_start, ...)

this function use torch.concat to replace a specific subset of variable channels in y_pred.

Module Contents#

credit.postblock.PI = 3.141592653589793#
credit.postblock.logger#
class credit.postblock.PostBlock(post_conf)#

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.

operations#
forward(x)#
class credit.postblock.TracerFixer(post_conf)#

Bases: torch.nn.Module

This module fixes tracer values by replacing their values to a given threshold (e.g., tracer[tracer<thres] = thres).

Parameters:

post_conf (dict) – config dictionary that includes all specs for the tracer fixer.

tracer_indices#
tracer_thres#
forward(x)#
class credit.postblock.GlobalMassFixer(post_conf)#

Bases: torch.nn.Module

This module applies global mass conservation fixes for both dry air and water budget. The output ensures that the global dry air mass and global water budgets are conserved through correction ratios applied during model runs. Variables specific total water and precipitation will be corrected to close the budget. All corrections are done using float32 PyTorch tensors.

Parameters:

post_conf (dict) – config dictionary that includes all specs for the global mass fixer.

q_ind_start#
q_ind_end#
forward(x)#
class credit.postblock.GlobalWaterFixer(post_conf)#

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.

q_ind_start#
q_ind_end#
precip_ind#
evapor_ind#
forward(x)#
class credit.postblock.GlobalEnergyFixer(post_conf)#

Bases: torch.nn.Module

This module applys global energy conservation fixes. The output ensures that the global sum of total energy in the atmosphere is balanced by radiantion and energy fluxes at the top of the atmosphere and the surface. Variables air temperature will be modified to close the budget. All corrections are done using float32 Pytorch tensors.

Parameters:

post_conf (dict) – config dictionary that includes all specs for the global energy fixer.

T_ind_start#
T_ind_end#
q_ind_start#
q_ind_end#
U_ind_start#
U_ind_end#
V_ind_start#
V_ind_end#
TOA_solar_ind#
TOA_OLR_ind#
surf_solar_ind#
surf_LR_ind#
surf_SH_ind#
surf_LH_ind#
forward(x)#
credit.postblock.concat_fix(y_pred, q_pred_correct, q_ind_start, q_ind_end, N_vars)#

this function use torch.concat to replace a specific subset of variable channels in y_pred.

Given q_pred = y_pred[:, ind_start:ind_end, …], and q_pred_correct this function does: y_pred[:, ind_start:ind_end, …] = q_pred_correct, but without using in-place modifications, so the graph of y_pred is maintained. It also handles q_ind_start == q_ind_end cases.

All input tensors must have 5 dims of batch, level-or-var, time, lat, lon

Parameters:

  • y_pred (torch.Tensor) – Original y_pred tensor of shape (batch, var, time, lat, lon).

  • q_pred_correct (torch.Tensor) – Corrected q_pred tensor.

  • q_ind_start (int) – Index where q_pred starts in y_pred.

  • q_ind_end (int) – Index where q_pred ends in y_pred.

  • N_vars (int) – Total number of variables in y_pred (i.e., y_pred.shape[1]).

Returns:

Concatenated y_pred with corrected q_pred.

Return type:

torch.Tensor

Contents