WindgustAFNO

GlobalMRF202440 GBNVIDIAPyTorch

Import path: earth2studio.models.dx.WindgustAFNO

View source on GitHub View install commands

Documentation

Bases: Module, AutoModelMixin

Wind gust AFNO diagnsotic model. Predicts the maximum wind gust during the preceding hour with the units m/s. This model uses an 17 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:

• catalog.ngc.nvidia.com/orgs/nvidia/teams/earth-2/models/afno_dx_wg-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 [17,1,1]

• scale (Tensor) –

  Model scale normalization tensor of size [17,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 windgustnet model package from Nvidia model registry

load_model classmethod

load_model(package: Package) -> DiagnosticModel

Load diagnostic from package