AtlasCRPS

GlobalMRF202680 GBNVIDIAPyTorch

Import path: earth2studio.models.px.AtlasCRPS

View source on GitHub View install commands

Documentation

Bases: Module, AutoModelMixin, PrognosticMixin

Atlas CRPS ensemble prognostic model for ERA5 variables on a 0.25 degree global lat-lon grid.

AtlasCRPS consumes two input lead times (t-6h and t) and predicts a single step at t+6h on a 721x1440 latitude-longitude grid. The model predicts a low-resolution residual which is decoded to full resolution by an autoencoder. Ensemble members are generated by calling the model repeatedly from the same initial condition, each forward pass draws a new noise vector that modulates every transformer block.

Note

For more information see the following references:

• huggingface.co/nvidia/atlas-era5

Parameters:

• model (Module) –

  CRPS latent transformer predicting a normalized low-resolution residual.

• model_processor (Module) –

  Processor providing the state and residual normalizers, the static invariant channels, the cosine zenith angle channel, the low-resolution downsampling and the low-resolution conditioning.

• autoencoder (Module) –

  Autoencoder decoding the low-resolution residual to full resolution.

• autoencoder_processor (Module) –

  Autoencoder processor mapping the decoded residual back to physical units.

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

Ensemble noise is drawn from the global PyTorch generator, use torch.manual_seed for reproducible members.

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 AtlasCRPS model.

Note

This package is the same merged Atlas-SI / Atlas-CRPS HuggingFace package used by earth2studio.models.px.Atlas.load_default_package, laid out with a crps/config.json manifest and the CRPS model and processor checkpoints under crps/genmodel, sharing the autoencoders/ directory with the Atlas (SI) si/ subtree.

load_model classmethod

load_model(package: Package) -> PrognosticModel

Instantiate and load AtlasCRPS from a package.

The autoencoder used to decode the low-resolution residual is shared with the Atlas model but described directly in the CRPS config's package.autoencoders entry, so it is loaded from the same package rather than a separate one.

Examples using earth2studio.models.px.AtlasCRPS

Running Atlas-CRPS InferenceDeterministic inference using the Atlas-CRPS prognostic model.