API

Computation engine

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

Kernels

class tinygp.kernels.Kernel
class tinygp.kernels.Sum(kernel1: tinygp.kernels.Kernel, kernel2: tinygp.kernels.Kernel)
class tinygp.kernels.Product(kernel1: tinygp.kernels.Kernel, kernel2: tinygp.kernels.Kernel)
class tinygp.kernels.Constant(value: tinygp.types.JAXArray)
class tinygp.kernels.SubspaceKernel(*, axis: Optional[Axis] = None)
class tinygp.kernels.DotProduct(*, axis: Optional[Axis] = None)
class tinygp.kernels.Linear(*, order: int, sigma: JAXArray, axis: Optional[Axis] = None)
class tinygp.kernels.Polynomial(*, order: int, sigma: JAXArray, axis: Optional[Axis] = None)
class tinygp.kernels.MetricKernel(metric: Optional[Union[Metric, JAXArray]] = None)
class tinygp.kernels.Exp(metric: Optional[Union[Metric, JAXArray]] = None)
class tinygp.kernels.ExpSquared(metric: Optional[Union[Metric, JAXArray]] = None)
class tinygp.kernels.Matern32(metric: Optional[Union[Metric, JAXArray]] = None)
class tinygp.kernels.Matern52(metric: Optional[Union[Metric, JAXArray]] = None)
class tinygp.kernels.Cosine(metric: Optional[Union[Metric, JAXArray]] = None)
class tinygp.kernels.ExpSineSquared(metric: Union[Metric, JAXArray], *, gamma: JAXArray)
class tinygp.kernels.RationalQuadratic(metric: Union[Metric, JAXArray], *, alpha: JAXArray)

Metrics

tinygp.metrics.unit_metric(r: tinygp.types.JAXArray) tinygp.types.JAXArray

Compute the squared norm of a vector

This corresponds to the trivial metric with a unit length scale.

Parameters

r (JAXArray) – The radial difference vector

Returns

The squared difference

Return type

JAXArray

tinygp.metrics.diagonal_metric(ell: tinygp.types.JAXArray) tinygp.metrics.Metric

A diagonal metric with a length scale per dimension

Parameters

ell (JAXArray) – The length scale for each dimension. If this is a scalar, the metric will be isotropic.

tinygp.metrics.dense_metric(cov: tinygp.types.JAXArray, *, lower: bool = True) tinygp.metrics.Metric

A full-rank general metric

The units of the covariance parameter are length^2, unlike the other metrics. Therefore,

dense_metric(jnp.diag(jnp.square(ell)))

and

diagonal_metric(ell)

are equivalent.

Parameters
  • cov (JAXArray) – The covariance matrix metric. This must be positive definite.

  • lower (bool, optional) – Should the lower triangular Cholesky factor be returned?

tinygp.metrics.cholesky_metric(chol: tinygp.types.JAXArray, *, lower: bool = True) tinygp.metrics.Metric

A general metric parameterized by its Cholesky factor

The units of the Cholesky factor are length, unlike the dense metric. Therefore,

cholesky_metric(jnp.diag(ell))

and

diagonal_metric(ell)

are equivalent.

Parameters
  • chol (JAXArray) – The covariance matrix metric. This must be positive definite.

  • lower (bool, optional) – Is chol lower triangular?

tinygp.metrics.compose(*functions: tinygp.metrics.Metric) tinygp.metrics.Metric

A helper function for composing metrics

Mean functions

tinygp.means.zero_mean(X: tinygp.types.JAXArray) tinygp.types.JAXArray
tinygp.means.constant_mean(value: tinygp.types.JAXArray) tinygp.means.Mean