PrecipitationAFNOv2

class earth2studio.models.dx.PrecipitationAFNOv2(

core_model,

landsea_mask,

orography,

center,

scale,

)

GlobalMRF202440 GB

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

Warning

PrecipitationAFNOv2 performs worse than PrecipitationAFNO v1 and will be deprecated in upcoming releases.

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