Skip to main content
Yanked

This release has been yanked by its maintainers, and will be ignored by installers, except when explicitly specified.
Consider using release 0.18.0 instead.
Reason given by maintainers: Jaxlib release issue on M1

GPJax's logo

codecov CodeFactor Documentation 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 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

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

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.

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

Note

We advise you create virtual environment before installing:

conda create -n gpjax_experimental python=3.10.0
conda activate gpjax_experimental

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

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.6.tar.gz (55.8 kB view details)

Uploaded Source

Built Distribution

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

gpjax-0.6-py3-none-any.whl (97.9 kB view details)

Uploaded Python 3

File details

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

File metadata

  • Download URL: gpjax-0.6.tar.gz
  • Upload date:
  • Size: 55.8 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: poetry/1.2.2 CPython/3.9.7 Darwin/22.3.0

File hashes

Hashes for gpjax-0.6.tar.gz
Algorithm Hash digest
SHA256 b0a40e89556e8d8c8165d5a904575895a9ec32b97c50a994f56a1aa6f1a4591a
MD5 4c489d237c703c32b7ca328697b4d207
BLAKE2b-256 b7996db1a922636a56eb36c6a58f7c6aaccaef026169a6b350c3726615a6cf89

See more details on using hashes here.

File details

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

File metadata

  • Download URL: gpjax-0.6-py3-none-any.whl
  • Upload date:
  • Size: 97.9 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: poetry/1.2.2 CPython/3.9.7 Darwin/22.3.0

File hashes

Hashes for gpjax-0.6-py3-none-any.whl
Algorithm Hash digest
SHA256 6a2147531d6753cb5b8b9f64222e53c1b6ddb16c2d3adbcc5d3503e7f4cf574d
MD5 16f65222b4d5b4a6658a07355a0dca0d
BLAKE2b-256 4c69a761bb8e9a79402ecfc41244cc33d75a08287a707a43c6111f4140fdcbba

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

0.9.2

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

This release

0.6 This release

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