# Import libraries.
using Turing, Random, LinearAlgebra
= 10
d @model function funnel()
~ Truncated(Normal(0, 3), -3, 3)
θ ~ MvNormal(zeros(d - 1), exp(θ) * I)
z return x ~ MvNormal(z, I)
end
funnel (generic function with 2 methods)
Turing
provides several wrapped samplers from external sampling libraries, e.g., HMC samplers from AdvancedHMC
. These wrappers allow new users to seamlessly sample statistical models without leaving Turing
However, these wrappers might only sometimes be complete, missing some functionality from the wrapped sampling library. Moreover, users might want to use samplers currently not wrapped within Turing
.
For these reasons, Turing
also makes running external samplers on Turing models easy without any necessary modifications or wrapping! Throughout, we will use a 10-dimensional Neal’s funnel as a running example::
# Import libraries.
using Turing, Random, LinearAlgebra
d = 10
@model function funnel()
θ ~ Truncated(Normal(0, 3), -3, 3)
z ~ MvNormal(zeros(d - 1), exp(θ) * I)
return x ~ MvNormal(z, I)
end
funnel (generic function with 2 methods)
Now we sample the model to generate some observations, which we can then condition on.
Users can use any sampler algorithm to sample this model if it follows the AbstractMCMC
API. Before discussing how this is done in practice, giving a high-level description of the process is interesting. Imagine that we created an instance of an external sampler that we will call spl
such that typeof(spl)<:AbstractMCMC.AbstractSampler
. In order to avoid type ambiguity within Turing, at the moment it is necessary to declare spl
as an external sampler to Turing espl = externalsampler(spl)
, where externalsampler(s::AbstractMCMC.AbstractSampler)
is a Turing function that types our external sampler adequately.
An excellent point to start to show how this is done in practice is by looking at the sampling library AdvancedMH
(AdvancedMH
’s GitHub) for Metropolis-Hastings (MH) methods. Let’s say we want to use a random walk Metropolis-Hastings sampler without specifying the proposal distributions. The code below constructs an MH sampler using a multivariate Gaussian distribution with zero mean and unit variance in d
dimensions as a random walk proposal.
MetropolisHastings{RandomWalkProposal{false, ZeroMeanIsoNormal{Tuple{Base.OneTo{Int64}}}}}(RandomWalkProposal{false, ZeroMeanIsoNormal{Tuple{Base.OneTo{Int64}}}}(ZeroMeanIsoNormal(
dim: 10
μ: Zeros(10)
Σ: [1.0 0.0 … 0.0 0.0; 0.0 1.0 … 0.0 0.0; … ; 0.0 0.0 … 1.0 0.0; 0.0 0.0 … 0.0 1.0]
)
))
Sampling is then as easy as:
Chains MCMC chain (10000×11×1 Array{Float64, 3}): Iterations = 1:1:10000 Number of chains = 1 Samples per chain = 10000 Wall duration = 3.97 seconds Compute duration = 3.97 seconds parameters = θ, z[1], z[2], z[3], z[4], z[5], z[6], z[7], z[8], z[9] internals = lp Summary Statistics parameters mean std mcse ess_bulk ess_tail rhat e ⋯ Symbol Float64 Float64 Float64 Float64 Float64 Float64 ⋯ θ -0.4492 0.7256 0.1145 41.9730 42.1825 1.0016 ⋯ z[1] -0.0956 0.5786 0.0577 99.4981 124.3028 1.0412 ⋯ z[2] 0.3996 0.6347 0.0814 59.6526 57.2347 1.0028 ⋯ z[3] -0.4232 0.7127 0.1063 43.4003 47.0137 1.0384 ⋯ z[4] -0.3469 0.7144 0.1058 42.3182 126.7066 1.0754 ⋯ z[5] 0.1390 0.7025 0.0917 59.4241 41.9321 1.0000 ⋯ z[6] -0.5742 0.6169 0.0645 84.3416 103.3731 1.0227 ⋯ z[7] -0.0662 0.6722 0.0705 88.7808 126.7438 1.0225 ⋯ z[8] -0.7906 0.7181 0.0899 68.6921 130.8598 1.0131 ⋯ z[9] -0.0301 0.6534 0.0820 66.2318 53.0874 1.0171 ⋯ 1 column omitted Quantiles parameters 2.5% 25.0% 50.0% 75.0% 97.5% Symbol Float64 Float64 Float64 Float64 Float64 θ -1.8790 -0.9878 -0.4017 0.0111 0.9557 z[1] -1.1404 -0.4890 -0.1671 0.3340 1.1722 z[2] -0.9593 0.0862 0.4185 0.7621 1.6356 z[3] -2.0399 -0.9046 -0.3574 0.0475 0.8702 z[4] -1.9963 -0.7026 -0.3361 0.1223 0.9517 z[5] -0.7854 -0.3842 0.0850 0.5289 1.7770 z[6] -1.9327 -0.8841 -0.4546 -0.2277 0.5470 z[7] -1.2160 -0.5477 -0.1121 0.3895 1.2614 z[8] -2.4609 -1.1767 -0.7001 -0.3970 0.7104 z[9] -1.3750 -0.5860 0.0768 0.3986 1.1421
As previously mentioned, the Turing wrappers can often limit the capabilities of the sampling libraries they wrap. AdvancedHMC
1 (AdvancedHMC
’s GitHub) is a clear example of this. A common practice when performing HMC is to provide an initial guess for the mass matrix. However, the native HMC sampler within Turing only allows the user to specify the type of the mass matrix despite the two options being possible within AdvancedHMC
. Thankfully, we can use Turing’s support for external samplers to define an HMC sampler with a custom mass matrix in AdvancedHMC
and then use it to sample our Turing model.
We can use the library Pathfinder
2 (Pathfinder
’s GitHub) to construct our estimate of mass matrix. Pathfinder
is a variational inference algorithm that first finds the maximum a posteriori (MAP) estimate of a target posterior distribution and then uses the trace of the optimization to construct a sequence of multivariate normal approximations to the target distribution. In this process, Pathfinder
computes an estimate of the mass matrix the user can access. You can see an example of how to use Pathfinder
with Turing in Pathfinder
’s docs.
So far we have used Turing’s support for external samplers to go beyond the capabilities of the wrappers. We want to use this support to employ a sampler not supported within Turing’s ecosystem yet. We will use the recently developed Micro-Cannoncial Hamiltonian Monte Carlo (MCHMC) sampler to showcase this. MCHMC[3,4] ((MCHMC’s GitHub)[https://github.com/JaimeRZP/MicroCanonicalHMC.jl]) is HMC sampler that uses one single Hamiltonian energy level to explore the whole parameter space. This is achieved by simulating the dynamics of a microcanonical Hamiltonian with an additional noise term to ensure ergodicity.
Using this as well as other inference methods outside the Turing ecosystem is as simple as executing the code shown below:
using MicroCanonicalHMC
# Create MCHMC sampler
n_adapts = 1_000 # adaptation steps
tev = 0.01 # target energy variance
mchmc = MCHMC(n_adapts, tev; adaptive=true)
# Sample
chain = sample(model, externalsampler(mchmc), 10_000)
[ Info: Tuning eps ⏳ [ Info: Tuning L ⏳ [ Info: Tuning sigma ⏳ Tuning: 0%|▏ | ETA: 0:06:34 ϵ: 1.4012862813109894 L: 3.1622776601683795 dE/d: -0.09757821953600399 Tuning: 1%|▍ | ETA: 0:04:23 ϵ: 0.7248143150401173 L: 3.2349719817128504 dE/d: 0.019166777701060057 Tuning: 100%|███████████████████████████████████████████| Time: 0:00:02 ϵ: 1.2126555226910607 L: 510.0171864995854 dE/d: 0.0011335833570463194
Chains MCMC chain (10000×11×1 Array{Float64, 3}): Iterations = 1:1:10000 Number of chains = 1 Samples per chain = 10000 Wall duration = 6.92 seconds Compute duration = 6.92 seconds parameters = θ, z[1], z[2], z[3], z[4], z[5], z[6], z[7], z[8], z[9] internals = lp Summary Statistics parameters mean std mcse ess_bulk ess_tail rhat ⋯ Symbol Float64 Float64 Float64 Float64 Float64 Float64 ⋯ θ -1.1937 1.1315 0.0350 975.9356 1006.6128 1.0030 ⋯ z[1] -0.0448 0.5681 0.0144 1642.2792 1735.8553 1.0000 ⋯ z[2] 0.1396 0.5372 0.0146 1427.0191 1552.8275 1.0013 ⋯ z[3] -0.4166 0.6182 0.0164 1629.4635 1517.6993 1.0020 ⋯ z[4] -0.2495 0.5375 0.0120 2204.4125 1992.0178 1.0033 ⋯ z[5] 0.0326 0.4742 0.0119 1735.1601 1682.1140 1.0010 ⋯ z[6] -0.5169 0.6576 0.0160 2051.9667 1712.7199 1.0014 ⋯ z[7] -0.0223 0.5086 0.0112 2191.3493 2160.9334 1.0002 ⋯ z[8] -0.5750 0.6418 0.0172 1727.3383 1617.6530 1.0000 ⋯ z[9] -0.0301 0.5099 0.0127 1716.5405 1729.4446 1.0106 ⋯ 1 column omitted Quantiles parameters 2.5% 25.0% 50.0% 75.0% 97.5% Symbol Float64 Float64 Float64 Float64 Float64 θ -2.9392 -2.1574 -1.2472 -0.3594 1.0409 z[1] -1.2566 -0.3672 -0.0320 0.2947 1.1045 z[2] -0.8938 -0.1928 0.1047 0.4376 1.3074 z[3] -1.9090 -0.7606 -0.2979 0.0077 0.5539 z[4] -1.5372 -0.5263 -0.1885 0.0885 0.6885 z[5] -0.9658 -0.2161 0.0238 0.2790 1.0704 z[6] -2.1539 -0.8403 -0.4013 -0.0672 0.4740 z[7] -1.1549 -0.2965 -0.0122 0.2654 1.0164 z[8] -2.1512 -0.9102 -0.4259 -0.1209 0.3286 z[9] -1.0668 -0.3271 -0.0294 0.2607 1.0454
The only requirement to work with externalsampler
is that the provided sampler
must implement the AbstractMCMC.jl-interface [INSERT LINK] for a model
of type AbstractMCMC.LogDensityModel
[INSERT LINK].
As previously stated, in order to use external sampling libraries within Turing
they must follow the AbstractMCMC
API. In this section, we will briefly dwell on what this entails. First and foremost, the sampler should be a subtype of AbstractMCMC.AbstractSampler
. Second, the stepping function of the MCMC algorithm must be made defined using AbstractMCMC.step
and follow the structure below:
# First step
function AbstractMCMC.step{T<:AbstractMCMC.AbstractSampler}(
rng::Random.AbstractRNG,
model::AbstractMCMC.LogDensityModel,
spl::T;
kwargs...,
)
[...]
return transition, sample
end
# N+1 step
function AbstractMCMC.step{T<:AbstractMCMC.AbstractSampler}(
rng::Random.AbstractRNG,
model::AbstractMCMC.LogDensityModel,
sampler::T,
state;
kwargs...,
)
[...]
return transition, sample
end
There are several characteristics to note in these functions:
There must be two step
functions:
state
, which carries the initialization information.The functions must follow the displayed signatures.
The output of the functions must be a transition, the current state of the sampler, and a sample, what is saved to the MCMC chain.
The last requirement is that the transition must be structured with a field θ
, which contains the values of the parameters of the model for said transition. This allows Turing
to seamlessly extract the parameter values at each step of the chain when bundling the chains. Note that if the external sampler produces transitions that Turing cannot parse, the bundling of the samples will be different or fail.
For practical examples of how to adapt a sampling library to the AbstractMCMC
interface, the readers can consult the following libraries:
Xu et al., AdvancedHMC.jl: A robust, modular and efficient implementation of advanced HMC algorithms, 2019↩︎
Zhang et al., Pathfinder: Parallel quasi-Newton variational inference, 2021↩︎
Robnik et al, Microcanonical Hamiltonian Monte Carlo, 2022↩︎
Robnik and Seljak, Langevine Hamiltonian Monte Carlo, 2023↩︎