---
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}
}
```