Aurora Atmospheric/Weather Forecasting Models #2823
Conversation
github-actions
bot
added
documentation
| model = aurora.Aurora(*args, **kwargs) | ||
| else: | ||
| model = getattr(aurora, weights.meta['model'])(*args, **kwargs) | ||
| model.load_checkpoint() |
There was a problem hiding this comment.
This still needs some work because the aurora package automatically downloads the checkpoint from huggingface however it downloads from main and we want to download from a commit hash.
| ) | ||
|
|
||
|
|
||
| def aurora_swin_unet( |
There was a problem hiding this comment.
Not sure how to name this because the package is named aurora so I can't also name the model aurora. Any suggestions on this @adamjstewart?
| @@ -0,0 +1,57 @@ | |||
| # Copyright (c) Microsoft Corporation. All rights reserved. | |||
There was a problem hiding this comment.
Don't review this yet because this entire script needs to be revamped to load dummy test data as xarray.
| @@ -106,6 +106,9 @@ datasets = [ | |||
| # xarray 0.12.3+ required for pandas 1.3.3 support | |||
| "xarray>=0.12.3", | |||
| ] | |||
| models = [ | |||
There was a problem hiding this comment.
@adamjstewart I added this new extras group but not sure if I'm missing something else in the repo that I need to update as well.
isaaccorley
requested review from
Copilot,
adamjstewart and
calebrob6
and removed request for
Copilot and
adamjstewart
There was a problem hiding this comment.
Pull Request Overview
This PR introduces the Aurora Climate Models into TorchGeo, providing pre-trained weights and a new aurora_swin_unet model function. Key changes include:
- A new file (torchgeo/models/aurora.py) defining Aurora_Weights and the aurora_swin_unet model with related metadata.
- Integration of Aurora models into the API (torchgeo/models/api.py, init.py, hubconf.py) and documentation, along with updates to tests and project dependencies.
- Addition of tests for the Aurora model in tests/models/test_aurora.py.
Reviewed Changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| torchgeo/models/aurora.py | New Aurora model and weights definitions with metadata and transforms. |
| torchgeo/models/api.py | Included aurora_swin_unet and Aurora_Weights in model registries. |
| torchgeo/models/init.py | Exported the new Aurora API components. |
| tests/models/test_aurora.py | Added tests to validate the Aurora model functionality. |
| pyproject.toml | Updated dependencies to include "microsoft-aurora". |
| hubconf.py | Exposed aurora_swin_unet in the hub configuration. |
| docs/api/weights/climate.csv | Documented Aurora weight configurations. |
| docs/api/models.rst | Updated API documentation to include Aurora models. |
Comments suppressed due to low confidence (1)
torchgeo/models/api.py:106
- [nitpick] Consider adding a comment to clarify why 'aurora_swin_unet' is entered both as a function key and a string key in the _model_weights dictionary. This will help future maintainers understand the intent behind the dual mapping.
'aurora_swin_unet': Aurora_Weights,
|
Last I knew, Aurora only supports medium-range weather forecasting (2 hrs–2 wks), not climate modeling (decadal scale). Would be good to confirm with the authors. We need to be very careful with naming here, as there is a huge distinction between data-driven stateful models and physics-driven statistical distribution models. |
You're correct, the paper specifically says "forecasting of weather and atmospheric processes". I'll make sure to not use the term "climate modeling". |
isaaccorley
changed the title
|
@adamjstewart, @isaaccorley, you're correct that Aurora only does medium-term forecasting (up to 10 days in 6 hour steps). Note, however, that it doesn't just do weather forecasting: we have models for ocean waves and air pollution too. |
|
@cnagda can you review? |
docs/api/weights/climate.csv
Outdated
There was a problem hiding this comment.
Filename: atmospheric.csv?
There was a problem hiding this comment.
Add column for nowcasting vs. medium-range weather forecasting vs. S2S vs. decadal-scale climate. Could use a key under the table and keys like N/MWF/S2S/C.
There was a problem hiding this comment.
I'm actually not 100% sure but I think these are all going to be nowcasting right since they are only forecasting less than a day out in hours? but you can do it autoregressively so technically you can also do MWF
There was a problem hiding this comment.
Also what do you think the column should be named? Task maybe?
There was a problem hiding this comment.
Yes, they're designed for up to 10 days at 6 hr increments, so definitely MWF. Task seems fine.
Planning to upstream this fix directly to the aurora repo microsoft/aurora#114 |
isaaccorley
force-pushed
the
models/aurora-climate
branch
from
ba9a06a to
9821bc1
Compare
Co-authored-by: Adam J. Stewart <[email protected]>
Co-authored-by: Adam J. Stewart <[email protected]>
Adding the Aurora Atmospheric/Weather Forecasting models from https://github.com/microsoft/aurora!
See their docs and paper for more info!
cc: @wesselb @crisbodnar @megstanley