Skip to content

ACE2ERA5

GlobalCM202440 GBAI2PyTorch

Import path: earth2studio.models.px.ACE2ERA5

View source on GitHub View install commands

Documentation

Bases: Module, AutoModelMixin, PrognosticMixin

ACE2-ERA5 prognostic model wrapper.

ACE2 (Ai2 Climate Emulator v2) is a 450M-parameter autoregressive emulator with 6-hour time steps, 1-degree horizontal resolution, and eight vertical layers that exactly conserves global dry air mass and moisture and can be stepped stably for arbitrarily many steps. ACE2-ERA5 was trained on the ERA5 dataset and requires forcing data during rollout (see forcing_data_source parameter). This wrapper makes use of the fme package to run model forward passes.

Parameters:

  • stepper (Stepper) –

    ACE2-ERA5 fme.ace.stepper.single_module.Stepper instance loaded from a checkpoint.

  • forcing_data_source (DataSource, default: ACE2ERA5Data(mode='forcing') ) –

    Data source providing forcing data during rollout. Must provide all forcing variables described in the ACE2-ERA5 paper, by default ACE2ERA5(mode="forcing").

  • dt (timedelta64, default: timedelta64(6, 'h') ) –

    Model timestep used to advance lead time coordinates, by default 6 hours.

References
Notes

For throughput-sensitive GPU inference, enabling TensorFloat-32 matmul kernels before importing PyTorch can improve performance on supported NVIDIA GPUs:

export TORCH_ALLOW_TF32_CUBLAS_OVERRIDE=1

For in-process control, this can also be enabled with:

torch.set_float32_matmul_precision("high")

Both settings trade some float32 matmul precision for faster matrix operations; the environment variable is a process-wide cuBLAS override.

Warning

This model may only be used with input data on the GPU device that the model was loaded on. Specifically, the data must be on the same device as whatever torch.cuda.current_device() was set to when the model package was loaded.

__call__ 

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

Runs one prognostic step using fme predict_paired API.

Parameters:

  • x (Tensor) –

    Input tensor

  • coords (CoordSystem) –

    Input coordinate system

Returns:

  • tuple[Tensor, CoordSystem] –

    Output tensor and coordinate system 6 hours in the future

create_iterator 

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

Creates an iterator to perform time-integration of ACE2ERA5.

Yields the first forecast step, then continues autoregressively by feeding previous outputs as the next prognostic state while fetching/using external forcings under the hood via _forward.

Parameters:

  • x (Tensor) –

    Input tensor

  • coords (CoordSystem) –

    Input coordinate system

Returns:

  • Iterator[tuple[Tensor, CoordSystem]] –

    Iterator of output tensors and coordinate systems

load_default_package classmethod 

load_default_package() -> Package

Load default ACE2-ERA5 package from HuggingFace.

load_model classmethod 

load_model(
    package: Package,
    forcing_data_source: DataSource = ACE2ERA5Data(
        mode="forcing", verbose=False
    ),
    dt: timedelta64 = timedelta64(6, "h"),
) -> PrognosticModel

Load ACE2-ERA5 prognostic model from a package.

Parameters:

  • package (Package) –

    Package to load the model checkpoint from.

  • forcing_data_source (DataSource, default: ACE2ERA5Data(mode='forcing', verbose=False) ) –

    External forcing data source. Must provide all forcing variables described in the ACE2-ERA5 paper, by default ACE2ERA5(mode="forcing").

  • dt (timedelta64, default: timedelta64(6, 'h') ) –

    Timestep for advancing lead time coordinates, by default 6 hours.

Returns:

  • PrognosticModel –

    ACE2-ERA5 prognostic model