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.

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 full output variable set (prognostics + diagnostics, in the same order as output_coords's variable axis), shape (1, 1, 1, num_output_variables, 1, 1, 1)

• scale (ndarray) –

  Standard deviations of the full output variable set, same shape and ordering as center

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

• ocean_output_times (ndarray) –

  Ocean output times, shape (num_ocean_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

• atmos_diagnostic_variables (list[str], default: None ) –

  Atmospheric diagnostic output variables. These are produced by the atmos model but, unlike atmos_variables, are not fed back in as input to the next autoregressive step, by default []

• ocean_diagnostic_variables (list[str], default: None ) –

  Ocean diagnostic output variables, analogous to atmos_diagnostic_variables, by default []

• use_cln (bool, default: False ) –

  Whether the atmos/ocean models use conditional layer norm, which requires sampling and passing noise to the model to produce ensemble variability from a single set of weights, by default False

• condition_shape (int, default: None ) –

  Dimension of the conditional layer norm noise vector. Required if use_cln is True, by default None

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

```python 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)

__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

The package's top-level config.yaml lists the available checkpoint versions; see load_model for how to select between them.

load_model classmethod

load_model(
    package: Package,
    atmos_model_idx: int = 0,
    ocean_model_idx: int = 0,
    version: Literal["v1.0", "v1.1"] = "v1.1",
) -> PrognosticModel

Load prognostic from package

Parameters:

• package (Package) –

  Package to load model from

• version (('v1.0', 'v1.1'), default: "v1.0" ) –

  Checkpoint version to load; see each version's entry in the package's config.yaml for details. v1.1 is the checkpoint submitted to the ECMWF AI Weather Quest competition (aiweatherquest.ecmwf.int/). v1.0 is the previous checkpoint, kept for reproducibility; loading it logs a deprecation warning, by default "v1.1"

• atmos_model_idx (int, default: 0 ) –

  Index of atmos model weights to load. Only meaningful for checkpoint versions that ship multiple atmos checkpoints (used to build ensembles without conditional layer norm), by default 0

• ocean_model_idx (int, default: 0 ) –

  Index of ocean model weights to load. Only meaningful for checkpoint versions that ship multiple ocean checkpoints, by default 0

Returns:

• PrognosticModel –

  Prognostic model

Examples using earth2studio.models.px.DLESyMLatLon

Running DLESyM InferenceBasic deterministic inference workflow for the DLESyM model.