Compatibility with MCMCChains (#309)

* Added MCMCChains connection

* Using bijector in test

* Apply suggestions from code review

Co-authored-by: Seth Axen <[email protected]>

* Apply suggestions from code review

* Update mcmcchains-connect.jl

* Update src/mcmcchains-connect.jl

Co-authored-by: Seth Axen <[email protected]>

Co-authored-by: Seth Axen <[email protected]>
Co-authored-by: Hong Ge <[email protected]>

Project.toml

@@ -33,3 +33,6 @@ StatsBase = "0.31, 0.32, 0.33"
 StatsFuns = "0.8, 0.9, 1"
 UnPack = "1"
 julia = "1"
+
+[extras]
+MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"

src/AdvancedHMC.jl

@@ -195,6 +195,10 @@ function __init__()
     @require CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" begin 
         include("contrib/cuda.jl")
     end
+
+    @require MCMCChains="c7f686f2-ff18-58e9-bc7b-31028e88f75d" begin
+        include("mcmcchains-connect.jl")
+    end
 end
 
 end # module

src/mcmcchains-connect.jl

@@ -0,0 +1,35 @@
+import .MCMCChains: Chains
+
+# A basic chains constructor that works with the Transition struct we defined.
+function AbstractMCMC.bundle_samples(
+    ts::Vector{<:Transition},
+    model::AbstractMCMC.AbstractModel,
+    sampler::AbstractMCMC.AbstractSampler,
+    state,
+    chain_type::Type{Chains};
+    discard_initial=0,
+    thinning=1,
+    param_names=missing,
+    bijector=identity,
+    kwargs...
+)
+    # Turn all the transitions into a vector-of-vectors.
+    t = ts[1]
+    tstat = merge((; lp = t.z.ℓπ.value), stat(t))
+    tstat_names = collect(keys(tstat))
+    vals = [vcat(bijector(t.z.θ), t.z.ℓπ.value, collect(values(stat(t)))) for t in ts]
+
+    # Check if we received any parameter names.
+    if ismissing(param_names)
+        param_names = [Symbol(:param_, i) for i in 1:length(keys(ts[1].z.θ))]
+    else
+        # Generate new array to be thread safe.
+        param_names = Symbol.(param_names)
+    end
+
+    # Bundle everything up and return a Chains struct.
+    return Chains(
+        vals, vcat(param_names, tstat_names), (parameters = param_names, internals = tstat_names,);
+        start=discard_initial + 1, thin=thinning,
+    )
+end

test/Project.toml

@@ -10,6 +10,7 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1"
+MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
@@ -21,4 +22,4 @@ UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
-LogDensityProblemsAD = "1.1.1"
+LogDensityProblemsAD = "1.1.1"

test/mcmcchains.jl

@@ -0,0 +1,35 @@
+using ReTest, Random, AdvancedHMC, ForwardDiff, AbstractMCMC, MCMCChains
+using Statistics: mean
+include("common.jl")
+
+@testset "MCMCChains w/ gdemo" begin
+    rng = MersenneTwister(0)
+
+    n_samples = 5_000
+    n_adapts = 5_000
+
+    θ_init = randn(rng, 2)
+
+    model = AdvancedHMC.LogDensityModel(LogDensityProblemsAD.ADgradient(Val(:ForwardDiff), ℓπ_gdemo))
+    init_eps = Leapfrog(1e-3)
+    κ = NUTS(init_eps)
+    metric = DiagEuclideanMetric(2)
+    adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, κ.τ.integrator))
+
+    samples = AbstractMCMC.sample(
+        rng, model, κ, metric, adaptor, n_adapts + n_samples;
+        nadapts = n_adapts,
+        init_params = θ_init,
+        chain_type = Chains,
+        progress=false,
+        bijector = invlink_gdemo,
+        verbose=false
+    );
+
+    # Transform back to original space.
+    # NOTE: We're not correcting for the `logabsdetjac` here since, but
+    # we're only interested in the mean it doesn't matter.
+    m_est = mean(samples[n_adapts + 1:end]) 
+
+    @test m_est[:,2] ≈ [49 / 24, 7 / 6] atol=RNDATOL
+end

test/runtests.jl

@@ -24,6 +24,7 @@ if GROUP == "All" || GROUP == "AdvancedHMC"
     include("demo.jl")
     include("models.jl")
     include("abstractmcmc.jl")
+    include("mcmcchains.jl")
 
     if CUDA.functional()
         include("cuda.jl")