README.md

metadata

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.

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

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