PrecipitationAFNOv2

GlobalMRF202440 GBNVIDIAPyTorch

Import path: earth2studio.models.dx.PrecipitationAFNOv2

View source on GitHub View install commands

Documentation

Bases: Module, AutoModelMixin

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:

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

Parameters:

• core_model (Module) –

  Core pytorch model

• landsea_mask (Tensor) –

  Land sea mask tensor of size [720,1440]

• orography (Tensor) –

  Surface geopotential (orography) tensor of size [720,1440]

• center (Tensor) –

  Model center normalization tensor of size [20,1,1]

• scale (Tensor) –

  Model scale normalization tensor of size [20,1,1]

__call__

__call__(
    x: Tensor, coords: CoordSystem
) -> tuple[Tensor, CoordSystem]

Forward pass of diagnostic

load_default_package classmethod

load_default_package() -> Package

Load prognostic package

load_model classmethod

load_model(package: Package) -> DiagnosticModel

Load diagnostic from package