ClimODE is a physics-informed neural ordinary differential equation based model that implements advection to precisely describe weather evolution with mass-conserving dynamics. It predicts the evolution of five ERA5 state variables, namely 10m horizontal winds, temperature at 850hPa, 2m temperature and geopotential at 500hPa,  in 48 hours at 6-hour time interval with their standard deviation, based on the initial state. 

In the paper, the authors trained both global and regional models, using not only the state variables from ERA5 but also the flow velocity derived from the corresponding state as the initial condition. This strategy together with the spatial-temporal embedding channels contributes to the stability and accuracy of the model. They also implement a probabilistic emission model so that ClimODE quantifies prediction uncertainties for each prediction it makes.

Their "approach outperforms existing data-driven methods in global and regional forecasting with an order of magnitude smaller parameterisation, establishing a new state of the art."

At NCI, we made the dk92 environment module NCI-ai-ml compatible to run this model, and prepared a notebook to show the training and the inference at global scale using NCI-Weatherbench. We used the released code to train the model with some variations in dataset construction.

  1. strategy f365: removing the last day in the leap years instead of 29 Feb so that there is no discontinuity in all the samples (each has 8 time steps in 2 days). 

  2. strategy ct: coding the timestamp in every 48 hours as 2-9 constantly instead of using their actual sequence in a year rolling in the range between 2 and 1459.

Use the Environment

$ module use /g/data/dk92/apps/Modules/modulefiles/
$ module load NCI-ai-ml/24.08
$ python3
 
>>> import torchdiffeq
>>> torchdiffeq.__version__
'0.2.4'
>>> import torch
>>> torch.__version__
'2.3.1'

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: analysis
    3. Compute Size: 1gpu
    4. Project:  <xy01>
    5. Storage: gdata/dk92+gdata/wb00+scratch/<xy01>
  3. Click Advanced options and fill in the following fields
    1. Module directories: /g/data/dk92/apps/Modules/modulefiles/
    2. Modules: NCI-ai-ml/24.08
    3. Jobfs size: 10GB  
  4. Launch the session to run the tested notebook

 

Note on the tested notebook

  • copy the tested notebook from "/g/data/dk92/notebooks/examples-aiml/climODE/climODE_global.ipynb" 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.
  • the notebook climODE_global.ipynb contains an example to
    1. train the model using the NCI-WeatherBench 5.625° dataset
    2. run inference using the checkpoint trained on Gadi 


Further Improvement Possibilities

At NCI, we tried two easy approaches to improve the model performance further, see details below for the comparison in three scalar metrics (ACC, CRPS, and RMSE) for all the five variables.

To try models trained using different strategies, currently we have three available, namely, original, ct vs ct+f365, modify the path to the corresponding checkpoint in your notebook.

  • No labels