Add InterpCRPSDiT for CRPS temporal interpolation

[gertln]

Description

Add InterpCRPSDiT, an endpoint-pinned CRPS temporal-interpolation model that refines a base model’s coarse trajectory to hourly and sub-hourly resolution. For a 6 h bracket, num_interp_steps=6, 24, and 36 produce 1 h, 15 min, and 10 min output, respectively. The requested cadence must be a whole number of minutes and evenly partition the coarse gap.

For each interior fraction tau the field is produced in a single forward as out(tau) = (1−tau)·x0 + tau·xT + sin(pi·tau)·f(x0, xT, cond, z) consisting of a linear base plus a DiT correction whose sin(pi·tau) envelope vanishes at both endpoints. The bracket endpoints are the base model's own frames and only the interior is learned. A regional sub-domain is available via set_domain. A base does not have to provide all 73 variables. If the only ones it lacks are the four optional channels (sp, u100m, v100m, tcwv), pass them to drop_variables. The model then no longer requires them from the base model but it fills them internally and omits them from the output.

Closes #986

Blocked by:

• PR Add DataReplay to drive the prognostic iterator with a DataSource #985 which is used by the example and tests.
• nvidia-physicsnemo 2.2.0 on PyPI — needed to resolve uv.lock (2.2.0 is the first release with natten2d_rope + proj_reshape_2d_conv, PR #1731). Until then uv lock cannot resolve the extra, so this PR stays a draft.
• Model weights hosting — load_default_package currently returns a placeholder hf://nvidia/earth2studio-interp-crps-dit; needs the bundle hosted

Model details

Property	Value
Architecture	PyTorch — DiT backbone with natten2d_rope neighborhood attention (physicsnemo Module)
Time step	Output = coarse gap / num_interp_steps (trained gaps 3–10 h; e.g. 6 h/6 = 1 h). Base gap auto-detected from the wrapped model.
Input variables	73 (ERA5 0.25° set); the base need only supply 73 − drop_variables
Output variables	73 minus any drop_variables
Spatial resolution	0.25° × 0.25° (720 × 1440); 721-lat bases are cropped to 720
Checkpoint source	HuggingFace — hf://nvidia/earth2studio-interp-crps-dit (TODO: host + @commit)
Checkpoint size	~382 MiB total (model CRPSModel.mdlus 376 MiB, grid set_phys.nc 5.6 MiB, plus norm .npy + config.json)

Dependencies added

The model has the interp-crps-dit optional-dependency extra.

Package	Version	License	Reason
natten	(py < 3.14)	MIT	Neighborhood attention kernels for the DiT natten2d_rope attention
nvidia-physicsnemo	>=2.2.0	Apache-2.0	DiT natten2d_rope layer + Module checkpoint format (first in 2.2.0 / PR #1731)

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes (example, models_px.rst, install.md).
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.
• Assess and address Greptile feedback (AI code review bot).
• PR 1 (DataReplay) merged — this PR is stacked on it (example + tests import DataReplay).
• Weights hosted on HuggingFace and load_default_package URL pinned with @commit.
• nvidia-physicsnemo 2.2.0 available on PyPI (unblocks uv.lock).

[greptile-apps]

Greptile Summary

This PR introduces two new prognostic model classes — InterpCRPSDiT and DataReplay — along with their tests, a worked example, dependency declarations, and documentation updates.

• InterpCRPSDiT (1110 lines): a one-shot, endpoint-pinned CRPS temporal interpolator backed by a DiT with neighborhood attention. The interpolation formula, coordinate bookkeeping, halo trimming, variable drop, and batch_func interaction are all correct. P2 issues include a misleading variance comment in _gaussian_blur_valid and a placeholder load_default_package URL.
• DataReplay (272 lines): wraps any DataSource as a prognostic model; batch-broadcast, grid-mismatch guard, finite-value check, and hook contracts are sound and well-tested.
• The pyproject.toml changes correctly add the interp-crps-dit extra and uv conflict declarations to prevent resolution failures against the da-healda physicsnemo version cap.

Confidence Score: 4/5

Safe to merge once the placeholder HuggingFace URL is pinned to real weights; the core interpolation logic, coordinate handling, and test coverage are solid.

The core math is correct throughout. The Gaussian blur normalizer has a misleading comment but reproduces training-time behavior faithfully. The only non-trivial concern before shipping to users is that load_default_package returns a non-existent URL, producing a hard runtime error for anyone following the example as written. The PR is already marked as a draft pending weight hosting.

interpcrpsdit.py — the _gaussian_blur_valid comment and the load_default_package placeholder URL; 07_crps_temporal_interpolation.py — documents sub.num_interp_steps mutation as an implicit API pattern without a docstring entry in set_domain.

Important Files Changed

Filename	Overview
earth2studio/models/px/interpcrpsdit.py	New 1110-line InterpCRPSDiT model; core interpolation logic, halo/sub-domain handling, and coordinate bookkeeping are correct, with P2 issues around a misleading variance comment in _gaussian_blur_valid and the placeholder load_default_package URL
earth2studio/models/px/datareplay.py	New DataReplay prognostic adapter; batch-dim handling, grid validation, hook contract, and step arithmetic are all correct and well-tested
test/models/px/test_interpcrpsdit.py	Comprehensive 1238-line test suite covering gap validation, endpoint pinning, drop_variables, set_domain, amp_dtype, seeded noise, and multi-time init; test math and dummy DiT contracts are correct
test/models/px/test_datareplay.py	Good coverage of step validation, grid-mismatch guard, finite-value guard, hook contracts, and coordinate ordering; all test assertions are mathematically correct
examples/02_medium_range/07_crps_temporal_interpolation.py	Clear worked example; post-construction mutation of sub.num_interp_steps is an undocumented API pattern, and the example will fail at runtime until weights are hosted
pyproject.toml	Correctly adds the interp-crps-dit extra and uv conflict declarations to prevent resolution failures with da-healda's physicsnemo cap
earth2studio/models/px/init.py	Adds DataReplay and InterpCRPSDiT imports in alphabetical order; no issues
test/conftest.py	Correctly registers test_interpcrpsdit.py with the interp-crps-dit optional-dependency guard; DataReplay has no optional deps so test_datareplay.py is correctly left unregistered

Comments Outside Diff (2)

1. earth2studio/models/px/interpcrpsdit.py, line 1163-1171 (link)

   _trim with bot=0 or right=0 — behavior is correct but a comment would help

   The slice x[..., top : h - bot, left : w - right] correctly handles bot=0 and right=0 via integer arithmetic (h - 0 = h, w - 0 = w), avoiding any negative-index issue. The early-return guard fires only when all four halo values are zero. This is fine as-is, but a short comment noting that bot=0/right=0 are intentionally handled by h-0=h arithmetic would help future readers distinguish the correct intent from an apparent off-by-one.

   Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

2. examples/02_medium_range/07_crps_temporal_interpolation.py, line 1678 (link)

   Direct attribute mutation post-set_domain as an undocumented API pattern

   sub.num_interp_steps = 24 mutates a public attribute after object construction. Because set_domain inherits num_interp_steps from the parent, users who want a different cadence for the sub-domain must mutate it afterwards. The example is the primary documentation of this workflow, implicitly establishing mutation as the intended API. If num_interp_steps is later cached or used to derive other state at construction time, this pattern could silently break. Consider accepting num_interp_steps as a parameter to set_domain, or documenting the mutation explicitly in the set_domain docstring.

_{Reviews (1): Last reviewed commit: "Add InterpCRPSDiT: one-shot endpoint-pin..." | Re-trigger Greptile}

[greptile-apps]

Misleading unit-variance claim in blur normalizer

The comment states "unit-variance white noise stays unit-variance", but this is mathematically incorrect. A separable 2D Gaussian blur (two successive 1D convolutions, each with normalized kernel k) multiplies the input variance by (sum(k_i²))², not sum(k_i²). Dividing by (k * k).sum() = sum(k_i²) only removes one factor, leaving output variance = sum(k_i²) — which equals 1 only for a Dirac kernel. For the largest noise scale (scale=64, σ=32, r=96), sum(k_i²) is roughly 1/(2·96+1) ≈ 0.005, so the "unit-variance" noise is actually much weaker than claimed.

This does not affect inference correctness if the training code used the identical normalization (as stated by "Reproduces the noise generator used at training"), but the incorrect claim could mislead future maintainers into "fixing" the normalization in a way that breaks the training/inference match.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[gertln]

For a normalized 1d kernel, each blur pass scales the marginal white-noise variance by the sum of the squared kernel weights. The two perpendicular passes therefore apply this factor twice. Dividing the output values by that factor applies the inverse factor twice to the variance, restoring unit marginal variance. This matches the training-time noise generator and is retained for consistency.

[greptile-apps]

You're correct, and my original comment was wrong. I conflated scalar value division with variance scaling: dividing the output values by c divides the variance by c², not c. So dividing by sum(k_i²) restores unit marginal variance exactly — two convolution passes each multiply variance by sum(k_i²), giving sum(k_i²)², and dividing values by sum(k_i²) divides variance by sum(k_i²)², which fully cancels it. The comment and normalization are correct.

Tip: You can customize Greptile's behavior for this repo with .greptile/rules.md and .greptile/config.json.

[greptile-apps]

Placeholder URL will cause a runtime failure for any user calling load_default_package

load_default_package currently returns hf://nvidia/earth2studio-interp-crps-dit, which does not yet exist. Any user who calls InterpCRPSDiT.load_default_package() (e.g. by following the example as written) will get a hard runtime error. The TODO comment makes this clear to developers, but the public example links directly to load_default_package() without a visible guard. Before removing the draft label, the URL should be pinned to the real commit hash, or the example and class docstring should prominently state that this method will fail until the weights are hosted.

[gertln]

fixed in 269f0fd