Package: gpjax-nightly
Quickstart|Install guide|Documentation|Slack CommunityGPJax aims to provide a low-level interface to Gaussian process (GP) models inJax, 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 supportGPJax was founded byThomas Pinder. Today, the maintenance of GPJax is undertaken byThomas PinderandDaniel Dodd.We would be delighted to receive contributions from interested individuals and groups. To learn how you can get involved, please read ourguide for contributing. If you have any questions, we encourage you toopen an issue. For broader conversations, such as best GP fitting practices or questions about the mathematics of GPs, we invite you toopen a discussion.Feel free to join ourSlack Channel, where we can discuss the development of GPJax and broader support for Gaussian process modelling.Supported methods and interfacesNotebook examplesConjugate InferenceClassification with MCMCSparse Variational InferenceBlackJax IntegrationLaplace ApproximationTensorFlow Probability IntegrationInference on Non-Euclidean SpacesInference on GraphsLearning Gaussian Process BarycentresDeep Kernel RegressionNatural GradientsGuides for customisationCustom kernelsUCI regressionConversion between.ipynband.pyAbove examples are stored inexamplesdirectory in the double percent (py:percent) format. Checkoutjupytext using-clifor more info.To convertexample.pytoexample.ipynb, run:jupytext--tonotebookexample.pyTo convertexample.ipynbtoexample.py, run:jupytext--topy:percentexample.ipynbSimple exampleLet us import some dependencies and simulate a toy dataset $\mathcal{D}$.importgpjaxasgpxfromjaximportgrad,jitimportjax.numpyasjnpimportjax.randomasjrimportoptaxasoxkey=jr.PRNGKey(123)f=lambdax:10*jnp.sin(x)n=50x=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)The function of interest here, $f(\cdot)$, is sinusoidal, but our observations of it have been perturbed by Gaussian noise. We aim to utilise a Gaussian process to try and recover this latent function.1. Constructing the prior and posteriorWe begin by defining a zero-mean Gaussian process prior with a radial basis function kernel and assume the likelihood to be Gaussian.prior=gpx.Prior(kernel=gpx.RBF())likelihood=gpx.Gaussian(num_datapoints=n)Similar to how we would write on paper, the posterior is constructed by the product of our prior with our likelihood.posterior=prior*likelihood2. Learning hyperparametersEquipped with the posterior, we seek to learn the model's hyperparameters through gradient-optimisation of the marginal log-likelihood. We this below, adding Jax'sjust-in-time (JIT)compilation to accelerate training.mll=jit(posterior.marginal_log_likelihood(D,negative=True))For purposes of optimisation, we'll use optax's Adam.opt = ox.adam(learning_rate=1e-3)We define an initial parameter state through theinitialisecallable.parameter_state=gpx.initialise(posterior,key=key)Finally, we run an optimisation loop using the Adam optimiser via thefitcallable.inference_state=gpx.fit(mll,parameter_state,opt,num_iters=500)3. Making predictionsUsing our learned hyperparameters, we can obtain the posterior distribution of the latent function at novel test points.learned_params,_=inference_state.unpack()xtest=jnp.linspace(-3.,3.,100).reshape(-1,1)latent_distribution=posterior(learned_params,D)(xtest)predictive_distribution=likelihood(learned_params,latent_distribution)predictive_mean=predictive_distribution.mean()predictive_cov=predictive_distribution.covariance()InstallationStable versionThe latest stable version of GPJax can be installed viapip:pipinstallgpjaxNoteWe recommend you check your installation version:python-c'import gpjax; print(gpjax.__version__)'Development versionWarningThis 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.gitclonehttps://github.com/JaxGaussianProcesses/GPJax.gitcdGPJax
pythonsetup.pydevelopNoteWe advise you create virtual environment before installing:conda create -n gpjax_experimental python=3.10.0
conda activate gpjax_experimentaland recommend you check your installation passes the supplied unit tests:python-mpytesttests/Citing GPJaxIf you use GPJax in your research, please cite ourJOSS 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}
}