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]]]