AIFS2ENS#

class earth2studio.models.px.AIFS2ENS(
model,
latitudes,
longitudes,
interpolation_matrix,
inverse_interpolation_matrix,
invariants,
preload_invariants=True,
seed=None,
)[source]#
GlobalMRF202640 GB

Artificial Intelligence Forecasting System Ensemble version 2 (AIFS ENS v2), a data driven ensemble forecast model developed by the European Centre for Medium-Range Weather Forecasts (ECMWF). AIFS ENS v2 is based on a graph neural network (GNN) encoder and decoder, and a sliding window transformer processor. It extends AIFS ENS v1 with wave forecasting capabilities, improved stratospheric representation (10 hPa level), and snow cover. The ensemble model incorporates inherent stochasticity through noise conditioning, producing different forecasts for the same initial conditions.

Note

Key features of AIFS ENS v2:

  • Ensemble model with inherent stochasticity (noise conditioning)

  • New wave component with 11 wave variables

  • New snow variable (snowc - snow coverage)

  • Extended pressure levels to 10 hPa

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:

Parameters:

  • model (torch.nn.Module) – Core PyTorch module with the pretrained AIFS ENS 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.

  • seed (int | None) – If specified, sets the random seed before each model forward pass for reproducible stochastic noise. The seed used is seed + step where step is the forecast step number (0, 1, 2, …). Use the same seed in both E2S and vanilla anemoi-inference to get identical outputs, by default None.

Warning

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

__call__(x, coords)[source]#

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)[source]#

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()[source]#

Load prognostic package

Return type:

Package

classmethod load_model(package, preload_invariants=True, seed=None)[source]#

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.

  • seed (int | None, optional) – If specified, sets the random seed before each model forward pass for reproducible stochastic noise. The seed used is seed + step where step is the forecast step number, by default None.

Returns:

Loaded AIFS2ENS model

Return type:

PrognosticModel