CorrDiffCosmoEra5SDA

GlobalDADS202680 GBNVIDIAPyTorch

Import path: earth2studio.models.da.CorrDiffCosmoEra5SDA

View source on GitHub View install commands

Documentation

Bases: Module, AutoModelMixin

CorrDiff-COSMO with score-based data assimilation (SDA) via diffusion posterior sampling (DPS). Wraps a diffusion-mode CorrDiffCosmoEra5 downscaler and, at inference, nudges each denoising step toward sparse point observations -- producing a high-resolution COSMO-REA analysis over Europe that is conditioned on the ERA5 driving state and guided toward the observations.

Takes as input:

• An ERA5 driving state (the same input the downscaler conditions on).
• Sparse point observations (a DataFrame of lat/lon/variable/observation).

CorrDiff-COSMO directly models p(y | ERA5), so observations are mapped into its normalized output space without residual mean subtraction or addition. It is a single-shot diagnostic downscaler with no propagated state: __call__ and create_generator produce an independent analysis for each requested time, conditioned on the corresponding ERA5 state.

Each observation is mapped to its nearest output-grid cell and normalized as (obs - center) / scale. This direct mapping supports only identity-transform, unit-scale channels, such as u10m / v10m and u3d_l47 / v3d_l47. Other channels require a custom observation operator and are rejected.

Note that the u3d_l* channels are terrain-following model levels: the above-ground height of u3d_l47 is a + b * elevation_norm (in normalized elevation; nominally ~119.5 m at zero elevation anomaly). The exact height varies with terrain, so pick the level whose height suits the observations.

Parameters:

• model (CorrDiffCosmoEra5) –

  A diffusion-mode downscaler (mode="diffusion"), already cropped to the target region with CorrDiffCosmoEra5.set_domain if desired.

• assimilate_variables (tuple[str, ...]) –

  Output variables to assimilate (required); must be identity-transform, unit-scale channels, e.g. ("u10m", "v10m") (surface wind, available in either resolution) or ("u3d_l47", "v3d_l47") (REA2 terrain-following model-level wind; above-ground height varies with elevation, see the note above). Choose the channel whose physical height matches your observations -- there is no default, since the choice depends on the observations being assimilated.

• time_tolerance (TimeTolerance, default: timedelta64(10, 'm') ) –

  Observations within this window of the analysis time are used. A single value is symmetric; the default is ±10 minutes.

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

  Posterior ensemble size (independent DPS draws, seeds seed + i); defaults to the wrapped model's number_of_samples.

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

  Diffusion sampler steps; defaults to the wrapped model's number_of_steps.

• sda_std_obs (float | Mapping[str, float], default: 0.5 ) –

  Observation-noise standard deviation for DPS guidance, by default 0.5. A scalar is broadcast to every assimilated variable (in that variable's physical units); a mapping sets it per variable and must contain exactly the assimilated variables (unknown keys are rejected) -- e.g. {"u10m": 0.5, "v10m": 0.5} in m/s. When temperature is also assimilated, its entry could be "t2m": 1.0 in K. Every value must be finite and > 0. It is the effective uncertainty per occupied grid cell (multiple observations in one cell are averaged, not reduced by sqrt(n)).

• sda_gamma (float, default: 5e-05 ) –

  SDA covariance scaling in the DPS likelihood, by default 5e-5. Positive values account for denoiser-estimate uncertainty across diffusion noise levels. Larger values weaken observation guidance, especially early in denoising, so the analysis stays closer to the unguided downscaler and may fit observations less closely. Set to 0 for classical DPS without this correction and the strongest guidance for a fixed sda_std_obs. Must be

  = 0.

• amp (bool, default: False ) –

  Run the guided diffusion sampler under torch.autocast bf16, by default False (full precision, matching the wrapped downscaler). Setting True can reduce runtime and peak GPU memory; it takes effect on CUDA only (CPU always runs full precision).

__call__

__call__(
    x: DataArray, obs: DataFrame | None = None
) -> DataArray

Assimilate obs into the downscaled analysis for each time in x.

Parameters:

• x (DataArray) –

  ERA5 driving state on the native input grid (dims include time, variable, lat, lon; a size-1 lead_time is squeezed).

• obs (DataFrame | None, default: None ) –

  Sparse observations (columns time/lat/lon/variable/ observation), or None for a free (unconstrained) downscaling.

Returns:

• DataArray –

  Posterior COSMO-REA analysis, dims (time, sample, variable, y, x) with 2D lat/lon coordinates.

Raises:

• ValueError –

  If x carries a lead_time dimension of size != 1, or if obs is missing a required column.

• RuntimeError –

  If the model is on CUDA but CuPy is not installed (the analysis cannot be returned on-device; install the da-cosmo extra or run on CPU).

create_generator

create_generator(
    x: DataArray,
) -> Generator[DataArray | None, DataFrame | None, None]

Generator of independent COSMO-REA analyses (a diagnostic re-analysis, not a propagated forecast). Primed with a no-compute yield; each send(obs) produces the obs-constrained analysis for the next time in x.

Parameters:

• x (DataArray) –

  ERA5 driving state (dims include time, variable, lat, lon; a size-1 lead_time is squeezed). One analysis is produced per time.

Yields:

• DataArray | None –

  None on the priming yield (before the first send); thereafter the obs-constrained analysis for the current time, dims (time, sample, variable, y, x) with a single time.

Receives:

• DataFrame | None –

  Sparse observations for the current time (columns time/lat/ lon/variable/observation), or None for a free analysis.

Raises:

• RuntimeError –

  On priming, if the model is on CUDA but CuPy is not installed (the analysis cannot be returned on-device; install the da-cosmo extra or run on CPU).

Example

>>> gen = model.create_generator(x)
>>> next(gen)                  # prime
>>> state = gen.send(obs_df)   # analysis for x's first time

load_default_package classmethod

load_default_package() -> Package

Load the CorrDiff-COSMO model package (shared with the downscaler).

Returns:

• Package –

  The default CorrDiff-COSMO model package.

load_model classmethod

load_model(
    package: Package,
    assimilate_variables: tuple[str, ...],
    resolution: str = "rea2",
    domain: dict | None = None,
    time_tolerance: TimeTolerance = timedelta64(10, "m"),
    number_of_samples: int | None = None,
    sampler_steps: int | None = None,
    sda_std_obs: float | Mapping[str, float] = 0.5,
    sda_gamma: float = 5e-05,
    amp: bool = False,
) -> AssimilationModel

Load the assimilation model from a CorrDiff-COSMO package.

Loads a diffusion-mode CorrDiffCosmoEra5 internally and wraps it.

domain (optional) crops the downscaler to a sub-region BEFORE wrapping, so the observation grid stays in sync; it is forwarded verbatim to CorrDiffCosmoEra5.set_domain (e.g. domain=dict(lat_min=52.5, lat_max=55.5, lon_min=5.5, lon_max=9.5)). Cropping the wrapped .model afterwards would leave a stale obs grid.

Parameters:

• package (Package) –

  CorrDiff-COSMO model package (see load_default_package).

• assimilate_variables (tuple[str, ...]) –

  Output channels to assimilate (required); must be identity-transform, unit-scale channels (e.g. ("u10m", "v10m") surface wind, available in either resolution, or ("u3d_l47", "v3d_l47") REA2 terrain-following model-level wind whose above-ground height varies with elevation). There is no default, since the choice depends on the observations being assimilated.

• resolution (str, default: 'rea2' ) –

  COSMO-REA resolution, "rea2" or "rea6", by default "rea2".

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

  Sub-region crop forwarded to CorrDiffCosmoEra5.set_domain (keys lat_min/lat_max/lon_min/lon_max), by default None (full native footprint).

• time_tolerance (TimeTolerance, default: timedelta64(10, 'm') ) –

  Window for matching observation times to the analysis time, by default 10 minutes.

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

  Posterior ensemble size; defaults to the wrapped model's value.

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

  Number of diffusion sampler steps; defaults to the wrapped model's number_of_steps.

• sda_std_obs (float | Mapping[str, float], default: 0.5 ) –

  Assumed observation-noise std (lower trusts obs more), by default 0.5. A scalar broadcasts to every variable; a mapping sets it per variable and must contain exactly the assimilated variables (e.g. {"u10m": 0.5, "v10m": 0.5}).

• sda_gamma (float, default: 5e-05 ) –

  SDA covariance scaling in the DPS likelihood, by default 5e-5. Positive values account for denoiser-estimate uncertainty across diffusion noise levels. Larger values weaken observation guidance, especially early in denoising, so the analysis stays closer to the unguided downscaler and may fit observations less closely. Set to 0 for classical DPS without this correction and the strongest guidance for a fixed sda_std_obs. Must be >= 0.

• amp (bool, default: False ) –

  Run the guided sampler under torch.autocast bf16 (CUDA only), by default False. Setting True can reduce runtime and GPU memory.

Returns:

• AssimilationModel –

  The wrapped CorrDiffCosmoEra5SDA assimilation model.