Kernels and means
Built-in kernel families
Standard kernels accept variance=1.0 and lengthscale=1.0 by default.
variance is the covariance amplitude. A scalar lengthscale gives isotropic
scaling, and a vector supplies one lengthscale per input coordinate (ARD).
import jax.numpy as jnp
from gpex.kernels import RBFKernel, Matern52Kernel
isotropic = RBFKernel(variance=1.0, lengthscale=0.7)
ard = Matern52Kernel(variance=1.0, lengthscale=jnp.array([0.7, 1.2]))
| Class | Kernel family | Supported GP derivative order |
|---|---|---|
Matern12Kernel |
Matérn, ν = 1/2 | 0 |
Matern32Kernel |
Matérn, ν = 3/2 | 1 |
Matern52Kernel |
Matérn, ν = 5/2 | 2 |
RBFKernel |
Squared exponential | Any finite order |
All four families support callable posterior paths.
Kernel sums and products, constructed with + and *, support conditioning,
prediction, and finite-query sampling. They do not currently provide features
for callable paths.
Full input metrics
The Generalized...Kernel variants of these four families accept a complete
lower-triangular factor L with a positive diagonal. It defines the precision
metric M = L @ L.T, with distance norm(L.T @ (x - y)):
from gpex.kernels import GeneralizedMatern52Kernel
L = jnp.array([[1.4, 0.0], [0.3, 0.8]])
kernel = GeneralizedMatern52Kernel(variance=1.0, L=L)
For a diagonal L, its diagonal entries are reciprocal lengthscales.
L = jnp.diag(1 / jnp.array([0.7, 1.2])) reproduces the corresponding ARD metric.
Generalized kernels also support callable paths.
Plain numbers and arrays remain fixed during fitting. Use PositiveParameter
for variance and lengthscales, or CholeskyParameter for L, to make them
trainable. Read values through kernel.variance, kernel.lengthscales.lengthscale,
or kernel.lengthscales.L. See Fitting hyperparameters
and parameters and priors.
Compose kernels and means
Use + and * to compose kernels or means. For example, this two-dimensional
model combines a trainable full input metric with a product of two fixed RBF
kernels and a linear mean:
from gpex import GaussianProcess
from gpex.kernels import GeneralizedRBFKernel
from gpex.means import ConstantMean, LinearMean
from gpex.parameters import CholeskyParameter
from gpex.parameters.priors import IndependentCholeskyPrior, LogNormalPrior, NormalPrior
factor = CholeskyParameter(
jnp.array([[1.0, 0.0], [0.2, 1.4]]),
prior=IndependentCholeskyPrior(
diagonal_prior=LogNormalPrior(log_standard_deviation=0.5),
off_diagonal_prior=NormalPrior(standard_deviation=0.4),
),
name="precision_factor",
)
local = GeneralizedRBFKernel(variance=1.0, L=factor)
broad = RBFKernel(variance=0.3, lengthscale=2.0)
modulation = RBFKernel(lengthscale=jnp.array([1.5, 0.8]))
mean = LinearMean(weights=jnp.array([0.15, -0.1])) + ConstantMean(0.2)
gp = GaussianProcess(kernel=local + broad * modulation, mean=mean)
Here only the precision factor is trainable. The kernel variances, fixed lengthscales, and mean coefficients remain fixed.
Extension interfaces
A Kernel maps two points of shape (d,) to a scalar covariance. A Mean maps
one point to a scalar mean. Both are Equinox modules. Custom kernels must be
positive semidefinite and declare max_derivative_order when the latent GP has
finite smoothness. Evaluation must support the JAX transformations required by
the chosen operators, including derivatives at coincident inputs.
RadialKernel supplies input scaling and evaluation from squared distance.
Subclasses define _correlation and _derivative_order. RadialFeatureKernel
also requires _sample_frequencies for its unit-scale correlation.
SeparableKernel.factorization(input_dimension) returns a
KernelFactorization(coefficient, factors). Each factor is a scalar-coordinate
kernel, and their product times the coefficient must equal the kernel and its
derivatives. Derive factors from current parameter values. Factors with array
state should be Equinox modules. This interface enables
factored tensor quadrature.
Pass the kernel object directly to GaussianProcess to expose this capability.
FeatureKernel.random_features returns a fixed callable mapping one point to
a feature vector. Feature inner products approximate the complete kernel,
including its variance. Store random draws as array fields so repeated calls
use the same features and the map can cross JIT boundaries. Exact finite feature
maps are supported as well. See the class contracts below.
Classes
gpex.kernels
Covariance kernels for scalar latent functions.
Kernel
FeatureKernel
Bases: Kernel
A covariance kernel that can construct a finite feature map.
Feature support is independent of radiality and coordinate separability.
Subclasses implement both covariance evaluation and random_features.
Ordinary GP inference and finite-query sampling require only Kernel.
random_features
abstractmethod
random_features(
key: PRNGKeyArray,
input_dimension: int,
*,
feature_count: int = 1024,
) -> Callable[[Point], Array]
Construct a fixed random feature map approximating this kernel.
The returned JAX-compatible callable maps a point of shape (d,) to
(feature_count,). Its feature inner product approximates the kernel,
including its variance. Reusing the callable reuses the same features.
Features may be random Fourier features or another finite basis. They need not be Fourier or radial. An exact finite feature map is also supported. The callable must store numerical state in dynamic PyTree leaves and support JIT, vmap, and differentiation up to the kernel's supported derivative order.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
key
|
PRNGKeyArray
|
JAX random key for feature construction. |
required |
input_dimension
|
int
|
Number of coordinates in each input point. |
required |
feature_count
|
int
|
Positive number of features. Defaults to 1024. |
1024
|
Source code in src/gpex/kernels/features.py
RadialKernel
RadialKernel(
lengthscales: Lengthscales | GeneralizedLengthscales,
*,
variance: PositiveParameter | Scalar | float,
)
Bases: Kernel
Kernel evaluated from squared distance after lengthscale transformation.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
lengthscales
|
Lengthscales | GeneralizedLengthscales
|
|
required |
variance
|
PositiveParameter | Scalar | float
|
Finite positive scalar amplitude or a scalar |
required |
Subclasses supply _correlation(squared_distance) as a static callable
and set _derivative_order to the largest supported latent derivative
order. The default assumes unlimited smoothness. Subclass
RadialFeatureKernel to additionally provide random Fourier features.
Radiality is with respect to the transformed coordinates and does not imply a product over the original coordinates. Evaluation and numerical state support JIT, vmap, and differentiation.
Source code in src/gpex/kernels/radial.py
variance_parameter
instance-attribute
variance_parameter: PositiveParameter | Scalar = (
variance
if isinstance(variance, PositiveParameter)
else error_if(
value,
jnp.any(~jnp.isfinite(value) | (value <= 0)),
"variance must be finite and positive",
)
)
lengthscales
instance-attribute
lengthscales: Lengthscales | GeneralizedLengthscales = (
lengthscales
)
RadialFeatureKernel
Bases: RadialKernel, FeatureKernel
A radial kernel with a sampler for its normalized spectral distribution.
Subclasses supply _sample_frequencies(key, shape) with
shape = (feature_count, input_dimension). The shared implementation
applies the same lengthscale convention to distances and frequency draws.
Source code in src/gpex/kernels/radial.py
random_features
random_features(
key: PRNGKeyArray,
input_dimension: int,
*,
feature_count: int = 1024,
) -> Callable[[Point], Array]
Construct Fourier features; see
FeatureKernel.random_features.
Source code in src/gpex/kernels/radial.py
Lengthscales
Lengthscales(
lengthscale: PositiveParameter | Lengthscale | float,
)
Bases: Module
Ordinary scalar or ARD lengthscales, stored once as model parameters.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
lengthscale
|
PositiveParameter | Lengthscale | float
|
Finite positive scalar or vector of shape |
required |
Displacements and rows of Fourier frequencies are divided by these values. Construction and numerical operations support JIT, vmap, and differentiation.
Source code in src/gpex/kernels/lengthscales.py
lengthscale_parameter
instance-attribute
lengthscale_parameter: PositiveParameter | Lengthscale = (
lengthscale
if isinstance(lengthscale, PositiveParameter)
else error_if(
value,
jnp.any(~jnp.isfinite(value) | (value <= 0)),
"lengthscale must be finite and positive",
)
)
validate_dimension
transform_displacement
GeneralizedLengthscales
GeneralizedLengthscales(
L: CholeskyParameter | CholeskyFactor,
)
Bases: Module
Generalized lengthscales represented by a lower-triangular precision factor.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
L
|
CholeskyParameter | CholeskyFactor
|
Finite lower-triangular array of shape |
required |
The metric is L @ L.T. Displacements are transformed by L.T,
and rows of Fourier frequencies by right-multiplication with L.T.
A diagonal L contains reciprocal ordinary lengthscales. Off-diagonal
entries allow directions of variation that are not aligned with the axes.
Construction and numerical operations support JIT, vmap, and differentiation.
Source code in src/gpex/kernels/lengthscales.py
L_parameter
instance-attribute
L_parameter: CholeskyParameter | CholeskyFactor = (
L
if isinstance(L, CholeskyParameter)
else validate_cholesky_factor(L)
)
validate_dimension
transform_displacement
SeparableKernel
Bases: Kernel
A kernel with an exact product of factors in the input coordinates.
Subclasses implement factorization(input_dimension). The default
point evaluation multiplies the factors; subclasses may supply an
equivalent direct implementation. A sum of separable kernels is generally
not itself separable and remains a SumKernel.
factorization
abstractmethod
factorization(input_dimension: int) -> KernelFactorization
Return one factor per coordinate and a scalar coefficient.
Validate the requested dimension against kernel parameter shapes. Derive factors from current numerical leaves so differentiation through this representation agrees with point evaluation.
Source code in src/gpex/kernels/separable.py
KernelFactorization
Bases: Module
Represent coefficient * product(factor_j(x_j, y_j)).
coefficient is a scalar array, and factors contains one scalar
two-argument callable per input coordinate, in coordinate order. Factors
need not be stationary. Their numerical state must remain dynamic PyTree
leaves and their evaluation must support JAX transformations.
Construct this representation from the kernel's current values; it does not own additional trainable parameters or approximate the kernel.
RBFKernel
RBFKernel(
*,
variance: PositiveParameter | Scalar | float = 1.0,
lengthscale: PositiveParameter
| Lengthscale
| float = 1.0,
)
Bases: RadialFeatureKernel, SeparableKernel
Squared-exponential kernel with isotropic or ARD lengthscales.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
variance
|
PositiveParameter | Scalar | float
|
Positive scalar covariance amplitude or scalar |
1.0
|
lengthscale
|
PositiveParameter | Lengthscale | float
|
Positive scalar for isotropic distances or vector of shape
|
1.0
|
Uses squared distance sum(((x - y) / lengthscale)**2).
Supports derivatives of every order. Provides random Fourier features.
Source code in src/gpex/kernels/rbf.py
factorization
factorization(input_dimension: int) -> KernelFactorization
Expose scalar RBF factors using the current scalar or ARD lengths.
Source code in src/gpex/kernels/rbf.py
Matern12Kernel
Matern12Kernel(
*,
variance: PositiveParameter | Scalar | float = 1.0,
lengthscale: PositiveParameter
| Lengthscale
| float = 1.0,
)
Bases: RadialFeatureKernel
Matérn-1/2 kernel with scalar or ARD lengthscales.
PositiveParameter shapes, meanings, and defaults match
RBFKernel.
Supports value observations, with no derivative observations.
Fourier features use multivariate Student-t frequencies with 1
degree of freedom.
Source code in src/gpex/kernels/matern.py
Matern32Kernel
Matern32Kernel(
*,
variance: PositiveParameter | Scalar | float = 1.0,
lengthscale: PositiveParameter
| Lengthscale
| float = 1.0,
)
Bases: RadialFeatureKernel
Matérn-3/2 kernel with scalar or ARD lengthscales.
PositiveParameter shapes, meanings, and defaults match
RBFKernel.
Supports values and first derivatives.
Fourier features use multivariate Student-t frequencies with 3
degrees of freedom.
Source code in src/gpex/kernels/matern.py
Matern52Kernel
Matern52Kernel(
*,
variance: PositiveParameter | Scalar | float = 1.0,
lengthscale: PositiveParameter
| Lengthscale
| float = 1.0,
)
Bases: RadialFeatureKernel
Matérn-5/2 kernel with scalar or ARD lengthscales.
PositiveParameter shapes, meanings, and defaults match
RBFKernel.
Supports values, first derivatives, and second derivatives.
Fourier features use multivariate Student-t frequencies with 5
degrees of freedom.
Source code in src/gpex/kernels/matern.py
GeneralizedRBFKernel
GeneralizedRBFKernel(
*,
L: CholeskyParameter | CholeskyFactor,
variance: PositiveParameter | Scalar | float = 1.0,
)
Bases: RadialFeatureKernel
Squared-exponential kernel with a full Mahalanobis metric.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
variance
|
PositiveParameter | Scalar | float
|
Positive scalar covariance amplitude or scalar |
1.0
|
L
|
CholeskyParameter | CholeskyFactor
|
Complete lower-triangular factor of shape |
required |
Uses squared distance ||(L.T @ (x - y))||², with precision matrix
M = L @ L.T. A diagonal L contains reciprocal lengthscales.
Supports derivatives of every order. Provides random Fourier features.
Source code in src/gpex/kernels/rbf.py
GeneralizedMatern12Kernel
GeneralizedMatern12Kernel(
*,
L: CholeskyParameter | CholeskyFactor,
variance: PositiveParameter | Scalar | float = 1.0,
)
Bases: RadialFeatureKernel
Matérn-1/2 kernel with a full precision metric.
PositiveParameter shapes, meanings, and defaults match
GeneralizedRBFKernel.
Supports value observations, with no derivative observations.
Fourier features use multivariate Student-t frequencies with 1
degree of freedom.
Source code in src/gpex/kernels/matern.py
GeneralizedMatern32Kernel
GeneralizedMatern32Kernel(
*,
L: CholeskyParameter | CholeskyFactor,
variance: PositiveParameter | Scalar | float = 1.0,
)
Bases: RadialFeatureKernel
Matérn-3/2 kernel with a full precision metric.
PositiveParameter shapes, meanings, and defaults match
GeneralizedRBFKernel.
Supports values and first derivatives.
Fourier features use multivariate Student-t frequencies with 3
degrees of freedom.
Source code in src/gpex/kernels/matern.py
GeneralizedMatern52Kernel
GeneralizedMatern52Kernel(
*,
L: CholeskyParameter | CholeskyFactor,
variance: PositiveParameter | Scalar | float = 1.0,
)
Bases: RadialFeatureKernel
Matérn-5/2 kernel with a full precision metric.
PositiveParameter shapes, meanings, and defaults match
GeneralizedRBFKernel.
Supports values, first derivatives, and second derivatives.
Fourier features use multivariate Student-t frequencies with 5
degrees of freedom.
Source code in src/gpex/kernels/matern.py
SumKernel
ZeroMean() has no configuration. Mean composition uses
SumMean(left: Mean, right: Mean) and ProductMean(left: Mean, right: Mean),
also available through + and *.
gpex.means
Scalar mean functions and their composition.
Mean
Bases: Module
Base class providing composition for scalar mean functions.
ZeroMean
Bases: Mean
ConstantMean
ConstantMean(value: Parameter | Scalar | float = 0.0)
Bases: Mean
A scalar real floating mean, fixed or explicitly trainable.
Supply values in JAX's default floating precision, such as 0.0.
Source code in src/gpex/means.py
LinearMean
Bases: Mean
Linear mean weights @ point + intercept with explicit parameters.