|
--- |
|
license: apache-2.0 |
|
language: |
|
- en |
|
tags: |
|
- Pytorch |
|
- gravity wave |
|
- Weather & Climate |
|
- Foundation model |
|
datasets: |
|
- Prithvi-WxC/Gravity_wave_Parameterization |
|
base_model: |
|
- Prithvi-WxC/prithvi.wxc.2300m.v1 |
|
--- |
|
This repository contains pretrained model for Gravity Wave Flux Parametrization downstream task. |
|
|
|
<img src="https://cdn-uploads.huggingface.co/production/uploads/6488f1d3e22a0081a561ec8f/lOFP_1dAVKCw90uLpj2vu.png" alt="Gravity Wave" width="1024"/> |
|
|
|
|
|
|
|
|
|
### Model |
|
The pretrained [Prithvi WxC](https://huggingface.co/Prithvi-WxC/prithvi.wxc.2300m.v1) parameter model is finetuned to predict momentum fluxes from |
|
the [Gravity Wave Parameterization dataset](https://huggingface.co/datasets/Prithvi-WxC/Gravity_wave_Parameterization). |
|
|
|
<b>Input:</b> 491 (3 + 4x122) channels. |
|
1. latitude (1) |
|
2. longitude (1) |
|
3. surface elevation (1) |
|
4. zonal winds \\(u\\) (122) |
|
5. meridional winds \\(v\\) (122) 6. |
|
6. temperature \\(T\\) (122) |
|
7. pressure \\(P\\) (122) |
|
|
|
<b>Output:</b> 366 (3x122) channels. |
|
1. potential temperature \\(\theta\\) (122) |
|
2. zonal flux of vertical momentum \\(u'\omega'\\) (122) |
|
3. meridional flux of vertical momentum \\(v'\omega'\\) (122) |
|
|
|
### Code |
|
Code for fine-tuning is available through [Github](https://github.com/NASA-IMPACT/gravity-wave-finetuning). |
|
|
|
### Results |
|
<img src="https://cdn-uploads.huggingface.co/production/uploads/6488f1d3e22a0081a561ec8f/Vk1EKgzf_j90ZPiw2hGHE.png" alt="Gravity Wave" width="1024"/> |
|
|
|
For the Andes (mountain waves) and the Southern Ocean (non-mountain waves), |
|
the fine-tuned model achieves correlation coefficients of 0.99 and 0.97, respectively, when compared to the observed fluxes. |
|
|
|
|
|
### Inference and demo |
|
The github repo includes an inference script that allows to run |
|
the [gravity_wave_model](https://huggingface.co/Prithvi-WxC/Gravity_wave_Parameterization/blob/main/magnet-flux-uvtp122-epoch-99-loss-0.1022.pt) model |
|
for inference on [sample dataset](https://huggingface.co/datasets/Prithvi-WxC/Gravity_wave_Parameterization/blob/main/wxc_input_u_v_t_p_output_theta_uw_vw_era5_training_data_hourly_2015_constant_mu_sigma_scaling05.nc). |
|
|
|
## Citation |
|
If you use this work, consider citing our paper |
|
|
|
|
|
``` |
|
@misc{schmude2024prithviwxcfoundationmodel, |
|
title={Prithvi WxC: Foundation Model for Weather and Climate}, |
|
author={Johannes Schmude and Sujit Roy and Will Trojak and Johannes Jakubik and Daniel Salles Civitarese and Shraddha Singh and Julian Kuehnert and Kumar Ankur and Aman Gupta and Christopher E Phillips and Romeo Kienzler and Daniela Szwarcman and Vishal Gaur and Rajat Shinde and Rohit Lal and Arlindo Da Silva and Jorge Luis Guevara Diaz and Anne Jones and Simon Pfreundschuh and Amy Lin and Aditi Sheshadri and Udaysankar Nair and Valentine Anantharaj and Hendrik Hamann and Campbell Watson and Manil Maskey and Tsengdar J Lee and Juan Bernabe Moreno and Rahul Ramachandran}, |
|
year={2024}, |
|
eprint={2409.13598}, |
|
archivePrefix={arXiv}, |
|
primaryClass={cs.LG}, |
|
url={https://arxiv.org/abs/2409.13598}, |
|
} |
|
``` |
|
|
|
``` |
|
@article{gupta2024machine, |
|
title={Machine learning global simulation of nonlocal gravity wave propagation}, |
|
author={Gupta, Aman and Sheshadri, Aditi and Roy, Sujit and Gaur, Vishal and Maskey, Manil and Ramachandran, Rahul}, |
|
journal={arXiv preprint arXiv:2406.14775}, |
|
year={2024} |
|
} |
|
``` |
