earth2studio.models.px
.GraphCastSmall#
- class earth2studio.models.px.GraphCastSmall(ckpt, diffs_stddev_by_level, mean_by_level, stddev_by_level)[source]#
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:
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)[source]#
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)[source]#
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]]]