yingkaisha / ViT_Diffusion_GEFS

Improving ensemble extreme precipitation forecasts using generative artificial intelligence

Yingkai Sha, Ryan A. Sobash, David John Gagne II

NSF National Center for Atmospheric Research, Boulder, Colorado, USA

Abstract

An ensemble post-processing method is developed to improve the probabilistic forecasts of extreme precipitation events across the conterminous United States (CONUS). The method combines a 3-D Vision Transformer (ViT) for bias correction with a Latent Diffusion Model (LDM), a generative Artificial Intelligence (AI) method, to post-process 6-hourly precipitation ensemble forecasts and produce an enlarged generative ensemble that contains spatiotemporally consistent precipitation trajectories. These trajectories are expected to improve the characterization of extreme precipitation events and offer skillful multi-day accumulated and 6-hourly forecast guidance. The method is tested using the Global Ensemble Forecast System (GEFS) precipitation forecasts out to day 6 and verified against the Climate-Calibrated Precipitation Analysis (CCPA) data. Verification results indicate that the method generated skillful ensemble members with improved Continuous Ranked Probabilistic Skill Scores (CRPSSs) and Brier Skill Scores (BSSs). It outperformed the operational GEFS and a multivariate statistical post-processing baseline, showing reliable calibrations for extreme precipitation events that are also distinguishable from climatological forecasts. Explainability studies were further conducted, which revealed the decision-making process of the method and confirmed its effectiveness on ensemble member generation. This work introduces a novel, generative-AI-based approach to address the limitation of small numerical ensembles and the need for larger ensembles for accessing extreme precipitation events.

Data

• The forecast to post-process:
  • GEFS version 12 6 hourly accumulated total precipitation (APCP), upto day 6
• Training and verification target:
  • Climatology-Calibrated Precipitation Analysis (CCPA) 6 hourly quantitative precipitation estimation.

Method

• Vector Quantisation Variational Autoencoder (VQ-VAE)
• 3-D Vision Transformer (ViT)
• Latent Diffusion Model (LDM)

Navigation

• Model builder: model_utiles.py
• Model summary: MODEL03_all_model_summary.ipynb
• Model inference pipeline: MODEL03_model_full_pipeline.ipynb
• Model training: VQ-VAE, VQ-VAE output section, 3-D ViT, LDM
• Explainability study: RESULT04_Latent_space_verif.ipynb, PLOT_Latent_space_vis.ipynb