Changes for DynamicHMC.jl and related packages

By: Tamás K. Papp

Re-posted from: https://tamaspapp.eu/post/dynamichmc-reorganization/

TL;DR: new wrapper types for AD (with support for Flux.jl), new transformation interface, see the examples.

DynamicHMC.jl is Julia package for Hamiltonian MCMC using the NUTS algorithm. Its interface and the related packages underwent some reorganization which involves API changes. This post is a brief summary, as I prefer to maintain the actual examples in a repository I can test. You are encouraged to read these, and the documentation for each package (some of which is WIP).

A summary of the changes:

  1. TransformVariables.jl replaces the previous DSL for transforming parameters (which was ContinuousTransformations.jl). It has a different syntax, and should make it easier to specify custom transformations in the future (this is work in progress).

    For example,

    as((μ = asℝ, σ = asℝ₊))

    will specify a transformation of 2 real numbers to a NamedTuple with fields μ (unconstrained) and σ (constrained to be positive, eg a standard deviation).

  2. LogDensityProblems.jl contains the wrapper types for log density functions \(\mathbb{R}^n \to \mathbb{R}\). The most common workflow is making your function map a tuple of parameters to a ::Real (which may be -Inf), define a transformation, and then obtain the gradient via automatic differentiation. Currently ForwardDiff.jl (suboptimal for gradients, but very robust) and Flux.jl (faster) are supported. Support for Zygote.jl and ReverseDiff.jl should be easy to add, just open an issue.

    For example, once you have specified a callable p that takes the transformed parameters which result from transformation t,

    P = TransformedLogDensity(t, p)
    ∇P = FluxGradientLogDensity(P)

    will wrap them so that Flux.jl is used for obtaining the gradient via AD.

  3. DynamicHMC.jl expects a callable that supports the dimension query function, takes a vector of reals, and returns a LogDensityProblems.ValueGradient type. The recommended way of obtaining this is using the wrapper libraries above, but you are of course free to code your own.

    To continue the example, obtaining samples is now as simple as

    chain, NUTS_tuned = NUTS_init_tune_mcmc(∇P, 1000);

    because ∇P now contains the dimension information.

All of the above packages are registered. The new version for DynamicHMC.jl is v1.0.0, as the API change is breaking. All of the packages require at least Julia v0.7, but v1.0 or later is recommended.

DynamicHMCExamples.jl has some examples on suggested usage. I am happy to provide support with coding specific models, just please open an issue there. As usual, bug reports, feature requests and PRs are welcome.

Finally, I would like to thank users for their patience and feedback. I hope that you will find new interface more convenient, and the improved speed useful.

PS.: I have removed the Disqus comments from the blog; they broke the style and hardly anyone was using them anyway.