Aurora1p5

class earth2studio.models.px.Aurora1p5(core_model, static_vars)

GlobalMRF202648 GB

Aurora v1.5 0.25 degree global forecast model. This model is the improved version of Aurora, featuring an expanded set of surface variables (18 vs 4) and a richer set of static fields. It consists of a single auto-regressive model with a base time-step of 6 hours, operating on a 0.25 degree lat-lon grid (720, 1440) with 5 atmospheric variables across 13 pressure levels and 18 surface variables plus 7 output-only surface variables.

This wrapper uses an hourly rollout by default: the underlying 6-hour auto-regressive step is queried at each integer lead time from t+1h to t+6h before advancing the AR state.

Note

This model uses the checkpoints from the microsoft/aurora HuggingFace repository. For additional information see the following resources:

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

• microsoft/aurora

• https://huggingface.co/microsoft/aurora

• https://microsoft.github.io/aurora/example_v1p5.html

Aurora v1.5 was pretrained on ERA5 and fine-tuned on IFS operational analyses and as such recommended to be initialized with IFS analyses. The open-data IFS does not publish sea ice concentration (sic). earth2studio.data.NCAR_ERA5 or earth2studio.data.ARCO (which provide all required variables) may be used instead. GFS is not supported due to missing surface variables.

Note

The iterator yields the initial condition (h=0) first, as required by the Earth2Studio convention. For the 7 output-only diagnostic variables (i10fg, blh, uvb1h, ssrd1h, ttr1h, tp1h, sf1h), the h=0 output contains NaN because the decoder has not run at that step. All subsequent outputs (h≥1) contain real model predictions.

Warning

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

Parameters:

• core_model (torch.nn.Module) – Core Aurora1p5 model

• static_vars (dict[str, torch.Tensor]) – Dictionary of static field tensors (e.g., lsm, z, slt_*, tvh_*, tvl_*, …). Each tensor should have shape (720, 1440).

__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 1 hour 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 containing 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)

Load prognostic from package

Parameters:

package (Package)

Return type:

PrognosticModel