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. Level-2 checkpoints preserve physical and latent states and the CPU and active CUDA RNG states. Construct the model inside the checkpoint context to resume from the next forecast step. For reproducible ensemble restarts, use earth2studio.perturbation.Zero; the model supplies its own ensemble noise.

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.