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

Home/Archive/arXiv

arXiv · 2610.35095

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

View original source ↗Read source document ↗Export metadata ↓

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.

Explore related subjects

physics.ao-ph

Keep this discovery

Explore connections, maps & timelines

BibTeX RIS

Shijin Yuan, Tianyu Shen, Bin Mu, Zhen Wei, Bo Qin, Guokun Dai. 2026-10-26. 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. https://arxiv.org/abs/2610.35095

Cite the original work for its findings. Save a collection to share your selection of sources.

KEEP EXPLORING

Related papers

arXiv2026

A-UTE: Advection Informed, Uncertainty Aware Temperature Emulator

Physics-based Earth system models (ESMs) are essential for attributing climate change and generating scenario projections, yet their reliance on high-resolution numerical integration makes multi-decadal experiments expensive. In parallel, deep learning has delivered strong gains in short-range weather forecasting; however, auto-regressive roll-outs can accumulate error and become unstable when extended to decade-scale climate emulation. We introduce A-UTE: Advection Informed, Uncertainty Aware Temperature Emulator, aimed at stable multi-year emulation across heterogeneous climate models and grid resolutions. A-UTE is trained on various physics-based models at varying spatial resolutions to emulate temperature fields over a 10-year horizon. A-UTE formulates climate emulation as a forward-time stochastic dynamical system. We propose an auto-regressive ODE-SDE surrogate in which transport dynamics are constrained by an advection consistent ODE component, while a learned neural SDE term models coarse-grained variability and cross-model discrepancy at monthly resolution. We train A-UTE under negative log-likelihood objective for principled uncertainty estimates and probabilistic evaluation. Experiments across 20 climate models show that A-UTE improves long roll-out stability and accuracy relative to relevant baselines, advancing data-driven climate emulation with explicit physical structure and uncertainty-aware predictions.

physics.ao-ph↗

arXiv2026

Bistability of ice-shelf cavity circulation and melt under coastal polynya forcing

Antarctic ice shelves play a key role in controlling the mass loss of the Antarctic Ice Sheet. Increasing evidence from regional studies suggests the possibility of abrupt transitions from low to high basal melt regimes under future climate forcing. The computational cost and complexity of these realistic simulations yet limit their ability to systematically investigate the steady-state response of ice shelves to incremental forcing, the reversibility of abrupt transitions and the underlying positive feedback loops. Here, we report a detailed analysis of abrupt regime changes in an ice-shelf cavity with idealized geometry, subject to prescribed temperature and salinity stratifications in the far field and near-ice-front surface salt fluxes representing the impact of sea-ice production in coastal polynyas. We show that a cold regime emerges when the buoyancy loss in the polynya is sufficient to trigger deep convection, while a warm regime develops otherwise. The steady response of the system exhibits hysteresis, characteristic of irreversible transitions, and scaling laws are proposed for both melt rates and transition thresholds. We demonstrate that bistability emerges when a gyre recirculates the meltwater within the ice-shelf front area. The bistable forcing range increases with far-field ocean temperature but decreases with background stratification. We also show that increasing horizontal viscosity limits bistability, which may explain the range of behaviors reported in previous studies.

physics.ao-ph↗

arXiv2026

Automated Detection and Climatological Analysis of Ripple-Scale Gravity Wave Instabilities Using a Squeeze-and-Excitation Convolutional Neural Network

All-sky OH airglow imaging provides two-dimensional observations of mesospheric wave morphology near the nominal OH emission peak altitude. Small, localized ripple-like structures are difficult to catalogue consistently through manual inspection and may represent either instability-related features or secondary gravity waves. We present an automated detection framework based on a squeeze-and-excitation convolutional neural network (SE-CNN) trained on 41 x 41-pixel patches extracted from 512 x 512-pixel all-sky images acquired at Yucca Ridge Field Station. Signed differences between consecutive images separated by approximately 2 min are normalized using spatial median-absolute-deviation scaling before classification. The trained model is applied to complete images through sliding-window inference, followed by score-map reconstruction, component filtering, and spatial matching. On the held-out patch test set, the classifier achieved an accuracy of 0.92, precision of 0.89, recall of 0.93, and F1 score of 0.92. Within a selected cohort of manually positive validation frames, 299 of 415 evaluable frames contained at least one reference-automated spatial overlap, and 324 of 559 eligible reference supports received a one-to-one automated-component assignment. In a separate blinded single-reviewer audit of disagreement patches, 411 of 543 model-candidate/reference-unmarked crops were classified as clear ripple-like, corresponding to a clear-ripple fraction of 75.7% with a night-cluster bootstrap 95% interval of 70.4%-80.2%. Because the audit was restricted to selected disagreement cases from manually positive frames, this value is not an archive-wide precision estimate.

physics.ao-ph↗