CorrDiffEra5Hrrr

NADS202680 GBNVIDIAPyTorch

Import path: earth2studio.models.dx.CorrDiffEra5Hrrr

View source on GitHub View install commands

Documentation

Bases: Module, AutoModelMixin

Generative downscaling from 0.25-degree ERA5 to 3 km HRRR over CONUS.

Note

For more information see the following references:

• huggingface.co/nvidia/corrdiff-era5-hrrr

Parameters:

• network (Module) –

  Conditioned network. For network_kind="rectified_flow" a ConcatConditionWrapper(DiT) whose output is the velocity or the clean data (see prediction_type); for network_kind="edm" an EDMPreconditioner(ConcatConditionWrapper(DiT)) (an x0-predictor).

• lat_input_grid (Tensor) –

  1-D regular ERA5 input grid (the native training footprint). Latitude may be ascending or descending; input longitudes may use either the [0, 360) or the [-180, 180) convention (compared modulo 360).

• lon_input_grid (Tensor) –

  1-D regular ERA5 input grid (the native training footprint). Latitude may be ascending or descending; input longitudes may use either the [0, 360) or the [-180, 180) convention (compared modulo 360).

• lat_output_grid (Tensor) –

  2-D HRRR latitude / longitude [H, W] of the output crop (degrees).

• lon_output_grid (Tensor) –

  2-D HRRR latitude / longitude [H, W] of the output crop (degrees).

• hrrr_y (Tensor) –

  1-D native HRRR projection coordinates (m) of the output crop.

• hrrr_x (Tensor) –

  1-D native HRRR projection coordinates (m) of the output crop.

• era5_center (Tensor) –

  ERA5 input normalization (mean / std), size [n_era5].

• era5_scale (Tensor) –

  ERA5 input normalization (mean / std), size [n_era5].

• out_center (Tensor) –

  Output normalization (mean / std), size [n_out].

• out_scale (Tensor) –

  Output normalization (mean / std), size [n_out].

• invariants (Tensor) –

  Normalized static invariant channels [n_inv, H, W], appended to the conditioning after the cosine-zenith channel.

• network_kind (Literal['rectified_flow', 'edm'], default: 'rectified_flow' ) –

  Generative formulation, by default "rectified_flow"

• era5_variables (Sequence[str], default: ERA5_VARIABLES ) –

  Input channel order, by default ERA5_VARIABLES

• output_variables (Sequence[str], default: OUTPUT_VARIABLES ) –

  Output channel order, by default OUTPUT_VARIABLES

• presence_flags (Sequence[str], default: () ) –

  ERA5 variables that training randomly dropped from the input; one scalar "present" flag (always 1 at inference) per name.

• day_of_year (bool, default: True ) –

  Whether the scalar conditioning carries [sin, cos] of the day-of-year phase at the validity time.

• prediction_type (Literal['x0', 'flow'], default: 'x0' ) –

  Rectified-flow output parameterization. Ignored for EDM, by default "x0"

• time_scale (float, default: 999.0 ) –

  Multiplier applied to the rectified-flow time t in [0, 1] before the network's timestep embedder (the examples train with 999.0).

• number_of_samples (int, default: 1 ) –

  Ensemble members per input; settable between calls.

• number_of_steps (int, default: 50 ) –

  ODE solver steps (each Heun step costs two network evaluations).

• solver ((heun, euler), default: "heun" ) –

  ODE solver.

• shift (float, default: 32.0 ) –

  SD3 resolution shift a applied to the rectified-flow time grid, t -> a t / (1 + (a - 1) t); 1.0 disables it. The examples found a of 12-32 optimal for full-domain sampling (the networks train on 256 x 256 patches). Ignored for EDM.

• t_max (float, default: 0.99 ) –

  Rectified-flow start time (the noise end); kept below 1.

• x0v_clip (float, default: 0.05 ) –

  Lower clamp of the denominator in the x-prediction to velocity conversion v = (x_t - x0_hat) / max(t, x0v_clip) near the data end (t -> 0).

• sigma_min (float, default: 0.01 ) –

  EDM Karras schedule parameters. Ignored for rectified flow.

• sigma_max (float, default: 0.01 ) –

  EDM Karras schedule parameters. Ignored for rectified flow.

• rho (float, default: 0.01 ) –

  EDM Karras schedule parameters. Ignored for rectified flow.

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

  Base RNG seed for the sampling latents; member i uses seed + i. None leaves sampling unseeded.

• amp (bool, default: True ) –

  Run network forwards under bf16 autocast while the ODE integration stays in fp32 (the examples' evaluation setting; roughly halves inference time).

__call__

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

Downscale ERA5 states to the HRRR CONUS crop.

Parameters:

• x (Tensor) –

  Input tensor [batch, time, variable, lat, lon] (the wrapper flattens leading batch dimensions automatically).

• coords (CoordSystem) –

  Input coordinate system.

Returns:

• tuple[Tensor, CoordSystem] –

  [batch, sample, time, variable, hrrr_y, hrrr_x] and its coordinates.

load_default_package classmethod

load_default_package() -> Package

Default pre-trained model package.

Returns:

• Package –

  Model package with default checkpoint location

load_model classmethod

load_model(
    package: Package,
    device: str | None = None,
    number_of_samples: int | None = None,
    number_of_steps: int | None = None,
    shift: float | None = None,
    seed: int | None = None,
    amp: bool | None = None,
    variant: Literal["x_pred"] = "x_pred",
) -> DiagnosticModel

Load the model from a package.

Parameters:

• package (Package) –

  Model package to load from.

• device (str | None, default: None ) –

  Device to place the model on, by default None (CPU).

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

  Ensemble members per input; defaults to the package metadata.

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

  ODE steps; defaults to the package metadata.

• shift (float | None, default: None ) –

  Rectified-flow resolution shift; defaults to the package metadata.

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

  Base RNG seed, by default None (unseeded).

• amp (bool | None, default: None ) –

  bf16 autocast for the network; defaults to the package metadata.

• variant (x_pred, default: "x_pred" ) –

  Sub-folder of the package to load, by default "x_pred" (the x-prediction rectified-flow model).

Returns:

• DiagnosticModel –

  Loaded model.