Skip to content

Aurora1p5Ensemble

GlobalMRF202648 GBMicrosoftPyTorch

Import path: earth2studio.models.px.Aurora1p5Ensemble

View source on GitHub View install commands

Documentation

Bases: Aurora1p5

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 uses an hourly rollout by default, leveraging the 6-hour base time-step to produce hourly lead times without additional model evaluations per AR cycle.

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