Repository navigation
feat(kernels): nonstationary kernels for climate data (amplitude, Gibbs/Paciorek, Gneiting) #812
Description
Feature Request
Describe the Feature Request
Climate fields are not stationary. Variability is larger over land than over sea, correlation lengths are shorter over mountains, and on the globe the structure changes with latitude. GPJax's nonstationary kernels (Linear, Polynomial, ArcCosine) are not the ones climate scientists use. The xarray example (#808) shows its largest errors at coastlines and over the Andes, which is a stationary kernel's typical failure.
Describe Preferred Solution
In order of priority:
-
Changing amplitude.
$k(x, y) = \sigma(x),\sigma(y),k_0(x, y)$ , where$\sigma(\cdot)$ is a positive function, for example a softplus or exp of a mean function or of a linear model in covariates. It is a wrapper around any kernel. -
Gibbs / Paciorek–Schervish kernel. The lengthscale changes with location. For an isotropic
$\ell(x)$ in$d$ dimensions:
$$k(x, y) = \left(\frac{2,\ell(x),\ell(y)}{\ell(x)^2 + \ell(y)^2}\right)^{d/2} k_0!\left(\frac{\sqrt{2},\lVert x - y\rVert}{\sqrt{\ell(x)^2 + \ell(y)^2}}\right).$$ - This is valid for any isotropic correlation that is positive definite on
$\mathbb{R}^d$ for all$d$ : RBF, the Matérn family, rational quadratic (Paciorek & Schervish, 2006). - An anisotropic version with a matrix
$\Sigma(x)$ could come later. -
$\log \ell(x)$ can be a parametric function of covariates (Risser & Calder, 2015), or a latent GP (Heinonen et al., 2016) as a later step.
- This is valid for any isotropic correlation that is positive definite on
-
Axially symmetric use on the sphere. On the
UnitSpherecolumns, make$\ell$ and$\sigma$ functions of latitude only (Jun & Stein, 2007, 2008). A kernel that is positive definite on$\mathbb{R}^3$ stays positive definite on the sphere. This could be documentation and a helper, not a new kernel. -
Gneiting space–time kernel (Gneiting, 2002). This is the standard nonseparable space–time covariance:
$$C(h, u) = \frac{\sigma^2}{\psi(u^2)^{d/2}},\varphi!\left(\frac{\lVert h\rVert^2}{\psi(u^2)}\right),$$
with$\varphi$ completely monotone and$\psi$ positive with a completely monotone derivative, for example$\varphi(t) = e^{-c t^{\gamma}}$ and$\psi(t) = (a t^{\alpha} + 1)^{\beta}$ . Strictly it is stationary, but it is often requested with the kernels above.
API sketch:
kernel = gpx.kernels.Gibbs(
base=gpx.kernels.Matern32(),
lengthscale_fn=gpx.mean_functions.Linear(...), # log-lengthscale over inputs
active_dims=[0, 1, 2], # where the distance is taken
)
kernel = gpx.kernels.Amplitude(base=kernel, amplitude_fn=...)Open design point: the lengthscale or amplitude function often depends on covariates (elevation, land fraction), while the distance is taken over spatial columns. The kernel must read different columns for these two roles, for example with active_dims for the distance and a separate column selection for the functions.
Describe Alternatives
- Input warping / deep kernels. Already possible, and it covers the spatial-deformation method (Sampson & Guttorp, 1992). It is less interpretable, and it is harder to constrain.
- SPDE/INLA nonstationary Matérn (Lindgren, Rue & Lindström, 2011). Very common in climate work, but it is a separate solver. Consider it after the Vecchia work.
- Mixtures of local stationary processes (Fuentes, 2001). This is possible with sums of kernels and weights, but it has boundary artefacts.
Related Code
gpjax/kernels/base.pygpjax/kernels/stationary/gpjax/kernels/computations/gpjax/mean_functions.pygpjax/xarray.py(UnitSphere)
Additional Context
Acceptance criteria:
- Positive-definiteness tests (hypothesis) for random lengthscale functions.
- Gibbs reduces to the base kernel when
$\ell$ is constant. - The kernels work with
jit,gradand the existing compute engines. - An example on real data: a covariate-driven lengthscale over complex terrain, or latitude-dependent structure on the globe.
References:
- Gibbs (1997), PhD thesis, University of Cambridge.
- Paciorek & Schervish (2006), Spatial modelling using a new class of nonstationary covariance functions, Environmetrics 17:483–506.
- Risser & Calder (2015), Regression-based covariance functions for nonstationary spatial modeling, Environmetrics 26:284–297.
- Heinonen et al. (2016), Non-stationary Gaussian process regression with Hamiltonian Monte Carlo, AISTATS.
- Jun & Stein (2007), An approach to producing space–time covariance functions on spheres, Technometrics 49:468–479; Jun & Stein (2008), Nonstationary covariance models for global data, Annals of Applied Statistics 2:1271–1289.
- Gneiting (2002), Nonseparable, stationary covariance functions for space–time data, JASA 97:590–600.
If the feature request is approved, would you be willing to submit a PR?
Yes
Activity
thomaspinder commentedon Oct 4, 2026
Agreed design
A design review settled the API below. The first PR implements all three kernels.
Findings from the code survey
__call__(x, y) -> scalaron single points) is sufficient. All compute engines vmap over__call__, so no existing code has to change.This matches a direct Paciorek–Schervish formula (ARD
Matern32, 300 points) to 1.3e-15, and ARD base kernels work.__call__runs only once per row (vmapbatches only the operations that depend on the mapped input. This is also true inside sum and product kernels. A special compute engine is therefore not necessary.gpjax.mean_functionshas noLinear. The API sketch in the issue body is replaced by the design below.Location functions —
gpjax/kernels/location_functions.pyA new abstract type, separate from mean functions (ADR 0001):
(D,)and returns a scalar on the log scale. The kernel appliesexp.active_dims.Constant: one log value. A constant function gives back the stationary kernel exactly.Linear:active_dimsis required, the weights start at zero (the model starts stationary), and there is no intercept by default (the base kernel holds the overall scale).Kernels
VaryingAmplitude(base, amplitude)inkernels/nonstationary/:Gibbs(base, lengthscale)inkernels/nonstationary/: input scaling as above. The base must have the new class flagisotropic_radial(RBF, Matern12/32/52, RationalQuadratic, PoweredExponential). Other bases raiseTypeError.active_dimsmust stayslice(None). The base kernel selects the distance columns, and the location function selects the covariate columns.Gneitinginkernels/stationary/, a subclass ofAbstractKernel(two lengthscales, no closed-form spectral density). Gneiting (2002, eq. 14):Fields:
space_dims: list[int]andtime_dim: int. They are column indices of any array, not only of xarray data.Other changes
sample_approxandRFFgive a clearer error for kernels that have no spectral density.PoweredExponential.poweris a separate issue.Tests
jit,gradandfitwork.VaryingAmplitudeandGibbsreduce to the base kernel with a constant function. Gibbs matches the direct Paciorek–Schervish formula. Invalid bases raiseTypeError.time_diminspace_dimsraises an error.Examples (Applied modelling)
thomaspinder commentedon Oct 4, 2026
Update: implementation in #815, with one change to the examples
The space–time example does not use the Irish wind data, as the design above said. On that data the Gneiting kernel gave no benefit over its separable version. The profile log marginal likelihood over β ∈ [0, 1] was flat (Δ ≈ 1.5 in two different 180-day windows), and the free fit went to β ≈ 0. The empirical correlations show an interaction, but it comes from westerly advection, which is asymmetric, and the Gneiting family is symmetric (Gneiting, Genton & Guttorp, 2007).
We tested other candidates:
The example uses the 2019 European winter. The physical reason is that large anomalies persist longer than small ones, which is a symmetric interaction that the Gneiting kernel can represent. The notebook states the advection limit.
The other parts of the design are implemented as agreed.