earth2studio.models.px.GraphCastSmall

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

GraphCast Small 1.0 degree model

A smaller, low-resolution version of GraphCast (1 degree resolution, 13 pressure levels and a smaller mesh), trained on ERA5 data from 1979 to 2015. This model is useful for running with lower memory and compute constraints while maintaining good forecast skill. The model operates on a 1-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