Closed JaimeRZP closed 1 year ago
I have now coded the changes into AdvancedHMC src. Users should be able to sample a Turing model using AdvancedHMC directly by using:
n_samples, n_adapts = 10_000, 1_000
sample(model, metric, proposal, initial_θ, n_samples, adaptor, n_adapts)
or even simpler
samples, stats = sample(model, 0.1, 0.95, n_samples, n_adapts; initial_θ=initial_θ)
which will use NUTS by defautl.
The unbounded to bounded space transforms are still missing.
Hi All!
My name is Jaime Ruiz Zapatero and I am 3rd year PhD in astrophysics at Oxford.
@yebai and I have been discussing possible ways of interfacing Turing and AdvancedHMC in a more general without the current glue code currently present inside Turing. This motivated by the limitations of the current interface already discussed in this PR by @sethaxen. The fundamental idea is to extract a
LogDensityProblem
object from a generic Turing model which then is used to build aHamiltonian
object for AdvancedHMC. This can be done as follows:Where model is a conditioned Turing model.
I have written a Neal's funnel example of how this interface can work which you can find in this notebook. Here I have also drafted some place holder functions to wrap the interface such that it is user-friendly. Essentially, I have added an additional signature to
sample
and created a wrapper structure for the sampler ingridients. To do so I have tried to follow the idea proposed in this issue by @yebai.This example is already fully functional. However, it is missing an important aspect. Turing models can have priors with hard boundaries. However, samplers work best in continuous spaces. The solution is map the bounded prior space to the unbounded sampling space. Turing does this internally but to avoid relying too much on Turing one can also do the following:
Then the generated samples can be transformed back to the prior space using:
Note that these already account for the jacobian of the transformation. Also, this might not be the most elegant way of doing the transformation.
Question: Where would you suggest writing these functions into the code? My guess would be within sampler but I don't want to mess around with the design idea.
Question: Where would you suggest incorporating the transformation functions?
Once we have interfaced with the internal sampling method, we should also be able to use
AbstractMCMC
to do the sample following what Turing does. I already have written something like this for a micro-canonical HMC sampler I have been working on, MicrocanonicalHMC.jl. The most involved step is to overload theAbstractMCMC.step
function with AdvancedHMC equivalent.All the best, Jaime