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

__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) -> PrognosticModel

Load prognostic from package