earth2studio.models.px.AIFS#

class earth2studio.models.px.AIFS(model, latitudes, longitudes, interpolation_matrix, inverse_interpolation_matrix)[source]#

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. For additional information see the following resources:

Warning

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

Parameters:
  • model (Module)

  • latitudes (Tensor)

  • longitudes (Tensor)

  • interpolation_matrix (Tensor)

  • inverse_interpolation_matrix (Tensor)

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

Load prognostic package

Return type:

Package

classmethod load_model(package)[source]#

Load prognostic from package

Parameters:

package (Package)

Return type:

PrognosticModel