PyTorch-powered data assimilation slashes air quality model startup times

Air quality forecasts live or die by their starting point. A numerical model that begins a run with biased initial concentrations of pollutants will carry that error into every subsequent hour of prediction, and the error typically grows fastest in the earliest hours of the simulation. Data assimilation is the statistical machinery that corrects this: it blends real-world observations from surface stations and satellites with a model’s background fields to produce an improved analysis, which then serves as the starting state for the forecast. For atmospheric chemistry, however, the traditional tools for doing this are showing their age. A team of Chinese researchers now reports a rebuilt assimilation system, written in the PyTorch deep learning framework, that runs more than twenty-five times faster than the established Fortran-based equivalent on graphics hardware while producing essentially identical results.

The new system, called PY3DVar, was developed by Hancheng Ye, Zengliang Zang, Wei You, Yiwen Hu, Ning Liu and Yi Li at the College of Meteorology and Oceanography of the National University of Defense Technology in Changsha, and is described in a development and technical paper published in Geoscientific Model Development on 28 October 2026. It performs three-dimensional variational data assimilation, or 3DVar, for the Weather Research and Forecasting model coupled to Chemistry, widely known as WRF-Chem. The work targets a persistent bottleneck: as observation volumes, model resolutions and the number of simulated chemical species have all grown, the state vectors and matrices inside atmospheric chemistry assimilation have expanded dramatically, pushing traditional Fortran implementations toward the limits of operational timeliness.

The core of 3DVar is a cost function that weighs the departure of a candidate analysis from the model background against its departure from the observations, each scaled by its respective error covariance. Minimizing that function yields the statistically optimal estimate of the atmospheric state. Doing so requires the gradient of the cost function with respect to every control variable, and in classical systems this gradient is derived analytically and coded by hand as an adjoint routine. PY3DVar replaces that manual derivation with automatic differentiation, PyTorch’s autograd engine, which records every tensor operation that depends on the control variable during the forward computation of the cost and then back-propagates along the recorded graph to produce the gradient without any hand-written adjoint code.

The system is far from a black-box neural network. The authors emphasize that PY3DVar inherits the complete classical variational optimization theory without modifying any underlying physical equations, retaining full mathematical traceability. The background error covariance matrix is factorized into a diagonal matrix of background error standard deviations and a correlation matrix, which is itself decomposed via Cholesky factorization and split into tensor products of longitudinal, latitudinal and vertical correlation matrices. Only surface in-situ observations are assimilated, so the observation operator reduces to horizontal bilinear interpolation from the model surface level to the observation locations. Both the new and the old systems use the same L-BFGS minimization algorithm with identical parameter settings: a memory size of six, a maximum of four inner iterations per optimization step, and convergence declared when the scaled gradient norm falls below a threshold of ten to the minus two.

The configuration targets six conventional air pollutants simultaneously: the gases sulfur dioxide, nitrogen dioxide, carbon monoxide and ozone, plus the particulate matter categories PM2.5 and PM10. Because WRF-Chem’s aerosol scheme, CMAQ-AERO5, tracks dozens of individual species, the researchers merged components by physicochemical properties and particle size, consolidating 27 aerosol species into 12 control variables and adding the four gases directly, for 16 control variables in total. Background error statistics were estimated with the National Meteorological Center method, using differences between pairs of 24-hour and 48-hour forecasts initiated daily at 00:00 UTC over 30 consecutive days. The experiments used a single domain covering most of China at 27-kilometer horizontal resolution with 168 by 210 grid points and 34 vertical layers, driven by NCEP final analysis meteorology, a 2013 multi-resolution emission inventory for China and WACCM chemical boundary conditions.

Testing proceeded in three stages. First, an idealized twin-observation experiment placed one station with extremely low prescribed concentrations and another with extremely high ones far apart on the 27-kilometer grid, verifying that both systems generate physically reasonable, Gaussian-shaped increment fields centered on the observation points, decaying with distance and weakening with height in ways governed by the estimated background error correlations. Second, twenty consecutive real-case assimilation cycles were run at 00:00 UTC each day from 10 to 29 October 2023, each followed by a 24-hour forecast, with observations from the China National Environmental Monitoring Centre split daily into assimilation and independent validation sets. Third, both systems were benchmarked across 1, 2, 4, 8, 16, 32 and 64 CPU cores on a dual AMD EPYC 7542 node, while PY3DVar was additionally deployed on a single NVIDIA RTX 3090 graphics card with 24 gigabytes of memory, all calculations in 32-bit floating-point precision.

The headline result is stark. On the GPU, PY3DVar completed an average assimilation cycle in 1.2 seconds, against 32.7 seconds for the Fortran system running on 64 CPU cores, a 26.2-fold speedup. Two factors combine to produce this. The mean per-iteration time fell from 0.94 seconds to 0.13 seconds, a 7.5-fold gain attributed chiefly to GPU-accelerated vectorized tensor operations. Meanwhile, the average number of iterations needed for convergence dropped from 34.6 to 9.9, which the authors attribute to the different line-search strategies inside the otherwise identical L-BFGS optimizers: PyTorch’s built-in Strong Wolfe line search achieves a larger reduction of the cost function per iteration than the classical Moré-Thuente scheme used in the Fortran code. The cost function descent curves of the two systems converge to closely matching final values.

The picture on CPU hardware is more nuanced and, the authors argue, more honest than a single headline number. Per iteration, the compiled Fortran implementation remains faster than the PyTorch one at every tested core count, with FT3DVar improving only modestly from 5.47 to 0.92 seconds as cores are added because its single-core baseline is already fast. PY3DVar scales more steeply, from 42.81 to 2.04 seconds per iteration, yet stays slower per iteration throughout. Nevertheless, because it needs far fewer iterations, PY3DVar achieves shorter total assimilation times at 8 cores, 73.9 versus 85.7 seconds; at 32 cores, 42.4 versus 51.1 seconds; and at 64 cores, 20.2 versus 32.7 seconds, a relative speedup of 1.61. At 1, 2, 4 and 16 cores the Fortran system retains the total-time advantage.

Accuracy, the other half of the equation, held up. In the real-case experiments, assimilation raised the correlation coefficient between analysis and observations to 0.58 for sulfur dioxide and 0.48 for carbon monoxide, above 0.8 for nitrogen dioxide, ozone and PM2.5, and to 0.79 for PM10, while reducing root mean square and mean absolute errors for all six pollutants relative to an unassimilated control run. The scatter distributions and evaluation metrics from PY3DVar and FT3DVar were essentially identical across every pollutant, and the 24-hour forecasts launched from their respective analysis fields showed fundamentally consistent metric curves. The assimilation’s positive effect on forecasts persisted for more than 24 hours for sulfur dioxide, carbon monoxide, PM2.5 and PM10, whose longer atmospheric lifetimes make them sensitive to large-scale transport, but only around 6 hours for nitrogen dioxide and 8 hours for ozone, whose daytime concentrations are dominated by rapid local photochemistry.

The authors position PY3DVar between two extremes: traditional Fortran systems with rigorous physical foundations but laborious manual coding and limited throughput, and fully data-driven neural assimilation methods that compute quickly but lack physical interpretability and have mostly been validated only on idealized low-dimensional problems. They note the closely related TorchDA package, which implements 3DVar and other algorithms in PyTorch but was demonstrated primarily on low-dimensional Lorenz models, leaving full-scale three-dimensional atmospheric chemistry assimilation untested before this work. With its source code, observation data and background error covariance modeling code released on Zenodo, PY3DVar offers operational air quality centers a path to faster analysis cycles on commodity graphics hardware, though its demonstrated gains rest on a single 27-kilometer domain over China, surface observations only, and a 20-day October 2023 test period, so broader operational adoption will depend on how the system holds up under wider conditions.

Subject of Research: A PyTorch-based, GPU-accelerated three-dimensional variational data assimilation system for the WRF-Chem atmospheric chemistry model

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

Article References: Ye, H., Zang, Z., You, W., Hu, Y., Liu, N., & Li, Y. (2026). GPU-accelerated three-dimensional variational data assimilation with automatic differentiation for WRF-Chem. Geoscientific Model Development, 19(20), 10263-10282. https://doi.org/10.5194/gmd-19-10263-2026

Image Credits: AI Generated

DOI: 10.5194/gmd-19-10263-2026

Keywords: data assimilation, WRF-Chem, 3DVar, PyTorch, GPU acceleration, automatic differentiation, air quality forecasting, atmospheric chemistry, background error covariance, L-BFGS, computational efficiency, Geoscientific Model Development