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

Abstract. 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 Deep Computing Unit (DCU) acceleration for high computational efficiency. Py3DVAR enables the simultaneous assimilation of gaseous pollutants (SO₂, NO₂, CO, O₃) and particulate matter (PM₂.₅, PM₁₀) and supports flexible deployment on both Central Processing Unit (CPU) and DCU computing platforms. To evaluate its performance and efficiency, idealized and real-case assimilation experiments (27 km and 9 km grid resolutions) were conducted, compared against a traditional CPU-parallelized Fortran-based 3DVAR system (Fortran-3DVAR). Idealized results show Py3DVAR effectively propagates observation information, generating increment fields consistent with Fortran-3DVAR. In real-case experiments, Py3DVAR substantially improves the model initial field quality: at 27 km resolution, correlation coefficients (CORR) for SO₂, NO₂, CO, O₃, PM₂.₅, and PM₁₀ increased by 0.77, 0.51, 0.71, 0.98, 0.60, and 0.69, respectively; corresponding improvements at 9 km resolution are 0.78, 0.98, 0.66, 0.96, 0.63, and 0.78. The root mean square error (RMSE) and mean absolute error (MAE) are also significantly reduced, with analysis field accuracy comparable to Fortran-3DVAR. In terms of computational efficiency, Py3DVAR shows remarkable advantages: on the same CPU platform, the total iteration time at 27 km resolution is only 7.1 s, approximately 8.8 times faster than Fortran-3DVAR (62.5 s); on the DCU platform, the speedup reaches 32.7 times at 27 km and 40.3 times at 9 km. A 24-hour forecast test shows that the improved initial fields have sustained positive effects on short-term forecasts: the improvements persist for over 24 hours for SO₂, CO, PM₂.₅, and PM₁₀, and for over 6 hours for NO₂ and O₃. This study confirms that Py3DVAR achieves order-of-magnitude gains in computational efficiency while maintaining accuracy equivalent to traditional assimilation algorithms, providing a flexible new technical pathway for operational atmospheric chemical data assimilation and future intelligent assimilation systems.