Aurora1p5Ensemble

class earth2studio.models.px.Aurora1p5Ensemble(core_model, static_vars, seed=None)

GlobalMRF202648 GB

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:

• 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. 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 (torch.nn.Module) – Core Aurora1p5Ensemble model (stochastic=True)

• 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).

• seed (int | None, optional) – 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__(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