credit.postblock

credit.postblock#

Submodules#

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
`GlobalEnergyFixerUpDown`	Global energy conservation fixer using explicit up/down flux decomposition.
`WetMaskBlock`	Post-processing layer that applies wet mask to ocean predictions

Package Contents#

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#

tracer_thres_max#

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

class credit.postblock.GlobalEnergyFixerUpDown(post_conf)#

Bases: torch.nn.Module

Global energy conservation fixer using explicit up/down flux decomposition.

Identical correction logic to GlobalEnergyFixer but uses separate downwelling and upwelling flux indices rather than pre-computed net fluxes. The net TOA and surface imbalances are formed as:

R_T  = (DSWRFtoa  - USWRFtoa  - ULWRFtoa) / N_seconds
F_S  = (FSDS_J    - FSUS      + FLDS_J    - FLUS - SHF - LHF) / N_seconds

where *_J variables are in J/m² (energy over the timestep) and SHF/LHF are positive-upward surface turbulent heat fluxes also in J/m².

Parameters:: post_conf (dict) – config dictionary. The sub-key global_energy_fixer_updown must be present and contain all specs listed below.

Config keys (under global_energy_fixer_updown):

activate / activate_outside_model / simple_demo
midpoint, denorm, surface_geopotential_name
T_inds, q_inds, U_inds, V_inds
sp_inds (required when grid_type == 'sigma')
TOA_down_solar_ind — DSWRFtoa index in y_pred
TOA_up_solar_ind — USWRFtoa index in y_pred
TOA_up_OLR_ind — ULWRFtoa index in y_pred
surf_down_solar_ind — FSDS_J index in y_pred
surf_up_solar_ind — FSUS index in y_pred
surf_down_LW_ind — FLDS_J index in y_pred
surf_up_LW_ind — FLUS index in y_pred
surf_SH_ind — SHF index in y_pred (positive-upward)
surf_LH_ind — LHF index in y_pred (positive-upward)

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

TOA_up_solar_ind#

TOA_up_OLR_ind#

surf_down_solar_ind#

surf_up_solar_ind#

surf_down_LW_ind#

surf_up_LW_ind#

surf_SH_ind#

surf_LH_ind#

forward(x)#

class credit.postblock.WetMaskBlock(conf)#

Bases: torch.nn.Module

Post-processing layer that applies wet mask to ocean predictions Zero trainable parameters, but mask influences gradients

Masks predictions so land points = 0, ocean points preserve values. This encourages the model to focus learning on ocean regions.

forward(predictions)#

Apply wet mask to predictions with gradient influence

Parameters:: predictions – tensor of shape (batch, n_vars, time, lat, lon)
Returns:: same shape, with land values set to zero
Return type:: masked_predictions