earth2studio.models.dx.PrecipitationAFNO

class earth2studio.models.dx.PrecipitationAFNO(core_model, center, scale)

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 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

• NVlabs/FourCastNet

Parameters:

• core_model (torch.nn.Module) – Core pytorch model

• 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()

Default pre-trained precipation model package from Nvidia model registry

Return type:

Package

classmethod load_model(package)

Load diagnostic from package

Parameters:

package (Package)

Return type:

DiagnosticModel

Examples using earth2studio.models.dx.PrecipitationAFNO

Running Diagnostic Inference

Running Diagnostic Inference