GPU-accelerated three-dimensional variational data assimilation with automatic differentiation for WRF-Chem

This study developed a PyTorch-based three-dimensional variational (3DVar) data assimilation system (PY3DVar) for the Weather Research and Forecasting model coupled to Chemistry (WRF-Chem), which integrates automatic differentiation (AD) to replace traditional manual gradient derivation and adopts Graphics Processing Unit (GPU) acceleration for high computational efficiency. PY3DVar enables the simultaneous assimilation of gaseous pollutants (SO ₂ , NO ₂ , CO, O ₃ ) and particulate matter (PM _2.5 , PM ₁₀ ) and supports flexible deployment on both Central Processing Unit (CPU) and GPU computing platforms. Idealized tests were first conducted to verify the physical consistency of the assimilation system. Twenty consecutive real-case assimilation experiments, each followed by a 24 h forecast, were then conducted to assess computational performance and efficiency, benchmarked against the classic Fortran-based 3DVar system (FT3DVar) using CPU parallelization. The results of the idealized tests show PY3DVar effectively propagates observation information, generating increment fields consistent with FT3DVar. Long-term real-case experiments further verify that the two assimilation systems achieve nearly identical analysis accuracy for chemical initial fields. On the GPU platform, PY3DVar achieves an average assimilation time of 1.2 s, representing a 26.2-fold speedup compared to FT3DVar on 64 CPU cores (32.7 s). Parallel scalability tests across 1, 2, 4, 8, 16, 32 and 64 CPU cores reveal that FT3DVar runs faster per iteration across all core counts. Consequently, FT3DVar achieves shorter total assimilation time on 1, 2, 4 and 16 cores, whereas PY3DVar outperforms it at 8, 32 and 64 cores, owing to its intrinsically fewer required iterations for convergence. The average 24 h forecast results show that the improved initial fields have sustained positive effects on short-term forecasts: the improvements persist for over 24 h for SO ₂ , CO, PM _2.5 , and PM ₁₀ , and for over 6 h for NO ₂ and O ₃ . These results demonstrate that PY3DVar delivers substantially faster assimilation than FT3DVar while maintaining equivalent analysis accuracy, offering an efficient and practical solution for operational atmospheric chemical data assimilation.