Installation & Requirements

wrf_gpu is a JAX/CUDA application: an Ampere-or-newer NVIDIA GPU (compute capability ≥8.0), a CUDA 13 JAX build, and Python. This page covers the hardware you need, the full and the minimal (source-only) install, how to verify the GPU is visible, and the runtime-table environment variable that table-using physics schemes require before a run.

GPU and host memory¶

Minimum: NVIDIA Ampere, compute capability 8.0, for the GPU kernels. Older cards (Pascal, Turing) are not supported: on a GTX 1080 Ti only a slow debug path ran, about 3.5 times slower than CPU-WRF on six CPU cores. Hardware requirements.

The tested card is an RTX 5090 (CUDA 13) with a Ryzen 9 9950X host; other NVIDIA GPUs are untested. How much memory a forecast needs depends on the grid, the physics and the number of nests: gpuwrf doctor estimates it for a case and gpuwrf quickstart --profiles gives rough grid sizes for 12 to 32 GB cards. Compiling a new set-up needs more host memory than running it. See GPU memory.

Prefer a container? The onboarding container builds wrf_gpu together with the original WPS and WRF tools for region set-up.

Full install¶

Use an isolated environment. Python 3.11 is the reference interpreter; the package requires Python ≥ 3.10. Install the CUDA build of JAX first, then the package.

git clone https://github.com/wrf-gpu/wrf_gpu.git
cd wrf_gpu

# Isolated environment (venv or conda both work):
python -m venv .venv && . .venv/bin/activate
# or:  conda create -n wrfgpu python=3.11 && conda activate wrfgpu

# 1. GPU build of JAX (CUDA 13 wheels, tested with JAX 0.10):
pip install "jax[cuda13]==0.10.*"

# 2. wrf_gpu itself (pulls numpy, netCDF4, xarray, PyYAML, zarr, ...):
pip install -e .

# 3. Verify JAX sees the GPU (should list a cuda device):
python -c "import jax; print(jax.devices())"

Equivalently, the cuda extra pins the CUDA 13 jaxlib for you:

pip install -e ".[cuda]"

Minimal source-only install¶

You do not need the whole repository (proofs, agent infrastructure, validation corpora) just to run wrf_gpu. A cone sparse-checkout of src plus the vendored runtime data tables is enough; this path was verified end-to-end on a clean machine.

# Shallow, no-checkout clone, then cone-sparse-checkout only what `run` needs:
git clone --depth 1 --no-checkout \
    https://github.com/wrf-gpu/wrf_gpu.git wrf_gpu && cd wrf_gpu
git sparse-checkout init --cone
# In cone mode the root files (pyproject.toml, README.md, LICENSE, ...) come
# automatically; add the source package and the vendored runtime data tables
# (small physics tables such as RRTMG live there -- `src` alone is not enough):
git sparse-checkout set src data/fixtures
git checkout

# Install (CPU import works without a GPU; add the CUDA jaxlib to run on GPU):
python -m venv .venv && . .venv/bin/activate
pip install -e .                  # CPU import-check; or:
pip install -e ".[cuda]"          # GPU execution (CUDA 13 jaxlib)

# Verify:
python -c "import gpuwrf; print('wrf_gpu', gpuwrf.__version__)"
gpuwrf run --help                # the command-line entry point

The data/fixtures/ tables (about 32 MiB) are required, so the sparse checkout must include them alongside src. Larger tables are downloaded on first use.

Prerequisite environment variables¶

Table-using schemes (Noah-MP and classic RRTM, plus WRF-table-dependent options) load lookup tables from a pristine WRF v4 install at runtime, so their paths must be set before you run a case that selects them. Without GPUWRF_WRF_ROOT, those schemes fail closed with a clear named error before any compute — they never silently substitute a scheme. Cases whose physics menu does not use those schemes run without it.

Two optional paths (both defaulted) tune performance and scratch: