earth2studio.models.px.DLESyMLatLon

class earth2studio.models.px.DLESyMLatLon(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 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)

Parameters:

• *args – Arguments for DLESyM

• **kwargs – Keyword arguments for DLESyM

• atmos_model (Module)

• ocean_model (Module)

• hpx_lat (ndarray)

• hpx_lon (ndarray)

• nside (int)

• center (ndarray)

• scale (ndarray)

• atmos_constants (ndarray)

• ocean_constants (ndarray)

• atmos_input_times (ndarray)

• ocean_input_times (ndarray)

• atmos_output_times (ndarray)

• ocean_output_times (ndarray)

• atmos_variables (list[str])

• ocean_variables (list[str])

• atmos_coupling_variables (list[str])

• ocean_coupling_variables (list[str])

__call__(x, coords)

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

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

Running DLESyM Inference

Running DLESyM Inference