Skip to main content

GPJax's logo

codecov CodeFactor Netlify Status PyPI version DOI Downloads Slack Invite

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 organisation

Contributions

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.

Another way you can contribute to GPJax is through issue triaging. This can include reproducing bug reports, asking for vital information such as version numbers and reproduction instructions, or identifying stale issues. If you would like to begin triaging issues, an easy way to get started is to subscribe to GPJax on CodeTriage.

As a contributor to GPJax, you are expected to abide by our code of conduct. If you feel that you have either experienced or witnessed behaviour that violates this standard, then we ask that you report any such behaviours through this form or reach out to one of the project's gardeners.

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

Governance

GPJax was founded by Thomas Pinder. Today, the project's gardeners are daniel-dodd@, henrymoss@, st--@, and thomaspinder@, listed in alphabetical order. The full governance structure of GPJax is detailed here. We appreciate all the contributors to GPJax who have helped to shape GPJax into the package it is today.

Supported methods and interfaces

Notebook examples

Guides for customisation

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}$.

from jax import config

config.update("jax_enable_x64", True)

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.key(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.gps.Prior(mean_function=meanf, kernel = kernel)

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

# Construct the posterior
posterior = prior * likelihood

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

# Obtain Type 2 MLEs of the hyperparameters
opt_posterior, history = gpx.fit(
    model=posterior,
    objective=gpx.objectives.conjugate_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
hatch env create
hatch shell

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

hatch run dev:test

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}
}

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

gpjax-0.9.2.tar.gz (4.7 MB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

gpjax-0.9.2-py3-none-any.whl (109.5 kB view details)

Uploaded Python 3

File details

Details for the file gpjax-0.9.2.tar.gz.

File metadata

  • Download URL: gpjax-0.9.2.tar.gz
  • Upload date:
  • Size: 4.7 MB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/5.1.1 CPython/3.10.12

File hashes

Hashes for gpjax-0.9.2.tar.gz
Algorithm Hash digest
SHA256 3e6903c37340a6218b95ac52fece85e7bcb2c5301db2c17f14fb0b8f62ff81bd
MD5 58d7cf278323721e5946f8741bfb5ea5
BLAKE2b-256 b52e09da88c48307fb2696c46af66b1889c12d8e2edc788fda7e79ad0e4bc649

See more details on using hashes here.

File details

Details for the file gpjax-0.9.2-py3-none-any.whl.

File metadata

  • Download URL: gpjax-0.9.2-py3-none-any.whl
  • Upload date:
  • Size: 109.5 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/5.1.1 CPython/3.10.12

File hashes

Hashes for gpjax-0.9.2-py3-none-any.whl
Algorithm Hash digest
SHA256 84856eab7ca720cf89551cd48abf69f086d3cfa2aa1e3f4af02975320925a8e7
MD5 dc48b8a3131b24136a8c85131fd2cc9e
BLAKE2b-256 9ad86315637ad61c574b2f235b835aed950d6977e179433e475d0fd63bda6795

See more details on using hashes here.

Release history Release notifications | RSS feed

0.18.0

2 files

0.17.0

2 files

0.15.0

2 files

0.14.0

2 files

0.13.6

2 files

0.13.5

2 files

0.13.4

2 files

0.13.3

2 files

0.13.2

2 files

0.13.1

2 files

0.13.0

2 files

0.12.2

2 files

0.12.0

2 files

0.11.2

2 files

0.11.1

2 files

0.11.0

2 files

0.10.2

2 files

0.10.1

2 files

0.10.0

2 files

0.9.5

2 files

0.9.4

2 files

0.9.3

2 files

This release

0.9.2 This release

2 files

0.9.1

2 files

0.9.0

2 files

0.8.2

2 files

0.8.1

2 files

0.8.0

2 files

0.7.3

2 files

0.7.2

2 files

0.7.1

2 files

0.7.0

2 files

0.6.9

2 files

0.6.8

2 files

0.6.7

2 files

0.6.6

2 files

0.6.5

2 files

0.6.4

2 files

0.6.3

2 files

0.6.2

2 files

0.6.1

2 files

0.6

2 files

0.5.9

2 files

0.5.8

2 files

0.5.7

2 files

0.5.6

2 files

0.5.5

1 file

0.5.4

1 file

0.5.3

1 file

0.5.2

1 file

0.5.1

1 file

0.5.0

1 file

0.4.13

1 file

0.4.12

1 file

0.4.11

1 file

0.4.10

1 file

0.4.9

1 file

0.4.8

1 file

0.4.7

2 files

0.4.6

2 files

0.4.5

2 files

0.4.3

2 files

0.4.2

2 files

0.4.1

2 files

0.4.0

2 files

0.3.8

2 files

0.3.7

2 files

0.3.6

2 files

0.3.5

2 files

0.3.4

2 files

0.3.3

2 files

0.3.2

2 files

0.3.1

2 files

0.3.0

2 files

0.2.0

2 files

0.1.6

2 files

0.1.5

2 files

0.1.4

2 files

0.1.3

2 files

0.1.2

2 files

0.1.1

2 files

0.1.0

2 files

0.0.0

2 files

Anthropic, PBC Visionary sponsor Bloomberg Visionary sponsor Hudson River Trading Visionary sponsor Meta Visionary sponsor NVIDIA Visionary sponsor Microsoft Sustainability sponsor Depot Continuous Integration AWS Cloud computing and Security Sponsor Datadog Monitoring Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page