Probabilistic weather forecasts from a pre-trained FourCastNet¶

Backend: TensorFlow, JAX or PyTorch
Dataset: ERA5 (WeatherBench 2)
Techniques: SplitConformalRegressor

Turn a pre-trained global weather model into 90% forecast intervals at every point on Earth, without any training.

FourCastNet is a neural weather model from NVIDIA, trained on 40 years of ERA5 reanalysis data. From the state of the atmosphere now (20 variables on a 0.25° grid, about 1 million points), it predicts the state 6 hours later. keras-climate provides it as a Keras model with NVIDIA's original weights.

Like most weather models, it outputs a single forecast, with no indication of how far off it might be. Here we add that, with no training at all:

Run the pre-trained model on past days, and compare its forecasts with what actually happened.
At every grid point, find the error size that 90% of these errors stay below (split conformal prediction).
Use it as the margin around every new forecast at that point.

Runtime: use a GPU (Runtime → Change runtime type → T4 GPU on Colab). The notebook downloads the model (0.9 GB) and 30 days of ERA5 data (about 3 GB), and runs the model 120 times.

# Install keras-proba and keras-climate, then choose a backend: "jax", "tensorflow" or "torch".
# PyTorch is also needed, to read NVIDIA's original checkpoint (it is preinstalled on Colab).
!pip install -q keras-proba "keras-climate>=0.1.1" ruamel.yaml xarray zarr gcsfs
import os

os.environ["KERAS_BACKEND"] = "jax"

Load the pre-trained model¶

fourcastnet_backbone downloads NVIDIA's checkpoint and loads it into the Keras model. The model works on standardized variables, so we also download the mean and standard deviation of each variable that it was trained with.

import io
import urllib.request

import matplotlib.pyplot as plt
import numpy as np
import xarray as xr

import keras
import keras_proba
from keras_climate.weights.pretrained import fourcastnet_backbone

model, _ = fourcastnet_backbone()

STATS_URL = "https://portal.nersc.gov/project/m4134/FCN_weights_v0/stats_v0/"
means = np.load(io.BytesIO(urllib.request.urlopen(STATS_URL + "global_means.npy").read()))[0, :20, 0, 0]
stds = np.load(io.BytesIO(urllib.request.urlopen(STATS_URL + "global_stds.npy").read()))[0, :20, 0, 0]

[keras_climate] converted 148/148 weights.

Load the weather data¶

WeatherBench 2 hosts ERA5 on Google Cloud, readable without an account. load_state reads the 20 variables FourCastNet expects, in its order, on its 720 × 1440 grid (latitudes 90° to -89.75°).

ERA5 = "gs://weatherbench2/datasets/era5/1959-2023_01_10-wb13-6h-1440x721_with_derived_variables.zarr"
era5 = xr.open_zarr(ERA5, storage_options={"token": "anon"})

# FourCastNet's 20 input variables: (ERA5 name, pressure level in hPa, or None at the surface).
VARIABLES = [
    ("10m_u_component_of_wind", None), ("10m_v_component_of_wind", None), ("2m_temperature", None),
    ("surface_pressure", None), ("mean_sea_level_pressure", None), ("temperature", 850),
    ("u_component_of_wind", 1000), ("v_component_of_wind", 1000), ("geopotential", 1000),
    ("u_component_of_wind", 850), ("v_component_of_wind", 850), ("geopotential", 850),
    ("u_component_of_wind", 500), ("v_component_of_wind", 500), ("geopotential", 500),
    ("temperature", 500), ("geopotential", 50), ("relative_humidity", 500), ("relative_humidity", 850),
    ("total_column_water_vapour", None),
]
T2M = 2  # Index of the 2 m temperature, the variable we forecast with intervals.


def load_state(time):
    snapshot = era5.sel(time=time)
    loaded, state = {}, []
    for name, level in VARIABLES:
        if name not in loaded:
            loaded[name] = snapshot[name].load()  # All its pressure levels, in one request.
        field = loaded[name] if level is None else loaded[name].sel(level=level)
        state.append(field.values[:720])
    state = np.stack(state, axis=-1)
    state[..., 17:19] *= 100  # Relative humidity: a fraction in WeatherBench 2, a percentage for FourCastNet.
    return state.astype("float32")


def load_t2m(time):
    return era5["2m_temperature"].sel(time=time).values[:720].astype("float32")

Forecast 24 hours ahead¶

Each model step advances the atmosphere by 6 hours, so 4 steps give a 24-hour forecast. We keep the 2 m temperature, in °C.

def forecast_t2m(state, steps=4):
    x = ((state - means) / stds)[None]
    for _ in range(steps):
        x = model.predict(x, verbose=0)
    return x[0, ..., T2M] * stds[T2M] + means[T2M] - 273.15


# Grid-cell areas shrink towards the poles: weight averages by cos(latitude).
latitudes = np.linspace(90, -89.75, 720)
area = np.cos(np.deg2rad(latitudes))[:, None] * np.ones((1, 1440))
area /= area.mean()

Run the model on past days¶

We forecast every day of January 2020 at 00 UTC: the first 20 days to calibrate the intervals, the last 10 to test them. For each forecast we also load the temperature that was actually observed 24 hours later. This is the slow part of the notebook.

days = np.arange("2020-01-01T00", "2020-01-31T00", np.timedelta64(1, "D"), dtype="datetime64[h]")
forecasts, actual, persistence = [], [], []
for day in days:
    state = load_state(day)
    forecasts.append(forecast_t2m(state))
    actual.append(load_t2m(day + np.timedelta64(24, "h")) - 273.15)
    persistence.append(state[..., T2M] - 273.15)  # Baseline: "tomorrow will be like today".
forecasts, actual, persistence = np.stack(forecasts), np.stack(actual), np.stack(persistence)

rmse = lambda prediction: np.sqrt(np.mean(area * (prediction - actual) ** 2))
print(f"24-hour forecast error (RMSE): FourCastNet {rmse(forecasts):.2f} °C, 'same as today' {rmse(persistence):.2f} °C")

24-hour forecast error (RMSE): FourCastNet 1.02 °C, 'same as today' 2.19 °C

Calibrate the intervals¶

SplitConformalRegressor treats each grid point as its own output: for each, it measures the absolute errors on the 20 calibration days and keeps a threshold that at least 90% of future errors should stay below. With 20 days, this is the second largest of the 20 errors (conformal prediction rounds up, so that the guarantee holds with few calibration examples).

CALIBRATION = 20
y_cal = actual[:CALIBRATION].reshape(CALIBRATION, -1)  # (days, grid points)
f_cal = forecasts[:CALIBRATION].reshape(CALIBRATION, -1)

conformal = keras_proba.postprocessing.SplitConformalRegressor(alpha=0.1)  # alpha = 1 - 90%.
conformal.fit(y_cal, f_cal)
half_width = keras.ops.convert_to_numpy(conformal.quantile_).reshape(720, 1440)

The map shows the margin around every forecast. It is small over the oceans, where the temperature changes slowly, and large over land: most of all over the northern continents in winter, where cold air masses come and go, and over the Antarctic ice sheet.

land = era5["land_sea_mask"].values[:720]
extent = [-180, 180, -89.75, 90]
centered = lambda field: np.roll(field, 720, axis=-1)  # Longitudes 0..360 -> -180..180.

plt.figure(figsize=(10, 4.5))
plt.imshow(centered(half_width), extent=extent, cmap="Blues", vmin=0, vmax=np.percentile(half_width, 99))
plt.colorbar(label="interval half-width (°C)", shrink=0.8)
plt.contour(np.linspace(-180, 180, 1440), latitudes, centered(land), levels=[0.5], colors="k", linewidths=0.4)
plt.title("90% interval: forecast ± this margin, for the 24-hour temperature forecast")
plt.xlabel("longitude")
plt.ylabel("latitude")
plt.show()

Output plot

Check the intervals on the test days¶

On the 10 test days, unseen during calibration, the intervals should contain the observed temperature about 90% of the time.

f_test = forecasts[CALIBRATION:].reshape(len(days) - CALIBRATION, -1)
y_test = actual[CALIBRATION:].reshape(len(days) - CALIBRATION, -1)
lower, upper = conformal.predict_interval(f_test)
covered = keras.ops.convert_to_numpy(keras_proba.postprocessing.interval_coverage(y_test, lower, upper))
covered = covered.reshape(720, 1440)  # Fraction of the test days covered, at each grid point.
print(f"The 90% intervals contain the observed temperature {np.mean(area * covered):.1%} of the time (area-weighted)")

The 90% intervals contain the observed temperature 90.2% of the time (area-weighted)

Finally, the forecasts and intervals at three cities, through the test days. The black dots are the temperatures that were actually observed.

cities = {"Moscow": (55.75, 37.6), "Chicago": (41.9, -87.6), "Singapore": (1.35, 103.8)}
lower = keras.ops.convert_to_numpy(lower).reshape(-1, 720, 1440)
upper = keras.ops.convert_to_numpy(upper).reshape(-1, 720, 1440)
test_days = days[CALIBRATION:] + np.timedelta64(24, "h")

fig, axes = plt.subplots(1, 3, figsize=(12, 3.5))
for ax, (city, (lat, lon)) in zip(axes, cities.items()):
    i, j = int(round((90 - lat) * 4)), int(round((lon % 360) * 4))
    dates = test_days.astype("datetime64[D]")
    ax.fill_between(dates, lower[:, i, j], upper[:, i, j], alpha=0.3, label="90% interval")
    ax.plot(dates, forecasts[CALIBRATION:, i, j], label="forecast")
    ax.plot(dates, actual[CALIBRATION:, i, j], "ko", ms=4, label="observed")
    ax.set(title=city, ylabel="temperature (°C)")
    ax.tick_params(axis="x", rotation=45)
axes[0].legend()
plt.tight_layout()
plt.show()

Output plot

What we learned¶

A pre-trained model's single forecast becomes a forecast with an interval, with no training: only its errors on a few past days are needed.
Each grid point gets its own margin, so the intervals are wide where the weather is hard to predict and narrow where it is not.
The guarantee holds when the calibration days resemble the forecast days. Weather depends on the season, so calibrate on recent days, for the same time of day and forecast lead time, and recalibrate as the season changes.
The guarantee is for each grid point on average over days, not for the whole map at once: on any given day, some regions will fall outside their intervals.