Prithvi.WxC is a 2.3 billion parameter foundation model developed as a collaboration between NASA and IBM Research. It is an encoder-decoder based transformer model. The pre-training task is a combination of forecasting and reconstruction from the randomly masked input. The model can take very sparse masked input, which enables the possibility to integrate observational data into the training/fine-tuning dataset for the future. 

The model takes 160 state variables with masking at two different timestamps within 12 hours from 40 years of NASA's MERRA-2 datasets and either reconstructs the full latter input state or predicts a future state within 24 hours. Besides the 160 dynamical input variables such as temperature and wind components, the model also takes static ones such as surface geopotential height, land, ocean and ice fractions, as well as the context information such as lead time and input time interval. Moreover, it adds climatology to the input so that the model learns the tendency normalised to the temporal mean and standard deviation.

Prithvi.WxC model has been fine-tuned for several downstream tasks, namely, autoregressive rollout forecasting, downscaling and gravity wave flux parametrisation, which addressed the following changes: 

  • temporal and/or spatial resolutions in dataset
  • input and output variables
  • domain coverage 

In the paper, the authors evaluate the models by the task-specific metrics and/or case study. For forecasting, they show the competitive zero-shot forecasting performance comparing to IFS HRES, IFS ENS, GraphCast and Pangu-Weather, together with the case study of 75 hurricane tracks of historical Category 3-5 hurricanes over the Atlantic Ocean from 2017 to 2023.  For downscaling, they present two fine-tuned models: a global model trained on 6x MERRA-2; and a regional model trained on 12x CORDEX-CMIP5-RCP8.5. The evaluation shows their capability to recover both the spatial fine features and the power spectrum of the ground truth.  For parametrisation, the fine-tuned model demonstrates its ability to predict the momentum fluxes in the upper troposphere on a coarser grid.

We have prepared the environment on gadi for this foundation model and the three fine-tuned models. Four released checkpoints have been downloaded from Hugging Face ( ibm-granite and Prithvi-WxC ) and saved to our notebook directory, together with a 3-month of MERRA-2 dataset and 1-month derived ERA5 dataset so that the following five notebooks, adapted from the developers' notebooks and scripts(https://github.com/NASA-IMPACT/Prithvi-WxC/tree/main/examples, https://github.com/IBM/granite-wxc/tree/main/examples/granitewxc_downscaling/notebooks, https://github.com/NASA-IMPACT/gravity-wave-finetuning/blob/main/inference.py and https://github.com/NASA-IMPACT/gravity-wave-finetuning/blob/main/finetune_gravity_wave.py), can run directly on Gadi. 

  • inference notebook for the reconstruction task
  • inference notebook for the autoregressive rollout forecasting task
  • inference notebook for the global downscaling task
  • inference notebook for the gravity wave parametrisation task
  • fine-tuning notebook to learn the new encoder and decoder for the gravity wave parametrisation task based on the core foundation model 

Use the Environment

$ module use /g/data/dk92/apps/Modules/modulefiles/
$ module load prithvi-wxc/2024.09
$ python3
 
>>> import PrithviWxC
>>> PrithviWxC.__version__
'0.1.0'
>>> import granitewxc
>>> granitewxc.__version__
'0.1.0'

Run the tested notebook

  1. Go to ARE site: are.nci.org.au
  2. Fill out the JupyterLab request form
    1. Walltime (hours): 1
    2. Queue: normalsr [or dgxa100, see note below]
    3. Compute Size: X-large (26 cpus, 124 G mem ) [or 1gpu, see note below]
    4. Project:  <xy01>
    5. Storage: gdata/dk92+scratch/<xy01>
  3. Click Advanced options and fill in the following fields
    1. Module directories: /g/data/dk92/apps/Modules/modulefiles/
    2. Modules:  prithvi-wxc/2024.09
    3. Jobfs size: 10GB  
  4. Launch the session to run the tested notebook

Note on the tested notebook

  • copy the tested notebook from either/all of the following paths to your own working directory. If your working directory is different from "/scratch/<xy01>", remember to change the storage directive in the JupyterLab request form accordingly.
    • /g/data/dk92/notebooks/examples-aiml/prithvi-WxC/inference.ipynb
    • /g/data/dk92/notebooks/examples-aiml/prithvi-WxC/autoreg.prediction.ipynb
    • /g/data/dk92/notebooks/examples-aiml/prithvi-WxC/downscaling.ipynb
    • /g/data/dk92/notebooks/examples-aiml/prithvi-WxC/parametrization.ipynb
    • /g/data/dk92/notebooks/examples-aiml/prithvi-WxC/finetuning_param.ipynb
  • to get a similar model used in the parametrization.ipynb  inference notebook, one need to create more samples to be used in the `finetuning_param.ipynb` notebook. The 1-month worth of data stored in dk92 is for demonstration purposes only. 
  • if to run  autoreg.prediction.ipynb or finetuning_param.ipynb, it is recommended to request a A100 GPU (job submitted to the dgxa100  queue with compute size of 1gpu) to run the compute intensive cells faster. 


  • No labels