Skip to content

AIFS

GlobalMRF202540 GBECMWFPyTorch

Import path: earth2studio.models.px.AIFS

View source on GitHub View install commands

Documentation

Bases: Module, AutoModelMixin, PrognosticMixin

Artificial Intelligence Forecasting System (AIFS), a data driven forecast model developed by the European Centre for Medium-Range Weather Forecasts (ECMWF). AIFS is based on a graph neural network (GNN) encoder and decoder, and a sliding window transformer processor, and is trained on ECMWF's ERA5 re-analysis and ECMWF's operational numerical weather prediction (NWP) analyses. Consists of a single model with a time-step size of 6 hours.

Note

This model uses the checkpoints provided by ECMWF. Multiple checkpoint versions are supported. Use:

  • AIFS.load_default_package() for the default (AIFS-Single v1.0)
  • AIFS.load_default_package(version="1.1") for AIFS-Single v1.1

The checkpoint metadata (ai-models.json) is used to derive the correct variable ordering and indices for each checkpoint version. For additional information see the following resources:

Parameters:

  • model (Module) –

    Core PyTorch module with the pretrained AIFS weights loaded.

  • latitudes (Tensor) –

    Latitude values for the native octahedral grid, registered as a buffer for interpolation.

  • longitudes (Tensor) –

    Longitude values for the native octahedral grid, registered as a buffer for interpolation.

  • interpolation_matrix (Tensor) –

    CSR sparse matrix mapping ERA5 lat/lon inputs onto the octahedral grid.

  • inverse_interpolation_matrix (Tensor) –

    CSR sparse matrix mapping outputs from the octahedral grid back to ERA5 lat/lon.

  • invariants (Tensor) –

    Tensor of shape [4, 721, 1440] containing the invariant fields "lsm", "sdor", "slor" and "z"

Warning

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

__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