earth2studio.models.dx.CBottleSR#

class earth2studio.models.dx.CBottleSR(core_model, lat_lon=True, output_resolution=(2161, 4320), super_resolution_window=None, sampler_steps=18, sigma_max=800, seed=None)[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 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.

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

  • lat_lon (bool, optional, by 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], optional) – High-resolution output dimensions for lat/lon output. Only used when lat_lon=True, 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 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.

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

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

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

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

Load diagnostic model from package

Parameters:
  • package (Package) – CBottle AI model package

  • lat_lon (bool, optional, by 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], optional) – High-resolution output dimensions for lat/lon output, by default (2161, 4320)

  • super_resolution_window (None | tuple[int, int, int, int], optional) – Super-resolution window for lat/lon output, 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

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

Returns:

Diagnostic model

Return type:

DiagnosticModel

Examples using earth2studio.models.dx.CBottleSR#

CBottle Super Resolution

CBottle Super Resolution