Package: gpjax
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 organisationContributionsWe 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.Another way you can contribute to GPJax is throughissue
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 tosubscribe to GPJax on CodeTriage.As a contributor to GPJax, you are expected to abide by ourcode 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 throughthis formor reach out to
one of the project'sgardeners.Feel free to join ourSlack
Channel,
where we can discuss the development of GPJax and broader support for Gaussian
process modelling.GovernanceGPJax was founded byThomas Pinder. Today, the
project's gardeners aredaniel-dodd@,henrymoss@,st--@, andthomaspinder@, listed in alphabetical order. The full
governance structure of GPJax is detailedhere. We appreciate allthe contributors to
GPJaxwho have
helped to shape GPJax into the package it is today.Supported methods and interfacesNotebook examplesConjugate InferenceClassification with MCMCSparse Variational InferenceStochastic Variational InferenceBlackJax IntegrationLaplace ApproximationInference on Non-Euclidean SpacesInference on GraphsPathwise SamplingLearning Gaussian Process BarycentresDeep Kernel RegressionPoisson RegressionBayesian OptimisationGuides 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}$.fromjaximportconfigconfig.update("jax_enable_x64",True)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)# Construct the priormeanf=gpx.mean_functions.Zero()kernel=gpx.kernels.RBF()prior=gpx.gps.Prior(mean_function=meanf,kernel=kernel)# Define a likelihoodlikelihood=gpx.likelihoods.Gaussian(num_datapoints=n)# Construct the posteriorposterior=prior*likelihood# Define an optimiseroptimiser=ox.adam(learning_rate=1e-2)# Define the marginal log-likelihoodnegative_mll=jit(gpx.objectives.ConjugateMLL(negative=True))# Obtain Type 2 MLEs of the hyperparametersopt_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 distributionxtest=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 deviationpred_mean=predictive_dist.mean()pred_std=predictive_dist.stddev()InstallationStable versionThe latest stable version of GPJax can be installed via
pip:pipinstallgpjaxNoteWe recommend you check your installation version:python-c'import gpjax; print(gpjax.__version__)'Development versionWarningThis version is possibly unstable and may contain bugs.NoteWe advise you create virtual environment before installing:conda create -n gpjax_experimental python=3.10.0
conda activate gpjax_experimentalClone a copy of the repository to your local machine and run the setup
configuration in development mode.gitclonehttps://github.com/JaxGaussianProcesses/GPJax.gitcdGPJax
poetryinstallWe recommend you check your installation passes the supplied unit tests:poetryrunpytestCiting 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}
}