Skip to content

GraphCastOperational

GlobalMRF202240 GBGoogleJAX

Import path: earth2studio.models.px.GraphCastOperational

View source on GitHub View install commands

Documentation

Bases: Module, AutoModelMixin, PrognosticMixin

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:

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

    Quater degree resolution [721x1440] land sea mask on lat-lon grid

  • geopotential_at_surface (array) –

    Quater degree resolution [721x1440] 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