Aurora Atmospheric/Weather Forecasting Models

[isaaccorley]

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

[Copilot]

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,

[adamjstewart]

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.

[isaaccorley]

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".

[wesselb]

@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.

[adamjstewart]

@cnagda can you review?

[cnagda]

Looks good to me. For load_checkpoint , can we override to use commit hash / have our own version of it in this file?

[isaaccorley]

Looks good to me. For load_checkpoint , can we override to use commit hash / have our own version of it in this file?

Planning to upstream this fix directly to the aurora repo microsoft/aurora#114

[adamjstewart]

Failing minimum tests due to huggingface/safetensors#641, for now maybe we can pin to older safetensors.

[isaaccorley]

Can we support older aurora? The aurora 1.6.0 version pins to a specific version of timm<1.0 which conflicts with our min version of timm being >1.0

[adamjstewart]

I believe aurora 1.6 includes your PR

[adamjstewart]

@isaaccorley I'll let you take over again. Now that the tests are actually running, there are a bunch of new errors.