Source code: Kriging in Python with PyKrige: IDW vs ordinary kriging of UK NO₂ monitoring data

"""Spatial interpolation in Python: IDW vs ordinary kriging of UK NO2 monitoring data.

Downloads the 2025 annual mean NO2 at every UK Automatic Urban and Rural Network (AURN) station from
DEFRA UK-AIR, keeps background sites with good data capture, fits and compares variogram models,
interpolates a 2 km surface with inverse distance weighting and ordinary kriging (PyKrige), maps the
kriging standard deviation, scores every method with leave-one-out cross-validation and, if the
Sentinel-5P tutorial has been run, tests the satellite NO2 column as a covariate (regression kriging).
Run with:  python kriging_no2_uk.py
"""
import platform
import time
from importlib.metadata import version
from pathlib import Path

import geopandas as gpd
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pykrige
import pyreadr
import rasterio
import requests
from pykrige.ok import OrdinaryKriging
from rasterio.features import geometry_mask
from rasterio.transform import from_origin
from scipy.stats import pearsonr

HERE = Path(__file__).parent
OUT, DATA = HERE / "output", HERE / "data"
for d in (OUT, DATA):
    d.mkdir(exist_ok=True)
plt.rcParams.update({"figure.dpi": 150, "savefig.bbox": "tight", "font.size": 9})
pd.set_option("display.width", 140)

print("=== versions ===")
print(f"python {platform.python_version()} | pykrige {pykrige.__version__} | numpy {np.__version__} | "
      f"pandas {pd.__version__} | geopandas {gpd.__version__} | pyreadr {version('pyreadr')} | "
      f"rasterio {rasterio.__version__}")


def section(title: str) -> None:
    print(f"\n=== {title} ===")


# --- Step 1: download the AURN annual means and site metadata -----------------------
YEAR = 2025
UKAIR = "https://uk-air.defra.gov.uk/openair/R_data/"


def fetch(name):
    dest = DATA / name
    if not dest.exists():
        for attempt in range(4):
            try:
                r = requests.get(UKAIR + name, timeout=60)
                r.raise_for_status()
                dest.write_bytes(r.content)
                break
            except requests.RequestException:
                if attempt == 3:
                    raise
                time.sleep(5 * (attempt + 1))
    return dest


annual = pyreadr.read_r(fetch(f"summary_annual_AURN_{YEAR}.rds"))[None]
meta = pyreadr.read_r(fetch("AURN_metadata.RData"))["AURN_metadata"]
meta = meta[meta["parameter"] == "NO2"].drop_duplicates("site_id")
aurn = (annual[["code", "site", "NO2.mean", "NO2.capture"]]
        .rename(columns={"NO2.mean": "no2", "NO2.capture": "capture"})
        .merge(meta[["site_id", "location_type", "latitude", "longitude", "ratified_to"]],
               left_on="code", right_on="site_id", how="left")
        .drop(columns="site_id"))
aurn = aurn[aurn["no2"].notna()]
section(f"AURN NO2 annual means, {YEAR}")
print(f"{len(aurn)} stations report NO2; {(aurn['capture'] >= 0.75).sum()} have >= 75% data capture")
print(f"ratified to: {aurn['ratified_to'].astype(str).str[:10].value_counts().head(2).to_dict()}")
print(aurn[aurn["capture"] >= 0.75].groupby("location_type")["no2"].agg(["count", "mean", "min", "max"]).round(1))

# Background sites only: traffic and industrial sites measure the street or the works next door,
# which changes over tens of metres and cannot be interpolated between stations kilometres apart.
BACKGROUND = ["Urban Background", "Suburban Background", "Rural Background"]
stations = aurn[(aurn["capture"] >= 0.75) & aurn["location_type"].isin(BACKGROUND)].copy()
roadside = aurn[(aurn["capture"] >= 0.75) & (aurn["location_type"] == "Urban Traffic")].copy()
CRS = "EPSG:27700"   # British National Grid, metres
gdf = gpd.GeoDataFrame(stations, geometry=gpd.points_from_xy(stations.longitude, stations.latitude),
                       crs="EPSG:4326").to_crs(CRS)
road = gpd.GeoDataFrame(roadside, geometry=gpd.points_from_xy(roadside.longitude, roadside.latitude),
                        crs="EPSG:4326").to_crs(CRS)
x, y, z = gdf.geometry.x.to_numpy() / 1000, gdf.geometry.y.to_numpy() / 1000, gdf["no2"].to_numpy()  # km
print(f"\nkept {len(gdf)} background stations: mean {z.mean():.1f}, sd {z.std(ddof=1):.1f}, "
      f"range {z.min():.1f} to {z.max():.1f} ug/m3")
dists = np.hypot(x[:, None] - x, y[:, None] - y)
np.fill_diagonal(dists, np.inf)
print(f"nearest-neighbour distance: median {np.median(dists.min(1)):.1f} km, max {dists.min(1).max():.1f} km")

# --- Step 2: experimental variogram and model fits ----------------------------------
NLAGS = 12
MODELS = ["spherical", "exponential", "gaussian"]
fits = {}
section("variogram fits (PyKrige, distances in km)")
for model in MODELS:
    ok = OrdinaryKriging(x, y, z, variogram_model=model, nlags=NLAGS, weight=True)
    fits[model] = ok
    psill, rng, nugget = ok.variogram_model_parameters     # PyKrige returns the partial sill
    print(f"{model:12s} partial sill {psill:6.2f}  range {rng:6.1f} km  "
          f"nugget {nugget:5.2f}  -> nugget/sill {nugget / (psill + nugget):.2f}")
print("lag centres (km):", np.round(fits["spherical"].lags, 0).astype(int).tolist())
print("semivariance:   ", np.round(fits["spherical"].semivariance, 1).tolist())

# --- Step 3: prediction grid over the UK --------------------------------------------
ne = DATA / "countries.gpkg"
if not ne.exists():
    gpd.read_file("https://naciscdn.org/naturalearth/10m/cultural/ne_10m_admin_0_countries.zip").to_file(ne)
uk = gpd.read_file(ne)
uk = uk[uk["ADM0_A3"] == "GBR"].to_crs(CRS)
uk = uk.clip((-1e5, 0, 7e5, 1.25e6))                      # drop overseas territories
CELL = 2000
X0, Y1 = 0, 1_220_000
GX, GY = 330, 610                                         # 660 km x 1220 km at 2 km
gx = (X0 + CELL * (np.arange(GX) + 0.5)) / 1000
gy = (Y1 - CELL * (np.arange(GY) + 0.5)) / 1000
transform = from_origin(X0, Y1, CELL, CELL)
on_land = ~geometry_mask(uk.geometry, (GY, GX), transform)
section("prediction grid")
print(f"{GX} x {GY} cells of {CELL / 1000:.0f} km, {on_land.sum():,} on land")

# --- Step 4: ordinary kriging and IDW ------------------------------------------------
MODEL = "spherical"
ok = fits[MODEL]
t0 = time.time()
z_ok, var_ok = ok.execute("grid", gx, gy, backend="vectorized")
z_ok, var_ok = np.asarray(z_ok), np.asarray(var_ok)
section("ordinary kriging")
print(f"{MODEL} model, {time.time() - t0:.1f} s for {GX * GY:,} cells")
sd_ok = np.sqrt(np.clip(var_ok, 0, None))
print(f"land cells: predicted NO2 {np.nanmin(z_ok[on_land]):.1f} to {np.nanmax(z_ok[on_land]):.1f} ug/m3; "
      f"kriging SD {sd_ok[on_land].min():.1f} to {sd_ok[on_land].max():.1f} (median {np.median(sd_ok[on_land]):.1f})")


def idw(xs, ys, zs, xq, yq, power=2.0):
    """Inverse distance weighting: a weighted mean of every station, with weights 1 / d^power."""
    d = np.hypot(xq[:, None] - xs[None, :], yq[:, None] - ys[None, :])
    d = np.maximum(d, 1e-6)
    w = 1.0 / d ** power
    return (w @ zs) / w.sum(1)


GXX, GYY = np.meshgrid(gx, gy)
z_idw = idw(x, y, z, GXX.ravel(), GYY.ravel(), 2).reshape(GY, GX)
section("IDW (power 2)")
print(f"land cells: {z_idw[on_land].min():.1f} to {z_idw[on_land].max():.1f} ug/m3")
diff = z_ok - z_idw
print(f"kriging minus IDW over land: mean {diff[on_land].mean():+.2f}, 5th to 95th percentile "
      f"{np.percentile(diff[on_land], 5):+.1f} to {np.percentile(diff[on_land], 95):+.1f} ug/m3")

# --- Step 5: leave-one-out cross-validation ------------------------------------------
n = len(z)


def loo(predict):
    pred = np.empty(n)
    for i in range(n):
        keep = np.arange(n) != i
        pred[i] = predict(keep, i)
    return pred


def ok_loo(model):
    psill, rng, nugget = fits[model].variogram_model_parameters   # variogram fixed from the full data set
    params = {"psill": psill, "range": rng, "nugget": nugget}

    def predict(keep, i):
        m = OrdinaryKriging(x[keep], y[keep], z[keep], variogram_model=model,
                            variogram_parameters=params)
        p, v = m.execute("points", x[i:i + 1], y[i:i + 1])
        ok_var[model][i] = v[0]
        return p[0]
    return predict


ok_var = {m: np.empty(n) for m in MODELS}
preds = {"Global mean": loo(lambda k, i: z[k].mean())}
for p in (1, 2, 3):
    preds[f"IDW p={p}"] = loo(lambda k, i, p=p: idw(x[k], y[k], z[k], x[i:i + 1], y[i:i + 1], p)[0])
for m in MODELS:
    preds[f"OK {m}"] = loo(ok_loo(m))


def scores(name, pred):
    err = pred - z
    # The leave-one-out mean is perfectly anti-correlated with the left-out value, so r means nothing there.
    r = np.nan if name == "Global mean" else pearsonr(pred, z)[0]
    return {"RMSE": np.sqrt(np.mean(err ** 2)), "MAE": np.mean(np.abs(err)), "bias": err.mean(), "r": r}


cv = pd.DataFrame({k: scores(k, v) for k, v in preds.items()}).T
section(f"leave-one-out cross-validation, {n} background stations (ug/m3)")
print(cv.round(2).sort_values("RMSE").to_string())
zscore = (preds[f"OK {MODEL}"] - z) / np.sqrt(ok_var[MODEL])
print(f"OK {MODEL}: mean standardised error {zscore.mean():+.2f}, SD of standardised error {zscore.std():.2f} "
      f"(1.0 = kriging variance is well calibrated)")
print(f"share of stations inside the 95% kriging interval: {(np.abs(zscore) <= 1.96).mean():.0%}")
worst = pd.DataFrame({"site": gdf["site"].to_numpy(), "type": gdf["location_type"].to_numpy(), "observed": z,
                      "OK": preds[f"OK {MODEL}"], "IDW p=2": preds["IDW p=2"]})
worst["OK error"] = worst["OK"] - worst["observed"]
print("\nlargest kriging errors:")
print(worst.reindex(worst["OK error"].abs().sort_values(ascending=False).index).head(6).round(1).to_string(index=False))

# Roadside stations were not used: how badly does the background surface under-predict them?
rx, ry = road.geometry.x.to_numpy() / 1000, road.geometry.y.to_numpy() / 1000
road_pred, _ = ok.execute("points", rx, ry)
print(f"\n{len(road)} roadside stations: observed mean {road['no2'].mean():.1f}, background kriging at the "
      f"same points {np.mean(road_pred):.1f} ug/m3 (under-prediction {np.mean(road['no2'] - road_pred):.1f})")

# --- Step 6: does the Sentinel-5P column help? (needs the Sentinel-5P tutorial output) ----
S5P = HERE.parent / "sentinel-5p-no2-air-pollution-python" / "output" / f"s5p_no2_{YEAR}_mean.tif"
rk = None
if S5P.exists():
    with rasterio.open(S5P) as src:
        lonlat = gdf.to_crs("EPSG:4326").geometry
        col = np.array([v[0] for v in src.sample(zip(lonlat.x, lonlat.y), indexes=1)])
        road_ll = road.to_crs("EPSG:4326").geometry
        col_road = np.array([v[0] for v in src.sample(zip(road_ll.x, road_ll.y), indexes=1)])
        grid_ll = gpd.GeoSeries(gpd.points_from_xy(GXX.ravel() * 1000, GYY.ravel() * 1000), crs=CRS).to_crs("EPSG:4326")
        col_grid = np.array([v[0] for v in src.sample(zip(grid_ll.x, grid_ll.y), indexes=1)]).reshape(GY, GX)
    section("Sentinel-5P column at the stations")
    r_bg = pearsonr(col, z)[0]
    slope, intercept = np.polyfit(col, z, 1)
    print(f"background stations: Pearson r = {r_bg:.2f}, NO2 = {intercept:.2f} + {slope:.3f} x column (umol/m2)")
    print(f"roadside stations:   Pearson r = {pearsonr(col_road, road['no2'])[0]:.2f}")

    def rk_predict(keep, i):
        b1, b0 = np.polyfit(col[keep], z[keep], 1)
        res = z[keep] - (b0 + b1 * col[keep])
        m = OrdinaryKriging(x[keep], y[keep], res, variogram_model="spherical", nlags=NLAGS, weight=True)
        return b0 + b1 * col[i] + m.execute("points", x[i:i + 1], y[i:i + 1])[0][0]

    preds["Regression kriging (S5P)"] = loo(rk_predict)
    preds["Regression only (S5P)"] = loo(lambda k, i: np.polyval(np.polyfit(col[k], z[k], 1), col[i]))
    cv = pd.DataFrame({k: scores(k, v) for k, v in preds.items()}).T
    print("\nLOOCV including the satellite covariate:")
    print(cv.round(2).sort_values("RMSE").to_string())
    res_all = z - (intercept + slope * col)
    rk_ok = OrdinaryKriging(x, y, res_all, variogram_model="spherical", nlags=NLAGS, weight=True)
    print("residual variogram: partial sill {:.2f}, range {:.1f} km, nugget {:.2f}".format(*rk_ok.variogram_model_parameters))
    res_grid, _ = rk_ok.execute("grid", gx, gy, backend="vectorized")
    rk = intercept + slope * col_grid + np.asarray(res_grid)
else:
    section("Sentinel-5P column at the stations")
    print(f"{S5P} not found: run s5p_no2_air_pollution.py from the Sentinel-5P tutorial first to add this step.")

# --- Step 7: save the surfaces ----------------------------------------------------------
profile = dict(driver="GTiff", width=GX, height=GY, count=1, dtype="float32", crs=CRS, transform=transform,
               nodata=np.nan, compress="deflate")
for name, arr in [("no2_ok", z_ok), ("no2_ok_sd", sd_ok), ("no2_idw", z_idw)] + ([("no2_rk", rk)] if rk is not None else []):
    with rasterio.open(OUT / f"{name}_{YEAR}.tif", "w", **profile) as dst:
        dst.write(np.where(on_land, arr, np.nan).astype("float32"), 1)
cv.round(3).to_csv(OUT / "loocv_scores.csv")
print(f"\nsaved GeoTIFFs and loocv_scores.csv to {OUT.name}")

# --- Figures ----------------------------------------------------------------------------
EXT = (X0 / 1000, (X0 + GX * CELL) / 1000, (Y1 - GY * CELL) / 1000, Y1 / 1000)
CMAP = "Spectral_r"
VMAX = 25
ukkm = uk.copy()
ukkm.geometry = ukkm.geometry.scale(1 / 1000, 1 / 1000, origin=(0, 0))
TYPE_MARK = {"Urban Background": "o", "Suburban Background": "s", "Rural Background": "^"}


def frame(ax):
    ukkm.boundary.plot(ax=ax, color="#333333", lw=0.4)
    ax.set_xlim(-10, 670)
    ax.set_ylim(0, 1220)
    ax.set_aspect("equal")
    ax.tick_params(labelsize=7)
    ax.set_xlabel("Easting (km)", fontsize=7)
    ax.set_ylabel("Northing (km)", fontsize=7)


def masked(arr):
    return np.where(on_land, arr, np.nan)


# Figure 1: stations.
fig, ax = plt.subplots(figsize=(6.4, 9))
frame(ax)
ax.scatter(road.geometry.x / 1000, road.geometry.y / 1000, s=10, marker="x", color="#777777", lw=0.8,
           label=f"Roadside, not used ({len(road)})")
for t, mk in TYPE_MARK.items():
    s = gdf["location_type"] == t
    sc = ax.scatter(x[s], y[s], c=z[s], cmap=CMAP, vmin=0, vmax=VMAX, marker=mk, s=38, edgecolor="black", lw=0.5,
                    label=f"{t} ({s.sum()})")
fig.colorbar(sc, ax=ax, shrink=0.6, label="Annual mean NO$_2$ (µg/m³)", extend="max")
ax.legend(loc="upper left", fontsize=7, frameon=True)
ax.set_ylabel("Northing (km)", fontsize=7)
ax.set_title(f"AURN NO$_2$ stations, {YEAR} annual mean (≥ 75% capture)", fontsize=10)
fig.savefig(OUT / "01-aurn-stations.png")
plt.close(fig)

# Figure 2: experimental variogram and the three model fits.
fig, ax = plt.subplots(figsize=(8, 4.6))
ref = fits["spherical"]
ax.plot(ref.lags, ref.semivariance, "o", color="black", label="Experimental semivariance")
h = np.linspace(0, ref.lags.max() * 1.05, 300)
for (m, f), colour in zip(fits.items(), ["#d53e4f", "#3288bd", "#66a61e"]):
    ps, rg, nu = f.variogram_model_parameters
    ax.plot(h, f.variogram_function(f.variogram_model_parameters, h), color=colour, lw=1.8,
            label=f"{m}: nugget {nu:.1f}, sill {ps + nu:.1f}, range {rg:.0f} km")
ax.axhline(z.var(), color="#999999", ls="--", lw=1)
ax.text(h.max(), z.var(), "sample variance", ha="right", va="bottom", fontsize=7, color="#555555")
ax.set_xlabel("Separation distance (km)")
ax.set_ylabel("Semivariance (µg/m³)²")
ax.set_title(f"Variogram of {YEAR} background NO$_2$ ({n} AURN stations, {NLAGS} lags)", fontsize=10)
ax.legend(frameon=False, fontsize=7.5)
ax.spines[["top", "right"]].set_visible(False)
ax.grid(alpha=0.3)
fig.savefig(OUT / "02-variogram.png")
plt.close(fig)

# Figure 3: IDW vs ordinary kriging.
fig, axes = plt.subplots(1, 2, figsize=(10, 8.6), constrained_layout=True)
for ax, arr, title in [(axes[0], z_idw, "IDW, power 2"), (axes[1], z_ok, f"Ordinary kriging ({MODEL})")]:
    im = ax.imshow(masked(arr), extent=EXT, cmap=CMAP, vmin=0, vmax=VMAX, interpolation="nearest")
    frame(ax)
    ax.scatter(x, y, s=5, color="black")
    ax.set_title(title, fontsize=10)
fig.colorbar(im, ax=axes, shrink=0.55, label="Background NO$_2$, annual mean (µg/m³)", extend="max")
fig.savefig(OUT / "03-idw-vs-kriging.png")
plt.close(fig)

# Figure 4: kriging standard deviation.
fig, ax = plt.subplots(figsize=(6.4, 9))
im = ax.imshow(masked(sd_ok), extent=EXT, cmap="viridis", interpolation="nearest")
frame(ax)
ax.scatter(x, y, s=8, color="white", edgecolor="black", lw=0.4)
fig.colorbar(im, ax=ax, shrink=0.6, label="Kriging standard deviation (µg/m³)")
ax.set_ylabel("Northing (km)", fontsize=7)
ax.set_title("Where the kriged map is least certain", fontsize=10)
fig.savefig(OUT / "04-kriging-standard-deviation.png")
plt.close(fig)

# Figure 5: cross-validation.
order = cv.sort_values("RMSE").index
fig, (a1, a2) = plt.subplots(1, 2, figsize=(11, 4.8), constrained_layout=True)
cols = ["#3288bd" if k.startswith("OK") else "#66a61e" if "S5P" in k else "#f46d43" if k.startswith("IDW")
        else "#999999" for k in order]
a1.barh(range(len(order)), cv.loc[order, "RMSE"], color=cols)
a1.set_yticks(range(len(order)), order)
a1.invert_yaxis()
for i, v in enumerate(cv.loc[order, "RMSE"]):
    a1.text(v + 0.05, i, f"{v:.2f}", va="center", fontsize=7.5)
a1.set_xlabel("Leave-one-out RMSE (µg/m³)")
a1.set_title("Cross-validation error by method (lower is better)", fontsize=10)
a1.spines[["top", "right"]].set_visible(False)
lim = (0, max(z.max(), preds[f"OK {MODEL}"].max(), preds["IDW p=2"].max()) * 1.08)
a2.plot(lim, lim, color="#999999", lw=1)
a2.scatter(z, preds["IDW p=2"], s=18, color="#f46d43", label=f"IDW p=2 (RMSE {cv.loc['IDW p=2', 'RMSE']:.2f})")
a2.scatter(z, preds[f"OK {MODEL}"], s=18, color="#3288bd", label=f"OK {MODEL} (RMSE {cv.loc[f'OK {MODEL}', 'RMSE']:.2f})")
a2.set_xlim(lim)
a2.set_ylim(lim)
a2.set_aspect("equal")
a2.set_xlabel("Observed annual mean (µg/m³)")
a2.set_ylabel("Predicted with the station left out (µg/m³)")
a2.legend(frameon=False, fontsize=7.5)
a2.set_title("Each station predicted from all the others", fontsize=10)
fig.savefig(OUT / "05-cross-validation.png")
plt.close(fig)

# Figure 6: satellite column vs ground stations, and the regression-kriging map.
if rk is not None:
    fig, (a1, a2) = plt.subplots(1, 2, figsize=(11, 6.4), constrained_layout=True,
                                 gridspec_kw={"width_ratios": [1.1, 1]})
    a1.scatter(col_road, road["no2"], s=16, marker="x", color="#777777", label="Roadside")
    a1.scatter(col, z, s=22, color="#3288bd", edgecolor="black", lw=0.4, label="Background")
    xx = np.linspace(0, max(col.max(), col_road.max()) * 1.05, 10)
    a1.plot(xx, intercept + slope * xx, color="#d53e4f", lw=1.5, label=f"Background fit (r = {r_bg:.2f})")
    a1.set_xlabel(f"Sentinel-5P tropospheric column, {YEAR} mean (µmol/m²)")
    a1.set_ylabel(f"AURN surface NO$_2$, {YEAR} mean (µg/m³)")
    a1.set_title("Column (satellite) vs surface (ground station)", fontsize=10)
    a1.legend(frameon=False, fontsize=7.5)
    a1.spines[["top", "right"]].set_visible(False)
    a1.grid(alpha=0.3)
    im = a2.imshow(masked(rk), extent=EXT, cmap=CMAP, vmin=0, vmax=VMAX, interpolation="nearest")
    frame(a2)
    a2.scatter(x, y, s=5, color="black")
    a2.set_title("Regression kriging with the Sentinel-5P column", fontsize=10)
    fig.colorbar(im, ax=a2, shrink=0.7, label="Background NO$_2$ (µg/m³)", extend="max")
    fig.savefig(OUT / "06-satellite-vs-ground.png")
    plt.close(fig)
print("figures written")