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