DLESyMv0_ISCCP_ERA5LatLon

GlobalCM202440 GBNVIDIAPyTorch

Import path: earth2studio.models.px.DLESyMv0_ISCCP_ERA5LatLon

View source on GitHub View install commands

Documentation

Bases: DLESyMv0_ISCCP_ERA5, DLESyMLatLon

Lat/lon convenience wrapper for DLESyMv0_ISCCP_ERA5.

Combines the DLESyMv0_ISCCP_ERA5 climate checkpoints (see DLESyMv0_ISCCP_ERA5) with the lat/lon regridding interface of [DLESyMLatLon][earth2studio.models.px.dlesym.DLESyMLatLon]. Inputs are accepted on the equiangular lat/lon grid (so any ERA5-compatible [DataSource][earth2studio.data.DataSource] works directly), regridded to HEALPix nside=64 internally, and the outputs are regridded back to lat/lon before being returned. This is the recommended entry point for most users of the upstream checkpoints, including climate-timescale rollouts.

Like DLESyMv0_ISCCP_ERA5, when use_ttr=True (default) the wrapper advertises ERA5 ttr as the radiative input variable and applies the per-day-of-year moment-matching TTR -> OLR transform internally. The transform is applied after the initial condition is regridded to HEALPix; subsequent rollout steps reuse the model's own OLR output, so it is only applied once. Derived variables (ws10m from u10m/v10m and tau300-700 from z300/z700) and SST NaN-interpolation are handled identically to [DLESyMLatLon][earth2studio.models.px.dlesym.DLESyMLatLon].

Parameters:

• *args (Any, default: () ) –

  Positional arguments forwarded to DLESyMv0_ISCCP_ERA5.

• **kwargs (Any, default: {} ) –

  Keyword arguments forwarded to DLESyMv0_ISCCP_ERA5 (including use_ttr and the climatology arrays).

Note

See DLESyMv0_ISCCP_ERA5 and [DLESyMLatLon][earth2studio.models.px.dlesym.DLESyMLatLon] for details. Model hooks applied during iteration operate on the HEALPix grid, as with [DLESyMLatLon][earth2studio.models.px.dlesym.DLESyMLatLon].

Example

pkg = DLESyMv0_ISCCP_ERA5LatLon.load_default_package()
model = DLESyMv0_ISCCP_ERA5LatLon.load_model(pkg, use_ttr=True)

# x, coords come straight from an ERA5 data source on the lat/lon grid
x, coords = fetch_data(...)
y, y_coords = model(x, coords)

atmos, atmos_coords = model.retrieve_valid_atmos_outputs(y, y_coords)
ocean, ocean_coords = model.retrieve_valid_ocean_outputs(y, y_coords)

__call__

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

Runs upstream DLESyM forward 1 coupled step, regridding to/from HEALPix.

Parameters:

• x (Tensor) –

  Input tensor on the lat/lon grid, with ttr in variable when use_ttr=True.

• coords (CoordSystem) –

  Input coordinate system (lat/lon).

Returns:

• tuple[Tensor, CoordSystem] –

  Output tensor and coordinates on the lat/lon grid (in model variable space: rlut rather than ttr).

create_iterator

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

Create a time-integration iterator (yields initial condition first).

Parameters:

• x (Tensor) –

  Input tensor.

• coords (CoordSystem) –

  Input coordinate system.

Yields:

• Iterator[tuple[Tensor, CoordSystem]] –

  (x, coords) at each step. The first yield is the initial condition in model variable space (rlut, post-transform).

load_default_package classmethod

load_default_package() -> Package

Default DLESyMv0_ISCCP_ERA5 model package on HuggingFace.

load_model classmethod

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

Load the DLESyMv0_ISCCP_ERA5 prognostic from a package.

Parameters:

• package (Package) –

  Model package containing config.yaml, atmos_model_*.mdlus, ocean_model_*.mdlus, the HEALPix lat/lon and constant fields, and (when use_ttr=True) era5_ttr_doy_stats_hpx64.nc and isccp_olr_doy_stats_hpx64.nc.

• use_ttr (bool, default: True ) –

  See DLESyMv0_ISCCP_ERA5. Defaults to True.

• atmos_model_idx (int, default: 0 ) –

  Index into cfg.models.atmos_model_checkpoints. Defaults to 0.

• ocean_model_idx (int, default: 0 ) –

  Index into cfg.models.ocean_model_checkpoints. Defaults to 0.

Returns:

• PrognosticModel –

  Loaded DLESyMv0_ISCCP_ERA5 instance.

Examples using earth2studio.models.px.DLESyMv0_ISCCP_ERA5LatLon

Running the DLESyM ISCCP-ERA5 ModelCoupled inference with the upstream DLESyM checkpoints, including precipitation.