O6.2 Create SOTA gradient-based Samplers

[odunbar]

Issue

Following work from a summer research project, we have determined that, though not always clear picture, using gradient-based samplers can leverage some performance gains when the gradient (w.r.t the current parameter) is fast to calculate. For example, if we utilize auto-differentiable emulator codes, this is possible. During the project some nice comparison tests were developed for this assessment though are currently stored on a private repository.

The idea of O6.2 is to incorporate the pipeline into src, and the assessements into tests/examples for CES.jl

Current state of the art

We assessed the following during the summer research project

• Samplers: Currently there are many different state-of-art MCMC approaches that are simple to implement and can leverage differentiation of the state
• Autodiff: There are several julia auto-diff packages in various states of use that could be used for these implementations (Zygote, Enzyme, ForwardDiff, BackwardDiff etc.). We
• Compatible emulators: The weakest of the three, many GP emulator codes are not compatible with any of the autodiff packages, often different packages would be compatible with only one of them.

For the summer we settled upon ForwardDiff (and BackwardDiff) compatible with AbstractGPs. The drawback is AbstractGPs does not include the interface for training the GP.

Plan

We wish to create a framework that is easy to extend to new emulators/samplers/autodiff frameworks as this area of software is still in flux
Completed

• Implement the Barker Gradient-based sampler in src/ using ForwardDiff #341
• Implement the useful diagnostics of the chains in src/ which are able to differentiate different methodologies performance (e.g. ESS, ESJD). #342
• Implement the AbstractGP + forward-diff emulator object from provided GP hyperparameters #345
• Add modularity in the choice of autodiff package for the samplers #343
• Add an end-to-end unit test with CES and autodiff emulators #344

TODO

• Implement/improve interface for AbstractGP training
• Comparison test for gradient vs non-gradient samplers into example in the CES repo
• Documentation of the samplers, and new emulators.
• Add non-Barker Variant samplers e.g. NUTS