How wrf_gpu Works

wrf_gpu is a clean-room JAX rewrite of the WRF-ARW forecast model that targets the GPU memory hierarchy from day one. It reproduces the same physics solver WRF implements, but its architecture, execution model, and validation strategy are built for the GPU — not inherited from the WRF Fortran source.

A rewrite, not a port¶

wrf_gpu is not a line-by-line translation of WRF Fortran and it is not a CUDA-annotated copy of the WRF codebase. It is an independent implementation of the WRF-ARW equations and physics, written in JAX and compiled to the GPU through XLA. It reads the same inputs as WRF and writes WRF-compatible output, but the solver in between is its own code, held to WRF's numerical behavior by validation rather than by shared source.

Why this matters for you

Because it is a rewrite validated against an oracle, wrf_gpu can be honest and precise about exactly which parts match WRF and how well. The sealed identity fixtures prove their dynamics/thermodynamics fields cell-for-cell identical to CPU-WRF v4; this is fixture-backed evidence, not universal control/configuration parity. The broader forecast-skill equivalence over 24–72 h is a separate, still-open gate. The guide never blurs those two. See Validation against WRF.

GPU-native architecture¶

The design follows a small set of immutable principles:

Whole-state device residency. After initialization the full model state lives on the GPU. There are no host↔device transfers inside the timestep loop (any exception requires an architecture decision record).
Fused, timestep-scale kernels. The integrator compiles into a few large fused programs rather than a storm of tiny kernel launches, so the GPU stays busy instead of waiting on launch overhead.
WRF's precision. Since v0.3.0 the default path runs fused fp32 kernels — WRF's own single precision — with double precision only where WRF itself uses it, revalidated against CPU-WRF. GPUWRF_FAST_DEFAULTS=0 turns the fast paths off; see Performance → fp32.
No masking on the result path. There is no clamping, no nan_to_num, and no silent substitution of a scheme. If state goes non-finite, a default-on guard aborts with the exact location instead of hiding it.

The forecast loop, physics dispatch, and I/O live under src/gpuwrf/runtime/, src/gpuwrf/coupling/, src/gpuwrf/physics/ and src/gpuwrf/io/; the frozen State, grid, and physics registry are in src/gpuwrf/contracts/.

Standalone initialization¶

In a standard WRF workflow you run WPS (geogrid → ungrib → metgrid) to produce met_em files, then real.exe to turn them into wrfinput/wrfbdy, then wrf.exe to integrate. The current public CLI keeps the WRF real-data artifacts (wrfinput_* and wrfbdy_d01) plus your namelist.input, and replaces wrf.exe with GPU integration. It does not need CPU-WRF wrfout history files for a standalone forecast.

The single entry point gpuwrf run auto-detects the case:

Input directory contains…	Mode	Meaning
`wrfinput_<domain>` + `wrfbdy_d01`, no CPU `wrfout`	Standalone native-init	Loads the initial state and boundary forcing from WRF real-data artifacts and integrates on the GPU — the normal operational path.
≥2 CPU-WRF `wrfout_<domain>` history files	Replay	Runs from a CPU-WRF case for direct GPU↔CPU comparison / validation.

See Input Data & Initialization for details.

Validation against WRF, not bitwise reproducibility¶

wrf_gpu does not aim to be bit-for-bit identical to a specific WRF build (which is impossible across compilers and core counts anyway). Instead it validates against WRF as an oracle through a tiered pyramid, strongest evidence first:

Micro-fixture / savepoint parity — individual physics kernels checked against pristine-WRF single-column oracles to machine precision.
Physical invariants — conservation and boundedness checks.
Short-run / convergence — behavior over the first forecast hours.
Cell-identity proof — a full GPU wrfout compared to a CPU-WRF wrfout over all grid cells, all forecast leads, and all core prognostics, against a frozen tolerance manifest that is read before comparison and never tuned.

Every claim in this guide has a proof object on disk in the repository (proofs/), and every architecture decision has a cross-model-reviewed ADR. Read the method and reproduce commands on the Validation page.

The priority order¶

When two goals conflict, wrf_gpu resolves them in a fixed order:

stability > identity > speed > memory

A result is never destabilized or altered to go faster or fit in less memory. Since v0.3.0 the default path runs the validated fp32 kernels (WRF's own single precision), revalidated against CPU-WRF; other validated fp32 fast paths are release defaults, with explicit opt-outs; performance changes retain their numerical scope. Proven WRF-fidelity fixes may intentionally move results, as v0.23.4 and v0.3.0 do; affected trajectories must then be revalidated. See Performance & GPU Modes.

Agent-driven development¶

wrf_gpu is built and extended by an AI manager agent that dispatches implementer and critic models, reviews every diff, runs the acceptance gates, and merges. The two AI systems credited on the logo — an Anthropic Claude model and an OpenAI GPT model — are the primary authors and reviewers. This is why the repository is unusually proof-dense: the development process itself requires a falsifiable artifact for every claim. See Credits.