earth2studio.models.px.GraphCastOperational

class earth2studio.models.px.GraphCastOperational(ckpt, diffs_stddev_by_level, mean_by_level, stddev_by_level)

GraphCast operational model

A high-resolution model (0.25 degree resolution, 13 pressure levels) pre-trained on ERA5 data from 1979 to 2017 and fine-tuned on HRES data from 2016 to 2021. This model can be initialized from HRES data (does not require precipitation inputs).

The model operates on a 0.25-degree lat-lon grid (south-pole including) equirectangular grid with 85 variables including:

• Surface variables (2m temperature, 10m winds, etc.)

• Pressure level variables (temperature, winds, geopotential, etc.)

• Static variables (land-sea mask, surface geopotential)

Note

This model and checkpoint are based on the GraphCast architecture from DeepMind. For more information see the following references:

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

• google-deepmind/graphcast

• https://www.science.org/doi/10.1126/science.adi2336

Warning

We encourage users to familiarize themselves with the license restrictions of this model’s checkpoints.

Parameters:

• ckpt (graphcast.CheckPoint) – Model checkpoint containing weights and configuration

• diffs_stddev_by_level (xr.Dataset) – Standard deviation of differences by level for normalization

• mean_by_level (xr.Dataset) – Mean values by level for normalization

• stddev_by_level (xr.Dataset) – Standard deviation by level for normalization

__call__(x, coords)

Runs prognostic model 1 step.

Parameters:

• x (torch.Tensor) – Input tensor

• coords (CoordSystem) – Input coordinate system

Returns:

Output tensor and coordinate system 6 hours in the future

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

Load prognostic package

Return type:

Package

classmethod load_model(package)

Load prognostic from package

Parameters:

package (Package) – Package to load model from

Returns:

Prognostic model

Return type:

PrognosticModel