Skip to content

Aurora1p5

GlobalMRF202648 GBMicrosoftPyTorch

Import path: earth2studio.models.px.Aurora1p5

View source on GitHub View install commands

Documentation

Bases: Module, AutoModelMixin, PrognosticMixin

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:

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 (Module) –

    Core Aurora1p5 model

  • static_vars (dict[str, Tensor]) –

    Dictionary of static field tensors (e.g., lsm, z, slt_, tvh_, tvl_*, ...). Each tensor should have shape (720, 1440).

__call__ 

__call__(
    x: Tensor, coords: CoordSystem
) -> tuple[Tensor, CoordSystem]

Runs prognostic model 1 step.

Parameters:

  • x (Tensor) –

    Input tensor

  • coords (CoordSystem) –

    Input coordinate system

Returns:

  • tuple[Tensor, CoordSystem] –

    Output tensor and coordinate system 1 hour in the future

create_iterator 

create_iterator(
    x: Tensor, coords: CoordSystem
) -> Iterator[tuple[Tensor, CoordSystem]]

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 (Tensor) –

    Input tensor

  • coords (CoordSystem) –

    Input coordinate system

Yields:

  • Iterator[tuple[Tensor, CoordSystem]] –

    Iterator that generates time-steps of the prognostic model containing the output data tensor and coordinate system dictionary.

load_default_package classmethod 

load_default_package() -> Package

Load prognostic package

load_model classmethod 

load_model(package: Package) -> PrognosticModel

Load prognostic from package