Skip to content

Pangu6

GlobalMRF202340 GBONNX

Import path: earth2studio.models.px.Pangu6

View source on GitHub View install commands

Documentation

Bases: PanguBase

Pangu Weather 6 hour model. This model consists of two underlying auto-regressive models with a time-step size of 24 hours and 6 hours. These two models are interweaved during prediction. Pangu Weather operates on 0.25 degree lat-lon grid (south-pole including) equirectangular grid with 69 atmospheric/surface variables.

Note

This model uses the ONNX checkpoints from the original publication. For additional information see the following resources:

Note

To avoid ONNX init session overhead of this model we recommend setting the default Pytorch device to the correct target prior to model construction.

Warning

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

Parameters:

  • ort_24hr (str) –

    Path to Pangu 24 hour onnx file

  • ort_6hr (str) –

    Path to Pangu 6 hour onnx file

  • eager_sessions (bool, default: False ) –

    Build the 24 hour session at construction instead of on first use in a rollout. Either way the session is built once and cached; eager keeps the build cost out of the first rollout step, by default False

__call__ 

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

Runs 6 hour 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, eager_sessions: bool = False
) -> PrognosticModel

Load prognostic from package