Skip to content

CorrDiffCosmoEra5

EUDS202680 GBNVIDIAPyTorch

Import path: earth2studio.models.dx.CorrDiffCosmoEra5

View source on GitHub View install commands

Documentation

Bases: Module, AutoModelMixin

COSMO-REA downscaling model: ERA5 -> high-resolution COSMO-REA.

Diagnostic model that downscales a global ERA5 state to high-resolution COSMO-REA regional reanalysis over Europe -- COSMO-REA6 (~6 km) or COSMO-REA2 (~2.2 km), selected with resolution. The input is an ERA5 state, so it can downscale an ERA5 analysis directly or run behind a global forecast model (e.g. SFNO -> CorrDiffCosmoEra5).

mode selects one of two networks trained per resolution:

  • "mean" -- a deterministic regression predicting the conditional mean E[y | x]: a single smooth field in one forward pass. Fast; a good first high-resolution look or a deterministic conditional-mean baseline.
  • "diffusion" -- a generative model sampling the conditional distribution p(y | x): an ensemble of number_of_samples realizations (seeded for reproducibility) that also captures the spread the mean cannot represent.

Both networks are PhysicsNeMo DiTs at a fixed grid resolution but crop-size agnostic, so set_domain returns an instance restricted to any lat/lon sub-region without retraining (bounded to the trained footprint). Each resolution ships an extended grid beyond its native footprint; for COSMO-REA2 -- whose native footprint is central-European -- the extended grid reaches a broad European domain, enabling 2.2 km downscaling across it. Outputs are surface and model-level (3D) fields -- winds, temperature, humidity, precipitation, cloud cover, fluxes, TKE, PBL height; variables with a canonical Earth2Studio name are relabelled via [CosmoLexicon][earth2studio.lexicon.CosmoLexicon] and COSMO-specific fields keep a descriptive name. Optionally emits derived hub-height wind components (see hub_heights).

Parameters:

  • era5_variables (Sequence[str]) –

    ERA5 input variable names (canonical Earth2Studio lexicon), in network order.

  • output_variables (Sequence[str]) –

    Output (COSMO-REA) variable names, in network output-channel order.

  • regression_model (Module | None) –

    Network for mode="mean" (an upstream physicsnemo RoPE-NATTEN DiT).

  • diffusion_model (Module | None) –

    Network for mode="diffusion" (an upstream physicsnemo EDMPreconditioner(ConcatConditionWrapper(DiT))).

  • resolution ((rea6, rea2), default: "rea6" ) –

    Which COSMO-REA resolution this package targets.

  • mode ((mean, diffusion), default: "mean" ) –

    Which model to run.

  • lat_input_grid (Tensor) –

    1D increasing regular ERA5 input grid (the native default footprint).

  • lon_input_grid (Tensor) –

    1D increasing regular ERA5 input grid (the native default footprint).

  • lat_output_grid (Tensor) –

    2D curvilinear native (rotated-pole) target grid [out_lat, out_lon].

  • lon_output_grid (Tensor) –

    2D curvilinear native (rotated-pole) target grid [out_lat, out_lon].

  • 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) in transformed space, size [n_out].

  • out_scale (Tensor) –

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

  • static_invariants (OrderedDict[str, Tensor]) –

    Native, already-normalized static invariant channels (non-position), each [out_lat, out_lon] — e.g. elevation_norm, land_fraction, z0_lu, slope_east, slope_north. Position channels (sin/cos lat/lon) and cos_zenith are computed, not stored here.

  • pre_invariant_variables (Sequence[str]) –

    Names of the invariant channels before / after cos_zenith in the background channel order.

  • post_invariant_variables (Sequence[str]) –

    Names of the invariant channels before / after cos_zenith in the background channel order.

  • channel_transforms (dict | None, default: None ) –

    Per-output-channel nonlinear transform spec (from the package metadata), used to invert after de-normalizing.

  • constraints (dict | None, default: None ) –

    Physical-constraint spec (metadata["constraints"]): per-channel bounds (min/max clamps) and a shortwave sza_gate (the dawn/dusk solar gate), applied in physical space in postprocess for both modes.

  • number_of_samples (int, default: 1 ) –

    Number of samples (diffusion); the sample dim is kept (size 1) for mode="mean" so both modes share an output contract. The constructor value is the default; it is settable at runtime between calls (model.number_of_samples = N).

  • physical_clamp (bool, default: True ) –

    Apply physical-bounds + the shortwave dawn/dusk solar gate to outputs.

  • number_of_steps (int, default: 18 ) –

    Diffusion sampler step count (EDM Heun). Defaults to 18.

  • sigma_min (float, default: 0.002 ) –

    Diffusion noise-schedule bounds. Default to 0.002 / 800.0.

  • sigma_max (float, default: 0.002 ) –

    Diffusion noise-schedule bounds. Default to 0.002 / 800.0.

  • rho (float, default: 7.0 ) –

    Karras noise-schedule exponent. Defaults to 7.0.

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

    Diffusion ODE solver: "heun" (2nd-order, default) or "euler" (1st-order).

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

    Base RNG seed for diffusion sampling (a per-sample offset is added). None (default) leaves sampling unseeded.

  • amp (bool, default: False ) –

    Run the network forward passes (regression patches and the diffusion denoiser) under torch.autocast bf16. Roughly halves inference time on Ampere+ GPUs at a negligible accuracy cost; the sampler bookkeeping stays in fp64. Defaults to False (full precision).

  • hub_heights (Sequence[float], default: None ) –

    Optional list of above-ground heights (m) at which to emit interpolated wind components u{H}m/v{H}m (e.g. [100] -> u100m, v100m), vertically interpolated from the model's 3D wind levels. Wind speed is obtained by composing the stock DerivedWS(levels=["{H}m"]) diagnostic on the output. None (default) -> no derived outputs and behavior/output_variables are unchanged. Requires wind_levels. Names follow the lexicon m-suffix convention (m above surface, as u100m); arbitrary H need not be pre-registered in the lexicon (outputs are produced, not fetched). A height outside the levels' range is clamped (not extrapolated) and warns.

  • hub_interp ((linear, log), default: "linear" ) –

    Interpolation method for the hub-height wind components: "linear" (in geometric height) or "log" (in ln(height)). Both are roughness-free between levels. Defaults to "linear".

  • wind_levels (dict, default: None ) –

    Package metadata describing the 3D wind levels for hub-height derivation: each level's {u, v} output channels + its a, b height coefficients (height = a + b*elevation_invariant) and the elevation invariant name. Supplied by the package; required when hub_heights is set.

Notes

The diffusion sampler controls (number_of_steps, sigma_min, sigma_max, rho, solver) default to the package metadata but are exposed for inference tuning rather than fixed: as with any diffusion sampler, these settings can shape the sampled output distribution (e.g. ensemble spread and the representation of extremes).

The models are trained on ERA5 (global input) paired with COSMO-REA regional reanalysis (high-resolution target) over Europe:

__call__ 

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

Run the model. x is [batch, time, variable, lat, lon]; coords carries a time axis (validity times) driving the solar-zenith channel.

load_default_package classmethod 

load_default_package() -> Package

Default pre-trained COSMO-REA downscaling package (combined rea6/rea2).

Returns the hosted package; pick the model within it via load_model(..., mode=, resolution=) (or rely on the mean/rea6 defaults through from_pretrained).

load_model classmethod 

load_model(
    package: Package,
    device: str | None = None,
    mode: Literal["mean", "diffusion"] = "mean",
    resolution: Literal["rea6", "rea2"] = "rea6",
    hub_heights: Sequence[float] | None = None,
    hub_interp: Literal["linear", "log"] = "linear",
) -> DiagnosticModel

Load a COSMO-REA downscaling model from a package.

mode ({"mean", "diffusion"}) and resolution ({"rea6", "rea2"}) select the checkpoint within one package. All other configuration (sampler, normalization, transforms) comes from the package metadata.

hub_heights/hub_interp opt into derived hub-height wind components (u{H}m/v{H}m; compose DerivedWS for speed); they require the package to carry a wind_levels metadata block (else requesting hub_heights raises). See the constructor.