earth2studio.models.dx.WindgustAFNO

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

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:

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

• scale (torch.Tensor) – Model scale normalization tensor of size [17,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()

Default pre-trained climatenet model package from Nvidia model registry

Return type:

Package

classmethod load_model(package)

Load diagnostic from package

Parameters:

package (Package)

Return type:

DiagnosticModel