Skip to content

feat(px): add StormCastEurope prognostic model - #1144

Draft
gertln wants to merge 1 commit into
NVIDIA:mainfrom
gertln:stormcast-cosmo-era5
Draft

gertln wants to merge 1 commit into
NVIDIA:mainfrom
gertln:stormcast-cosmo-era5

Conversation

@gertln

@gertln gertln commented Sep 8, 2026

Copy link
Copy Markdown
Collaborator

Description

Adds StormCastEurope, an autoregressive prognostic downscaling model for the European domain. Given the current state and the target-time ERA5 fields, it predicts the next hourly state, p(REA_{t+1} | REA_t, ERA5_{t+1}) and rolls out autoregressively. It is available at two resolutions: rea6 (COSMO-REA6, roughly 6 km) and rea2 (COSMO-REA2, roughly 2.2 km). It shares its variable channels with the CorrDiff-COSMO Europe downscaler, so the two can be used on the same fields, CorrDiff for single-shot downscaling and StormCastEurope for hourly temporal consitent rollouts.

ERA5 conditioning is supplied through any DataSource or ForecastSource; the wrapper regrids and normalizes ERA5 onto the COSMO grid and concatenates it with the previous state as the diffusion condition. Physical constraints are optional and applied in physical space each step (per-variable bounds together with a shortwave solar gate), and passing constraints=None runs the model unconstrained. Extended static invariants (elevation, land fraction, and latitude/longitude encodings) are loaded from the package.
Hub-height wind is an optional, output-only derivation that is emitted but never fed back into the rollout. Sub-domain inference is available through set_domain(...), which crops to a latitude/longitude bounding box with patch alignment and a NATTEN minimum-domain check.

image

Checklist

  • I am familiar with the Contributing Guidelines.
  • New or existing tests cover these changes.
  • The documentation is up to date with these changes.
  • 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 for guidance; use discretion, addressing all feedback is not required).

Dependencies

Adds the stormcast-europe extra (earth2studio[cosmo] + nvidia-physicsnemo>=2.2.0a0 + scipy>=1.15.2), added to all. Like cosmo/stormcast-conus, PhysicsNeMo is git-pinned via [tool.uv.sources] until the natten2d_rope (RoPE) operator ships on PyPI.

Add one-hour autoregressive COSMO-REA forecasts at rea6 and rea2,
conditioned on the previous state and target-time ERA5 data. Configure the
default package location at hf://nvidia/stormcast-cosmo-era5.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
@copy-pr-bot

copy-pr-bot Bot commented Sep 8, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant