`CBottleInfill`¶

GlobalCM202580 GBNVIDIAPyTorch

Import path: earth2studio.models.dx.CBottleInfill

View source on GitHub View install commands

Documentation¶

Bases: Module, AutoModelMixin

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 (Module) –

Core Pytorch model
sst_ds (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, 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 | None, default: None ) –

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

call ¶

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

Forward pass of diagnostic

load_default_package `classmethod` ¶

load_default_package() -> Package

Default pre-trained cBottle model package from Nvidia model registry

load_model `classmethod` ¶

load_model(
    package: Package,
    input_variables: list[str] | VariableArray = [
        "u10m",
        "v10m",
    ],
    sampler_steps: int = 18,
    sigma_max: float = 200,
) -> DiagnosticModel

Load AI datasource from package

Parameters:

package (Package) –

CBottle AI model package
input_variables (list[str] | VariableArray, default: ['u10m', 'v10m'] ) –

List of input variables that will be provided for conditioning the output generation, by default ["u10m", "v10m"]
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