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 dynamical core of the Weather Research and Forecasting (WRF) model version 4.5 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 zonal wind, meridional wind, dry potential temperature perturbation, and water vapour mixing ratio (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.