Skip to content

CBottleTCGuidance

GlobalCM202580 GBNVIDIAPyTorch

Import path: earth2studio.models.dx.CBottleTCGuidance

View source on GitHub View install commands

Documentation

Bases: Module, AutoModelMixin

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 (Module) –

    Core Pytorch diffusion model

  • classifier_model (Module) –

    Pytorch classifier model

  • sst_ds (Dataset) –

    Sea surface temperature xarray dataset

  • lat_lon (bool, default: True ) –

    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, default: 18 ) –

    Number of diffusion steps, by default 18

  • sigma_max (float, default: 200.0 ) –

    Noise amplitude used to generate latent variables, by default 200

  • batch_size (int, default: 4 ) –

    Batch size to generate time samples at, consider adjusting based on hardware being used, by default 4

  • seed (int, default: None ) –

    Random generator seed for latent variables. If None will use no seed, by default None

  • dataset_modality (DatasetModality, default: ERA5 ) –

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

__call__ 

__call__(
    x: Tensor, coords: CoordSystem
) -> tuple[Tensor, CoordSystem]

Forward pass of diagnostic

load_default_package classmethod 

load_default_package() -> Package

Default pre-trained cBottle model package from Nvidia model registry

load_model classmethod 

load_model(
    package: Package,
    lat_lon: bool = True,
    sampler_steps: int = 18,
    sigma_max: float = 200,
    seed: int | None = None,
    allow_second_order_derivatives: bool = False,
) -> DiagnosticModel

Load diagnostic from package

Parameters:

  • package (Package) –

    CBottle AI model package

  • lat_lon (bool, default: True ) –

    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, default: 18 ) –

    Number of diffusion steps, by default 18

  • sigma_max (float, default: 200 ) –

    Noise amplitude used to generate latent variables, by default 80

  • seed (int, default: None ) –

    Random generator seed for latent variables. If None, no seed will be used, by default None

  • allow_second_order_derivatives (bool, default: False ) –

    Enable checkpoint/model loading path required for second-order autodiff (needed for odds-ratio computations). Keep False for faster standard guided inference, by default False.

Returns:

  • DiagnosticModel –

    Diagnostic model

Examples using earth2studio.models.dx.CBottleTCGuidance

Preview of CBottle Tropical Cyclone GuidanceCBottle Tropical Cyclone GuidanceGuided tropical cyclone sampling with cBottle and odds-ratio diagnostics.