No glue code

Project.toml

@@ -6,10 +6,12 @@ version = "0.4.6"
 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"
 LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1"
+MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
@@ -19,6 +21,16 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
 
+[weakdeps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
+OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
+
+[extensions]
+AdvancedHMCCUDAExt = "CUDA"
+AdvancedHMCMCMCChainsExt = "MCMCChains"
+AdvancedHMCOrdinaryDiffEqExt = "OrdinaryDiffEq"
+
 [compat]
 AbstractMCMC = "4.2"
 ArgCheck = "1, 2"
@@ -37,17 +49,7 @@ StatsBase = "0.31, 0.32, 0.33, 0.34"
 StatsFuns = "0.8, 0.9, 1"
 julia = "1.6"
 
-[extensions]
-AdvancedHMCCUDAExt = "CUDA"
-AdvancedHMCMCMCChainsExt = "MCMCChains"
-AdvancedHMCOrdinaryDiffEqExt = "OrdinaryDiffEq"
-
 [extras]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
-
-[weakdeps]
-CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
-MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
-OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"

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
 
 struct DiffEqIntegrator{T<:AbstractScalarOrVec{<:AbstractFloat},DiffEqSolver} <:

src/abstractmcmc.jl

@@ -25,6 +25,15 @@ struct HMCSampler{K,M,A} <: AbstractMCMC.AbstractSampler
 end
 HMCSampler(kernel, metric) = HMCSampler(kernel, metric, Adaptation.NoAdaptation())
 
+# Convinience constructor
+function NUTSSampler(ϵ::Float64, TAP::Float64, d::Int)
+    metric =  DiagEuclideanMetric(d)
+    integrator = Leapfrog(ϵ)
+    kernel = AdvancedHMC.NUTS{MultinomialTS, GeneralisedNoUTurn}(integrator)
+    adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(TAP, integrator))
+    return HMCSampler(kernel, metric, adaptor)
+end    
+
 """
     HMCState
 
@@ -53,6 +62,73 @@ struct HMCState{
     adaptor::TAdapt
 end
 
+################
+# No glue code #
+################
+function AbstractMCMC.sample(
+    model::DynamicPPL.Model, 
+    sampler::AbstractMCMC.AbstractSampler,
+    N::Integer;
+    progress = true,
+    verbose = false,
+    callback = nothing,
+    kwargs...,
+)
+    return AbstractMCMC.sample(
+        Random.GLOBAL_RNG,
+        model,
+        sampler,
+        N;
+        progress = progress,
+        verbose = verbose,
+        callback = callback,
+        kwargs...,
+    )
+end
+
+function AbstractMCMC.sample(
+    rng::Random.AbstractRNG,
+    model::DynamicPPL.Model,
+    sampler::AbstractMCMC.AbstractSampler,
+    N::Integer;
+    progress = true,
+    verbose = false,
+    callback = nothing,
+    kwargs...,
+)
+    # obtain dimensions of the model
+    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))
+
+    dists = _get_dists(vi)
+    dist_lengths = [length(dist) for dist in dists]
+    vsyms = _name_variables(vi, dist_lengths)
+    d = LogDensityProblems.dimension(ℓ)
+
+    if callback === nothing
+        callback = HMCProgressCallback(N, progress = progress, verbose = verbose)
+        progress = false # don't use AMCMC's progress-funtionality
+    end
+
+    return AbstractMCMC.mcmcsample(
+        rng,
+        AbstractMCMC.LogDensityModel(ℓ),
+        sampler,
+        N;
+        param_names = vsyms,
+        progress = progress,
+        verbose = verbose,
+        callback = callback,
+        kwargs...,
+    )
+end
+###
+
 """
     $(TYPEDSIGNATURES)
 

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