Skip to content

FuXi

GlobalMRF202340 GBONNX

Import path: earth2studio.models.px.FuXi

View source on GitHub View install commands

Documentation

Bases: Module, AutoModelMixin, PrognosticMixin

FuXi weather model consists of three auto-regressive U-net transfomer models with a time-step size of 6 hours. The three models are trained to predict short (5days), medium (10 days) and longer (15 days) forecasts respectively. FuXi operates on 0.25 degree lat-lon grid (south-pole including) equirectangular grid with 70 atmospheric/surface variables. This model uses two time-steps as an input.

Note

This model uses the ONNX checkpoint from the original publication repository. 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.

Parameters:

  • ort_short (str) –

    Path to FuXi short model onnx file

  • ort_medium (str) –

    Path to FuXi medium model onnx file

  • ort_long (str) –

    Path to FuXi long model onnx file

__call__ 

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

Runs short 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