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

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