CorrDiff is a stochastic downscaling model trained by the high-resolution 2-km weather model over Taiwan conditioned on 25-km ERA5 reanalysis. It uses 20 input variables from ERA5, namely tcwv, t2m, u10, v10, (t,z,u,v)x(500, 700, 850, 925) hPa to generate high resolution output for four variables, namely, maximum radar reflectivity, temperature at 2 meter, eastward wind at 10 meter and northward wind at 10 meter.

In the paper, the authors trained the model using a two stage approach: the first U-Net model learns to generate a high resolution mean field of the output variables and then the second diffusion model learns the variance between the generated mean field and the ground truth in the target dataset. This approach contributes to the skillful RMSE and CRPS and the faithful recovery of spectra and distributions even for extremes.

At NCI, we prepared the environment nci-corrdiff for this model and wrote an inference notebook using the checkpoints trained on NCI Gadi. The local training follows the released code and uses cropped BARRA2-C2, a high-resolution 4.4-km dataset over Australian region, as the target while uses the same condition of the 25-km ERA5 reanalysis. The NCI trained model generates four variables, namely tas, uas, vas, and maxcolrefl from the low resolution cropped ERA5 input in the same region. 

Use the Environment

$ module use /g/data/dk92/apps/Modules/modulefiles/
$ module load nci-corrdiff/24.07
$ python3
 
>>> import modulus
>>> modulus.__version__
'0.6.0a0'
>>> import torch
>>> torch.__version__
'2.2.1+cu118'

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: gpuvolta
    3. Compute Size: 1gpu
    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: nci-corrdiff/24.07
    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/corrdiff/inference/CorrDiff_inference.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 CorrDiff_inference.ipynb contains an example to run inference using the checkpoint trained on Gadi.


  • No labels