earth2studio.models.px.CBottleVideo

class earth2studio.models.px.CBottleVideo(core_model, sst_ds, lat_lon=True, sampler_steps=18, sigma_max=1000.0, sigma_min=0.02, seed=None)

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:

• https://arxiv.org/abs/2505.06474v1

• NVlabs/cBottle

• https://catalog.ngc.nvidia.com/orgs/nvidia/teams/earth-2/models/cbottle

Parameters:

• core_model (torch.nn.Module) – Core Pytorch model

• sst_ds (xr.Dataset) – Sea surface temperature xarray dataset

• lat_lon (bool, optional) – 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, optional) – Number of diffusion steps, by default 18

• sigma_max (float, optional) – Maximum supported noise level during sampling, by default 1000

• sigma_min (float, optional) – Minimum supported noise level during sampling, by default 0.02

• seed (int | None, optional) – If set, will fix the seed of the random generator for latent variables, by default None

__call__(x, coords)

Runs prognostic model 1 step.

Parameters:

• x (torch.Tensor) – Input conditional tensor for first frame, if all NaNs model will not use any conditioning.

• coords (CoordSystem) – Input coordinate system

Returns:

Output tensor and coordinate system 6 hours in the future

Return type:

tuple[torch.Tensor, CoordSystem]

create_iterator(x, coords)

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 (torch.Tensor) – Input tensor

• coords (CoordSystem) – Input coordinate system

Yields:

Iterator[tuple[torch.Tensor, CoordSystem]] – Iterator that generates time-steps of the prognostic model container the output data tensor and coordinate system dictionary.

Return type:

Iterator[tuple[Tensor, OrderedDict[str, ndarray]]]

classmethod load_default_package()

Default pre-trained CBottle3D model package from Nvidia model registry

Return type:

Package

classmethod load_model(package, lat_lon=True, seed=None)

Load AI datasource from package

Parameters:

• package (Package) – CBottle AI model package

• lat_lon (bool, optional) – 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

• seed (int | None, optional) – If set, will fix the seed of the random generator for latent variables, by default None

Returns:

Prognostic Model

Return type:

PrognosticModel