`DLESyM`¶

GlobalS2S202540 GBNVIDIAPyTorch

Import path: earth2studio.models.px.DLESyM

View source on GitHub View install commands

Documentation¶

Bases: Module, AutoModelMixin, PrognosticMixin

DLESyM-V1-ERA5 prognostic model. This is an ensemble forecast model for global earth system modeling. This model includes an atmosphere and ocean component, using atmospheric variables as well as the sea-surface temperature on a HEALPix nside=64 (approximately 1 degree) resolution grid. The model architecture is a U-Net with padding operations modified to support using the HEALPix grid. Because the atmosphere and ocean models are predicted at different times, not all entries in the output tensor are valid. As a result, we provide convenience methods for retrieving the valid atmospheric and oceanic outputs.

Parameters:

atmos_model (Module) –

Atmosphere model
ocean_model (Module) –

Ocean model
hpx_lat (ndarray) –

HEALPix latitude coordinates, shape (12, nside, nside)
hpx_lon (ndarray) –

HEALPix longitude coordinates, shape (12, nside, nside)
nside (int) –

HEALPix nside
center (ndarray) –

Means of the input data, shape (1, 1, 1, num_variables, 1, 1, 1)
scale (ndarray) –

Standard deviations of the input data, shape (1, 1, 1, num_variables, 1, 1, 1)
atmos_constants (ndarray) –

Constants for the atmosphere model, shape (12, num_atmos_constants, nside, nside)
ocean_constants (ndarray) –

Constants for the ocean model, shape (12, num_ocean_constants, nside, nside)
atmos_input_times (ndarray) –

Atmospheric input times, shape (num_atmos_input_times,)
ocean_input_times (ndarray) –

Ocean input times, shape (num_ocean_input_times,)
atmos_output_times (ndarray) –

Atmospheric output times, shape (num_atmos_output_times,)
atmos_variables (list[str]) –

Atmospheric variables
ocean_variables (list[str]) –

Ocean variables
atmos_coupling_variables (list[str]) –

Atmospheric coupling variables
ocean_coupling_variables (list[str]) –

Ocean coupling variables

Note

For more information about this model see:

For more information about the HEALPix grid see:

NVlabs/earth2grid

Example

pkg = DLESyM.load_default_package()
model = DLESyM.load_model(pkg)

# Create iterator
iterator = model.create_iterator(x, coords)

for step, (x, coords) in enumerate(iterator):
    if step > 0:
        # Valid atmos and ocean predictions with their respective coordinates extracted below
        atmos_outputs, atmos_coords = model.retrieve_valid_atmos_outputs(x, coords)
        ocean_outputs, ocean_coords = model.retrieve_valid_ocean_outputs(x, coords)
        ```pycon
        ...
        ```

call ¶

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

Runs coupled DLESyM model forward 1 step.

Parameters:

x (Tensor) –

Input tensor
coords (CoordSystem) –

Input coordinate system

Returns:

tuple[Tensor, CoordSystem] –

Output tensor and coordinate system for the prediction

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 DLESyM model package on NGC

load_model `classmethod` ¶

load_model(
    package: Package,
    atmos_model_idx: int = 0,
    ocean_model_idx: int = 0,
) -> PrognosticModel

Load prognostic from package

Parameters:

package (Package) –

Package to load model from
atmos_model_idx (int, default: 0 ) –

Index of atmos model weights in package to load, by default 0
ocean_model_idx (int, default: 0 ) –

Index of ocean model weights in package to load, by default 0