earth2studio.models.dx.PrecipitationAFNOv2

class earth2studio.models.dx.PrecipitationAFNOv2(core_model, landsea_mask, orography, center, scale)

Improved Precipitation AFNO diagnostic model. Predicts the total precipitation for the past 6 hours [t-6h, t] with the units m. This model uses an 20 atmospheric inputs and outputs one on a 0.25 degree lat-lon grid (south-pole excluding) [720 x 1440].

Note

For more information on the model, please refer to:

• https://catalog.ngc.nvidia.com/orgs/nvidia/teams/earth-2/models/afno_dx_tp-v1-era5

Parameters:

• core_model (torch.nn.Module) – Core pytorch model

• landsea_mask (torch.Tensor) – Land sea mask tensor of size [720,1440]

• orography (torch.Tensor) – Surface geopotential (orography) tensor of size [720,1440]

• center (torch.Tensor) – Model center normalization tensor of size [20,1,1]

• scale (torch.Tensor) – Model scale normalization tensor of size [20,1,1]

__call__(x, coords)

Forward pass of diagnostic

Parameters:

• x (Tensor)

• coords (OrderedDict[str, ndarray])

Return type:

tuple[Tensor, OrderedDict[str, ndarray]]

classmethod load_default_package()

Load prognostic package

Return type:

Package

classmethod load_model(package)

Load diagnostic from package

Parameters:

package (Package)

Return type:

DiagnosticModel