DLESyMLatLon

GlobalS2S202540 GBNVIDIAPyTorch

Import path: earth2studio.models.px.DLESyMLatLon

View source on GitHub View install commands

Documentation

Bases: DLESyM

DLESyM prognostic model supporting lat/lon input and output coordinates. This model still uses the HEALPix grid internally, but the first input is regridded from lat/lon and the outputs are regridded back to lat/lon upon returning from the model. Regridding is done using the earth2grid package. For convenience, we expose regridding methods that are accessible as .to_hpx and .to_ll.

Note

See DLESyM for more information about the prognostic model. Due to the internal regridding, model hooks applied during iteration will need to operate on the HEALPix grid.

Example

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

# x and coords are data defined on appropriate lat/lon grid
x, coords = fetch_data(...)

# Run model
x, coords = model(x, coords)

# Lat-lon outputs
atmos_outputs, atmos_coords = model.retrieve_valid_atmos_outputs(x, coords)
ocean_outputs, ocean_coords = model.retrieve_valid_ocean_outputs(x, coords)

# HEALPix outputs
atmos_outputs_hpx, atmos_coords_hpx = model.to_hpx(atmos_outputs), model.coords_to_hpx(atmos_coords)
ocean_outputs_hpx, ocean_coords_hpx = model.to_hpx(ocean_outputs), model.coords_to_hpx(ocean_coords)

Args

args Arguments for DLESyM *kwargs Keyword arguments for DLESyM

__call__

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

Runs coupled DLESyM model forward 1 step, regridding to/from HEALPix grid

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

Returns:

• PrognosticModel –

  Prognostic model

Examples using earth2studio.models.px.DLESyMLatLon

Running DLESyM InferenceBasic deterministic inference workflow for the DLESyM model.