Skip to content

Aurora1p5Ensemble_6h

GlobalMRF202648 GBMicrosoftPyTorch

Import path: earth2studio.models.px.Aurora1p5Ensemble_6h

View source on GitHub View install commands

Documentation

Bases: _Aurora

Aurora v1.5 ensemble global forecast model with six-hourly output.

Uses the stochastic checkpoint of Aurora1p5Ensemble, querying only t+6h per AR cycle. Inputs and variables match Aurora1p5. Diagnostic variables suffixed 1h retain their one-hour accumulation windows; they are not six-hour totals. The iterator uses a single-entry noise cache, so seeds do not give matching trajectories with the hourly ensemble variant.

See Aurora1p5Ensemble for checkpoint references, data source recommendations and license information. Initial-condition diagnostic handling is described in Aurora1p5.

Parameters:

  • core_model (Module) –

    Core Aurora1p5Ensemble model (stochastic=True)

  • static_vars (dict[str, Tensor]) –

    Static field tensors, each with shape (720, 1440).

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

    Seed applied at the start of each iterator for reproducible stochastic noise, by default None

__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