InterpCRPSDiT

GlobalMrf202680 GB

Import path: earth2studio.models.px.InterpCRPSDiT

View source on GitHub View install commands

Documentation

Bases: Module, AutoModelMixin, PrognosticMixin

One-shot, endpoint-pinned CRPS temporal interpolation with a DiT backbone (73 variables at 0.25 degree resolution, 721x1440).

Upsamples the coarse trajectory of a base prognostic model to num_interp_steps sub-steps per coarse interval. For interior fraction tau and per-member latent z each frame is produced in a single forward pass as out(tau) = (1 - tau) x0 + tau xT + sin(pi tau) f(x0, xT, cond, z), a linear blend plus a DiT correction f that vanishes at both endpoints. Endpoints are copied from the base model, only interior frames are modelled. Ensemble spread comes entirely from the per-member latent z (trained with a CRPS objective).

The model wraps a base prognostic model (px_model) so the pair acts as a single prognostic with higher temporal resolution. The coarse gap is taken from px_model (trained on 3-10 h) and the output resolution is gap / num_interp_steps, e.g. a 6 h gap with num_interp_steps=36 gives 10 min. num_interp_steps must divide the gap in minutes. A regional sub-domain is available via set_domain.

The DiT runs on 720 latitudes (90 to -89.75). The -90 row of the 721x1440 output is filled with the zonal mean of the -89.75 row.

Note

Trained and validated only on ERA5. The interior is a learned statistical interpolation with no physical conservation guarantee, so quality may degrade with other base data (e.g. GFS). A gap outside 3-10 h only warns. The DiT runs in fp32 by default, see amp_dtype.

Parameters:

• dit (Module) –

  DiT network predicting the correction f (a physicsnemo Module)

• center (Tensor) –

  Per-variable normalization, shape (1, 73, 1, 1)

• scale (Tensor) –

  Per-variable normalization, shape (1, 73, 1, 1)

• static_inv (Tensor) –

  Time-invariant conditioning maps (12, H, W): sin/cos lat, sin/cos lon, lsm, orog and 6 physical fields

• lat2d (Tensor) –

  (H, W) latitude/longitude grids in degrees, defines the output grid

• lon2d (Tensor) –

  (H, W) latitude/longitude grids in degrees, defines the output grid

• px_model (PrognosticModel) –

  Base model producing the coarse trajectory

• num_interp_steps (int, default: 6 ) –

  Sub-steps per coarse interval, must divide the coarse gap in minutes, by default 6

• lon_pad (int, default: 128 ) –

  Circular longitude pad (columns) to reduce the date-line seam, 0 disables, by default 128

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

  If set, noise is drawn from a per-device generator seeded with this value. Results depend on the number of ensemble members and init times, by default None (global RNG)

• drop_variables (list[OptionalVariable] | None, default: None ) –

  Variables the base model does not supply, a subset of sp, u100m, v100m, tcwv. Dropped channels are flagged absent to the DiT and omitted from the output, by default None

• amp_dtype (dtype | None, default: None ) –

  None runs the DiT in fp32, torch.bfloat16 / torch.float16 run it under autocast. Endpoint pinning stays in fp32, by default None

__call__

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

Run one step, returning the initial condition.

Parameters:

• x (Tensor) –

  Input tensor.

• coords (CoordSystem) –

  Input coordinate system.

Returns:

• tuple[Tensor, CoordSystem] –

  Initial-condition frame and its coordinate system

create_iterator

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

Yield the initial condition, then interpolated sub-steps and coarse endpoints.

Parameters:

• x (Tensor) –

  Input tensor.

• coords (CoordSystem) –

  Input coordinate system.

Yields:

• tuple[Tensor, CoordSystem] –

  Each emitted frame and its coordinate system

load_default_package classmethod

load_default_package() -> Package

Load the default model package.

Returns:

• Package –

  Default model package

load_model classmethod

load_model(
    package: Package,
    px_model: PrognosticModel,
    num_interp_steps: int = 6,
    drop_variables: list[OptionalVariable] | None = None,
    amp_dtype: dtype | None = None,
) -> PrognosticModel

Load from a package with the DiT checkpoint, normalization stats and static maps.

Parameters:

• package (Package) –

  Package holding CRPSModel.mdlus, set_phys.nc, global_means.npy, global_stds.npy.

• px_model (PrognosticModel) –

  Base model producing the coarse trajectory

• num_interp_steps (int, default: 6 ) –

  Sub-steps per coarse interval, must divide the coarse gap in minutes, by default 6

• drop_variables (list[OptionalVariable] | None, default: None ) –

  Variables the base does not supply, see the class docstring, by default None

• amp_dtype (dtype | None, default: None ) –

  DiT inference precision, see the class docstring, by default None

Returns:

• PrognosticModel –

  The loaded InterpCRPSDiT.

Raises:

• ValueError –

  If set_phys.nc latitude is not descending or has non-finite values, or the normalization arrays are not len(VARIABLES) finite channels with positive scales.

Examples using earth2studio.models.px.InterpCRPSDiT

Temporal InterpolationTemporal Interpolation inference using InterpCRPSDiT model.