Skip to content

Add WeatherNext-style ONNX conversion demo - #469

Draft
justinchuby with Copilot wants to merge 16 commits into
mainfrom
copilot/convert-weather-next
Draft

Add WeatherNext-style ONNX conversion demo#469
justinchuby with Copilot wants to merge 16 commits into
mainfrom
copilot/convert-weather-next

Conversation

Copilot AI commented Aug 10, 2026

Copy link
Copy Markdown
Contributor

WeatherNext uses a JAX/Haiku + xarray architecture that does not map to the existing HuggingFace model registry flow. This PR adds a runnable ONNX demo that sketches the WeatherNext one-step forecast contract and graph workflow under examples/.

  • Demo workflow

    • Adds examples/weathernext.py.
    • Builds a small deterministic grid→mesh→grid forecast graph.
    • Saves the result as a standard Mobius ModelPackage.
    • Supports optional ONNX Runtime validation with --validate.
  • WeatherNext-style graph contract

    • Inputs:
      • input_state: weather variables on a lat/lon grid
      • forcings: known future forcing variables
      • sample_noise: stochastic FGN-style noise channels
    • Output:
      • next_state: one-step forecast grid
  • Coverage

    • Adds tests/weathernext_example_test.py to verify the exported ONNX graph I/O contract.

Example:

PYTHONPATH=src python examples/weathernext.py output/weathernext-mini --validate

@CLAassistant

Copy link
Copy Markdown

CLA assistant check
Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
You have signed the CLA already but the status is still pending? Let us recheck it.

Copilot AI linked an issue Aug 10, 2026 that may be closed by this pull request
Copilot AI and others added 7 commits August 10, 2026 06:22
Signed-off-by: GitHub <noreply@github.com>

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Signed-off-by: GitHub <noreply@github.com>

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Signed-off-by: GitHub <noreply@github.com>

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Signed-off-by: GitHub <noreply@github.com>

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Signed-off-by: GitHub <noreply@github.com>

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Signed-off-by: GitHub <noreply@github.com>

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Signed-off-by: GitHub <noreply@github.com>

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Comment thread examples/weathernext.py Fixed
Copilot AI changed the title [WIP] Explore conversion of weather next to onnx Add WeatherNext-style ONNX conversion demo Aug 10, 2026
Copilot AI requested a review from justinchuby August 10, 2026 06:33
@justinchuby

Copy link
Copy Markdown
Member

@copilot what I wanted was a formal support of this model in the modeling/integation files, and a demo script that can run this model on real data.

Copilot AI and others added 8 commits August 10, 2026 13:57
Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>

Copilot AI commented Aug 10, 2026

Copy link
Copy Markdown
Contributor Author

@copilot what I wanted was a formal support of this model in the modeling/integation files, and a demo script that can run this model on real data.

Implemented in f588c1a (with the WeatherNext support commits immediately before it): the WeatherNext graph is now formalized under mobius config/model/task/integration modules, and examples/weathernext.py uses that path with NPZ/xarray real-data inputs plus optional converted weights.

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.

Convert weather next

3 participants