GPJax's logo

Quickstart | Install guide | Documentation | Slack Community

GPJax aims to provide a low-level interface to Gaussian process (GP) models in Jax, structured to give researchers maximum flexibility in extending the code to suit their own needs. The idea is that the code should be as close as possible to the maths we write on paper when working with GP models.

Package support

GPJax was founded by Thomas Pinder. Today, the maintenance of GPJax is undertaken by Thomas Pinder and Daniel Dodd.

We would be delighted to receive contributions from interested individuals and groups. To learn how you can get involved, please read our guide for contributing. If you have any questions, we encourage you to open an issue. For broader conversations, such as best GP fitting practices or questions about the mathematics of GPs, we invite you to open a discussion.

Feel free to join our Slack Channel, where we can discuss the development of GPJax and broader support for Gaussian process modelling.

Supported methods and interfaces

Notebook examples

Conjugate Inference

Classification with MCMC

Sparse Variational Inference

Stochastic Variational Inference

BlackJax Integration

Laplace Approximation

Inference on Non-Euclidean Spaces

Inference on Graphs

Pathwise Sampling

Learning Gaussian Process Barycentres

Deep Kernel Regression

Poisson Regression

Guides for customisation

Custom kernels

UCI regression

Conversion between `.ipynb` and `.py`

Above examples are stored in examples directory in the double percent (py:percent) format. Checkout jupytext using-cli for more info.

To convert example.py to example.ipynb, run:

jupytext --to notebook example.py

To convert example.ipynb to example.py, run:

jupytext --to py:percent example.ipynb

Simple example

Let us import some dependencies and simulate a toy dataset $\mathcal{D}$.

import gpjax as gpx
from jax import grad, jit
import jax.numpy as jnp
import jax.random as jr
import optax as ox

key = jr.PRNGKey(123)

f = lambda x: 10 * jnp.sin(x)

n = 50
x = jr.uniform(key=key, minval=-3.0, maxval=3.0, shape=(n,1)).sort()
y = f(x) + jr.normal(key, shape=(n,1))
D = gpx.Dataset(X=x, y=y)

# Construct the prior
meanf = gpx.mean_functions.Zero()
kernel = gpx.kernels.RBF()
prior = gpx.Prior(mean_function=meanf, kernel = kernel)

# Define a likelihood
likelihood = gpx.Gaussian(num_datapoints = n)

# Construct the posterior
posterior = prior * likelihood

# Define an optimiser
optimiser = ox.adam(learning_rate=1e-2)

# Define the marginal log-likelihood
negative_mll = jit(gpx.objectives.ConjugateMLL(negative=True))

# Obtain Type 2 MLEs of the hyperparameters
opt_posterior, history = gpx.fit(
    model=posterior,
    objective=negative_mll,
    train_data=D,
    optim=optimiser,
    num_iters=500,
    safe=True,
    key=key,
)

# Infer the predictive posterior distribution
xtest = jnp.linspace(-3., 3., 100).reshape(-1, 1)
latent_dist = opt_posterior(xtest, D)
predictive_dist = opt_posterior.likelihood(latent_dist)

# Obtain the predictive mean and standard deviation
pred_mean = predictive_dist.mean()
pred_std = predictive_dist.stddev()

Installation

Stable version

The latest stable version of GPJax can be installed via pip:

pip install gpjax

Note

We recommend you check your installation version:
python -c 'import gpjax; print(gpjax.__version__)'

Development version

Warning

This version is possibly unstable and may contain bugs.

Note

We advise you create virtual environment before installing:
conda create -n gpjax_experimental python=3.10.0
conda activate gpjax_experimental

Clone a copy of the repository to your local machine and run the setup configuration in development mode.

git clone https://github.com/JaxGaussianProcesses/GPJax.git
cd GPJax
poetry install

We recommend you check your installation passes the supplied unit tests:
poetry run pytest

Citing GPJax

If you use GPJax in your research, please cite our JOSS paper.

@article{Pinder2022,
  doi = {10.21105/joss.04455},
  url = {https://doi.org/10.21105/joss.04455},
  year = {2022},
  publisher = {The Open Journal},
  volume = {7},
  number = {75},
  pages = {4455},
  author = {Thomas Pinder and Daniel Dodd},
  title = {GPJax: A Gaussian Process Framework in JAX},
  journal = {Journal of Open Source Software}
}

Metadata

Release files for gpjax 0.6.5

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for gpjax 0.6.5
File	Size	Uploaded
gpjax-0.6.5.tar.gz	59.2 kB	Jun 14, 2023	Details

Built distribution (wheel)

Table of built distributions (wheels) for gpjax 0.6.5
File	Interpreter	ABI	Platform	Reset
gpjax-0.6.5-py3-none-any.whl 100.9 kB Jun 14, 2023	Python 3	none	any	Details

Total release size: 160.1 kB

Release files / gpjax-0.6.5.tar.gz

Download URL	gpjax-0.6.5.tar.gz
Size	59.2 kB
Tags	Source
SHA-256 checksum How to use checksums	`34928252b9118a9f8955d35a42379cd37b0e0c0a969f037424cdb0a61c98fd8c`
BLAKE2b-256 checksum How to use checksums	`9ea58de7e4e1016bf3b2ecd714c9b9d4e3e71bc0791fe0bf0d7748aade9e85ab`
Upload date	Jun 14, 2023
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`poetry/1.2.2 CPython/3.9.7 Darwin/22.3.0`

Release files / gpjax-0.6.5-py3-none-any.whl

Download URL	gpjax-0.6.5-py3-none-any.whl
Size	100.9 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`f532462872ee5ef42969688b9c2d02d000eaf8bd8292904f2303e5b88e20a554`
BLAKE2b-256 checksum How to use checksums	`962992241d6f9e554a0bbc0187e8b042736e6112a76e374989149075e69abda9`
Upload date	Jun 14, 2023
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`poetry/1.2.2 CPython/3.9.7 Darwin/22.3.0`