GPs

AbstractPrior

AbstractPrior(kernel, mean_function, jitter=1e-06)

Bases: Module, Generic[M, K]

Abstract Gaussian process prior.

Parameters:

• kernel (K) –

  kernel object inheriting from AbstractKernel.

• mean_function (M) –

  mean function object inheriting from AbstractMeanFunction.

__call__

__call__(*args, **kwargs)

Evaluate the Gaussian process at the given points.

The output of this function is a TensorFlow probability distribution from which the the latent function's mean and covariance can be evaluated and the distribution can be sampled.

Under the hood, __call__ is calling the objects predict method. For this reasons, classes inheriting the AbstractPrior class, should not overwrite the __call__ method and should instead define a predict method.

Parameters:

• *args (Any, default: () ) –

  The arguments to pass to the GP's predict method.

• **kwargs (Any, default: {} ) –

  The keyword arguments to pass to the GP's predict method.

Returns:

• GaussianDistribution ( GaussianDistribution ) –

  A multivariate normal random variable representation of the Gaussian process.

predict abstractmethod

predict(*args, **kwargs)

Evaluate the predictive distribution.

Compute the latent function's multivariate normal distribution for a given set of parameters. For any class inheriting the AbstractPrior class, this method must be implemented.

Parameters:

• *args (Any, default: () ) –

  Arguments to the predict method.

• **kwargs (Any, default: {} ) –

  Keyword arguments to the predict method.

Returns:

• GaussianDistribution ( GaussianDistribution ) –

  A multivariate normal random variable representation of the Gaussian process.

Prior

Prior(kernel, mean_function, jitter=1e-06)

Bases: AbstractPrior[M, K]

A Gaussian process prior object.

The GP is parameterised by a mean and kernel function.

A Gaussian process prior parameterised by a mean function \(m(\cdot)\) and a kernel function \(k(\cdot, \cdot)\) is given by \(p(f(\cdot)) = \mathcal{GP}(m(\cdot), k(\cdot, \cdot))\).

To invoke a Prior distribution, a kernel and mean function must be specified.

Example:

    >>> import gpjax as gpx
    >>> kernel = gpx.kernels.RBF()
    >>> meanf = gpx.mean_functions.Zero()
    >>> prior = gpx.gps.Prior(mean_function=meanf, kernel = kernel)

Parameters:

• kernel (K) –

  kernel object inheriting from AbstractKernel.

• mean_function (M) –

  mean function object inheriting from AbstractMeanFunction.

__call__

__call__(*args, **kwargs)

Evaluate the Gaussian process at the given points.

The output of this function is a TensorFlow probability distribution from which the the latent function's mean and covariance can be evaluated and the distribution can be sampled.

Under the hood, __call__ is calling the objects predict method. For this reasons, classes inheriting the AbstractPrior class, should not overwrite the __call__ method and should instead define a predict method.

Parameters:

• *args (Any, default: () ) –

  The arguments to pass to the GP's predict method.

• **kwargs (Any, default: {} ) –

  The keyword arguments to pass to the GP's predict method.

Returns:

• GaussianDistribution ( GaussianDistribution ) –

  A multivariate normal random variable representation of the Gaussian process.

__mul__

__mul__(other)

Combine the prior with a likelihood to form a posterior distribution.

The product of a prior and likelihood is proportional to the posterior distribution. By computing the product of a GP prior and a likelihood object, a posterior GP object will be returned. Mathematically, this can be described by:

$$p(f(\cdot) \mid y) \propto p(y \mid f(\cdot))p(f(\cdot)),$$

where \(p(y | f(\cdot))\) is the likelihood and \(p(f(\cdot))\) is the prior.

Example:

    >>> import gpjax as gpx
    >>> meanf = gpx.mean_functions.Zero()
    >>> kernel = gpx.kernels.RBF()
    >>> prior = gpx.gps.Prior(mean_function=meanf, kernel = kernel)
    >>> likelihood = gpx.likelihoods.Gaussian(num_datapoints=100)
    >>> prior * likelihood

Args: other (Likelihood): The likelihood distribution of the observed dataset.

Returns Posterior: The relevant GP posterior for the given prior and likelihood. Special cases are accounted for where the model is conjugate.

predict

predict(test_inputs)

Compute the predictive prior distribution for a given set of parameters. The output of this function is a function that computes a TFP distribution for a given set of inputs.

In the following example, we compute the predictive prior distribution and then evaluate it on the interval :math:[0, 1]:

Example:

    >>> import gpjax as gpx
    >>> import jax.numpy as jnp
    >>> kernel = gpx.kernels.RBF()
    >>> mean_function = gpx.mean_functions.Zero()
    >>> prior = gpx.gps.Prior(mean_function=mean_function, kernel=kernel)
    >>> prior.predict(jnp.linspace(0, 1, 100)[:, None])

Parameters:

• test_inputs (Float[Array, 'N D']) –

  The inputs at which to evaluate the prior distribution.

Returns:

• GaussianDistribution ( GaussianDistribution ) –

  A multivariate normal random variable representation of the Gaussian process.

sample_approx

sample_approx(num_samples, key, num_features=100)

Approximate samples from the Gaussian process prior.

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

In particular, we approximate the Gaussian processes' prior as the finite feature approximation \(\hat{f}(x) = \sum_{i=1}^m\phi_i(x)\theta_i\) where \(\phi_i\) are \(m\) features sampled from the Fourier feature decomposition of the model's kernel and \(\theta_i\) are samples from a unit Gaussian.

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.

In the following example, we build 10 such samples and then evaluate them over the interval \([0, 1]\):

For a prior distribution, the following code snippet will build and evaluate an approximate sample.

Example:

    >>> import gpjax as gpx
    >>> import jax.numpy as jnp
    >>> import jax.random as jr
    >>> key = jr.PRNGKey(123)
    >>>
    >>> meanf = gpx.mean_functions.Zero()
    >>> kernel = gpx.kernels.RBF(n_dims=1)
    >>> prior = gpx.gps.Prior(mean_function=meanf, kernel = kernel)
    >>>
    >>> sample_fn = prior.sample_approx(10, key)
    >>> sample_fn(jnp.linspace(0, 1, 100).reshape(-1, 1))

Parameters:

• num_samples (int) –

  The desired number of samples.

• key (KeyArray) –

  The random seed used for the sample(s).

• num_features (int, default: 100 ) –

  The number of features used when approximating the kernel.

Returns:

• FunctionalSample ( FunctionalSample ) –

  A function representing an approximate sample from the Gaussian process prior.

AbstractPosterior

AbstractPosterior(prior, likelihood, jitter=1e-06)

Bases: Module, Generic[P, L]

Abstract Gaussian process posterior.

The base GP posterior object conditioned on an observed dataset. All posterior objects should inherit from this class.

Parameters:

• prior (AbstractPrior) –

  The prior distribution.

• likelihood (AbstractLikelihood) –

  The likelihood distribution.

• jitter (float, default: 1e-06 ) –

  A small constant added to the diagonal of the covariance matrix to ensure numerical stability.

__call__

__call__(*args, **kwargs)

Evaluate the Gaussian process posterior at the given points.

The output of this function is a TFP distribution from which the the latent function's mean and covariance can be evaluated and the distribution can be sampled.

Under the hood, __call__ is calling the objects predict method. For this reasons, classes inheriting the AbstractPrior class, should not overwrite the __call__ method and should instead define a predict method.

Parameters:

• *args (Any, default: () ) –

  The arguments to pass to the GP's predict method.

• **kwargs (Any, default: {} ) –

  The keyword arguments to pass to the GP's predict method.

Returns:

• GaussianDistribution ( GaussianDistribution ) –

  A multivariate normal random variable representation of the Gaussian process.

predict abstractmethod

predict(*args, **kwargs)

Compute the latent function's multivariate normal distribution for a given set of parameters. For any class inheriting the AbstractPrior class, this method must be implemented.

Parameters:

• *args (Any, default: () ) –

  Arguments to the predict method.

• **kwargs (Any, default: {} ) –

  Keyword arguments to the predict method.

Returns:

• GaussianDistribution ( GaussianDistribution ) –

  A multivariate normal random variable representation of the Gaussian process.

ConjugatePosterior

ConjugatePosterior(prior, likelihood, jitter=1e-06)

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{align} 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{align}$$

where

$$\begin{align} \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{align}$$

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:

• prior (AbstractPrior) –

  The prior distribution.

• likelihood (AbstractLikelihood) –

  The likelihood distribution.

• jitter (float, default: 1e-06 ) –

  A small constant added to the diagonal of the covariance matrix to ensure numerical stability.

__call__

__call__(*args, **kwargs)

Evaluate the Gaussian process posterior at the given points.

The output of this function is a TFP distribution from which the the latent function's mean and covariance can be evaluated and the distribution can be sampled.

Under the hood, __call__ is calling the objects predict method. For this reasons, classes inheriting the AbstractPrior class, should not overwrite the __call__ method and should instead define a predict method.

Parameters:

• *args (Any, default: () ) –

  The arguments to pass to the GP's predict method.

• **kwargs (Any, default: {} ) –

  The keyword arguments to pass to the GP's predict method.

Returns:

• GaussianDistribution ( GaussianDistribution ) –

  A multivariate normal random variable representation of the Gaussian process.

predict

predict(test_inputs, train_data)

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 $$ 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}} $$ where $$ \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}). $$

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.

Returns:

• GaussianDistribution ( GaussianDistribution ) –

  A function that accepts an input array and returns the predictive distribution as a GaussianDistribution.

sample_approx

sample_approx(num_samples, train_data, key, num_features=100, solver_algorithm=Cholesky())

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

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) 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, default: 100 ) –

  The number of features used when approximating the kernel.

• solver_algorithm (Optional[Algorithm], default: Cholesky() ) –

  The algorithm to use for the solves of the inverse of the covariance matrix. See the CoLA documentation for which solver to pick. For PSD matrices, CoLA currently recommends Cholesky() for small matrices and CG() for larger matrices. Select Auto() to let CoLA decide. Defaults to Cholesky().

Returns:

• FunctionalSample ( FunctionalSample ) –

  A function representing an approximate sample from the Gaussian

• FunctionalSample –

  process prior.

NonConjugatePosterior

NonConjugatePosterior(prior, likelihood, latent=None, jitter=1e-06, key=jr.PRNGKey(42))

Bases: AbstractPosterior[P, NGL]

A non-conjugate Gaussian process posterior object.

A Gaussian process posterior object for models where the likelihood is non-Gaussian. Unlike the ConjugatePosterior object, the NonConjugatePosterior object does not provide an exact marginal log-likelihood function. Instead, the NonConjugatePosterior object represents the posterior distributions as a function of the model's hyperparameters and the latent function. Markov chain Monte Carlo, variational inference, or Laplace approximations can then be used to sample from, or optimise an approximation to, the posterior distribution.

Parameters:

• prior (AbstractPrior) –

  The prior distribution.

• likelihood (AbstractLikelihood) –

  The likelihood distribution.

• jitter (float, default: 1e-06 ) –

  A small constant added to the diagonal of the covariance matrix to ensure numerical stability.

__call__

__call__(*args, **kwargs)

Evaluate the Gaussian process posterior at the given points.

The output of this function is a TFP distribution from which the the latent function's mean and covariance can be evaluated and the distribution can be sampled.

Under the hood, __call__ is calling the objects predict method. For this reasons, classes inheriting the AbstractPrior class, should not overwrite the __call__ method and should instead define a predict method.

Parameters:

• *args (Any, default: () ) –

  The arguments to pass to the GP's predict method.

• **kwargs (Any, default: {} ) –

  The keyword arguments to pass to the GP's predict method.

Returns:

• GaussianDistribution ( GaussianDistribution ) –

  A multivariate normal random variable representation of the Gaussian process.

predict

predict(test_inputs, train_data)

Query the predictive posterior distribution.

Conditional on a set of training data, 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. Note, to gain predictions on the scale of the original data, the returned distribution will need to be transformed through the likelihood function's inverse link function.

Parameters:

• train_data (Dataset) –

  A gpx.Dataset object that contains the input and output data used for training dataset.

Returns:

• GaussianDistribution ( GaussianDistribution ) –

  A function that accepts an input array and returns the predictive distribution as a dx.Distribution.

construct_posterior

construct_posterior(prior, likelihood)

Utility function for constructing a posterior object from a prior and likelihood. The function will automatically select the correct posterior object based on the likelihood.

Parameters:

• prior (Prior) –

  The Prior distribution.

• likelihood (AbstractLikelihood) –

  The likelihood that represents our beliefs around the distribution of the data.

Returns

AbstractPosterior: A posterior distribution. If the likelihood is
    Gaussian, then a `ConjugatePosterior` will be returned. Otherwise,
    a `NonConjugatePosterior` will be returned.