earth2studio.models.px.AIFSENS

class earth2studio.models.px.AIFSENS(model, latitudes, longitudes, interpolation_matrix, inverse_interpolation_matrix)

Artificial Intelligence Forecasting System Ensemble (AIFS ENS v1.0), a probabilistic, ensemble-based forecast model from the European Centre for Medium-Range Weather Forecasts (ECMWF). AIFS ENS uses a GNN encoder/decoder with a sliding-window transformer processor, trained on ERA5 reanalysis and operational NWP analyses, and is run four times daily with a 6-hour time step. The model is trained with a CRPS objective over a small ensemble to provide calibrated probabilistic output.

Note

This model uses the checkpoints provided by ECMWF. For additional information see the following resources:

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

• https://huggingface.co/ecmwf/aifs-ens-1.0

• ecmwf/anemoi-core

Parameters:

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

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

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

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

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

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). :param x: Input tensor :type x: torch.Tensor :param coords: Input coordinate system :type coords: CoordSystem

Yields:

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

Parameters:

• x (Tensor)

• coords (OrderedDict[str, ndarray])

Return type:

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

classmethod load_default_package()

Load prognostic package

Return type:

Package

classmethod load_model(package)

Load prognostic from package

Parameters:

package (Package)

Return type:

PrognosticModel