earth2studio.models.dx.CBottleInfill#

class earth2studio.models.dx.CBottleInfill(core_model, sst_ds, input_variables, sampler_steps=18, sigma_max=80, seed=0)[source]#

Climate in a bottle infill diagnostic 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 infill diagnostic enables users to generate all variables supported by cBottle from just a subset ontop of existing monthly average sea surface temperatures and solar conditioning. The cBottle infill model uses the a globally-trained coarse-resolution image diffusion model that generates 100km (50k-pixel) fields

Note

Unlike other diagnostics that have a fixed input, this diagnostic allows users to specify the input to be any subset of cBottle’s output variables. Namely, the model can be adapted to the data present in the inference pipeline and will always generate all output fields based on the information provided.

Note

For more information see the following references:

Parameters:

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

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

  • input_variables (list[str] | VariableArray) – List of input variables that will be provided for conditioning the output generation. Must be a subset of the cBottle output / supported variables. See cBottle lexicon for full list.

  • sampler_steps (int, optional) – Number of diffusion steps, by default 18

  • sigma_max (float, optional) – Noise amplitude used to generate latent variables, by default 80

  • batch_size (int, optional) – Batch size to generate time samples at, consider adjusting based on hardware being used, by default 4

  • seed (int, optional) – Random generator seed for latent variables, by default 0

__call__(x, coords)[source]#

Forward pass of diagnostic

Parameters:

  • x (Tensor)

  • coords (OrderedDict[str, ndarray])

Return type:

tuple[Tensor, OrderedDict[str, ndarray]]

classmethod load_default_package()[source]#

Default pre-trained cBottle model package from Nvidia model registry

Return type:

Package

classmethod load_model(package, input_variables=['u10m', 'v10m'], sampler_steps=18, sigma_max=80, seed=0)[source]#

Load AI datasource from package

Parameters:

  • package (Package) – CBottle AI model package

  • input_variables (list[str] | VariableArray) – List of input variables that will be provided for conditioning the output generation, by default [“u10m”, “v10m”]

  • sampler_steps (int, optional) – Number of diffusion steps, by default 18

  • sigma_max (float, optional) – Noise amplitude used to generate latent variables, by default 80

  • seed (int, optional) – Random generator seed for latent variables, by default 0

Returns:

Diagnostic model

Return type:

DiagnosticModel

Examples using earth2studio.models.dx.CBottleInfill#

CBottle Data Generation and Infilling

CBottle Data Generation and Infilling

CBottle Super Resolution

CBottle Super Resolution