Skip to content

Aurora1p5Ensemble

GlobalMRF202648 GBMicrosoftPyTorch

Import path: earth2studio.models.px.Aurora1p5Ensemble

View source on GitHub View install commands

Documentation

Bases: _Aurora

Aurora v1.5 ensemble 0.25 degree global forecast model. Identical to Aurora1p5 except it uses the stochastic ensemble checkpoint, where each forward pass injects fresh Gaussian noise into the backbone conditioning context. Calling the model N times (or with a batch of N copies of the same initial condition) therefore produces N statistically independent members.

Like Aurora1p5, this wrapper produces hourly output using six lead-time queries per 6-hour AR cycle. Use Aurora1p5Ensemble_6h for six-hourly output. The two cadences consume the RNG stream differently, so the same seed does not produce matching trajectories between variants.

Note

This model uses the ensemble checkpoint 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. See Aurora1p5 for data source recommendations.

Warning

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

Parameters:

  • core_model (Module) –

    Core Aurora1p5Ensemble model (stochastic=True)

  • static_vars (dict[str, Tensor]) –

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

  • seed (int | None, default: None ) –

    If specified, sets the random seed via set_rng at the start of each create_iterator call for reproducible stochastic noise. By default None (non-reproducible).

__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 one output time-step 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

Parameters:

  • package (Package) –

    Package to load model from

Returns:

  • PrognosticModel –

    Prognostic model