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
Authors: Shijin Yuan, Tianyu Shen, Bin Mu, Zhen Wei, Bo Qin, Guokun Dai
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\times700$ 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.
The papers around it, by date and closeness
Higher is closer to this paper; left is older. Lines join papers close to each other: groups of lines are sub-topics. Select a point to see why it is there. Back to the summary of this paper
Papers around this one
Ranked by how close their title and abstract are to this paper’s. Select some to read them together.
Closest papers
-
reGRAF: a global MPAS reforecast data set with convection-allowing refinements over the US and EuropeT. Hamill, Akshay Subramaniam, Montgomery L. Flora et al. · 2026 · Not summarized yet
reGRAF provides GPU-generated global MPAS reforecasts with convection-allowing regional mesh refinement, whereas WRF-JAX reimplements a regional WRF model.
Abstract
NVIDIA and The Weather Company (TWC) have generated a data set of reforecasts from TWC's GRAF (Global high-Resolution Atmospheric Forecasting) model, a version of the National Center for Atmospheric Research (NCAR) Model for Predictions Across Scales (\href{https://ncar.ucar.edu/what-we-offer/models/model-prediction-across-scales-mpas}{\underline{MPAS}}). GRAF is global, but the configuration for this reforecast had a mesh refinement to \textasciitilde4 km over the US, Caribbean Basin, and Europe, and 15 km elsewhere. This model was designed to run much of the computation on graphical processing units (GPUs), with this development assisted by NVIDIA. The 1836 reforecast cases (\textasciitild…
-
Toward GPU-Resident Climate Models: A Feasibility Study on Lossy Compression for the Spherical Harmonic Transform's Communication BottleneckL. Breschi, Flavio Vella · 2026 · Not summarized yet
This study tests lossy compression to reduce GPU-resident climate models’ spherical-harmonic communication bottleneck, rather than reimplementing WRF dynamics and physics.
Abstract
Operational pseudospectral atmospheric models such as the ECMWF Integrated Forecasting System (IFS) run today almost exclusively on CPUs; GPU ports are under active development but not yet used in production. These models rely on the Spherical Harmonic Transform (SHT). Each time-step requires forward and inverse SHTs, and both passes depend on global pencil transposition that redistribute multi-dimensional arrays across compute nodes. At large node counts these global collectives dominate wall-clock time. We investigate GPU-resident lossy compression, using representative fields from the DYAMOND high-resolution operational dataset as input, and combining measured GPU compression throughput w…
-
GPU-accelerated three-dimensional variational data assimilation with automatic differentiation for WRF-ChemHan-Cheng Ye, Zeng-Liang Zang, Wei You et al. · 2026 · Geoscientific Model Development · Not summarized yet
PY3DVar uses GPU-accelerated automatic differentiation for WRF-Chem data assimilation, while WRF-JAX differentiates through coupled regional forecasts.
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 Graphics Processing Unit (GPU) acceleration for high computational efficiency. PY3DVar enables the simultaneous assimilation of gaseous pollutants (SO 2 , NO 2 , CO, O 3 ) and particulate matter (PM 2.5 , PM 10 ) 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 assimilat…
-
A Description of the Advanced Research WRF Model Version 4C. Skamarock, B. Klemp, J. Dudhia et al. · 2019 · cited 1,712 times · Not summarized yet
The WRF version 4 description documents the model whose dynamics and physics WRF-JAX reimplements.
-
THE VERY-SHORT-RANGE NUMERICAL WEATHER AND THREATS OF HAZARD FORECAST BASED ON CONVECTIVE-RESOLVING MODELING WITH CONTINUOUS DATA ASSIMILATIOND. V. Blinov, G. Rivin, A. Bundel et al. · 2026 · Proceedings of the YSU C: Geological and Geographical Sciences · Not summarized yet
This system produces short-range regional forecasts with continuous data assimilation, while WRF-JAX focuses on a differentiable GPU implementation.
Abstract
The regional atmospheric modeling systems, which refine the results of global models, are currently actively used as a means of numerical weather forecasting for the coming days and hours. Using non-hydrostatic approach and grid spacing of several kilometers or less, these systems not only take into account complex geographical features of the terrain in detail but can also realistically predict complex weather formation mechanisms, including those associated with convection. Thе presentation is devoted to the description and results analysis of a continuously operating system for regional data assimilation and hourly numerical forecast updates for the next 12 h. It is based on the convectiv…
-
Technical note: A physically based correction of WRF temperature using atmospheric lapse rates for high-resolution hydrological simulations in a glacierized catchmentF. Hofmeister, Xin-Yang Fan, M. Pfeiffer et al. · 2026 · Hydrology and Earth System Sciences · Not summarized yet
This study uses WRF for high-resolution regional downscaling, applying its output to glacierized-catchment hydrology rather than differentiable GPU forecasting.
Abstract
This technical note presents a physically based post-processing workflow for correcting coarse-resolution regional climate model output used to force high-resolution hydrological models in high-elevation, glacierized catchments. The Weather Research and Forecasting (WRF) model was used to dynamically downscale the 20th Century Reanalysis Version 3 (20CRv3; 1° × 1°, 3-hourly) to a spatial resolution of 2 km and an hourly temporal resolution for the period 1969–2015. Evaluation of WRF output in the upper Kauner Valley (Austria) revealed a pronounced cold bias in the diagnosed 2 m temperature at elevations above 2600 m above sea level, which was attributed to the coarse representation of snow- …
-
Optimizer-Dependent Trade-Offs in Variational Multiple-Doppler Wind Retrieval: A Comparison of Adam, L-BFGS, and Newton-CGSeong-Jun Hwang, Ji-Sun Lee, Dong-In Lee et al. · 2026 · Atmosphere · Not summarized yet
This study applies GPU automatic differentiation to variational Doppler-wind retrieval, rather than differentiating through WRF’s coupled forecast time steps.
Abstract
The Wind Synthesis System using Doppler Measurements (WISSDOM) retrieves three-dimensional winds by minimizing a variational objective that combines Doppler-radar observations and physical constraints. Using a common Python implementation with PyTorch (2.9.1+cu130) automatic differentiation on an NVIDIA GPU, this study compared Adam, limited-memory Broyden–Fletcher–Goldfarb–Shanno (L-BFGS), and trust-region Newton conjugate-gradient (Newton-CG) for the same Typhoon Bavi case, initial wind field, and objective. At solver-specific stopping points, all three methods substantially reduced the objective and produced broadly similar principal horizontal- and vertical-wind structures. L-BFGS and Ne…
-
SD3: The First Operationally-Proven, Fully Heterogeneous-Accelerated Global Kilometer-Scale Model with Inline Physics ArchitectureYuan Liang, X. Hu, Xi Chen et al. · 2026 · Not summarized yet
SD3 builds a GPU-accelerated weather model with coupled physics, applying that approach to global kilometer-scale forecasting rather than regional WRF.
Abstract
We present Super Dynamics on the Cube (SD3), the first operationally-proven, fully heterogeneous-accelerated global kilometer-scale model. Through a full-stack reconstruction in native CUDA C, SD3 introduces an"Inline Physics"architecture that deeply couples its non-hydrostatic cubed-sphere dynamical core with fast physical processes. This design maintains high physical consistency, ensures energy and mass conservation, and mitigates"Grey Zone"mismatches. Five years of continuous operational forecasting, together with validation against diverse observations, demonstrates the scientific fidelity of SD3. Recently, SD3 achieved an unprecedented 205.7 simulated days per day at global 1.5 km reso…
Prior work it builds on
-
Explicit Forecasts of Winter Precipitation Using an Improved Bulk Microphysics Scheme. Part II: Implementation of a New Snow ParameterizationG. Thompson, P. Field, R. Rasmussen et al. · 2008 · Monthly Weather Review · cited 2,935 times · Not summarized yet
This paper describes a snow parameterization relevant to the WRF physics suite that WRF-JAX implements.
-
Important factors in the tracking of tropical cyclones in operational modelsT. Marchok · 2021 · Journal of Applied Meteorology and Climatology · cited 39 times · Not summarized yet
This study evaluates tropical-cyclone tracking configurations in operational forecasts, a downstream forecasting topic distinct from WRF-JAX’s model implementation.
Abstract
Multiple configurations of the Geophysical Fluid Dynamics Laboratory vortex tracker are tested to determine a setup that produces the best representation of a model forecast tropical cyclone center fix for the purpose of providing track guidance with the highest degree of accuracy and availability. Details of the tracking algorithms are provided, including descriptions of both the Barnes analysis used for center-fixing most variables and a separate scheme used for center-fixing wind circulation. The tracker is tested by running multiple configurations on all storms from the 2015-2017 hurricane seasons in the Atlantic and eastern Pacific Basins using forecasts from two operational National We…
-
More accuracy with less precisionS. Lang, A. Dawson, M. Diamantakis et al. · 2021 · Quarterly Journal of the Royal Meteorological Society · cited 47 times · Not summarized yet
This paper evaluates single-precision forecasts in ECMWF’s IFS, providing context for WRF-JAX’s numerical-precision and performance comparisons.
Abstract
Reducing the numerical precision of the forecast model of the Integrated Forecasting System (IFS) of the European Centre for Medium‐Range Weather Forecasts (ECMWF) from double to single precision results in significant computational savings without negatively affecting forecast accuracy. The computational savings allow to increase the vertical resolution of the operational ensemble forecasts from 91 to 137 levels earlier than anticipated and before the next upgrade of ECMWF's high‐performance computing facility. This upgrade to 137 levels harmonises the vertical resolution of the medium‐range deterministic forecasts and the medium‐range and extended‐range ensemble forecasts. Increasing the v…
-
Improved Representation of Boundary Layer Clouds over the Southeast Pacific in ARW-WRF Using a Modified Tiedtke Cumulus Parameterization Scheme*Chunxi Zhang, Yu-Qing Wang, K. Hamilton · 2011 · Monthly Weather Review · cited 486 times · Not summarized yet
This study modifies WRF’s Tiedtke cumulus scheme, a model-physics component within the scope of WRF-JAX’s physics implementation.
-
A new ensemble-based consistency test for the Community Earth System Model (pyCECT v1.0)A. Baker, D. Hammerling, M. Levy et al. · 2015 · Geoscientific Model Development · cited 69 times · Not summarized yet
pyCECT provides an ensemble-based software-consistency test for CESM, a verification resource for climate models rather than a weather-model reimplementation.
Abstract
Abstract. Climate simulation codes, such as the Community Earth System Model (CESM), are especially complex and continually evolving. Their ongoing state of development requires frequent software verification in the form of quality assurance to both preserve the quality of the code and instill model confidence. To formalize and simplify this previously subjective and computationally expensive aspect of the verification process, we have developed a new tool for evaluating climate consistency. Because an ensemble of simulations allows us to gauge the natural variability of the model's climate, our new tool uses an ensemble approach for consistency testing. In particular, an ensemble of CESM cl…
Go further
How are these papers found?
We compare the title and abstract of this paper with those of every paper summarized on this site, and with the papers Semantic Scholar recommends, those it cites and those citing it (for a much-cited paper, its most recent citations and, from OpenAlex, the most cited ones). Among the prior work and the follow-ups, the papers that use each other in depth, a method it uses and the much-cited follow-ups come a little higher. The levels (very close, close, related) come from the similarity of the texts, not from a judgment of quality; * marks a paper without an abstract, compared by its title. The sentence under a paper is written by an AI from the two abstracts: check it in the papers. A paper marked not relevant by three readers leaves the list.
Updated on 2 November 2026.