AIFS2

class earth2studio.models.px.AIFS2(

model,

latitudes,

longitudes,

interpolation_matrix,

inverse_interpolation_matrix,

invariants,

preload_invariants=True,

)

GlobalMRF202640 GB

Artificial Intelligence Forecasting System version 2 (AIFS v2), a data driven forecast model developed by the European Centre for Medium-Range Weather Forecasts (ECMWF). AIFS v2 is based on a graph neural network (GNN) encoder and decoder, and a sliding window transformer processor. It extends AIFS v1 with wave forecasting capabilities, improved stratospheric representation (10 hPa level), and snow cover. Consists of a single model with a time-step size of 6 hours.

Note

Key differences from AIFS v1:

• New wave component with 11 wave variables

• New snow variable (snowc - snow coverage)

• Extended pressure levels to 10 hPa

• Vertical velocity (W) changed from prognostic to diagnostic

It is recommended to use the IFS data source to prepare model inputs given the variable set required.

Note

For additional information see the following resources:

• https://arxiv.org/abs/2509.18994

• https://huggingface.co/ecmwf/aifs-single-2.0

• ecmwf/anemoi-core

Parameters:

• model (torch.nn.Module) – Core PyTorch module with the pretrained AIFS v2 weights loaded.

• latitudes (torch.Tensor) – Latitude values for the native octahedral grid, registered as a buffer for interpolation.

• longitudes (torch.Tensor) – Longitude values for the native octahedral grid, registered as a buffer for interpolation.

• interpolation_matrix (torch.Tensor) – CSR sparse matrix mapping ERA5 lat/lon inputs onto the octahedral grid.

• inverse_interpolation_matrix (torch.Tensor) – CSR sparse matrix mapping outputs from the octahedral grid back to ERA5 lat/lon.

• invariants (torch.Tensor) – Tensor of shape [5, 721, 1440] containing the invariant fields “lsm”, “sdor”, “slor”, “z”, and “wmb” (model bathymetry). Empty if preload_invariants=False.

• preload_invariants (bool) – If True, invariant fields are fetched from IFS at load time and cached. If False, invariant fields must be provided as input variables, enabling use of invariants from alternative sources for exact reproducibility, by default True.

Warning

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

__call__(x, coords)

Runs prognostic model 1 step.

Parameters:

• x (torch.Tensor) – Input tensor

• coords (CoordSystem) – Input coordinate system

Returns:

Output tensor and coordinate system 6 hours in the future

Return type:

tuple[torch.Tensor, CoordSystem]

create_iterator(x, coords)

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 (torch.Tensor) – Input tensor

• coords (CoordSystem) – Input coordinate system

Yields:

Iterator[tuple[torch.Tensor, CoordSystem]] – Iterator that generates time-steps of the prognostic model container the output data tensor and coordinate system dictionary.

Return type:

Iterator[tuple[Tensor, OrderedDict[str, ndarray]]]

classmethod load_default_package()

Load prognostic package

Return type:

Package

classmethod load_model(package, preload_invariants=True)

Load prognostic from package

Parameters:

• package (Package) – Model package to load from

• preload_invariants (bool, optional) – If True, invariant fields (lsm, sdor, slor, z, wmb) are fetched from IFS at load time and cached. If False, these fields must be provided as input variables, allowing use of invariants from alternative sources (e.g., ECMWF Open Data) for exact reproducibility with reference implementations, by default True.

Returns:

Loaded AIFS2 model

Return type:

PrognosticModel