earth2studio.models.dx.CBottleSR#
- class earth2studio.models.dx.CBottleSR(sr_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.
Note
For more information see the following references:
- Parameters:
sr_model (torch.nn.Module) – Core cBottle super-resolution helper module implementing the diffusion process
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 (Tuple[int, int, int, int] | None, 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, 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
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:
- classmethod load_model(package, lat_lon=True, output_resolution=(2161, 4320), super_resolution_window=None, sampler_steps=18, sigma_max=800, seed=None, distilled_model=False, device='cpu')[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 (tuple[int, int, int, int] | None, 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
distilled_model (bool, optional) – Whether to use the distilled model, If True, the distilled helper is used, enabling generation with fewer sampler steps, by default False
device (str) – Device to load model onto, by default cpu
- Returns:
Diagnostic model
- Return type:
DiagnosticModel
