earth2studio.models.dx.CBottleTCGuidance#

class earth2studio.models.dx.CBottleTCGuidance(core_model, classifier_model, sst_ds, lat_lon=True, sampler_steps=18, sigma_max=200.0, batch_size=4, seed=None, dataset_modality=DatasetModality.ERA5)[source]#

Climate in a Bottle tropical cyclone guidance diagnostic. This model for Climate in a Bottle (cBottle) allows users to provide an cyclone guidance map on a lat-lon grid and synthesis global climate realizations at that given time. The tropical cyclone guidance field is regridded to HPX Level 3, which is then used during the sampling process.

Note

For more information see the following references:

Note

This model provides the function model.create_guidance_tensor() as a utility to create the input guidance tensor.

Parameters:

  • core_model (torch.nn.Module) – Core Pytorch diffusion model

  • classifier_model (torch.nn.Module) – Pytorch classifier model

  • sst_ds (xr.Dataset) – Sea surface temperature xarray dataset

  • lat_lon (bool, optional) – Lat/lon toggle, if true the model will return output on a 0.25 deg lat/lon grid. If false, the native nested HealPix grid will be returned, by default True

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

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

  • batch_size (int, optional) – Batch size to generate time samples at, consider adjusting based on hardware being used, by default 4

  • seed (int, optional) – Random generator seed for latent variables. If None will use no seed, by default None

  • dataset_modality (DatasetModality, optional) – Dataset modality label to use when sampling (0=ICON, 1=ERA5), by default DatasetModality.ERA5

__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, sampler_steps=18, sigma_max=200, seed=None)[source]#

Load diagnostic from package

Parameters:

  • package (Package) – CBottle AI model package

  • lat_lon (bool, optional) – Lat/lon toggle, if true prognostic input/output on a 0.25 deg lat/lon grid. If false, the native nested HealPix grid will be returned, by default True

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

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

  • seed (int, optional) – Random generator seed for latent variables. If None, no seed will be used, by default None

Returns:

Diagnostic model

Return type:

DiagnosticModel