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 CBottle3D model package from Nvidia model registry
- Return type:
- 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
