earth2studio.models.dx.CBottleSR#

class earth2studio.models.dx.CBottleSR(core_model, output_resolution=(2161, 4320), super_resolution_window=None, sampler_steps=18, sigma_max=800)[source]#

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 are then regridded to a lat/lon grid. Suggested output dimensions are (2161, 4320) which corresponds to 10km resolution at the equator or (4321, 8640) which corresponds to 5km resolution at the equator. The model can also be used to generate results on a smaller region of the globe by specifying a super-resolution window. This can often be desirable as full global results are extremely expensive to generate.

Note

For more information see the following references: - https://arxiv.org/abs/2505.06474v1 - NVlabs/cBottle - https://catalog.ngc.nvidia.com/orgs/nvidia/teams/earth-2/models/cbottle

Parameters:

  • core_model (torch.nn.Module) – Core pytorch model implementing the diffusion-based super-resolution

  • output_resolution (Tuple[int, int], optional) – High-resolution output dimensions, by default (2161, 4320).

  • super_resolution_window (Union[None, Tuple[int, int, int, int]], optional) – 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 return results for the specified window, by default None

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

  • sigma_max (int, optional) – Maximum noise level for diffusion process, by default 800

__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, output_resolution=(2161, 4320), super_resolution_window=None, sampler_steps=18, sigma_max=800)[source]#

Load AI datasource from package

Parameters:

  • package (Package) – CBottle AI model package

  • output_resolution (tuple[int, int], optional) – High-resolution output dimensions, by default (2161, 4320)

  • super_resolution_window (None | tuple[int, int, int, int], optional) – Super-resolution window, by default None

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

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

Returns:

Diagnostic model

Return type:

DiagnosticModel

Examples using earth2studio.models.dx.CBottleSR#

CBottle Super Resolution

CBottle Super Resolution