No glue code

Project.toml

@@ -6,6 +6,7 @@ version = "0.4.4"
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
+DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8"
 InplaceOps = "505f98c9-085e-5b2c-8e89-488be7bf1f34"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"

src/AdvancedHMC.jl

@@ -23,6 +23,7 @@ using LogDensityProblemsAD: LogDensityProblemsAD
 
 import AbstractMCMC
 using AbstractMCMC: LogDensityModel
+using DynamicPPL
 
 import StatsBase: sample
 
@@ -222,6 +223,9 @@ function Hamiltonian(metric::AbstractMetric, ℓπ, kind::Union{Symbol,Val}; kwa
     return Hamiltonian(metric, ℓ)
 end
 
+### Turing Interface
+include("turing_utils.jl")
+
 ### Init
 
 using Requires

src/sampler.jl

@@ -131,6 +131,68 @@ sample(
     (pm_next!) = pm_next!,
 )
 
+###
+# Allows to pass Turing model to build Hamiltonian
+
+function sample(
+    model::DynamicPPL.Model,
+    metric::AbstractMetric,
+    κ::AbstractMCMCKernel,
+    θ::AbstractVecOrMat{<:AbstractFloat},
+    n_samples::Int,
+    adaptor::AbstractAdaptor = NoAdaptation(),
+    n_adapts::Int = min(div(n_samples, 10), 1_000);
+    drop_warmup = false,
+    verbose::Bool = true,
+    progress::Bool = false,
+    (pm_next!)::Function = pm_next!,
+)
+    ctxt = model.context
+    vi = DynamicPPL.VarInfo(model, ctxt)
+
+    # We will need to implement this but it is going to be
+    # Interesting how to plug the transforms along the sampling
+    # processes
+
+    #vi_t = Turing.link!!(vi, model)
+
+    ℓ = LogDensityProblemsAD.ADgradient(DynamicPPL.LogDensityFunction(vi, model, ctxt))
+    h = Hamiltonian(metric, ℓ)
+    return sample(
+        GLOBAL_RNG,
+        h,
+        κ,
+        θ,
+        n_samples,
+        adaptor,
+        n_adapts;
+        drop_warmup = drop_warmup,
+        verbose = verbose,
+        progress = progress,
+        (pm_next!) = pm_next!,
+    )
+end
+
+function sample(model::DynamicPPL.Model, ϵ::Number, TAP::Number, n_samples::Int, n_adapts::Int;
+                initial_θ=initial_θ, progress=true, kwargs...)
+    ctxt = model.context
+    vi = VarInfo(model, ctxt)
+
+    dists = _get_dists(vi)
+    dist_lengths = [length(dist) for dist in dists]
+    vsyms = _name_variables(vi, dist_lengths)
+    d = length(vsyms)
+
+    metric =  DiagEuclideanMetric(d)
+    integrator = Leapfrog(ϵ)
+    proposal = NUTS{MultinomialTS, GeneralisedNoUTurn}(integrator)
+    adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(TAP, integrator))
+    return sample(model, metric, proposal, initial_θ, n_samples, adaptor, n_adapts;
+                  progress=progress, kwargs...)
+end
+
+###
+
 """
     sample(
         rng::AbstractRNG,

src/turing_utils.jl

@@ -0,0 +1,19 @@
+function _get_dists(vi::VarInfo)
+    mds = values(vi.metadata)
+    return [md.dists[1] for md in mds]
+end
+
+function _name_variables(vi::VarInfo, dist_lengths::AbstractVector)
+    vsyms = keys(vi)
+    names = []
+    for (vsym, dist_length) in zip(vsyms, dist_lengths)
+        if dist_length==1
+            name = [vsym]
+            append!(names, name)
+        else
+            name = [DynamicPPL.VarName(Symbol(vsym, i,)) for i in 1:dist_length]
+            append!(names, name)
+         end
+    end
+    return names
+end