Rapid simulations of atmospheric data assimilation of hourly-scale phenomena with modern neural networks
Yiyuan Li,Xiting Ju,Yi Xiao,Qilong Jia,Yongxiao Zhou,Simeng Qian,Rongfen Lin,Bin Yang,Shupeng Shi,Xin Liu,Jie Gao,Zhen Wang,Sha Liu,Jian Tan,Xuan Wang,Zhengding Hu,Limin Yan,Wei Xue
DOI: https://doi.org/10.1145/3581784.3607031
2023-01-01
Abstract:Atmospheric data assimilation is essential for numerical weather prediction. Ensemble data assimilation connects multiple instances of an atmospheric model through a Kalman filter-based algorithm, which is regarded as a challenging computing task today. In this work, we build a fast, low-cost, and scalable atmospheric data assimilation prototype, DIDA, for the new-generation Sunway supercomputer, including: (1) a framework that enables flexible deployment of components, and manages and optimizes data communication among modules, achieving maximum resource efficiency; (2) an accurate, robust, UNet-based surrogate model for atmospheric dynamic simulation to generate the background ensemble; (3) a batch-LETKF algorithm with high-performance eigenvalue decomposition, which is up to 7.37 times faster than existing numerical libraries while exhibiting almost linear scalability. Experimental evaluations show that our AI-integrated ensemble data assimilation prototype can complete hour-cycle assimilation in minutes, maintain linear scalability, and save an order of magnitude of computing resources, compared with the traditional method.