keisler-2022

This repo contains code to run the machine learning weather model described in Forecasting Global Weather with Graph Neural Networks (Keisler 2022).

The model uses a three-stage graph neural network: an Encoder maps ERA5 lat/lon fields onto an H3 hexagonal mesh, a Processor runs 9 rounds of message passing on the mesh, and a Decoder maps back to lat/lon to produce a 6-hour forecast update. Autoregressive rollout produces multi-day forecasts.

Installation

Prerequisites: Python 3.10+ and uv installed.

git clone git@github.com:rkeisler/keisler-2022.git
cd keisler-2022

# CPU-only (default)
uv sync

# GPU with CUDA 12
uv sync --extra cuda12

# Optional: run tests
uv run pytest

Running Forecasts

A 10-day forecast should take about one minute total (load initial conditions, run forecast, save output) on a GPU machine. It will take a bit longer on a CPU machine, e.g. 2 minutes on a 8-vCPU machine.

The model can be initialized from two data sources:

ERA5 reanalysis (via Google ARCO) — historical dates, good for evaluation:

uv run forecast.py --init 2020-01-01T00 --steps 40

ECMWF IFS analysis (via ECMWF Open Data on AWS) — recent dates, good for near real-time forecasting:

uv run --extra opendata forecast.py --init 2026-02-15T00 --steps 20 --input opendata

Use --help to see all available options:

uv run forecast.py --help

GPU Setup

uv sync --extra cuda12

Verify GPU access:

uv run python -c "import jax; print('Devices:', jax.devices())"

You should see [CudaDevice(id=0)] instead of [CpuDevice(id=0)].

Troubleshooting: If JAX falls back to CPU, make sure LD_LIBRARY_PATH is not set. A pre-existing LD_LIBRARY_PATH can cause JAX to find incompatible system CUDA libraries instead of its own pip-bundled versions. To fix, run this command or put it in your .bashrc:

unset LD_LIBRARY_PATH

See the JAX installation docs for more details.

Scripts

The scripts/ directory contains example analysis and visualization scripts.

01_evaluation.py — ERA5 Evaluation

Runs a forecast initialized from ERA5 reanalysis, computes area-weighted RMSE at each 6-hour lead time, and produces a figure of specific humidity at 850 hPa comparing ERA5 truth vs. model forecast.

uv run --extra scripts scripts/01_evaluation.py --init 2020-01-01T00 --steps 12

02_sensitivity.py — Forecast Sensitivity Maps

Computes d(forecast)/d(initial_conditions) using JAX autodiff for a chosen target location, field, and lead time, then visualizes the sensitivity maps.

uv run --extra scripts scripts/02_sensitivity.py --init 2026-01-03T00 --steps 12

03_hurricane.py — Hurricane Sandy Tracking

Runs an 8-day forecast initialized from ERA5 at 2012-10-23T00 (2012 was a test-set year), tracks Hurricane Sandy's center via the Z1000 minimum, and compares the predicted track against the actual track.

uv run --extra scripts scripts/03_hurricane.py