ConjugatePosterior#

class gpjax.gps.ConjugatePosterior(prior, likelihood, jitter=1e-06)[source]#

Bases: AbstractPosterior[P, GL]

A Conjuate Gaussian process posterior object.

A Gaussian process posterior distribution when the constituent likelihood function is a Gaussian distribution. In such cases, the latent function values \(f\) can be analytically integrated out of the posterior distribution. As such, many computational operations can be simplified; something we make use of in this object.

For a Gaussian process prior \(p(\mathbf{f})\) and a Gaussian likelihood \(p(y | \mathbf{f}) = \mathcal{N}(y\mid \mathbf{f}, \sigma^2))\) where \(\mathbf{f} = f(\mathbf{x})\), the predictive posterior distribution at a set of inputs \(\mathbf{x}\) is given by

\[\begin{split}\begin{aligned} p(\mathbf{f}^{\star}\mid \mathbf{y}) & = \int p(\mathbf{f}^{\star}, \mathbf{f} \mid \mathbf{y})\\ & =\mathcal{N}(\mathbf{f}^{\star} \boldsymbol{\mu}_{\mid \mathbf{y}}, \boldsymbol{\Sigma}_{\mid \mathbf{y}} \end{aligned}\end{split}\]

where

\[\begin{split}\begin{aligned} \boldsymbol{\mu}_{\mid \mathbf{y}} & = k(\mathbf{x}^{\star}, \mathbf{x})\left(k(\mathbf{x}, \mathbf{x}')+\sigma^2\mathbf{I}_n\right)^{-1}\mathbf{y} \\ \boldsymbol{\Sigma}_{\mid \mathbf{y}} & =k(\mathbf{x}^{\star}, \mathbf{x}^{\star\prime}) -k(\mathbf{x}^{\star}, \mathbf{x})\left( k(\mathbf{x}, \mathbf{x}') + \sigma^2\mathbf{I}_n \right)^{-1}k(\mathbf{x}, \mathbf{x}^{\star}). \end{aligned}\end{split}\]

Example

>>> import gpjax as gpx
>>> import jax.numpy as jnp
>>>
>>> prior = gpx.gps.Prior(
...     mean_function = gpx.mean_functions.Zero(),
...     kernel = gpx.kernels.RBF()
... )
>>> likelihood = gpx.likelihoods.Gaussian(num_datapoints=100)
>>>
>>> posterior = prior * likelihood
Parameters:
predict(test_inputs, train_data, *, return_covariance_type='dense')[source]#

Query the predictive posterior distribution.

Conditional on a training data set, compute the GP’s posterior predictive distribution for a given set of parameters. The returned function can be evaluated at a set of test inputs to compute the corresponding predictive density.

The predictive distribution of a conjugate GP is given by

\[\begin{split} p(\mathbf{f}^{\star}\mid \mathbf{y}) & = \int p(\mathbf{f}^{\star} \mathbf{f} \mid \mathbf{y})\\ & =\mathcal{N}(\mathbf{f}^{\star} \boldsymbol{\mu}_{\mid \mathbf{y}}, \boldsymbol{\Sigma}_{\mid \mathbf{y}} \end{split}\]
where
\[\begin{split} \boldsymbol{\mu}_{\mid \mathbf{y}} & = k(\mathbf{x}^{\star}, \mathbf{x})\left(k(\mathbf{x}, \mathbf{x}')+\sigma^2\mathbf{I}_n\right)^{-1}\mathbf{y} \\ \boldsymbol{\Sigma}_{\mid \mathbf{y}} & =k(\mathbf{x}^{\star}, \mathbf{x}^{\star\prime}) -k(\mathbf{x}^{\star}, \mathbf{x})\left( k(\mathbf{x}, \mathbf{x}') + \sigma^2\mathbf{I}_n \right)^{-1}k(\mathbf{x}, \mathbf{x}^{\star}). \end{split}\]

The conditioning set is a GPJax Dataset object, whilst predictions are made on a regular Jax array.

Example

>>> import gpjax as gpx
>>> import jax.numpy as jnp
>>>
>>> xtrain = jnp.linspace(0, 1).reshape(-1, 1)
>>> ytrain = jnp.sin(xtrain)
>>> D = gpx.Dataset(X=xtrain, y=ytrain)
>>> xtest = jnp.linspace(0, 1).reshape(-1, 1)
>>>
>>> prior = gpx.gps.Prior(mean_function = gpx.mean_functions.Zero(), kernel = gpx.kernels.RBF())
>>> posterior = prior * gpx.likelihoods.Gaussian(num_datapoints = D.n)
>>> predictive_dist = posterior(xtest, D)
Parameters:

  • test_inputs (Num[Array, "N D"]) – A Jax array of test inputs at which the predictive distribution is evaluated.

  • train_data (Dataset) – A gpx.Dataset object that contains the input and output data used for training dataset.

  • return_covariance_type (Literal['dense', 'diagonal']) – Literal denoting whether to return the full covariance of the joint predictive distribution at the test_inputs (dense) or just the the standard-deviation of the predictive distribution at the test_inputs.

Returns:

A function that accepts an input array and

returns the predictive distribution as a GaussianDistribution.

Return type:

GaussianDistribution

sample_approx(num_samples, train_data, key, num_features=100)[source]#

Draw approximate samples from the Gaussian process posterior.

Build an approximate sample from the Gaussian process posterior. This method provides a function that returns the evaluations of a sample across any given inputs.

Unlike when building approximate samples from a Gaussian process prior, decompositions based on Fourier features alone rarely give accurate samples. Therefore, we must also include an additional set of features (known as canonical features) to better model the transition from Gaussian process prior to Gaussian process posterior. For more details see [Wilson et. al. (2020)](https://arxiv.org/abs/2002.09309).

In particular, we approximate the Gaussian processes’ posterior as the finite feature approximation \(\hat{f}(x) = \sum_{i=1}^m \phi_i(x)\theta_i + \sum{j=1}^N v_jk(.,x_j)\) where \(\phi_i\) are m features sampled from the Fourier feature decomposition of the model’s kernel and \(k(., x_j)\) are N canonical features. The Fourier weights \(\theta_i\) are samples from a unit Gaussian. See [Wilson et. al. (2020)](https://arxiv.org/abs/2002.09309) for expressions for the canonical weights \(v_j\).

A key property of such functional samples is that the same sample draw is evaluated for all queries. Consistency is a property that is prohibitively costly to ensure when sampling exactly from the GP prior, as the cost of exact sampling scales cubically with the size of the sample. In contrast, finite feature representations can be evaluated with constant cost regardless of the required number of queries.

Parameters:

  • num_samples (int) – The desired number of samples.

  • key (KeyArray) – The random seed used for the sample(s).

  • num_features (int) – The number of features used when approximating the kernel.

  • train_data (Dataset)

Returns:

A function representing an approximate sample from the Gaussian

process prior.

Return type:

FunctionalSample