feat(kernels): nonstationary (VaryingAmplitude, Gibbs) and space–time (Gneiting) kernels

[thomaspinder]

Summary

Three kernels for climate fields, as agreed in the design on #812. Fixes #812.

 gpjax/kernels/
+├── location_functions.py        # AbstractLocationFunction, Constant, Linear
 ├── nonstationary/
+│   ├── varying_amplitude.py     # k = σ(x) σ(y) k0(x, y)
+│   └── gibbs.py                 # lengthscale ℓ(x) via input scaling of k0
 ├── stationary/
+│   ├── gneiting.py              # nonseparable space–time kernel
 │   └── base.py                  # + isotropic_radial: ClassVar[bool]
 └── approximations/rff.py        # clearer error for kernels with no spectral density
+GLOSSARY.md
+docs/adr/0001-location-functions.md
+docs/examples/nonstationary_terrain.py
+docs/examples/spacetime_temperature.py

A location function gives a kernel parameter that changes with location. It evaluates one point, selects its own columns with active_dims, and returns a log value. It is a separate type from mean functions; ADR 0001 gives the reason.

kernel = gpx.kernels.VaryingAmplitude(
    gpx.kernels.Gibbs(
        gpx.kernels.Matern32(active_dims=[0, 1]),                    # distance columns
        lengthscale=location_functions.Linear(active_dims=[2]),      # covariate column
    ),
    amplitude=location_functions.Linear(active_dims=[2]),
)
kernel = gpx.kernels.Gneiting(space_dims=[1, 2], time_dim=0)

Gibbs needs no refactor of the stationary kernels:

gibbs(x, y)
  a, b   = ℓ-function(x), ℓ-function(y)           # log lengthscales
  c      = sqrt(2 / (e^{2a} + e^{2b}))
  return  (2 e^{a+b} / (e^{2a} + e^{2b}))^{d/2} * base_kernel(c·x, c·y)

vmap batches only the operations that depend on the mapped input, so a location function runs once for each row ($N + M$ times for an $N \times M$ matrix), also inside sum and product kernels.

Evidence

Tests: uv run poe test gives 3304 passed and 1 skipped. Lint and doctests pass. The CI docs build (-E -W) gives 0 warnings.

• Gibbs matches a direct Paciorek–Schervish formula with an ARD Matérn-3/2 base, to 1e-15.
• Hypothesis checks that the Gram matrices are positive definite for random weights and parameters.
• A constant location function gives exactly the base kernel.
• Gneiting with β = 0 equals the separable product. At lag 0 it equals a powered exponential kernel.
• Invalid bases and invalid columns raise errors. jit, grad and fit work.

Colorado precipitation normals, 247 stations (5-fold cross-validation):

Model	Log ML	Held-out NLPD	90% coverage, mountains / plains	RMSE (log)
Before: stationary	−283.1	1.095	0.82 / 0.99	0.196
After: Gibbs	−217.9	0.876	0.87 / 0.82	0.174

Europe, daily temperature anomalies, January–February 2019 (held-out gap of 140 points):

Model	Log ML	Held-out NLPD	RMSE (K)
Before: separable (β = 0)	−1845.0	1.033	2.17
After: nonseparable (β = 1.00)	−1799.4	0.996	2.09

On Irish wind (Haslett & Raftery, 1989), the profile likelihood over β was flat (Δ ≈ 1.5), because the interaction there comes from advection, which a symmetric kernel cannot model. The notebook states this limit.

Merge Danger

Door: two-way

All the changes add new API. The only change to existing behaviour is the text of the RFF TypeError, and RFF's base_kernel parameter annotation, which is now AbstractKernel so that the clear message is reached under beartype.

Blast Radius: small

• No existing kernel computes anything differently. Stationary kernels gain only a class flag.
• feat(xarray): input transforms, chunked GridSpec.predict, and a surface-temperature infilling example #808 also adds functions to docs/examples/data/_pull_reference_datasets.py, so the second PR to merge gets a small conflict there.
• The docs build runs the two new notebooks: about 20 s and 1 min 40 s.

Follow-up: #813 (constrain PoweredExponential.power).

🤖 Generated with Claude Code

[github-actions]

📖 Docs preview: https://pr-815--endearing-crepe-c2d5fe.netlify.app

Smoke render — the expensive notebooks run with reduced budgets, so
figures are not publication fidelity. /render-mode.txt says smoke.