Prithvi-WxC/Gravity_wave_Parameterization

This repository contains pretrained model for Gravity Wave Flux Parametrization downstream task.

Model

The pretrained Prithvi WxC parameter model is finetuned to predict momentum fluxes from the Gravity Wave Parameterization dataset.

Input: 491 (3 + 4x122) channels.

1. latitude (1)
2. longitude (1)
3. surface elevation (1)
4. zonal winds $uu$ (122)
5. meridional winds $vv$ (122) 6.
6. temperature $TT$ (122)
7. pressure $PP$ (122)

Output: 366 (3x122) channels.

1. potential temperature $\theta \backslash theta$ (122)
2. zonal flux of vertical momentum $u^{\mathrm{\prime }}{\omega }^{\mathrm{\prime }}u\text{'}\backslash omega\text{'}$ (122)
3. meridional flux of vertical momentum $v^{\mathrm{\prime }}{\omega }^{\mathrm{\prime }}v\text{'}\backslash omega\text{'}$ (122)

Code

Code for fine-tuning is available through Github.

Results

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 model for inference on sample dataset.

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