Skip to content

CBottleVideo

GlobalCM202540 GBNVIDIAPyTorch

Import path: earth2studio.models.px.CBottleVideo

View source on GitHub View install commands

Documentation

Bases: Module, AutoModelMixin, PrognosticMixin

Climate in a bottle video prognostic Climate in a Bottle (cBottle) is an AI model for emulating global km-scale climate simulations and reanalysis on the equal-area HEALPix grid. The cBottle video prognostic model uses the video diffusion checkpoint of CBottle trained to predict 12 frames (initial state including) at a time.

Note

This wrapper allows users to provide an input condition for the first frame of the model. If this tensor is all NaNs no variable conditioning will be used running the network outside of time-stamp and respective SST.

Warning

Default model package has SST data from January 1940 to December 2022, expanded SST data should be provided out of this range.

Note

For more information see the following references:

Parameters:

  • core_model (Module) –

    Core Pytorch model

  • sst_ds (Dataset) –

    Sea surface temperature xarray dataset

  • lat_lon (bool, default: True ) –

    Lat/lon toggle, if true data source will return output on a 0.25 deg lat/lon grid. If false, the native nested HealPix grid will be returned, by default True

  • sampler_steps (int, default: 18 ) –

    Number of diffusion steps, by default 18

  • sigma_max (float, default: 1000.0 ) –

    Maximum supported noise level during sampling, by default 1000

  • sigma_min (float, default: 0.02 ) –

    Minimum supported noise level during sampling, by default 0.02

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

    If set, will fix the seed of the random generator for latent variables, by default None

  • dataset_modality (DatasetModality, default: ERA5 ) –

    Dataset modality label to use when sampling (0=ICON, 1=ERA5), by default DatasetModality.ERA5

  • time_stepper (TimeStepperFunction, default: HEUN ) –

    Sampler function used to denoise, by default TimeStepperFunction.HEUN

__call__ 

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

Runs prognostic model 1 step.

Parameters:

  • x (Tensor) –

    Input conditional tensor for first frame, if all NaNs model will not use any conditioning.

  • 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 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 container the output data tensor and coordinate system dictionary.

load_default_package classmethod 

load_default_package() -> Package

Default pre-trained CBottle3D model package from Nvidia model registry

load_model classmethod 

load_model(
    package: Package,
    lat_lon: bool = True,
    sampler_steps: int = 18,
    sigma_max: float = 1000,
    seed: int | None = None,
) -> PrognosticModel

Load prognostic from package

Parameters:

  • package (Package) –

    CBottle AI model package

  • lat_lon (bool, default: True ) –

    Lat/lon toggle, if true prognostic input/output on a 0.25 deg lat/lon grid. If false, the native nested HealPix grid will be returned, by default True

  • sampler_steps (int, default: 18 ) –

    Number of diffusion steps, by default 18

  • sigma_max (float, default: 1000 ) –

    Noise amplitude used to generate latent variables, by default 200

  • seed (int, default: None ) –

    Random generator seed for latent variables. If None, no seed will be used, by default None

Returns:

  • PrognosticModel –

    Prognostic Model