Skip to content

CBottleSR

GlobalDS202580 GBNVIDIAPyTorch

Import path: earth2studio.models.dx.CBottleSR

View source on GitHub View install commands

Documentation

Bases: Module, AutoModelMixin

Climate in a Bottle Super-Resolution (CBottleSR) model.

CBottleSR is a diffusion-based super-resolution model that learns mappings between low- and high-resolution climate data with high fidelity. This model generates results at 5km resolution on a healpix grid with 10 levels of resolution (1024x1024). The results can be output in either HEALPix format or regridded to a lat/lon grid. If lat/lon is used for output, the results will be regridded to the specified output resolution. Suggested output resolutions are (2161, 4320) for ~10km equatorial resolution and (4321, 8640) for ~5km equatorial resolution. The model can also be used to generate results for a smaller region of the globe by specifying a super-resolution window. This is often desirable as full global results are computationally expensive.

Note

For more information see the following references:

Parameters:

  • sr_model (Module) –

    Core cBottle super-resolution helper module implementing the diffusion process

  • lat_lon (bool, optional, by default True, default: True ) –

    Lat/lon toggle, if true the model will expect a lat/lon grid as input and output a lat/lon grid. If false, the native nested HealPix grid will be used for input and output. Input HEALPix is level 6 and output HEALPix is level 10 with NEST pixel ordering

  • output_resolution (Tuple[int, int], default: (2161, 4320) ) –

    High-resolution output dimensions for lat/lon output. Only used when lat_lon=True, by default (2161, 4320)

  • super_resolution_window (Tuple[int, int, int, int] | None, default: None ) –

    Super-resolution window. If None, super-resolution is done on the entire global grid If provided, the super-resolution window is a tuple of (lat_south, lon_west, lat_north, lon_east) and will only apply super-resolution to the specified window. For lat/lon output, the result will just be returned for the specified window with the specified output resolution, by default None

  • sampler_steps (int, default: 18 ) –

    Number of diffusion steps, by default 18

  • sigma_max (int, default: 800 ) –

    Maximum noise level for diffusion process, by default 800

  • seed (int, default: None ) –

    Random generator seed for latent variables, 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,
    lat_lon: bool = True,
    output_resolution: tuple[int, int] = (2161, 4320),
    super_resolution_window: (
        tuple[int, int, int, int] | None
    ) = None,
    sampler_steps: int = 18,
    sigma_max: int = 800,
    seed: int | None = None,
    distilled_model: bool = False,
) -> DiagnosticModel

Load diagnostic model from package

Parameters:

  • package (Package) –

    CBottle AI model package

  • lat_lon (bool, optional, by default True, default: True ) –

    Lat/lon toggle, if true the model will expect a lat/lon grid as input and output a lat/lon grid. If false, the native nested HEALPix grid will be used for input and output Input HEALPix is level 6 and output HEALPix is level 10 with NEST pixel ordering

  • output_resolution (tuple[int, int], default: (2161, 4320) ) –

    High-resolution output dimensions for lat/lon output, by default (2161, 4320)

  • super_resolution_window (tuple[int, int, int, int] | None, default: None ) –

    Super-resolution window for lat/lon output, by default None

  • sampler_steps (int, default: 18 ) –

    Number of diffusion steps, by default 18

  • sigma_max (float, default: 800 ) –

    Noise amplitude used to generate latent variables, by default 800

  • seed (int, default: None ) –

    Random generator seed for latent variables, by default None

  • distilled_model (bool, default: False ) –

    Whether to use the distilled model, If True, the distilled helper is used, enabling generation with fewer sampler steps, by default False

Returns:

  • DiagnosticModel –

    Diagnostic model

Examples using earth2studio.models.dx.CBottleSR

Preview of CBottle Super ResolutionCBottle Super ResolutionClimate in a Bottle (cBottle) super resolution workflows for global weather.