API¶

Computation engine¶

class tinygp.GaussianProcess(kernel: Kernel, X: JAXArray, *, diag: Union[JAXArray, float] = 0.0, mean: Optional[Union[Callable[[JAXArray], JAXArray], JAXArray]] = None)¶

An interface for designing a Gaussian Process regression model

Parameters

kernel (Kernel) – The kernel function X (JAXArray): The input
is (coordinates. This can be any PyTree that) – compatible with kernel where the zeroth dimension is N_data, the size of the data set.
diag (JAXArray, optional) – The value to add to the diagonal of the covariance matrix, often used to capture measurement uncertainty. This should be a scalar or have the shape (N_data,).
mean – (Callable, optional): A callable or constant mean function that will be evaluated with the X as input: mean(X)

condition(y: tinygp.types.JAXArray) → tinygp.types.JAXArray¶

Condition the process on observed data

Parameters: y (JAXArray) – The observed data. This should have the shape (N_data,), where N_data was the zeroth axis of the X data provided when instantiating this object.
Returns: The marginal likelihood of this model, evaluated at y.

condition_and_predict(y: JAXArray, X_test: Optional[JAXArray] = None, *, kernel: Optional[Kernel] = None, include_mean: bool = True, return_var: bool = False, return_cov: bool = False) → Tuple[JAXArray, Union[JAXArray, Tuple[JAXArray, JAXArray]]]¶

Condition on observed data and return the predictive process

This combines GaussianProcess.condition() and GaussianProcess.predict() into a single operation which will be somewhat more efficient than calling them both separately. See those docstrings for a description of all the arguments.

numpyro_dist(**kwargs)¶: Get the numpyro MultivariateNormal distribution for this process

predict(y: JAXArray, X_test: Optional[JAXArray] = None, *, kernel: Optional[Kernel] = None, include_mean: bool = True, return_var: bool = False, return_cov: bool = False) → Union[JAXArray, Tuple[JAXArray, JAXArray]]¶

Predict the GP model at new test points conditioned on observed data

Parameters

y (JAXArray) – The observed data. This should have the shape (N_data,), where N_data was the zeroth axis of the X data provided when instantiating this object.
X_test (JAXArray, optional) – The coordinates where the prediction should be evaluated. This should have a data type compatible with the X data provided when instantiating this object. If it is not provided, X will be used by default, so the predictions will be made.
include_mean (bool, optional) – If True (default), the predicted values will include the mean function evaluated at X_test.
return_var (bool, optional) – If True, the variance of the predicted values at X_test will be returned.
return_cov (bool, optional) – If True, the covariance of the predicted values at X_test will be returned. If return_var is True, this flag will be ignored.

Returns

The mean of the predictive model evaluated at X_test, with shape (N_test,) where N_test is the zeroth dimension of X_test. If either return_var or return_cov is True, the variance or covariance of the predicted process will also be returned with shape (N_test,) or (N_test, N_test) respectively.

sample(key: jax._src.prng.PRNGKeyArray, shape: Optional[Sequence[int]] = None) → tinygp.types.JAXArray¶

Generate samples from the prior process

Parameters

key – A jax random number key array. shape (tuple, optional): The
to (number and shape of samples) – generate.

Returns

The sampled realizations from the process with shape (N_data,) + shape where N_data is the zeroth dimension of the X coordinates provided when instantiating this process.

Kernels¶

class tinygp.kernels.Kernel¶

The base class for all kernel implementations

This subclass provides default implementations to add and multiply kernels. Subclasses should accept parameters in their __init__ and then override Kernel.evaluate() with custom behavior.

evaluate(X1: tinygp.types.JAXArray, X2: tinygp.types.JAXArray) → tinygp.types.JAXArray¶

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 using Matern32(1.5)(x1, x2), for arrays of input coordinates x1 and x2.
When implementing a custom kernel, this method should treat X1 and X2 as single datapoints. In other words, these inputs will typically either be scalars of have shape n_dim, where n_dim is the number of input dimensions, rather than n_data or (n_data, n_dim), and you should let the Kernel vmap magic handle all the broadcasting for you.

class tinygp.kernels.Custom(function: Callable[[Any, Any], Any])¶

A custom kernel class implemented as a callable

Parameters: function – A callable with a signature and behavior that matches Kernel.evaluate().

class tinygp.kernels.Constant(value: tinygp.types.JAXArray)¶

This kernel returns the constant

\[k(\mathbf{x}_i,\,\mathbf{x}_j) = c\]

where \(c\) is a parameter.

Parameters: c – The parameter \(c\) in the above equation.

class tinygp.kernels.DotProduct¶

The dot product kernel

\[k(\mathbf{x}_i,\,\mathbf{x}_j) = \mathbf{x}_i \cdot \mathbf{x}_j\]

with no parameters.

class tinygp.kernels.Polynomial(*, order: tinygp.types.JAXArray, scale: tinygp.types.JAXArray = DeviceArray(1.0, dtype=float32), sigma: tinygp.types.JAXArray = DeviceArray(0.0, dtype=float32))¶

A polynomial kernel

\[k(\mathbf{x}_i,\,\mathbf{x}_j) = [(\mathbf{x}_i / \ell) \cdot (\mathbf{x}_j / \ell) + \sigma^2]^P\]

Parameters

order – The power \(P\).
scale – The parameter \(\ell\).
sigma – The parameter \(\sigma\).

class tinygp.kernels.Exp(scale: tinygp.types.JAXArray = DeviceArray(1.0, dtype=float32))¶

The exponential kernel

\[k(\mathbf{x}_i,\,\mathbf{x}_j) = \exp(-r)\]

where

\[r = ||(\mathbf{x}_i - \mathbf{x}_j) / \ell||_1\]

Parameters: scale – The parameter \(\ell\).

class tinygp.kernels.ExpSquared(scale: tinygp.types.JAXArray = DeviceArray(1.0, dtype=float32))¶

The exponential squared or radial basis function kernel

\[k(\mathbf{x}_i,\,\mathbf{x}_j) = \exp(-r^2 / 2)\]

where

\[r^2 = ||(\mathbf{x}_i - \mathbf{x}_j) / \ell||_2^2\]

Parameters: scale – The parameter \(\ell\).

class tinygp.kernels.Matern32(scale: tinygp.types.JAXArray = DeviceArray(1.0, dtype=float32))¶

The Matern-3/2 kernel

\[k(\mathbf{x}_i,\,\mathbf{x}_j) = (1 + \sqrt{3}\,r)\,\exp(-\sqrt{3}\,r)\]

where

\[r = ||(\mathbf{x}_i - \mathbf{x}_j) / \ell||_1\]

Parameters: scale – The parameter \(\ell\).

class tinygp.kernels.Matern52(scale: tinygp.types.JAXArray = DeviceArray(1.0, dtype=float32))¶

The Matern-5/2 kernel

\[k(\mathbf{x}_i,\,\mathbf{x}_j) = (1 + \sqrt{5}\,r + 5\,r^2/\sqrt{3})\,\exp(-\sqrt{5}\,r)\]

where

\[r = ||(\mathbf{x}_i - \mathbf{x}_j) / \ell||_1\]

Parameters: scale – The parameter \(\ell\).

class tinygp.kernels.Cosine(period: tinygp.types.JAXArray)¶

The cosine kernel

\[k(\mathbf{x}_i,\,\mathbf{x}_j) = \cos(2\,\pi\,r)\]

where

\[r = ||(\mathbf{x}_i - \mathbf{x}_j) / P||_1\]

Parameters: period – The parameter \(P\).

class tinygp.kernels.ExpSineSquared(*, period: tinygp.types.JAXArray, gamma: tinygp.types.JAXArray)¶

The exponential sine squared or quasiperiodic kernel

\[k(\mathbf{x}_i,\,\mathbf{x}_j) = \exp(-\Gamma\,\sin^2 \pi r)\]

where

\[r = ||(\mathbf{x}_i - \mathbf{x}_j) / P||_1\]

Parameters

period – The parameter \(P\).
gamma – The parameter \(\Gamma\).

class tinygp.kernels.RationalQuadratic(*, alpha: JAXArray, scale: Optional[JAXArray] = None)¶

The rational quadratic

\[k(\mathbf{x}_i,\,\mathbf{x}_j) = (1 + r^2 / 2\,\alpha)^{-\alpha}\]

where

\[r^2 = ||(\mathbf{x}_i - \mathbf{x}_j) / \ell||_2^2\]

Parameters

scale – The parameter \(\ell\).
alpha – The parameter \(\alpha\).

Transforms¶

In tinygp, a “transform” is any callable that takes an input coordinate and returns a transformed coordinate. There are some built in implementations for standard linear transformations that can be used to handle multivariate vector inputs.

class tinygp.transforms.Transform(transform: Callable[[Any], Any], kernel: tinygp.kernels.Kernel)¶

Apply a transformation to the input coordinates of the kernel

Parameters

transform – (Callable): A callable object that accepts coordinates as inputs and returns transformed coordinates.
kernel (Kernel) – The kernel to use in the transformed space.

class tinygp.transforms.Linear(scale: tinygp.types.JAXArray, kernel: tinygp.kernels.Kernel)¶

Apply a linear transformation to the input coordinates of the kernel

For example, the following transformed kernels are all equivalent, but the second supports more flexible transformations:

>>> import numpy as np
>>> from tinygp import kernels, transforms
>>> kernel0 = kernels.Matern32(4.5)
>>> kernel1 = transforms.Linear(1.0 / 4.5, kernels.Matern32())
>>> np.testing.assert_allclose(
...     kernel0.evaluate(0.5, 0.1), kernel1.evaluate(0.5, 0.1)
... )

Parameters

scale (JAXArray) – A 0-, 1-, or 2-dimensional array specifying the scale of this transform.
kernel (Kernel) – The kernel to use in the transformed space.

class tinygp.transforms.Cholesky(factor: tinygp.types.JAXArray, kernel: tinygp.kernels.Kernel, *, lower: bool = True)¶

Apply a Cholesky transformation to the input coordinates of the kernel

For example, the following transformed kernels are all equivalent, but the second supports more flexible transformations:

>>> import numpy as np
>>> from tinygp import kernels, transforms
>>> kernel0 = kernels.Matern32(4.5)
>>> kernel1 = transforms.Cholesky(4.5, kernels.Matern32())
>>> np.testing.assert_allclose(
...     kernel0.evaluate(0.5, 0.1), kernel1.evaluate(0.5, 0.1)
... )

Parameters

factor (JAXArray) – A 0-, 1-, or 2-dimensional array specifying the Cholesky factor. If 2-dimensional, this must be a lower or upper triangular matrix as specified by lower, but this is not checked.
kernel (Kernel) – The kernel to use in the transformed space.
lower – (bool, optional): Is factor lower (vs upper) triangular.

classmethod from_parameters(diagonal: tinygp.types.JAXArray, off_diagonal: tinygp.types.JAXArray, kernel: tinygp.kernels.Kernel) → Cholesky¶

Build a Cholesky transform with a sensible parameterization

Parameters

diagonal (JAXArray) – An (ndim,) array with the diagonal elements of factor. These must be positive, but this is not checked.
off_diagonal (JAXArray) – An ((ndim - 1) * ndim,) array with the off-diagonal elements of factor.
kernel (Kernel) – The kernel to use in the transformed space.

class tinygp.transforms.Subspace(axis: Union[Sequence[int], int], kernel: tinygp.kernels.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.

tinygp

API

Contents

API¶

Computation engine¶

Kernels¶

Transforms¶