Dependencies

python = "^3.11"torch = "2.1.0"timm = "0.9.10"numpy = "1.23.5"

Model-zoo

Pangu-Weather

Model Architecture

Example

# Pangufromweatherlearn.modelsimportPanguimporttorchif__name__=='__main__':B=1# batch_sizesurface=torch.randn(B,4,721,1440)# B, C, Lat, Lonsurface_mask=torch.randn(3,721,1440)# topography mask, land-sea mask, soil-type maskupper_air=torch.randn(B,5,13,721,1440)# B, C, Pl, Lat, Lonpangu_weather=Pangu()output_surface,output_upper_air=pangu_weather(surface,surface_mask,upper_air)

# Pangu_litefromweatherlearn.modelsimportPangu_liteimporttorchif__name__=='__main__':B=1# batch_sizesurface=torch.randn(B,4,721,1440)# B, C, Lat, Lonsurface_mask=torch.randn(3,721,1440)# topography mask, land-sea mask, soil-type maskupper_air=torch.randn(B,5,13,721,1440)# B, C, Pl, Lat, Lonpangu_lite=Pangu_lite()output_surface,output_upper_air=pangu_lite(surface,surface_mask,upper_air)

References

@article{bi2023accurate,  title={Accurate medium-range global weather forecasting with 3D neural networks},  author={Bi, Kaifeng and Xie, Lingxi and Zhang, Hengheng and Chen, Xin and Gu, Xiaotao and Tian, Qi},  journal={Nature},  volume={619},  number={7970},  pages={533--538},  year={2023},  publisher={Nature Publishing Group}}

@article{bi2022pangu,  title={Pangu-Weather: A 3D High-Resolution Model for Fast and Accurate Global Weather Forecast},  author={Bi, Kaifeng and Xie, Lingxi and Zhang, Hengheng and Chen, Xin and Gu, Xiaotao and Tian, Qi},  journal={arXiv preprint arXiv:2211.02556},  year={2022}}

Fuxi

Model Architecture

Example

fromweatherlearn.modelsimportFuxiimporttorchif__name__=='__main__':B=1# batch_sizein_chans=out_chans=70# number of input channels or output channelsinput=torch.randn(B,in_chans,2,721,1440)# B C T Lat Lonfuxi=Fuxi()# patch_size : Default: (2, 4, 4)# embed_dim : Default: 1536# num_groups : Default: 32# num_heads : Default: 8# window_size : Default: 7output=fuxi(input)# B C Lat Lon

References

License

TODO