CorrDiffTaiwan

ASDS202340 GBNVIDIAPyTorch

Import path: earth2studio.models.dx.CorrDiffTaiwan

View source on GitHub View install commands

Documentation

Bases: Module, AutoModelMixin

CorrDiff is a Corrector Diffusion model that learns mappings between low- and high-resolution weather data with high fidelity. This particular model was trained over a particular region near Taiwan.

Note

This model and checkpoint are from Mardani, Morteza, et al. 2023. For more information see the following references:

• arxiv.org/abs/2309.15214
• catalog.ngc.nvidia.com/orgs/nvidia/teams/modulus/models/corrdiff_inference_package

Parameters:

• residual_model (Module) –

  Core pytorch model

• regression_model (Module) –

  Core pytorch model

• in_center (Tensor) –

  Model input center normalization tensor of size [20,1,1]

• in_scale (Tensor) –

  Model input scale normalization tensor of size [20,1,1]

• out_center (Tensor) –

  Model output center normalization tensor of size [4,1,1]

• out_scale (Tensor) –

  Model output scale normalization tensor of size [4,1,1]

• out_lat (Tensor) –

  Output latitude grid of size [448, 448]

• out_lon (Tensor) –

  Output longitude grid of size [448, 448]

• number_of_samples (int, default: 1 ) –

  Number of high resolution samples to draw from diffusion model. Default is 1

• number_of_steps (int, default: 8 ) –

  Number of langevin diffusion steps during sampling algorithm. Default is 8

• solver (Literal['euler', 'heun'], default: 'euler' ) –

  Discretization of diffusion process. Only 'euler' and 'heun' are supported. Default is 'euler'

• seed (int | None, default: None ) –

  Random seed for reproducibility. Default is None.

__call__

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

Forward pass of diagnostic

load_default_package classmethod

load_default_package() -> Package

Default pre-trained corrdiff model package from Nvidia model registry

load_model classmethod

load_model(
    package: Package, device: str | None = None
) -> DiagnosticModel

Load diagnostic from package

Examples using earth2studio.models.dx.CorrDiffTaiwan

Generative DownscalingGenerative downscaling over Taiwan using CorrDiff diffusion model. Ensemble Forecasting with DownscalingCustom ensembling workflow with generative downscaling using CorrDiff.