DLWP

GlobalS2S202140 GBNVIDIAPyTorch

Import path: earth2studio.models.px.DLWP

View source on GitHub View install commands

Documentation

Bases: Module, AutoModelMixin, PrognosticMixin

Deep learning weather prediction (DLWP) prognostic model. This is a parsimonious global forecast model with a time-step size of 6 hours. The core model is a convolutional encoder-decoder trained on [64,64] cubed sphere data that has an input of 18 fields (2x7 atmos variables + 4 prescriptive) and outputs 14 fields (2x7 atmos variables). This implementation provides a wrapper that accepts [721,1440] lat-lon equirectangular grid of just the atmospheric varaibles as an input for better compatability with common data sources. Prescriptive fields are added inside the model wrapper.

Note

For more information about this model see:

• agupubs.onlinelibrary.wiley.com/doi/epdf/10.1029/2021MS002502
• catalog.ngc.nvidia.com/orgs/nvidia/teams/modulus/models/modulus_dlwp_cubesphere

Parameters:

• core_model (Module) –

  Core cubed-sphere DLWP model.

• landsea_mask (Tensor) –

  Land sea mask in cubed sphere form [6,64,64]

• orography (Tensor) –

  Surface geopotential (orography) in cubed sphere form [6,64,64]

• latgrid (Tensor) –

  Cubed sphere latitude coordinates [6,64,64]

• longrid (Tensor) –

  Cubed sphere longitude coordinates [6,64,64]

• cubed_sphere_transform (Tensor) –

  Sparse pytorch tensor to transform equirectangular fields to cubed sphere of size [24576, 1038240]

• cubed_sphere_inverse (Tensor) –

  Sparse pytorch tensor to transform cubed sphere fields to equirectangular of size [1038240, 24576]

• center (Tensor) –

  Model atmospheric variable center normalization tensor of size [1,7,1,1]

• scale (Tensor) –

  Model atmospheric variable scale normalization tensor of size [1,7,1,1]

__call__

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

Runs prognostic model 1 step.

Parameters:

• x (Tensor) –

  Input tensor

• coords (CoordSystem) –

  Input coordinate system

Returns:

• tuple[Tensor, CoordSystem] –

  Output tensor and coordinate system 6 hours in the future

create_iterator

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

Creates a iterator which can be used to perform time-integration of the prognostic model. Will return the initial condition first (0th step).

Parameters:

• x (Tensor) –

  Input tensor

• coords (CoordSystem) –

  Input coordinate system

Yields:

• Iterator[tuple[Tensor, CoordSystem]] –

  Iterator that generates time-steps of the prognostic model container the output data tensor and coordinate system dictionary.

load_default_package classmethod

load_default_package() -> Package

Default DLWP model package on NGC

load_model classmethod

load_model(package: Package) -> PrognosticModel

Load prognostic from package

Examples using earth2studio.models.px.DLWP

Single Variable Perturbation MethodIntermediate ensemble inference using a custom perturbation method. Model Hook Injection: PerturbationAdding model noise by using custom hooks. Distributed Manager InferenceSetting up distributed manager for parallel inference. IO Backend PerformanceLeverage different IO backends for storing inference results. Extending Diagnostic ModelsImplementing a custom diagnostic model