Regional weather diffusion dataloader updates

[pzharrington]

PhysicsNeMo Pull Request

Description

Modifies the recipe dataset contract to also support physicsnemo.datapipes style data loaders, along with some functionality allowing multi-region training (e.g. CONUS+EU). Also modifies the policy to only load the target under "state" key, and require users to always put the condition in the "background" key (including preceding state timesteps).

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.
• If I am implementing a new model or modifying any existing model, I have followed the Models Implementation Coding Standards.

Dependencies

Review Process

All PRs are reviewed by the PhysicsNeMo team before merging.

Depending on which files are changed, GitHub may automatically assign a maintainer for review.

We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score.
This score reflects the AI’s assessment of merge readiness and is not a qualitative judgment of your work, nor is
it an indication that the PR will be accepted / rejected.

AI-generated feedback should be reviewed critically for usefulness.
You are not required to respond to every AI comment, but they are intended to help both authors and reviewers.
Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.

[copy-pr-bot]

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.

[github-actions]

CODEOWNERS review map

Current for commit a2e28568fef9. An approval covers every file listed for that owner; one owner is sufficient for shared files.

No CODEOWNER

• examples/weather/regional_weather_diffusion/README.md
• examples/weather/regional_weather_diffusion/config/dataset/hrrr_era5.yaml
• examples/weather/regional_weather_diffusion/config/dataset/mock.yaml
• examples/weather/regional_weather_diffusion/config/inference/stormcast.yaml
• examples/weather/regional_weather_diffusion/config/model/stormcast.yaml
• examples/weather/regional_weather_diffusion/config/model/stormscope.yaml
• examples/weather/regional_weather_diffusion/config/test_diffusion_unet.yaml
• examples/weather/regional_weather_diffusion/config/test_regression_unet.yaml
• examples/weather/regional_weather_diffusion/datasets/init.py
• examples/weather/regional_weather_diffusion/datasets/data_loader_hrrr_era5.py
• examples/weather/regional_weather_diffusion/datasets/dataset.py
• examples/weather/regional_weather_diffusion/datasets/mock.py
• examples/weather/regional_weather_diffusion/inference.py
• examples/weather/regional_weather_diffusion/test_training.py
• examples/weather/regional_weather_diffusion/utils/config.py
• examples/weather/regional_weather_diffusion/utils/io.py
• examples/weather/regional_weather_diffusion/utils/nn.py
• examples/weather/regional_weather_diffusion/utils/parallel.py
• examples/weather/regional_weather_diffusion/utils/plots.py
• examples/weather/regional_weather_diffusion/utils/trainer.py

Comment /codeowners-info to refresh.

[greptile-apps]

The PR appears safe to merge; no actionable new issues remain, and both previous findings are fixed in the current code.

Summary

This PR expands the regional weather diffusion dataset contract to support configurable PyTorch and PhysicsNeMo datapipe loaders, grouped samples, and multi-region sharding. It also consolidates all per-sample conditioning—including previous state—under background, updates training and inference around that contract, and restores automatic custom-dataset discovery.

• Adds loader specifications and data-source abstractions for interchangeable loading backends.
• Adds rank-aware and segment-aware sampling with explicit rejection of empty rank-local shards.
• Updates HRRR/ERA5 and mock datasets to emit a single target state and combined conditioning tensor.
• Updates configurations, inference behavior, tests, and documentation for the revised contract.
• The two previous findings are resolved: custom dataset modules are discovered dynamically, and empty shards now raise before entering the infinite sampler loop.

_{Reviews (2) · Last reviewed commit: "greptile feedback"}

[pzharrington]

@greptileai