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

Shijin Yuan; Tianyu Shen; Bin Mu; Zhen Wei; Bo Qin; Guokun Dai

AI summary · arXiv: 2610.35095v1 · physics.ao-ph

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

Shijin Yuan, Tianyu Shen, Bin Mu, Zhen Wei, Bo Qin, Guokun Dai

Original paper Submitted to arXiv on 26 Oct. 2026 License: CC BY 4.0 Also in: fr

How to read this summary? Illustrated guide

The essentials

WRF-JAX reproduces WRF’s coupled regional simulations near its internal numerical-difference baseline, matches measured 128-process CPU throughput on one GPU, and provides approximate gradients through fully coupled physics.

Question

Can a GPU reimplementation of WRF preserve coupled dynamics and physics while achieving numerical consistency, practical integration speed, and differentiation across time steps?

Method

The authors reimplemented WRF 4.5 dynamics and the CONUS physics suite in JAX, using progressive component translation and numerical checks against identical-input Fortran calculations. They evaluated 27 ten-day integrations across nine regions and three initialization months on a 700 × 700 grid at 3 km spacing, comparing against WRF parallel-decomposition and perturbed-initial-condition baselines and measuring single-A800 throughput against CPU WRF. A separate 100-step, fully coupled backpropagation experiment checked initial wind, potential-temperature, and water-vapour gradients using 500 full-field directions per variable at three perturbation amplitudes.

Findings

Result Wind, temperature, and moisture differences approached WRF’s own numerical-difference baseline, with median case-wise RMSE ratios of 1.05–1.07 at 24 hours and 1.02–1.04 at 240 hours.
See the passage in the paper
The WRF internal numerical-difference baseline established for all cases shows that the wind, temperature, and moisture errors in WRF-JAX are of similar magnitude to the reference model’s own long-term numerical differences: the median case-wise RMSE ratios are 1.05–1.07 at 24 h and 1.02–1.04 at 240 h.
Result One A800 sustained 175.0 seconds per simulated hour, slightly outperforming the measured 128-process single-precision WRF configurations.
See the passage in the paper
The mean sustained integration time on a single A800 is 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.
Result The 100-step gradient evaluation produced median directional-derivative discrepancies of 2.58–10.06%, with matching signs in all 6000 checks.
See the passage in the paper
The evaluation of 500 full-field directions per variable at three amplitudes provides 6000 finite-difference comparisons: median relative differences across variables and amplitudes range from 2.58 to 10.06 %, and all derivative signs agree.
Authors’ interpretation The authors interpret the approximate gradients as potential update signals for initial-field inversion and trainable model corrections, rather than demonstrating completed optimization applications.
See the passage in the paper
The resulting approximate gradients can serve as update signals for initial-field inversion and trainable correction modules; through the chain rule for initial-field or parameter mappings, the final-state objective can be connected to upstream trainable variables, providing the differentiation interface needed for end-to-end optimization.

Limitations

The implementation covers the selected CONUS physics configuration rather than all WRF parameterization options.
See the passage in the paper
The scope of the WRF-JAX port comprises the WRF 4.5 ARW dynamical core and the CONUS physics suite; other, unselected parameterization schemes are outside the scope of this study.
Reported throughput excludes initialization, compilation, and output costs and therefore does not measure total application runtime.
See the passage in the paper
Performance is evaluated using sustained integration time, defined as the wall-clock seconds required for integration divided by the corresponding number of model hours, excluding initialization, compilation, and output costs.
Gradient agreement is limited by approximate derivative rules, non-smooth physics, and numerical errors in finite-difference references.
See the passage in the paper
Differences in gradient checks are therefore influenced jointly by approximate differentiation rules, model non-smoothness, and finite-difference numerical errors; they cannot all be attributed to errors in backpropagation itself.

Why it matters

The authors argue that WRF-JAX combines fully coupled regional forecasting, practical single-GPU execution, and approximate gradient computation within one framework. This provides a foundation for gradient-based initial-field adjustment, conditional nonlinear optimal perturbation problems, and future end-to-end training.

Written by AI from the paper. Passages are checked word for word against the text we analysed; the rest may contain errors.

AI-generated keywords: WRF-JAX; GPU acceleration; Automatic differentiation; Numerical consistency; Regional weather modelling

Read further

Comprehensive Summary
Layman's Summary
Blog article
Original abstract

This paper examines whether a regional weather model can combine WRF’s coupled dynamics and physical parameterizations with practical GPU performance and automatic differentiation while maintaining numerical consistency with the original implementation. The authors introduce WRF-JAX v1.1.28, a JAX reimplementation of the WRF 4.5 dynamical core and complete CONUS physics suite. Its scope encompasses Thompson microphysics, RRTMG radiation, surface exchange, Noah land-surface processes, Mellor–Yamada–Janjic boundary-layer mixing, and modified Tiedtke convection, but not every parameterization available in WRF. Development followed a progressive, verification-driven approach. Runtime coverage and call-dependency analysis identified relevant calculations and guided their translation. Temporary Fortran interfaces maintained full physics coupling during development and were removed once the corresponding components had been translated. Comparisons at matching computational stages assessed inputs, tendencies, persistent states, numerical precision, branching, and update timing. GPU acceleration relied on staged compilation, parallel processing of independent atmospheric columns, reusable executables, and controlled memory management. Reverse-mode automatic differentiation propagated sensitivities through the coupled forecast, with checkpointing and recomputation used to manage storage requirements. Numerical evaluation comprised 27 ten-day integrations across nine geographically diverse regions, initialized in March, July, and November 2021. Each integration used 3 km horizontal spacing, a 700-by-700 horizontal grid, 44 atmospheric mass levels, and an 18-second time step. Both implementations received identical initial and lateral boundary data. Crucially, the authors evaluated differences between WRF-JAX and a WRF reference against WRF’s own numerical variability. The principal baseline consisted of two 128-process WRF runs with different domain decompositions. Additional six-hour experiments used ten-member ensembles generated by perturbing initial potential temperature by one unit in the last place. Differences between WRF-JAX and WRF closely tracked these internal variability baselines. For three-dimensional wind, potential temperature, and water vapour, median case-wise root-mean-square error (RMSE) ratios relative to the domain-decomposition baseline were approximately 1.05–1.07 at 24 hours and 1.02–1.04 at 240 hours. At day ten, median WRF-JAX RMSEs were 1.035 and 0.980 metres per second for the two horizontal wind components, 0.396 K for potential temperature, and 0.457 grams per kilogram for water vapour. Near-surface, radiation, and precipitation comparisons generally supported similar numerical consistency, although accumulated explicit precipitation had a median RMSE ratio of 1.091. Six additional tropical-cyclone cases produced broadly consistent simulated tracks across the two implementations. On one NVIDIA A800 GPU, the mean sustained integration time was 175.0 seconds per simulated hour. This corresponded to speedups of 1.175 and 1.052 relative to 128-process single-precision WRF runs compiled with gfortran and ifort, respectively. These measurements excluded initialization, compilation, and input/output and therefore describe sustained integration performance rather than total workflow time. Differentiability was demonstrated over 100 steps, equivalent to 30 minutes, with all selected physics enabled. A single reverse-mode calculation produced gradients for the complete initial wind, potential-temperature, and water-vapour fields. Checks using 500 directions per variable and three perturbation amplitudes yielded median directional-derivative discrepancies of 2.58–5.22% for zonal wind, 2.94–6.49% for meridional wind, 7.86–10.06% for temperature, and 3.51–3.91% for moisture. All 6,000 comparisons agreed in sign. The authors conclude that WRF-JAX provides a practical foundation for gradient-based initial-condition adjustment and end-to-end optimization. Important limitations remain. Gradients are approximate because of thresholds, discrete lookups, custom derivative rules, and floating-point effects. Differentiation was validated only over a short window and for a single objective, while numerical agreement with WRF does not establish forecast skill against observations. Comparisons with other GPU implementations also involve differing configurations. Replacing conventional physics with smooth AI parameterizations and demonstrating effective optimization remain priorities for future work.

WRF-JAX: Bringing Regional Weather Modeling to GPUs and Automatic Differentiation

Modern weather models combine atmospheric dynamics with detailed representations of clouds, radiation, land surfaces, and turbulence. Making these systems run on new hardware is challenging. Making them differentiable—able to calculate how forecast outcomes respond to changes in their inputs—adds another layer of difficulty.

In “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,” Shijin Yuan, Tianyu Shen, Bin Mu, Zhen Wei, Bo Qin, and Guokun Dai investigate whether these goals can be achieved together. Their results offer a promising foundation, while clearly identifying limits in performance measurement, gradient accuracy, and validation.

What WRF-JAX Reimplements

WRF-JAX v1.1.28 translates the WRF 4.5 dynamical core and complete CONUS physics suite into JAX, a framework supporting accelerated computation and automatic differentiation. This is not a replacement for every parameterization available in WRF. Instead, it preserves a specific, fully coupled configuration.

The implemented physics includes Thompson microphysics, RRTMG radiation, surface exchange, Noah land-surface processes, Mellor–Yamada–Janjic boundary-layer mixing, and modified Tiedtke convection. Keeping these components coupled matters: each influences the evolving atmosphere, and a differentiable dynamical core alone would not capture the full chain of interactions.

A Verification-Driven Translation

The authors approached development progressively. Runtime coverage and call-dependency analysis identified calculations that needed translation. Temporary Fortran interfaces kept the physics connected during development, then were removed as the corresponding components became available in JAX.

Verification went beyond comparing final forecasts. The team checked inputs, tendencies, persistent states, numerical precision, branching, and update timing at matching computational stages. This helped distinguish implementation errors from differences arising through floating-point arithmetic and the atmosphere’s sensitivity to small perturbations.

Testing Agreement Against WRF’s Own Variability

The main numerical evaluation covered 27 ten-day integrations across nine geographically diverse regions, initialized in March, July, and November 2021. Each used a 700-by-700 horizontal grid, 3-kilometre spacing, 44 atmospheric mass levels, and an 18-second time step. Both implementations received identical initial and lateral boundary data.

Crucially, the benchmark was not exact agreement. WRF itself can produce different results when calculations are distributed differently across processors. The principal baseline therefore compared two 128-process WRF runs using different domain decompositions. Additional six-hour experiments examined ten-member ensembles with initial potential-temperature perturbations of one unit in the last place.

WRF-JAX’s differences from WRF closely tracked these internal variability baselines. For three-dimensional wind, potential temperature, and water vapour, median case-wise root-mean-square error ratios relative to the decomposition baseline were approximately 1.05–1.07 after 24 hours and 1.02–1.04 after ten days.

At day ten, median WRF-JAX errors relative to WRF were 1.035 and 0.980 metres per second for the horizontal wind components, 0.396 K for potential temperature, and 0.457 grams per kilogram for water vapour. Near-surface, radiation, and precipitation comparisons generally supported similar consistency, although accumulated explicit precipitation had a somewhat higher median error ratio of 1.091. Six additional tropical-cyclone cases also produced broadly consistent tracks.

GPU Performance: Competitive, With Important Boundaries

GPU acceleration combined staged compilation, parallel processing of independent atmospheric columns, reusable executables, and controlled memory management. On one NVIDIA A800 GPU, WRF-JAX averaged 175.0 seconds of integration time per simulated hour.

That represented speedups of 1.175 over 128-process single-precision WRF compiled with gfortran and 1.052 over the corresponding ifort configuration. The result is best understood as competitive sustained performance, rather than a dramatic acceleration over both CPU baselines.

These timings excluded initialization, compilation, and input/output. They therefore do not establish total workflow speed. Comparisons with other GPU implementations also require caution because configurations differ.

What Differentiability Adds

Automatic differentiation offers a way to calculate how a forecast objective responds to changes throughout the initial atmospheric state. This could support gradient-based initial-condition adjustment and end-to-end optimization.

The authors demonstrated reverse-mode differentiation through 100 model steps, equivalent to 30 minutes, with all selected physics enabled. A single calculation produced gradients for the complete initial wind, potential-temperature, and water-vapour fields. Checkpointing and recomputation helped manage storage requirements.

Checks using 500 directions per variable and three perturbation amplitudes yielded median directional-derivative discrepancies ranging from 2.58% to 10.06%, depending on the variable and amplitude. All 6,000 comparisons agreed in sign. This supports useful sensitivity information, but not exact gradients.

A Foundation, Not a Finished Optimization System

Thresholds, discrete lookups, custom derivative rules, and floating-point effects make the gradients approximate. Differentiation was validated only over a short window and for one objective. Likewise, numerical agreement with WRF does not demonstrate forecast skill against observations.

WRF-JAX’s contribution is bringing coupled regional weather physics, practical GPU execution, and automatic differentiation into one verified implementation. Demonstrating effective optimization—and exploring smoother AI-based physics parameterizations—remain important next steps.

Created on 02 Nov. 2026

Was this synthesis useful?

Score: 0

Ask questions about this paper to our AI assistant

How to chat with this paper? Illustrated guide

Compare with the PDF

You can also chat with multiple papers at once here.

AI assistant instructions?

Tools for this paper

Explain it differently 3 free rewrites left today

Save to a collection Group papers in collections to compare them, chat with them and export them. Create a free account Log in

Export BibTeX · RIS (Zotero, Mendeley, EndNote) · Obsidian (Markdown) · Notion (Markdown)

Advanced tools Premium

For researchers who follow a field closely: citation alerts, journal-club slides, version comparison and Zotero export. Discover Premium

Zotero

Citation alerts Citation alerts guide

Who works on it? The people and institutions who publish on its topic

Weekly letter Presented in "Climate and Earth this week" of 2 November 2026. Climate and Earth this week: Protection is not the same as resilience

Journal club Premium+ Slides to present it Propose for the journal club

Versions What changed between versions

Talk about it in public Premium+ Press kit: release, questions, hype check Press kit guide

Disclaimer: The AI-based summarization tool and virtual assistant provided on this website may not always provide accurate and complete summaries or responses. We encourage you to carefully review and evaluate the generated content to ensure its quality and relevance to your needs.

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

The essentials

Question

Method

Findings

Limitations

Why it matters

Read further

Definitions

WRF-JAX: Bringing Regional Weather Modeling to GPUs and Automatic Differentiation

What WRF-JAX Reimplements

A Verification-Driven Translation

Testing Agreement Against WRF’s Own Variability

GPU Performance: Competitive, With Important Boundaries

What Differentiability Adds

A Foundation, Not a Finished Optimization System

Ask questions about this paper to our AI assistant

Tools for this paper

Advanced tools Premium

Papers around this one

Closest papers

Prior work it builds on

Go further

Cite this paper (BibTeX)