Skip to content

GraphCastSmall

GlobalMRF202240 GBGoogleJAX

Import path: earth2studio.models.px.GraphCastSmall

View source on GitHub View install commands

Documentation

Bases: Module, AutoModelMixin, PrognosticMixin

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

    Model checkpoint containing weights and configuration

  • diffs_stddev_by_level (Dataset) –

    Standard deviation of differences by level for normalization

  • mean_by_level (Dataset) –

    Mean values by level for normalization

  • stddev_by_level (Dataset) –

    Standard deviation by level for normalization

  • land_sea_mask (array) –

    One degree resolution [181x360] land sea mask on lat-lon grid

  • geopotential_at_surface (array) –

    One degree resolution [181x360] geopotential at surface on lat-lon grid

__call__ 

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

Runs prognostic model 1 step.

Parameters:

  • x (Tensor) –

    Input tensor

  • coords (CoordSystem) –

    Input coordinate system

Returns:

  • tuple[Tensor, CoordSystem] –

    Output tensor and coordinate system 6 hours in the future

create_iterator 

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

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

    Input tensor

  • coords (CoordSystem) –

    Input coordinate system

Yields:

  • Iterator[tuple[Tensor, CoordSystem]] –

    Iterator that generates time-steps of the prognostic model container the output data tensor and coordinate system dictionary.

load_default_package classmethod 

load_default_package() -> Package

Load prognostic package

load_model classmethod 

load_model(package: Package) -> PrognosticModel

Load prognostic from package

Parameters:

  • package (Package) –

    Package to load model from

Returns:

  • PrognosticModel –

    Prognostic model