Weather Swarm - Pytorch
Project description
Aurora
Community and Open Source Implementation of the paper: "Aurora: A Foundation Model of the Atmosphere" in PyTorch: Paper link
Install
pip3 install aurora-torch
Example
import torch
from aurora_torch.main import SwinTransformerUNet3D
from loguru import logger
# Test with random input tensor of shape (B, D, H, W, C)
B, D, H, W, C = 2, 16, 64, 64, 32
model = SwinTransformerUNet3D(input_dim=C, output_dim=C)
input_tensor = torch.rand(B, D, H, W, C)
# Forward pass through the model
output = model(input_tensor)
logger.info(f"Output shape: {output.shape}")
License
MIT
Bibtex
@misc{bodnar2024aurora,
title={Aurora: A Foundation Model of the Atmosphere},
author={Cristian Bodnar and Wessel P. Bruinsma and Ana Lucic and Megan Stanley and Johannes Brandstetter and Patrick Garvan and Maik Riechert and Jonathan Weyn and Haiyu Dong and Anna Vaughan and Jayesh K. Gupta and Kit Tambiratnam and Alex Archibald and Elizabeth Heider and Max Welling and Richard E. Turner and Paris Perdikaris},
year={2024},
eprint={2405.13063},
archivePrefix={arXiv},
primaryClass={physics.ao-ph}
}
References
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
aurora_torch-0.0.9.tar.gz
(4.7 kB
view hashes)
Built Distribution
Close
Hashes for aurora_torch-0.0.9-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 868ee299337992912db5bee2da910eee356800795f79a703818a2c23bd72cfa7 |
|
MD5 | e405ae63f83d2021bd9911d84ea560c5 |
|
BLAKE2b-256 | a7f517f5139fde08df3eb9edf53c250811d635082af7661ddb12abee20abbd5e |