eztaox.kernels.direct

Kernels evaluated using a direct apporach, where the likelihood computation follows O(N^3) scaling.

Classes

MultibandLowRank	A multiband kernel implementating a low-rank Kronecker covariance structure.
MultibandFullRank	A multiband kernel implementating the full-rank Kronecker covariance structure.

Module Contents

class MultibandLowRank

Bases: tinygp.kernels.Kernel

A multiband kernel implementating a low-rank Kronecker covariance structure.

The specific form of the cross-band Kronecker covariance matrix is given by Equation 13 of Gordon et al. (2020). The implementation is inspired by this tinygp tutorial.

amplitudes: jax.numpy.ndarray

kernel: tinygp.kernels.Kernel

coord_to_sortable(X) → tinygp.helpers.JAXArray

evaluate(X1, X2) → tinygp.helpers.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:

1. 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.

2. 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 MultibandFullRank(kernel, diagonal, off_diagonal)

Bases: tinygp.kernels.Kernel

A multiband kernel implementating the full-rank Kronecker covariance structure.

The specific form of the cross-band Kronecker covariance matrix is given by Equation 18-20 of Gordon et al. (2020). The implementation is inspired by this tinygp tutorial.

Note

This kernel is still in development, please use with caution.

core_kernel: tinygp.kernels.Kernel

band_kernel: jax.numpy.ndarray

coord_to_sortable(X) → tinygp.helpers.JAXArray

evaluate(X1, X2) → tinygp.helpers.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:

1. 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.

2. 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.