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)[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 CBottleTCGuidance.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 data source 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, by default 0

__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=0)[source]#

Load AI datasource 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

  • input_variables (list[str] | VariableArray) – List of input variables that will be provided for conditioning the output generation, by default [“u10m”, “v10m”]

  • 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, by default 0

Returns:

Diagnostic model

Return type:

DiagnosticModel