WRF-JAX v1.1.28: a differentiable regional weather model with fully coupled physics, numerically consistent with WRF and matching 128-process CPU performance on a single GPU
Abstract
Porting a regional weather model in its entirety to a graphics processing unit (GPU) requires complete implementation of the dynamical core and all physical processes, and improved computational efficiency. This paper introduces WRF-JAX v1.1.28, a JAX reimplementation of the WRF 4.5 dynamical core and the complete CONUS physics suite. We conducted 27 ten-day integrations across nine regions and three initialization months, with a horizontal grid spacing of 3 km and 700×700 grid points, and established two numerical-difference baselines. The full-field RMSEs of wind, temperature, and moisture in WRF-JAX and the WRF internal numerical-difference baseline showed similar long-term evolution: the medians of the case-wise RMSE ratios were 1.05-1.07 at 24 h and 1.02-1.04 at 240 h. The mean sustained integration time on a single NVIDIA A800 was 175.0 s per model hour, corresponding to speedups of 1.17 and 1.05 relative to 128-process single-precision WRF runs compiled with gfortran and ifort, respectively. The experiments further backpropagated through 100 steps with full physics coupling, obtaining all three-dimensional gradient components of the initial UVTQ fields in a single calculation. The full-field Gaussian random-direction evaluation used 500 directions and three perturbation amplitudes per variable; the median relative differences in directional derivatives for zonal wind, meridional wind, dry potential temperature, and water vapour were 2.58-5.22%, 2.94-6.49%, 7.86-10.06%, and 3.51-3.91%, respectively. WRF-JAX combines fully coupled forecasting with differentiation across time steps, providing computable approximate gradients for regional weather simulation and end-to-end optimization, with accuracy comparable to that of the original WRF model and a forward integration speed equivalent to that of the original WRF model running on at least 128 cores.
Related
- DiffSpecMatch: Differentiable Waveform Inverse Design with Configurable Optimization Modules for Response-Spectrum Matching of Ground Motions
- DiffSWE2d: a differentiable Shallow Water Equations solver for end-to-end flood and tsunami modelling
- Suêtes: An end-to-end differentiable non-hydrostatic limited-area dynamical core
- A library for differentiable signal processing and machine learning on the sphere
- Differentiable land model reveals global environmental controls on latent ecological functions
BibTeX
@misc{yuan2026wrfjax,
title = {{WRF-JAX v1.1.28: a differentiable regional weather model with fully coupled physics, numerically consistent with WRF and matching 128-process CPU performance on a single GPU}},
author = {Shijin Yuan and Tianyu Shen and Bin Mu and Zhen Wei and Bo Qin and Guokun Dai},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.35095},
eprint = {2610.35095},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.35095},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - WRF-JAX v1.1.28: a differentiable regional weather model with fully coupled physics, numerically consistent with WRF and matching 128-process CPU performance on a single GPU
AU - Shijin Yuan
AU - Tianyu Shen
AU - Bin Mu
AU - Zhen Wei
AU - Bo Qin
AU - Guokun Dai
PY - 2026
DA - 2026/10/26
PB - arXiv
DO - 10.48550/arxiv.2610.35095
UR - https://arxiv.org/abs/2610.35095
AB - Porting a regional weather model in its entirety to a graphics processing unit (GPU) requires complete implementation of the dynamical core and all physical processes, and improved computational efficiency. This paper introduces WRF-JAX v1.1.28, a JAX reimplementation of the WRF 4.5 dynamical core and the complete CONUS physics suite. We conducted 27 ten-day integrations across nine regions and three initialization months, with a horizontal grid spacing of 3 km and 700×700 grid points, and established two numerical-difference baselines. The full-field RMSEs of wind, temperature, and moisture in WRF-JAX and the WRF internal numerical-difference baseline showed similar long-term evolution: the medians of the case-wise RMSE ratios were 1.05-1.07 at 24 h and 1.02-1.04 at 240 h. The mean sustained integration time on a single NVIDIA A800 was 175.0 s per model hour, corresponding to speedups of 1.17 and 1.05 relative to 128-process single-precision WRF runs compiled with gfortran and ifort, respectively. The experiments further backpropagated through 100 steps with full physics coupling, obtaining all three-dimensional gradient components of the initial UVTQ fields in a single calculation. The full-field Gaussian random-direction evaluation used 500 directions and three perturbation amplitudes per variable; the median relative differences in directional derivatives for zonal wind, meridional wind, dry potential temperature, and water vapour were 2.58-5.22%, 2.94-6.49%, 7.86-10.06%, and 3.51-3.91%, respectively. WRF-JAX combines fully coupled forecasting with differentiation across time steps, providing computable approximate gradients for regional weather simulation and end-to-end optimization, with accuracy comparable to that of the original WRF model and a forward integration speed equivalent to that of the original WRF model running on at least 128 cores.
ER -