Support multinomial sampling for NUTS, and more

[xukai92]

This PR resolves / will resolve:

• HMC: Multinomial Sampling Instead of Slice Sampling HMC: Multinomial Sampling Instead of Slice Sampling #62
• Get more comments and docstrings Get more comments and docstrings #75
• More robust and modular way of detecting divergence More robust and modular way of detecting divergence  #16
• Return more information for each step Return more information for each step #59
• Add more informative error messages. Add more informative error messages. #71
• Identify divergent transition Identify divergent transition #50
• RFC find_good_eps RFC find_good_eps #27
• Update gradient function to return both value and gradient Update gradient function to return both value and gradient #56

[xukai92]

@yebai Please take a look over the changes. So far I didn't change the transition interface (which means I don't touch the previous AbstractTrajectory stuff) yet. The only main change I did is introducing a tree struct called FullBinaryTree that contains a sampler::AbstractTreeSampler. With the FullBinaryTree I refactored the build_tree function with two tree sampling methods (slic and multinomial dispatched by the sampler.

[xukai92]

Just check the code is still type stable.

[yebai]

Please take a look over the changes. So far I didn't change the transition interface (which means I don't touch the previous AbstractTrajectory stuff) yet. The only main change I did is introducing a tree struct called FullBinaryTree that contains a sampler::AbstractTreeSampler. With the FullBinaryTree I refactored the build_tree function with two tree sampling methods (slic and multinomial dispatched by the sampler.

Thanks, Kai. Below are some issues I found so far.

The current PR contains 2 types for trajectory sampling method (slice/multinomial):

• TreeSampling
• TreeSampler

These 2 types have very similar functionality IMO. Maybe we can consider unify them to slightly reduce clutter and improve readability. Also, the following field names can be unified since they both indicate a trajectory sampling method (i.e., slice/multinomial)

• NUTS.sampling # maybe change this to sampler
• FullBinaryTree.sampler

[xukai92]

The reason I have NUTS.sampling is I don't want to initalize any sampler (which has some parameters based on the actual methods) when constructing NUTS.

One solution is to change NUTS.sampling to NUTS.samplerType::Type{T} where {T<:TreeSampler}. How do you think?

[yebai]

One solution is to change NUTS.sampling to NUTS.samplerType::Type{T} where {T<:TreeSampler}. How do you think?

Yeah, that could work.

[xukai92]

@yebai

Breaking changes

• ∂ℓπ∂θ is now assumed to retunr a (value, graident) tuple
• ∂H∂θ now returns (logprob, -∂ℓπ∂θ) tuple

Almost 2x speed-up by using the cache.

[yebai]

@yebai Also did you see my comments starting with REVIEW in the code? Those are thing that needs double check for the algorithm correctness.

I checked - it looks correct!

[yebai]

We need to update the example in README (and probably Turing at some point) since ∂ℓπ∂θ needs to return both value and gradient now.

[xukai92]

Sure. I will update it.

[xukai92]

Just merged it. Thanks for the help! @yebai

I will finish the rest issues for v0.2 in another PR in the next few days. Those are mostly interface changes and RFC so I hope it won't take as long as this one.