Subspace#
- class tinygp.transforms.Subspace(axis: Union[Sequence[int], int], kernel: Kernel)[source]#
Bases:
Kernel
A kernel transform that selects a subset of the input dimensions
For example, the following kernel only depends on the coordinates in the second (1-th) dimension:
>>> import numpy as np >>> from tinygp import kernels, transforms >>> kernel = transforms.Subspace(1, kernels.Matern32()) >>> np.testing.assert_allclose( ... kernel.evaluate(np.array([0.5, 0.1]), np.array([-0.4, 0.7])), ... kernel.evaluate(np.array([100.5, 0.1]), np.array([-70.4, 0.7])), ... )
- Parameters:
axis – (Axis, optional): An integer or tuple of integers specifying the axes to select.
kernel (Kernel) – The kernel to use in the transformed space.
- evaluate(X1: tinygp.helpers.JAXArray, X2: tinygp.helpers.JAXArray) tinygp.helpers.JAXArray [source]#
Evaluate the kernel at a pair of input coordinates
This should be overridden be subclasses to return the kernel-specific value. Two things to note:
Users shouldn’t generally call
Kernel.evaluate()
. Instead, always “call” the kernel instance directly; for example, you can evaluate the Matern-3/2 kernel usingMatern32(1.5)(x1, x2)
, for arrays of input coordinatesx1
andx2
.When implementing a custom kernel, this method should treat
X1
andX2
as single datapoints. In other words, these inputs will typically either be scalars of have shapen_dim
, wheren_dim
is the number of input dimensions, rather thann_data
or(n_data, n_dim)
, and you should let theKernel
vmap
magic handle all the broadcasting for you.
- evaluate_diag(X: tinygp.helpers.JAXArray) tinygp.helpers.JAXArray #
Evaluate the kernel on its diagonal
The default implementation simply calls
Kernel.evaluate()
withX
as both arguments, but subclasses can use this to make diagonal calcuations more efficient.