`data.CBottle3D`¶

Global

Import path: earth2studio.data.CBottle3D

Documentation¶

Bases: Module, AutoModelMixin

Climate in a bottle data source 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 data source uses the a globally-trained coarse-resolution image diffusion model that generates 100km (50k-pixel) fields given monthly average sea surface temperatures and solar conditioning.

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: 200.0 ) –

Noise amplitude used to generate latent variables, by default 200
batch_size (int, default: 4 ) –

Batch size to generate time samples at, consider adjusting based on hardware being used, by default 4
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
cache (bool, default: False ) –

Does nothing at the moment, by default False
verbose (bool, default: True ) –

Print generation progress, by default True

call ¶

__call__(
    time: datetime | list[datetime] | TimeArray,
    variable: str | list[str] | VariableArray,
) -> DataArray

Function to get data.

Parameters:

time (datetime | list[datetime] | TimeArray) –

Timestamps to return data for (UTC).
variable (str | list[str] | VariableArray) –

String, list of strings or array of strings that refer to variables to return. Must be in CBottle3D lexicon.

Returns:

DataArray –

Generated data from CBottle

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 = 200,
    seed: int | None = None,
    batch_size: int = 4,
    verbose: bool = True,
) -> DataSource

Load AI datasource 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: 200 ) –

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
batch_size (int, default: 4 ) –

Batch size to generate time samples at, consider adjusting based on hardware being used, by default 4
verbose (bool, default: True ) –

Print generation progress, by default True