earth2studio.models.px.DLESyM

class earth2studio.models.px.DLESyM(atmos_model, ocean_model, hpx_lat, hpx_lon, nside, center, scale, atmos_constants, ocean_constants, atmos_input_times, ocean_input_times, atmos_output_times, ocean_output_times, atmos_variables, ocean_variables, atmos_coupling_variables, ocean_coupling_variables)

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.

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

Note

For more information about this model see:

• https://arxiv.org/abs/2409.16247

• https://arxiv.org/abs/2311.06253v2

For more information about the HEALPix grid see:

• NVlabs/earth2grid

Parameters:

• atmos_model (torch.nn.Module) – Atmosphere model

• ocean_model (torch.nn.Module) – Ocean model

• hpx_lat (np.ndarray) – HEALPix latitude coordinates, shape (12, nside, nside)

• hpx_lon (np.ndarray) – HEALPix longitude coordinates, shape (12, nside, nside)

• nside (int) – HEALPix nside

• center (np.ndarray) – Means of the input data, shape (1, 1, 1, num_variables, 1, 1, 1)

• scale (np.ndarray) – Standard deviations of the input data, shape (1, 1, 1, num_variables, 1, 1, 1)

• atmos_constants (np.ndarray) – Constants for the atmosphere model, shape (12, num_atmos_constants, nside, nside)

• ocean_constants (np.ndarray) – Constants for the ocean model, shape (12, num_ocean_constants, nside, nside)

• atmos_input_times (np.ndarray) – Atmospheric input times, shape (num_atmos_input_times,)

• ocean_input_times (np.ndarray) – Ocean input times, shape (num_ocean_input_times,)

• atmos_output_times (np.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

• ocean_output_times (ndarray)

__call__(x, coords)

Runs coupled DLESyM model forward 1 step.

Parameters:

• x (torch.Tensor) – Input tensor

• coords (CoordSystem) – Input coordinate system

Returns:

Output tensor and coordinate system for the prediction

Return type:

tuple[torch.Tensor, CoordSystem]

create_iterator(x, coords)

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 (torch.Tensor) – Input tensor

• coords (CoordSystem) – Input coordinate system

Yields:

Iterator[tuple[torch.Tensor, CoordSystem]] – Iterator that generates time-steps of the prognostic model container the output data tensor and coordinate system dictionary.

Return type:

Iterator[tuple[Tensor, OrderedDict[str, ndarray]]]

classmethod load_default_package()

Default DLESyM model package on NGC

Return type:

Package

classmethod load_model(package, atmos_model_idx=0, ocean_model_idx=0)

Load prognostic from package

Parameters:

• package (Package) – Package to load model from

• atmos_model_idx (int) – Index of atmos model weights in package to load

• ocean_model_idx (int) – Index of ocean model weights in package to load

Returns:

Prognostic model

Return type:

PrognosticModel

Examples using earth2studio.models.px.DLESyM

Running DLESyM Inference

Running DLESyM Inference