PrecipitationAFNO

GlobalMRF202240 GBNVIDIAPyTorch

Import path: earth2studio.models.dx.PrecipitationAFNO

View source on GitHub View install commands

Documentation

Bases: Module, AutoModelMixin

Precipitation AFNO diagnsotic model. Predicts the total precipation parameter which is the accumulated amount of liquid and frozen water (rain or snow) with units m. This model was trained on ERA5 data and predicts the 6-hourly accumulated total precipitation at the validity date and time. 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: This checkpoint is from Parthik et al. 2022:

- https://arxiv.org/abs/2202.11214
- https://github.com/NVlabs/FourCastNet

Parameters:

• core_model (Module) –

  Core pytorch model

• 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

Default pre-trained precipation model package from Nvidia model registry

load_model classmethod

load_model(package: Package) -> DiagnosticModel

Load diagnostic from package

Examples using earth2studio.models.dx.PrecipitationAFNO

Running Diagnostic InferenceBasic prognostic + diagnostic inference workflow.