DirectSolver#

class tinygp.solvers.DirectSolver(X: tinygp.helpers.JAXArray, variance_value: tinygp.helpers.JAXArray, covariance_value: tinygp.helpers.JAXArray, scale_tril: tinygp.helpers.JAXArray)[source]#

Bases: tinygp.solvers.solver.Solver

A direct solver that uses jax’s built in Cholesky factorization

You generally won’t instantiate this object directly but, if you do, you’ll probably want to use the DirectSolver.init() method instead of the usual constructor.

condition(kernel: kernels.Kernel, X_test: Optional[JAXArray], diag: Optional[JAXArray]) → Any[source]#

Compute the covariance matrix for a conditional GP

Parameters

  • kernel – The kernel for the covariance between the observed and predicted data.

  • X_test – The coordinates of the predicted points. Defaults to the input coordinates.

  • diag – Any extra variance to add to the diagonal of the predicted model.

covariance() → tinygp.helpers.JAXArray[source]#

The evaluated covariance matrix

dot_triangular(y: tinygp.helpers.JAXArray) → tinygp.helpers.JAXArray[source]#

Compute a matrix product with the lower triangular linear system

If the covariance matrix is K = L @ L.T for some lower triangular matrix L, this method returns L @ y for some y.

classmethod init(kernel: tinygp.kernels.base.Kernel, X: tinygp.helpers.JAXArray, diag: tinygp.helpers.JAXArray, *, covariance: Optional[Any] = None) → DirectSolver[source]#

Build a DirectSolver for a given kernel and coordinates

Parameters

  • kernel – The kernel function.

  • X – The input coordinates.

  • diag – An extra diagonal component to add to the covariance matrix.

  • covariance – Optionally, a pre-computed array with the covariance matrix. This should be equal to the result of calling kernel and adding diag, but that is not checked.

normalization() → tinygp.helpers.JAXArray[source]#

The multivariate normal normalization constant

This should be (log_det + n*log(2*pi))/2, where n is the size of the covariance matrix, and log_det is the log determinant of the matrix.

solve_triangular(y: tinygp.helpers.JAXArray, *, transpose: bool = False) → tinygp.helpers.JAXArray[source]#

Solve the lower triangular linear system defined by this solver

If the covariance matrix is K = L @ L.T for some lower triangular matrix L, this method solves L @ x = y for some y. If the transpose parameter is True, this instead solves L.T @ x = y.

variance() → tinygp.helpers.JAXArray[source]#

The diagonal of the covariance matrix