Atlas

GlobalMRF202680 GBNVIDIAPyTorch

Import path: earth2studio.models.px.Atlas

View source on GitHub View install commands

Documentation

Bases: Module, AutoModelMixin, PrognosticMixin

Atlas prognostic model for ERA5 variables on a 0.25° global lat-lon grid.

Atlas consumes two input lead times (t-6h and t) and predicts a single step at t+6h on a 721x1440 latitude-longitude grid.

Parameters:

• autoencoders (ModuleList) –

  List of autoencoders for the full-resolution physical state.

• autoencoder_processors (ModuleList) –

  List of autoencoder processors for the full-resolution physical state.

• model (Module) –

  Model for the full-resolution physical state.

• model_processor (Module) –

  Model processor for the full-resolution physical state.

• sinterpolant (Module) –

  Stochastic interpolant for the low-resolution latent state.

• sinterpolant_sample_steps (int, default: 60 ) –

  Number of steps to sample for the stochastic interpolant.

Warning

This model is expected to use the iterator interface for autoregressive rollouts longer than one step. Iteratively using the __call__ and prep_next_input methods will not produce correct results, since the model performs autoregressive timestepping using a full-resolution physical state and an internal low-resolution latent state.

Note

For best inference performance, set the environment variable TORCH_ALLOW_TF32_CUBLAS_OVERRIDE=1. This is on by default in NGC containers, but other environments may need to set it manually.

__call__

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

Forward pass of the prognostic model, integrating a single 6h step.

Parameters:

• x (Tensor) –

  Input tensor of shape (..., lead_time, variable, lat, lon) corresponding to the coordinate system. Lead times expected: [-6h, 0h].

• coords (CoordSystem) –

  Coordinate dictionary describing x.

Returns:

• tuple[Tensor, CoordSystem] –

  Output tensor advanced to t+6h and its coordinate system.

create_iterator

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

Create an iterator that yields the initial state then successive 6h steps.

Parameters:

• x (Tensor) –

  Initial data tensor on device representing the initial condition.

• coords (CoordSystem) –

  Coordinate system for the initial data tensor.

Yields:

• Iterator[tuple[Tensor, CoordSystem]] –

  Iterator yielding successive model outputs and their coordinates.

load_default_package classmethod

load_default_package() -> Package

Load the default package for the Atlas model.

load_model classmethod

load_model(package: Package) -> PrognosticModel

Instantiate and load Atlas from a package.

Examples using earth2studio.models.px.Atlas

Running Atlas InferenceDeterministic inference using the Atlas prognostic model.